Update main.py

main.py

@@ -2,7 +2,7 @@ from fastapi import FastAPI
 from pydantic import BaseModel
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import StreamingResponse
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
 import torch
 
 app = FastAPI()
@@ -15,7 +15,7 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# Load model and tokenizer (do this once at startup)
+# Load model and tokenizer
 model_name = "Qwen/Qwen2.5-7B-Instruct-1M"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(
@@ -40,30 +40,31 @@ def generate_response_chunks(prompt: str):
             add_generation_prompt=True,
             return_tensors="pt"
         ).to(model.device)
-        
-        # Generate streamingly
+
+        # Set up streamer
+        streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+
+        # Generate response with streaming
         with torch.no_grad():
-            for outputs in model.generate(
+            model.generate(
                 inputs,
                 max_new_tokens=512,
                 do_sample=True,
                 temperature=0.7,
                 top_p=0.9,
-                streamer=None,  # We'll implement manual streaming
-                stopping_criteria=None
-            ):
-                chunk = outputs[0, inputs.shape[1]:]
-                text = tokenizer.decode(chunk, skip_special_tokens=True)
+                streamer=streamer
+            )
+            # Since TextStreamer handles printing, we yield chunks from the streamer
+            for text in streamer:
                 if text:
                     yield text
-                    
+
     except Exception as e:
-        yield f"Error occurred: {e}"
+        yield f"Error occurred: {str(e)}"
 
 @app.post("/ask")
 async def ask(question: Question):
     return StreamingResponse(
         generate_response_chunks(question.question),
         media_type="text/plain"
-    )
-    
+    )