Update app.py

app.py

@@ -39,20 +39,25 @@ def generate_response(model_id, conversation, user_message, max_length=512, temp
     load_time = time.time() - start_time
     print(f"✅ Model loaded in {load_time:.2f}s")
     
-    # Build conversation history
-    conversation_history = []
+    # Build messages in proper chat format
+    messages = []
+    
+    # Add conversation history
     for user_msg, assistant_msg in conversation:
         if user_msg:
-            conversation_history.append(f"User: {user_msg}")
+            messages.append({"role": "user", "content": user_msg})
         if assistant_msg:
-            conversation_history.append(f"Athena: {assistant_msg}")
+            messages.append({"role": "assistant", "content": assistant_msg})
     
     # Add current user message
-    conversation_history.append(f"User: {user_message}")
-    conversation_history.append("Athena:")
+    messages.append({"role": "user", "content": user_message})
     
-    # Create prompt
-    prompt = "\n".join(conversation_history)
+    # Apply chat template
+    prompt = tokenizer.apply_chat_template(
+        messages, 
+        tokenize=False, 
+        add_generation_prompt=True
+    )
     
     # Tokenize and move to GPU
     inputs = tokenizer(prompt, return_tensors="pt")