Update app.py

app.py

@@ -57,15 +57,12 @@ async def chat(req: ChatRequest):
     global stop_generation
     stop_generation = False
 
-    # 🔥 FORCE SHORT ANSWERS
-    prompt = f"""
-You are a concise assistant.
-Answer VERY SHORT (1-2 lines max).
-No long explanations.
+    prompt = f"""You are a helpful assistant.
+Reply in ONE short sentence only.
+Finish your answer and stop.
 
 User: {req.message}
-Assistant:
-"""
+Assistant:"""
 
     inputs = tokenizer(prompt, return_tensors="pt")
 
@@ -75,30 +72,48 @@ Assistant:
         skip_special_tokens=True
     )
 
+    # ✅ Define stop tokens
+    stop_tokens = ["User:", "\n\n"]
+
+    stop_token_ids = [
+        tokenizer.encode(token, add_special_tokens=False)
+        for token in stop_tokens
+    ]
+
     def generate():
         model.generate(
             **inputs,
             streamer=streamer,
-            max_new_tokens=40,   # 🔥 short output
-            temperature=0.6,
+            max_new_tokens=100,
+            temperature=0.5,
             do_sample=True,
-            eos_token_id=tokenizer.eos_token_id
+            eos_token_id=tokenizer.eos_token_id,
+            pad_token_id=tokenizer.eos_token_id
         )
 
-    thread = Thread(target=generate)
-    thread.start()
+    Thread(target=generate).start()
 
     async def stream():
         global stop_generation
 
+        buffer_ids = []
+
         for token in streamer:
             if stop_generation:
                 break
 
-            yield token  # 🔥 real-time streaming
+            # convert token → ids
+            token_ids = tokenizer.encode(token, add_special_tokens=False)
+            buffer_ids.extend(token_ids)
 
-    return StreamingResponse(stream(), media_type="text/plain")
+            # 🔥 STOP TOKEN CHECK (clean, not hacky)
+            for stop_seq in stop_token_ids:
+                if buffer_ids[-len(stop_seq):] == stop_seq:
+                    return
+
+            yield token
 
+    return StreamingResponse(stream(), media_type="text/plain")
 
 # -----------------------
 # START SERVER