feat: Add streaming translation endpoint and implementation.

main.py

@@ -1,6 +1,7 @@
 from fastapi import FastAPI
+from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
-from model_loader import translate
+from model_loader import translate, stream_translate
 
 app = FastAPI(title="NLLB Translation API")
 
@@ -10,3 +11,8 @@ class TranslationRequest(BaseModel):
 @app.post("/translate")
 def translate_text(req: TranslationRequest):
     return {"translated_text": translate(req.text)}
+
+@app.post("/translate/stream")
+def translate_stream(req: TranslationRequest):
+    generator = stream_translate(req.text)
+    return StreamingResponse(generator, media_type="text/plain")

model_loader.py

@@ -1,5 +1,10 @@
 import torch
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import threading
+from transformers import (
+    AutoTokenizer,
+    AutoModelForSeq2SeqLM,
+    TextIteratorStreamer,
+)
 
 MODEL_NAME = "facebook/nllb-200-distilled-600M"
 
@@ -12,15 +17,48 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 model.to(device)
 model.eval()
 
+# -------- normal (non-streaming) ----------
 def translate(text: str):
     inputs = tokenizer(text, return_tensors="pt").to(device)
 
     lang_id = tokenizer.convert_tokens_to_ids("arb_Arab")
 
-    outputs = model.generate(
-        inputs["input_ids"],
+    with torch.no_grad():
+        outputs = model.generate(
+            inputs["input_ids"],
+            forced_bos_token_id=lang_id,
+            max_length=300,
+        )
+
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+
+# -------- streaming ----------
+def stream_translate(text: str):
+    inputs = tokenizer(text, return_tensors="pt").to(device)
+
+    lang_id = tokenizer.convert_tokens_to_ids("arb_Arab")
+
+    streamer = TextIteratorStreamer(
+        tokenizer,
+        skip_special_tokens=True,
+        skip_prompt=True,
+    )
+
+    generation_kwargs = dict(
+        input_ids=inputs["input_ids"],
         forced_bos_token_id=lang_id,
-        max_length=300
+        max_length=300,
+        streamer=streamer,
     )
 
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+    # Run generation in background thread
+    thread = threading.Thread(
+        target=model.generate,
+        kwargs=generation_kwargs,
+    )
+    thread.start()
+
+    # Yield tokens as they are generated
+    for token in streamer:
+        yield token