Update app.py

app.py

@@ -4,49 +4,77 @@ import gradio as gr
 
 MODEL_NAME = "facebook/nllb-200-3.3B"
 
-# Pick device
+# Sélection du device
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load tokenizer + model (consider float16 for GPU to save memory)
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, src_lang="fra_Latn")
+# Chargement du modèle et du tokenizer
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
 model.to(device)
 
-# Language codes used by NLLB
+# Dictionnaire des langues supportées
 LANG_CODES = {
     "fr->ee": ("fra_Latn", "ewe_Latn"),
     "ee->fr": ("ewe_Latn", "fra_Latn"),
 }
 
 def translate(text: str, direction: str, max_length: int = 256) -> str:
-    if not text:
+    if not text.strip():
         return ""
-    src, tgt = LANG_CODES[direction]
 
-    # Tokenize and move to device
-    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True).to(device)
+    src_lang, tgt_lang = LANG_CODES[direction]
 
-    # Force target language token id as BOS so model generates target language
-    forced_bos_token_id = tokenizer.lang_code_to_id[tgt]
+    # Tokenization avec la langue source explicitement définie
+    inputs = tokenizer(
+        text,
+        return_tensors="pt",
+        padding=True,
+        truncation=True,
+        src_lang=src_lang
+    ).to(device)
 
-    generated = model.generate(
+    # On force la génération dans la langue cible
+    forced_bos_token_id = tokenizer.lang_code_to_id[tgt_lang]
+
+    # Génération
+    generated_tokens = model.generate(
         **inputs,
         forced_bos_token_id=forced_bos_token_id,
         max_length=max_length,
-        num_beams=4,
+        num_beams=4
     )
-    return tokenizer.batch_decode(generated, skip_special_tokens=True)[0]
 
-# Gradio UI
+    # Décodage
+    translation = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
+    return translation.strip()
+
+
+# === Interface Gradio ===
 with gr.Blocks() as demo:
-    gr.Markdown("## French ↔ Ewe translator (facebook/nllb-200-3.3B)")
+    gr.Markdown("## 🌍 French ↔ Ewe Translator (facebook/nllb-200-3.3B)")
     with gr.Row():
-        inp = gr.Textbox(lines=6, placeholder="Enter text to translate...")
-        out = gr.Textbox(lines=6, interactive=False)
-    direction = gr.Radio(choices=["fr->ee", "ee->fr"], value="fr->ee", label="Direction")
-    max_len = gr.Slider(minimum=32, maximum=1024, value=256, step=32, label="Max output tokens")
-    translate_btn = gr.Button("Translate")
+        inp = gr.Textbox(lines=6, label="Texte à traduire", placeholder="Entrez le texte ici...")
+        out = gr.Textbox(lines=6, label="Traduction", interactive=False)
+
+    direction = gr.Radio(
+        choices=["fr->ee", "ee->fr"],
+        value="fr->ee",
+        label="Direction de traduction"
+    )
+    max_len = gr.Slider(
+        minimum=32,
+        maximum=1024,
+        value=256,
+        step=32,
+        label="Longueur maximale de sortie"
+    )
 
-    translate_btn.click(lambda t, d, m: translate(t, d, m), inputs=[inp, direction, max_len], outputs=[out])
+    translate_btn = gr.Button("🔁 Traduire")
+
+    translate_btn.click(
+        fn=translate,
+        inputs=[inp, direction, max_len],
+        outputs=[out]
+    )
 
 demo.launch()