add fixes for mbart

app.py

@@ -1,8 +1,13 @@
 from fastapi import FastAPI, Request
-from transformers import MarianMTModel, MarianTokenizer
+from transformers import (
+    MarianMTModel,
+    MarianTokenizer,
+    MBartForConditionalGeneration,
+    MBart50TokenizerFast
+)
 import torch
 
-# import chunking service
+# import your chunking helpers
 from chunking import get_max_word_length, chunk_text
 
 app = FastAPI()
@@ -43,19 +48,30 @@ MODEL_MAP = {
 MODEL_CACHE = {}
 
 # ✅ Load Hugging Face model (Helsinki or Small100)
-def load_model(model_id):
+def load_model(model_id: str):
+    """
+    Load & cache either:
+      - MBart50 (facebook/mbart-*) 
+      - MarianMT otherwise
+    """
     if model_id not in MODEL_CACHE:
-        tokenizer = MarianTokenizer.from_pretrained(model_id)
-        model = MarianMTModel.from_pretrained(model_id).to("cpu")
+        if model_id.startswith("facebook/mbart"):
+            tokenizer = MBart50TokenizerFast.from_pretrained(model_id)
+            model     = MBartForConditionalGeneration.from_pretrained(model_id)
+        else:
+            tokenizer = MarianTokenizer.from_pretrained(model_id)
+            model     = MarianMTModel.from_pretrained(model_id)
+        model.to("cpu")
         MODEL_CACHE[model_id] = (tokenizer, model)
     return MODEL_CACHE[model_id]
 
+
 # ✅ POST /translate
 @app.post("/translate")
 async def translate(request: Request):
-    data = await request.json()
-    text = data.get("text")
-    target_lang = data.get("target_lang")
+    payload     = await request.json()
+    text        = payload.get("text")
+    target_lang = payload.get("target_lang")
 
     if not text or not target_lang:
         return {"error": "Missing 'text' or 'target_lang'"}
@@ -64,31 +80,25 @@ async def translate(request: Request):
     if not model_id:
         return {"error": f"No model found for target language '{target_lang}'"}
 
-    # Facebook/mbart placeholder check
-    if model_id.startswith("facebook/"):
-        return {"translation": f"[{target_lang}] uses model '{model_id}', which is not supported in this Space yet."}
-
     try:
-        # 1. figure out your safe word limit for this language
+        # chunk to safe length
         safe_limit = get_max_word_length([target_lang])
+        chunks     = chunk_text(text, safe_limit)
 
-        # 2. break the input up into chunks
-        chunks = chunk_text(text, safe_limit)
-
-        # 3. translate each chunk and collect results
         tokenizer, model = load_model(model_id)
         full_translation = []
+
         for chunk in chunks:
-            inputs = tokenizer(chunk, return_tensors="pt", padding=True, truncation=True).to(model.device)
+            inputs  = tokenizer(chunk, return_tensors="pt", padding=True, truncation=True)
+            inputs  = {k: v.to(model.device) for k, v in inputs.items()}
             outputs = model.generate(**inputs, num_beams=5, length_penalty=1.2, early_stopping=True)
             full_translation.append(tokenizer.decode(outputs[0], skip_special_tokens=True))
 
-        # 4. re-join the translated pieces
-        joined = " ".join(full_translation)
-        return {"translation": joined}
+        return {"translation": " ".join(full_translation)}
 
     except Exception as e:
-        return {"error": f"Translation failed: {str(e)}"}
+        return {"error": f"Translation failed: {e}"}
+
 
 # ✅ GET /languages
 @app.get("/languages")