Add NLLB-200 Translation API

app.py

@@ -1,7 +1,62 @@
 import gradio as gr
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+import torch
 
-def greet(name):
-    return "Hello " + name + "!!"
+# Load NLLB-200 (distilled for speed)
+MODEL_NAME = "facebook/nllb-200-distilled-600M"
+print(f"Loading {MODEL_NAME}...")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model = model.to(device)
+print(f"Model loaded on {device}")
+
+# Language codes
+LANGS = {
+    "English": "eng_Latn",
+    "French": "fra_Latn",
+    "Arabic": "arb_Arab",
+    "Moroccan Arabic": "ary_Arab",
+    "Spanish": "spa_Latn",
+    "German": "deu_Latn",
+    "Italian": "ita_Latn",
+    "Portuguese": "por_Latn",
+    "Chinese": "zho_Hans",
+    "Japanese": "jpn_Jpan",
+    "Korean": "kor_Hang",
+    "Russian": "rus_Cyrl",
+    "Turkish": "tur_Latn",
+    "Dutch": "nld_Latn",
+    "Hindi": "hin_Deva",
+}
+
+def translate(text, src_lang, tgt_lang):
+    if not text.strip():
+        return ""
+    
+    src_code = LANGS.get(src_lang, "eng_Latn")
+    tgt_code = LANGS.get(tgt_lang, "fra_Latn")
+    
+    tokenizer.src_lang = src_code
+    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512).to(device)
+    
+    forced_bos_token_id = tokenizer.convert_tokens_to_ids(tgt_code)
+    
+    with torch.no_grad():
+        outputs = model.generate(**inputs, forced_bos_token_id=forced_bos_token_id, max_length=512, num_beams=5)
+    
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+demo = gr.Interface(
+    fn=translate,
+    inputs=[
+        gr.Textbox(label="Text to translate", lines=3),
+        gr.Dropdown(list(LANGS.keys()), label="Source Language", value="English"),
+        gr.Dropdown(list(LANGS.keys()), label="Target Language", value="French"),
+    ],
+    outputs=gr.Textbox(label="Translation", lines=3),
+    title="NLLB-200 Translation API",
+    description="200 languages including Moroccan Arabic!",
+)
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
 demo.launch()

requirements.txt

@@ -0,0 +1,5 @@
+transformers>=4.30.0
+torch>=2.0.0
+sentencepiece
+protobuf
+gradio