Update app.py

app.py

@@ -1,7 +1,9 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, Request
+from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from threading import Thread
 import uvicorn
 
 # -----------------------
@@ -24,33 +26,82 @@ torch.set_num_threads(2)
 # -----------------------
 app = FastAPI()
 
+# stop flag (global)
+stop_generation = False
+
 class ChatRequest(BaseModel):
     message: str
 
+
 @app.get("/")
 def home():
-    return {"status": "API running 🚀"}
+    return {"status": "Streaming API running 🚀"}
+
 
+# -----------------------
+# STOP ENDPOINT
+# -----------------------
+@app.post("/stop")
+def stop():
+    global stop_generation
+    stop_generation = True
+    return {"status": "stopping"}
+
+
+# -----------------------
+# STREAMING CHAT
+# -----------------------
 @app.post("/chat")
-def chat(req: ChatRequest):
-    prompt = f"User: {req.message}\nAssistant:"
+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.
+
+User: {req.message}
+Assistant:
+"""
 
     inputs = tokenizer(prompt, return_tensors="pt")
 
-    outputs = model.generate(
-        **inputs,
-        max_new_tokens=80,
-        temperature=0.7,
-        do_sample=True
+    streamer = TextIteratorStreamer(
+        tokenizer,
+        skip_prompt=True,
+        skip_special_tokens=True
     )
 
-    reply = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    reply = reply.split("Assistant:")[-1].strip()
+    def generate():
+        model.generate(
+            **inputs,
+            streamer=streamer,
+            max_new_tokens=40,   # 🔥 short output
+            temperature=0.6,
+            do_sample=True,
+            eos_token_id=tokenizer.eos_token_id
+        )
+
+    thread = Thread(target=generate)
+    thread.start()
+
+    async def stream():
+        global stop_generation
+
+        for token in streamer:
+            if stop_generation:
+                break
+
+            yield token  # 🔥 real-time streaming
+
+    return StreamingResponse(stream(), media_type="text/plain")
 
-    return {"response": reply}
 
 # -----------------------
-# START SERVER DIRECTLY
+# START SERVER
 # -----------------------
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)