Update app.py

Use Streamlit interface

app.py

@@ -1,99 +1,125 @@
-from threading import Thread
-from huggingface_hub import login
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
-
-import gradio as gr
+import streamlit as st
 import os
 
-# Log In
-HF_TOKEN = os.environ.get("HF_TOKEN", None)
-login(token=HF_TOKEN)
-
-TITLE = "<h1><center>Chat with lianghsun/Llama-3.2-Taiwan-3B</center></h1>"
-
-DESCRIPTION = "<h3><center>Visit <a href='https://huggingface.co/lianghsun/Llama-3.2-Taiwan-3B' target='_blank'> the model page</a> for details.</center></h3>"
-
-DEFAULT_SYSTEM = "你是一個說中文的聊天機械人, 使用正體中文回答問題."
-
-CSS = """
-.duplicate-button {
-  margin: auto !important;
-  color: white !important;
-  background: green !important;
-  border-radius: 100vh !important;
-}
-"""
-
-
-tokenizer = AutoTokenizer.from_pretrained("shenzhi-wang/Gemma-2-9B-Chinese-Chat")
-model = AutoModelForCausalLM.from_pretrained("shenzhi-wang/Gemma-2-9B-Chinese-Chat", torch_dtype="auto", device_map="auto")
-
-def stream_chat(message: str, history: list, system: str, temperature: float, max_new_tokens: int):
-    conversation = [{"role": "system", "content": system or DEFAULT_SYSTEM}]
-    for prompt, answer in history:
-        conversation.extend([{"role": "user", "content": prompt}, {"role": "assistant", "content": answer}])
-
-    conversation.append({"role": "user", "content": message})
+from typing import Iterator
+from huggingface_hub import InferenceClient
 
-    input_ids = tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt").to(
-        model.device
-    )
-    streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
-
-    generate_kwargs = dict(
-        input_ids=input_ids,
-        streamer=streamer,
-        max_new_tokens=max_new_tokens,
-        temperature=temperature,
-        do_sample=True,
-    )
-    if temperature == 0:
-        generate_kwargs["do_sample"] = False
-
-    t = Thread(target=model.generate, kwargs=generate_kwargs)
-    t.start()
-
-    output = ""
-    for new_token in streamer:
-        output += new_token
-        yield output
-
-
-chatbot = gr.Chatbot(height=450)
+HF_TOKEN = os.environ.get("HF_TOKEN", None)
 
-with gr.Blocks(css=CSS) as demo:
-    gr.HTML(TITLE)
-    gr.HTML(DESCRIPTION)
-    gr.ChatInterface(
-        fn=stream_chat,
-        chatbot=chatbot,
-        fill_height=True,
-        additional_inputs_accordion=gr.Accordion(label="⚙️ Parameters", open=False, render=False),
-        additional_inputs=[
-            gr.Text(
-                value="",
-                label="System",
-                render=False,
-            ),
-            gr.Slider(
-                minimum=0,
-                maximum=1,
-                step=0.1,
-                value=0.8,
-                label="Temperature",
-                render=False,
-            ),
-            gr.Slider(
-                minimum=128,
-                maximum=4096,
-                step=1,
-                value=1024,
-                label="Max new tokens",
-                render=False,
-            ),
-        ],
+# Configure page settings
+st.set_page_config(
+    page_title="LLM Taiwan Chat",
+    page_icon="💬",
+    layout="centered"
+)
+
+# Initialize session state for chat history and system prompt
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+if "system_prompt" not in st.session_state:
+    st.session_state.system_prompt = "你是一個產自台灣的聊天機械人, 你以台灣本地人的身份, 使用正體中文回答問題."
+if "temperature" not in st.session_state:
+    st.session_state.temperature = 0.2
+if "top_p" not in st.session_state:
+    st.session_state.top_p = 0.95
+
+## model="lianghsun/Llama-3.2-Taiwan-3B" to meta-llama/Llama-3.2-3B-Instruct
+
+def stream_chat(prompt: str) -> Iterator[str]:
+    """Stream chat responses from the LLM API"""
+
+    client = InferenceClient(model="meta-llama/Llama-3.2-3B-Instruct", timeout=30, token=HF_TOKEN)
+
+    
+    messages = []
+    if st.session_state.system_prompt:
+        messages.append({"role": "system", "content": st.session_state.system_prompt})
+    messages.extend(st.session_state.messages)
+    
+    stream = client.chat.completions.create(
+        messages=messages,
+        model="meta-llama/Llama-3.2-3B-Instruct",
+        stream=True,
+        temperature=st.session_state.temperature,
+        top_p=st.session_state.top_p
     )
-
+    
+    for chunk in stream:
+        if chunk.choices[0].delta.content is not None:
+            yield chunk.choices[0].delta.content
+
+def clear_chat_history():
+    """Clear all chat messages and reset system prompt"""
+    st.session_state.messages = []
+    st.session_state.system_prompt = ""
+
+def main():
+    st.title("💬 LLM Taiwan Chat")
+    
+    # Add a clear chat button with custom styling
+    col1, col2 = st.columns([6, 1])
+    with col2:
+        if st.button("🗑️", type="secondary", use_container_width=True):
+            clear_chat_history()
+            st.rerun()
+    
+    # Advanced options in expander
+    with st.expander("進階選項 ⚙️", expanded=False):
+        # System prompt input
+        system_prompt = st.text_area(
+            "System Prompt 設定:",
+            value=st.session_state.system_prompt,
+            help="設定 system prompt 來定義 AI 助理的行為和角色。開始對話後將無法修改。",
+            height=100,
+            disabled=len(st.session_state.messages) > 0  # 當有對話時設為唯讀
+        )
+        if not st.session_state.messages and system_prompt != st.session_state.system_prompt:
+            st.session_state.system_prompt = system_prompt
+        
+        st.session_state.temperature = st.slider(
+            "Temperature",
+            min_value=0.0,
+            max_value=2.0,
+            value=st.session_state.temperature,
+            step=0.1,
+            help="較高的值會使輸出更加隨機，較低的值會使其更加集中和確定。"
+        )
+        st.session_state.top_p = st.slider(
+            "Top P",
+            min_value=0.1,
+            max_value=1.0,
+            value=st.session_state.top_p,
+            step=0.05,
+            help="控制模型輸出的多樣性，較低的值會使輸出更加保守。"
+        )
+    
+    # Display chat messages
+    for message in st.session_state.messages:
+        with st.chat_message(message["role"]):
+            st.write(message["content"])
+    
+    # Chat input
+    if prompt := st.chat_input("輸入您的訊息..."):
+        # Add user message to chat history
+        st.session_state.messages.append({"role": "user", "content": prompt})
+        
+        # Display user message
+        with st.chat_message("user"):
+            st.write(prompt)
+        
+        # Display assistant response with streaming
+        with st.chat_message("assistant"):
+            response_placeholder = st.empty()
+            full_response = ""
+            
+            # Stream the response
+            for response_chunk in stream_chat(prompt):
+                full_response += response_chunk
+                response_placeholder.markdown(full_response + "▌")
+            response_placeholder.markdown(full_response)
+        
+        # Add assistant response to chat history
+        st.session_state.messages.append({"role": "assistant", "content": full_response})
 
 if __name__ == "__main__":
-    demo.launch()
+    main()