Update app.py

app.py

@@ -1,6 +1,6 @@
 import spaces
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
 
 # Load the model and tokenizer locally
 model_name = "kz919/QwQ-0.5B-Distilled-SFT"
@@ -11,10 +11,21 @@ model = AutoModelForCausalLM.from_pretrained(model_name).to("cuda")
 @spaces.GPU
 def respond(message, history: list[tuple[str, str]], system_message, max_tokens, temperature, top_p):
     # Prepare the prompt by combining history and system messages
-    prompt = system_message + "\n"
+    msg = [
+        {"role": "system", "content": system_message}
+    ]
     for user_input, assistant_response in history:
-        prompt += f"User: {user_input}\nAssistant: {assistant_response}\n"
-    prompt += f"User: {message}\nAssistant:"
+        msg.extend(
+            {"role": "user", "content": user_input},
+            {"role": "assistant", "content": assistant_response}
+        )
+    msg.append({"role": "user", "content": message})
+
+    prompt = tokenizer.apply_chat_template(
+        msg,
+        tokenize=False,
+        add_generation_prompt=True
+    )
 
     # Tokenize the input prompt
     inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
@@ -26,11 +37,9 @@ def respond(message, history: list[tuple[str, str]], system_message, max_tokens,
         temperature=temperature,
         top_p=top_p,
         pad_token_id=tokenizer.eos_token_id,
+        streamer = TextStreamer(tokenizer)
     )
-
-    # Decode the generated tokens and yield the response
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    yield response.split("Assistant:")[-1].strip()
+    yield outputs
 
 
 # Create the Gradio interface