Update app.py

app.py

@@ -1,5 +1,6 @@
 from fastapi import FastAPI, Request
 from transformers import AutoModelForCausalLM, AutoTokenizer
+from fastapi.responses import StreamingResponse
 import torch
 
 app = FastAPI()
@@ -23,8 +24,17 @@ async def predict(request: Request):
     # Tokenize the input
     inputs = tokenizer(prompt, return_tensors="pt").to("cpu")  # Use "cuda" if GPU is enabled
 
-    # Generate tokens
-    outputs = model.generate(inputs.input_ids, max_length=40, num_return_sequences=1)
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-    return {"response": response}
+    # Generator function to stream tokens
+    def token_generator():
+        outputs = model.generate(
+            inputs.input_ids,
+            max_length=40,
+            do_sample=True,
+            num_return_sequences=1
+        )
+        for token_id in outputs[0]:
+            token = tokenizer.decode(token_id, skip_special_tokens=True)
+            yield f"{token} "
+    
+    # Return StreamingResponse
+    return StreamingResponse(token_generator(), media_type="text/plain")