Update main.py

main.py

@@ -1,22 +1,23 @@
 import os
+from fastapi import FastAPI
+from pydantic import BaseModel
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import torch
+
 # 2️⃣ Optional: force cache to writable directory
 os.environ["HF_HOME"] = "/tmp/hf"
 os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf"
 os.environ["HF_DATASETS_CACHE"] = "/tmp/hf"
 os.makedirs("/tmp/hf", exist_ok=True)
 
-
-from fastapi import FastAPI
-from pydantic import BaseModel
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-import torch
-
 app = FastAPI()
 
 # Load model once on startup
-MODEL_NAME = "facebook/nllb-200-1.3B"  # 3B version
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME,use_fast=False)
-model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float16).to("cuda" if torch.cuda.is_available() else "cpu")
+MODEL_NAME = "facebook/nllb-200-1.3B"
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=False)
+model = AutoModelForSeq2SeqLM.from_pretrained(
+    MODEL_NAME, torch_dtype=torch.float16
+).to("cuda" if torch.cuda.is_available() else "cpu")
 
 class TranslationRequest(BaseModel):
     text: str
@@ -30,9 +31,12 @@ def translate(req: TranslationRequest):
         return_tensors="pt",
     ).to(model.device)
 
+    # ✅ add "__" around lang codes
+    tgt_lang = "__" + req.tgt_lang + "__"
+
     outputs = model.generate(
         **inputs,
-        forced_bos_token_id=tokenizer.lang_code_to_id[req.tgt_lang],
+        forced_bos_token_id=tokenizer.lang_code_to_id[tgt_lang],
         max_length=512
     )
     translation = tokenizer.decode(outputs[0], skip_special_tokens=True)