correct issue with small100

app.py

@@ -3,7 +3,9 @@ from transformers import (
     MarianMTModel,
     MarianTokenizer,
     MBartForConditionalGeneration,
-    MBart50TokenizerFast
+    MBart50TokenizerFast,
+    AutoTokenizer,
+    AutoModelForSeq2SeqLM
 )
 import torch
 
@@ -25,16 +27,16 @@ MODEL_MAP = {
     "fr": "Helsinki-NLP/opus-mt-en-fr",
     "hr": "facebook/mbart-large-50-many-to-many-mmt",
     "hu": "Helsinki-NLP/opus-mt-tc-big-en-hu",
-    "is": "mkorada/opus-mt-en-is-finetuned-v4", #Manas's fine-tuned model
+    "is": "mkorada/opus-mt-en-is-finetuned-v4",  # Manas's fine-tuned model
     "it": "Helsinki-NLP/opus-mt-tc-big-en-it",
-    "lb": "alirezamsh/small100",
+    "lb": "alirezamsh/small100",                 # small100
     "lt": "Helsinki-NLP/opus-mt-tc-big-en-lt",
     "lv": "facebook/mbart-large-50-many-to-many-mmt",
-    "me": "Helsinki-NLP/opus-mt-tc-base-en-sh"
+    "me": "Helsinki-NLP/opus-mt-tc-base-en-sh",
     "mk": "Helsinki-NLP/opus-mt-en-mk",
-    "nb": "facebook/mbart-large-50-many-to-many-mmt", #place holder!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
+    "nb": "facebook/mbart-large-50-many-to-many-mmt",
     "nl": "facebook/mbart-large-50-many-to-many-mmt",
-    "no": "Confused404/eng-gmq-finetuned_v2-no", #Alex's fine-tuned model
+    "no": "Confused404/eng-gmq-finetuned_v2-no", # Alex's fine-tuned model
     "pl": "Helsinki-NLP/opus-mt-en-sla",
     "pt": "facebook/mbart-large-50-many-to-many-mmt",
     "ro": "facebook/mbart-large-50-many-to-many-mmt",
@@ -45,29 +47,32 @@ MODEL_MAP = {
     "tr": "Helsinki-NLP/opus-mt-tc-big-en-tr"
 }
 
-
+# Cache loaded models/tokenizers
 MODEL_CACHE = {}
 
-# ✅ Load Hugging Face model (Helsinki or Small100)
 def load_model(model_id: str):
     """
-    Load & cache either:
-      - MBart50 (facebook/mbart-*) 
-      - MarianMT otherwise
+    Load & cache:
+      - facebook/mbart-* via MBart50TokenizerFast & MBartForConditionalGeneration
+      - alirezamsh/small100 via AutoTokenizer & AutoModelForSeq2SeqLM
+      - all others via MarianTokenizer & MarianMTModel
     """
     if model_id not in MODEL_CACHE:
         if model_id.startswith("facebook/mbart"):
             tokenizer = MBart50TokenizerFast.from_pretrained(model_id)
             model     = MBartForConditionalGeneration.from_pretrained(model_id)
+        elif model_id == "alirezamsh/small100":
+            tokenizer = AutoTokenizer.from_pretrained(model_id)
+            model     = AutoModelForSeq2SeqLM.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]
 
+    return MODEL_CACHE[model_id]
 
-# ✅ POST /translate
 @app.post("/translate")
 async def translate(request: Request):
     payload     = await request.json()
@@ -100,18 +105,15 @@ async def translate(request: Request):
     except Exception as e:
         return {"error": f"Translation failed: {e}"}
 
-
-# ✅ GET /languages
 @app.get("/languages")
 def list_languages():
     return {"supported_languages": list(MODEL_MAP.keys())}
 
-# ✅ GET /health
 @app.get("/health")
 def health():
     return {"status": "ok"}
 
-# ✅ Uvicorn startup (required by Hugging Face)
-import uvicorn
+# Uvicorn startup for local testing
 if __name__ == "__main__":
-    uvicorn.run("app:app", host="0.0.0.0", port=7860)
+    import uvicorn
+    uvicorn.run("app:app", host="0.0.0.0", port=7860)