Update chat.py

chat.py

@@ -3,15 +3,11 @@ from model_loader import load_model
 from config import DEFAULT_MODEL, MAX_TOKENS, TEMPERATURE
 import torch
 
-# Conversation memory per model
 conversation_memory = {}
-
-# Load default model
 current_model_name = DEFAULT_MODEL
 tokenizer, model = load_model(current_model_name)
 
 def switch_model(model_name):
-    """Switch to a different model"""
     global tokenizer, model, current_model_name
     tokenizer, model = load_model(model_name)
     current_model_name = model_name
@@ -20,39 +16,35 @@ def switch_model(model_name):
     return f"Switched to model: {model_name}"
 
 def generate_response_stream(prompt, max_length=MAX_TOKENS, temperature=TEMPERATURE):
-    """Stream response token by token for typing effect"""
+    """Streaming response with memory and optimized memory usage"""
     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)
 
-    # Generate with streaming
-    output_ids = model.generate(
+    # Tokenize in small batches to save memory
+    inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
+    outputs = model.generate(
         **inputs,
-        max_length=max_length,
+        max_new_tokens=max_length,
         do_sample=True,
         temperature=temperature,
-        pad_token_id=tokenizer.eos_token_id
+        pad_token_id=tokenizer.eos_token_id,
+        streamer=None
     )[0]
 
-    # Decode token by token
     decoded_text = ""
-    for token_id in output_ids[len(inputs["input_ids"][0]):]:  # Skip prompt tokens
+    for token_id in outputs[len(inputs["input_ids"][0]):]:
         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"""
     global conversation_memory
     conversation_memory[current_model_name] = []
     return "Conversation reset."