Update app.py

app.py

@@ -517,6 +517,14 @@ class ModelWrapper:
             token_text = self.tokenizer.decode([token_id])
             response_text += token_text
             yield response_text
+    
+    def generate(self, prompt: str, max_new_tokens: int = 200, 
+                 temperature: float = 0.8, top_k: int = 50, top_p: float = 0.9):
+        """Non-streaming generation (returns full response)"""
+        response = ""
+        for partial_response in self.generate_stream(prompt, max_new_tokens, temperature, top_k, top_p):
+            response = partial_response
+        return response
 
 # ==============================================================================
 # GRADIO INTERFACE
@@ -545,16 +553,15 @@ def chat_fn(message, history, temperature, top_k, top_p, max_tokens):
     # Add current message
     conversation += f"User: {message} Sam:"
     
-    # Generate response
-    response = model.generate(
+    # Stream response token by token
+    for response in model.generate_stream(
         conversation,
         max_new_tokens=max_tokens,
         temperature=temperature,
         top_k=top_k,
         top_p=top_p
-    )
-    
-    return response
+    ):
+        yield response
 
 # Create Gradio interface
 with gr.Blocks(theme=gr.themes.Soft()) as demo: