Update chat.py

chat.py

@@ -1,26 +1,35 @@
 # chat.py
-import torch
 from model_loader import load_model
-from config import MAX_TOKENS, TEMPERATURE
+from config import DEFAULT_MODEL, MAX_TOKENS, TEMPERATURE
 
-tokenizer, model = load_model()
+# Conversation memory per model
+conversation_memory = {}
 
-# Conversation memory: stores previous exchanges
-conversation_history = []
+# 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
+    if model_name not in conversation_memory:
+        conversation_memory[model_name] = []
+    return f"Switched to model: {model_name}"
 
 def generate_response(prompt, max_length=MAX_TOKENS, temperature=TEMPERATURE):
-    """
-    Generate a response from the model with memory of previous conversation.
-    """
-    global conversation_history
+    """Generate response with conversation memory per model"""
+    global conversation_memory
+    
+    if current_model_name not in conversation_memory:
+        conversation_memory[current_model_name] = []
     
-    # Add user prompt to conversation
-    conversation_history.append(f"User: {prompt}")
+    history = conversation_memory[current_model_name]
     
-    # Combine conversation history for context
-    full_prompt = "\n".join(conversation_history) + "\nAI:"
+    history.append(f"User: {prompt}")
+    full_prompt = "\n".join(history) + "\nAI:"
     
-    # Tokenize and generate
     inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
     outputs = model.generate(
         **inputs,
@@ -30,19 +39,14 @@ def generate_response(prompt, max_length=MAX_TOKENS, temperature=TEMPERATURE):
         pad_token_id=tokenizer.eos_token_id
     )
     
-    # Decode response
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    
-    # Extract AI's latest message
     response = response.split("AI:")[-1].strip()
     
-    # Add AI response to conversation history
-    conversation_history.append(f"AI: {response}")
-    
+    history.append(f"AI: {response}")
     return response
 
 def reset_conversation():
-    """Reset conversation memory"""
-    global conversation_history
-    conversation_history = []
+    """Reset memory for current model"""
+    global conversation_memory
+    conversation_memory[current_model_name] = []
     return "Conversation reset."