Update app.py

app.py

@@ -462,14 +462,14 @@ class ModelWrapper:
         print(f"✅ Model loaded: {self.d_model}d × {self.n_layers}L × {self.n_heads}H")
     
     def generate_stream(self, prompt: str, max_new_tokens: int = 200, 
-                        temperature: float = 0.8, top_k: int = 50, top_p: float = 0.9):
+                    temperature: float = 0.8, top_k: int = 50, top_p: float = 0.9):
         """Generator that yields tokens one at a time for streaming"""
-        # Format prompt correctly (NO newline between User: and Sam:)
-        if not prompt.startswith("User:"):
-            prompt = f"User: {prompt} Sam:"
+        # Format prompt in ChatML format
+        if not prompt.startswith("<|im_start|>"):
+            prompt = f"<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n"
         else:
-            if " Sam:" not in prompt:
-                prompt = prompt + " Sam:"
+            if "<|im_start|>assistant" not in prompt:
+                prompt = prompt + "<|im_start|>assistant\n"
         
         # Tokenize
         encoding = self.tokenizer.encode(prompt)
@@ -509,8 +509,11 @@ class ModelWrapper:
             # Decode the new token
             token_id = int(next_token[0, 0])
             
-            # Stop on EOS
-            if token_id == self.tokenizer.token_to_id("<|endoftext|>"):
+            # Stop on EOS or end tokens
+            if token_id in [
+                self.tokenizer.token_to_id("<|endoftext|>"),
+                self.tokenizer.token_to_id("<|im_end|>")
+            ]:
                 break
             
             # Decode and yield the token
@@ -518,6 +521,7 @@ class ModelWrapper:
             response_text += token_text
             yield response_text
     
+    
     def generate(self, prompt: str, max_new_tokens: int = 200, 
                  temperature: float = 0.8, top_k: int = 50, top_p: float = 0.9):
         """Non-streaming generation (returns full response)"""
@@ -544,14 +548,15 @@ print(f"✅ Model downloaded to: {model_path}")
 # Load model
 model = ModelWrapper(model_path)
 
+
 def chat_fn(message, history, temperature, top_k, top_p, max_tokens):
-    # Build conversation context with proper template
+    # Build conversation context in ChatML format
     conversation = ""
     for user_msg, bot_msg in history:
-        conversation += f"User: {user_msg} Sam: {bot_msg} "
+        conversation += f"<|im_start|>user\n{user_msg}<|im_end|>\n<|im_start|>assistant\n{bot_msg}<|im_end|>\n"
     
     # Add current message
-    conversation += f"User: {message} Sam:"
+    conversation += f"<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n"
     
     # Stream response token by token
     partial_response = ""