Update app.py

app.py

@@ -1,25 +1,23 @@
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
 MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained(
+    MODEL_NAME, trust_remote_code=True
+)
+
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_NAME,
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
     device_map="auto",
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
     trust_remote_code=True
 )
 
 def chat(user_input, history):
-    messages = []
-
-    for h in history:
-        messages.append({"role": "user", "content": h[0]})
-        messages.append({"role": "assistant", "content": h[1]})
-
-    messages.append({"role": "user", "content": user_input})
+    # history is already in messages format
+    messages = history + [{"role": "user", "content": user_input}]
 
     prompt = tokenizer.apply_chat_template(
         messages,
@@ -32,9 +30,9 @@ def chat(user_input, history):
     outputs = model.generate(
         **inputs,
         max_new_tokens=512,
-        do_sample=True,
         temperature=0.7,
-        top_p=0.9
+        top_p=0.9,
+        do_sample=True
     )
 
     response = tokenizer.decode(
@@ -42,14 +40,15 @@ def chat(user_input, history):
         skip_special_tokens=True
     )
 
-    history.append((user_input, response))
-    return history, ""
+    messages.append({"role": "assistant", "content": response})
+
+    return messages, ""
 
 with gr.Blocks() as demo:
     gr.Markdown("## 🤖 Qwen Chatbot")
 
-    chatbot = gr.Chatbot()
-    msg = gr.Textbox(label="Your message")
+    chatbot = gr.Chatbot(type="messages")
+    msg = gr.Textbox(label="Your message", autofocus=True)
     clear = gr.Button("Clear")
 
     msg.submit(chat, [msg, chatbot], [chatbot, msg])