Update chat.py

chat.py

@@ -1,6 +1,7 @@
 # chat.py
 from model_loader import load_model
 from config import DEFAULT_MODEL, MAX_TOKENS, TEMPERATURE
+import torch
 
 # Conversation memory per model
 conversation_memory = {}
@@ -18,32 +19,37 @@ def switch_model(model_name):
         conversation_memory[model_name] = []
     return f"Switched to model: {model_name}"
 
-def generate_response(prompt, max_length=MAX_TOKENS, temperature=TEMPERATURE):
-    """Generate response with conversation memory per model"""
+def generate_response_stream(prompt, max_length=MAX_TOKENS, temperature=TEMPERATURE):
+    """Stream response token by token for typing effect"""
     global conversation_memory
     
     if current_model_name not in conversation_memory:
         conversation_memory[current_model_name] = []
     
     history = conversation_memory[current_model_name]
-    
     history.append(f"User: {prompt}")
     full_prompt = "\n".join(history) + "\nAI:"
     
     inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
-    outputs = model.generate(
+
+    # Generate with streaming
+    output_ids = model.generate(
         **inputs,
         max_length=max_length,
         do_sample=True,
         temperature=temperature,
         pad_token_id=tokenizer.eos_token_id
-    )
-    
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    response = response.split("AI:")[-1].strip()
-    
-    history.append(f"AI: {response}")
-    return response
+    )[0]
+
+    # Decode token by token
+    decoded_text = ""
+    for token_id in output_ids[len(inputs["input_ids"][0]):]:  # Skip prompt tokens
+        decoded_token = tokenizer.decode(token_id)
+        decoded_text += decoded_token
+        yield decoded_text
+
+    # Save to conversation memory
+    conversation_memory[current_model_name].append(f"AI: {decoded_text}")
 
 def reset_conversation():
     """Reset memory for current model"""