Update chat.py

chat.py

@@ -5,11 +5,23 @@ from config import MAX_TOKENS, TEMPERATURE
 
 tokenizer, model = load_model()
 
+# Conversation memory: stores previous exchanges
+conversation_history = []
+
 def generate_response(prompt, max_length=MAX_TOKENS, temperature=TEMPERATURE):
     """
-    Generate a response from the model
+    Generate a response from the model with memory of previous conversation.
     """
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    global conversation_history
+    
+    # Add user prompt to conversation
+    conversation_history.append(f"User: {prompt}")
+    
+    # Combine conversation history for context
+    full_prompt = "\n".join(conversation_history) + "\nAI:"
+    
+    # Tokenize and generate
+    inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
     outputs = model.generate(
         **inputs,
         max_length=max_length,
@@ -17,5 +29,20 @@ def generate_response(prompt, max_length=MAX_TOKENS, temperature=TEMPERATURE):
         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}")
+    
     return response
+
+def reset_conversation():
+    """Reset conversation memory"""
+    global conversation_history
+    conversation_history = []
+    return "Conversation reset."