Update app.py

app.py

@@ -11,17 +11,22 @@ model = AutoModelForCausalLM.from_pretrained(checkpoint).to(device)
 def predict(message, history, temperature, top_p):
     history.append({"role": "user", "content": message})
     input_text = tokenizer.apply_chat_template(history, tokenize=False, add_generation_prompt=True)
-    inputs = tokenizer.encode(input_text, return_tensors="pt").to(device)  
-    outputs = model.generate(
+    inputs = tokenizer.encode(input_text, return_tensors="pt").to(device)
+    
+    streamer = gr.TelegramStreamer()  # Use Gradio's built-in streamer
+    
+    # Generate with streaming
+    model.generate(
         inputs, 
         max_new_tokens=1024, 
         temperature=float(temperature), 
         top_p=float(top_p), 
-        do_sample=True
+        do_sample=True,
+        streamer=streamer
     )
-    decoded = tokenizer.decode(outputs[0])
-    response = decoded.split("<|start_header_id|>assistant<|end_header_id|>\n\n")[-1]
-    return response
+    
+    # The streamer will handle returning the tokens
+    return streamer
 
 with gr.Blocks() as demo:
     chatbot = gr.ChatInterface(