Update app.py

app.py

@@ -205,16 +205,35 @@ async def perform_core_search(raw_query: str):
     return [{"match_score": round(s, 2), **doc} for s, doc in scored_results[:100]], prediction, best_score
 
 def _generate_text(messages, max_tokens: int):
-    prompt = "\n".join([f"{m['role'].capitalize()}: {m['content']}" for m in messages]) + "\nAssistant:"
+    tokenizer = ai_generator.tokenizer
     
-    # 🔥 GREEDY DECODING (do_sample=False) FOR 3X FASTER AI RESPONSES
+    # 1. Use the model's Native Chat Template (Prevents random autocomplete/echoing)
+    try:
+        prompt = tokenizer.apply_chat_template(
+            messages, 
+            tokenize=False, 
+            add_generation_prompt=True
+        )
+    except Exception:
+        # Fallback if model doesn't have a strict template
+        prompt = "\n".join([f"{m['role'].capitalize()}: {m['content']}" for m in messages]) + "\nAssistant: "
+
+    # 2. return_full_text=False ensures the prompt isn't returned in the output!
     result = ai_generator(
         prompt, 
         max_new_tokens=max_tokens, 
         do_sample=False, 
-        truncation=True
+        truncation=True,
+        return_full_text=False, # 🔥 THIS FIXES THE REPEATING QUERY ISSUE
+        pad_token_id=tokenizer.eos_token_id if tokenizer.eos_token_id is not None else 0
     )
-    text_out = result[0]['generated_text'].split("Assistant:")[-1].strip()
+    
+    text_out = result[0]['generated_text'].strip()
+    
+    # 3. Extra safety cleanup for small AI models
+    if text_out.lower().startswith("assistant:"):
+        text_out = text_out[10:].strip()
+        
     return text_out
 
 async def run_ai_task(messages, max_tokens: int, query_key: str):