Update app.py

app.py

@@ -1,9 +1,7 @@
-from fastapi import FastAPI, Request
-from fastapi.responses import StreamingResponse
+from fastapi import FastAPI
 from pydantic import BaseModel
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
-from threading import Thread
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import uvicorn
 
 # -----------------------
@@ -12,6 +10,7 @@ import uvicorn
 MODEL_ID = "microsoft/phi-2"
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
     device_map="cpu",
@@ -26,36 +25,20 @@ torch.set_num_threads(2)
 # -----------------------
 app = FastAPI()
 
-# stop flag (global)
-stop_generation = False
-
 class ChatRequest(BaseModel):
     message: str
 
 
 @app.get("/")
 def home():
-    return {"status": "Streaming API running 🚀"}
-
-
-# -----------------------
-# STOP ENDPOINT
-# -----------------------
-@app.post("/stop")
-def stop():
-    global stop_generation
-    stop_generation = True
-    return {"status": "stopping"}
+    return {"status": "API running 🚀"}
 
 
 # -----------------------
-# STREAMING CHAT
+# CHAT (NO STREAMING)
 # -----------------------
 @app.post("/chat")
-async def chat(req: ChatRequest):
-
-    global stop_generation
-    stop_generation = False
+def chat(req: ChatRequest):
 
     prompt = f"""You are a concise assistant.
 Return plain text only.
@@ -70,58 +53,34 @@ Assistant:"""
 
     inputs = tokenizer(prompt, return_tensors="pt")
 
-    streamer = TextIteratorStreamer(
-        tokenizer,
-        skip_prompt=True,
-        skip_special_tokens=True
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=80,
+        temperature=0.5,
+        do_sample=True,
+        eos_token_id=tokenizer.eos_token_id,
+        pad_token_id=tokenizer.eos_token_id
     )
 
-    # ✅ 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=100,
-            temperature=0.5,
-            do_sample=True,
-            eos_token_id=tokenizer.eos_token_id,
-            pad_token_id=tokenizer.eos_token_id
-        )
+    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    Thread(target=generate).start()
+    # 🔥 CLEAN OUTPUT
+    if "Assistant:" in text:
+        text = text.split("Assistant:")[-1]
 
-    async def stream():
-        global stop_generation
+    if "User:" in text:
+        text = text.split("User:")[0]
 
-        buffer_ids = []
+    text = text.strip()
 
-        for token in streamer:
-            if stop_generation:
-                break
+    # remove unwanted formatting
+    text = text.replace("\n", " ")
+    text = text.replace("  ", " ")
 
-            # convert token → ids
-            token_ids = tokenizer.encode(token, add_special_tokens=False)
-            buffer_ids.extend(token_ids)
+    return {
+        "response": text
+    }
 
-            # 🔥 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",
-        headers={"Transfer-Encoding": "identity"}
-    )
 
 # -----------------------
 # START SERVER