Update app.py

app.py

@@ -1,80 +1,110 @@
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import gradio as gr
 
+# --------------------------------------------------
+# Chargement du modèle NLLB-200-3.3B
+# --------------------------------------------------
 MODEL_NAME = "facebook/nllb-200-3.3B"
 
-# Sélection du device
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# Chargement du modèle et du tokenizer
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+print("🚀 Chargement du modèle :", MODEL_NAME)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, src_lang="fra_Latn")
 model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
+device = "cuda" if torch.cuda.is_available() else "cpu"
 model.to(device)
 
-# Dictionnaire des langues supportées
-LANG_CODES = {
-    "fr->ee": ("fra_Latn", "ewe_Latn"),
-    "ee->fr": ("ewe_Latn", "fra_Latn"),
+# --------------------------------------------------
+# Langues supportées (tu peux en ajouter d’autres)
+# --------------------------------------------------
+LANGUAGES = {
+    "Français": "fra_Latn",
+    "Anglais": "eng_Latn",
+    "Ewe": "ewe_Latn",
+    "Fon": "fon_Latn",
+    "Espagnol": "spa_Latn",
+    "Allemand": "deu_Latn",
+    "Swahili": "swh_Latn",
+    "Lingala": "lin_Latn",
+    "Portugais": "por_Latn"
 }
 
-def translate(text: str, direction: str, max_length: int = 256) -> str:
+# --------------------------------------------------
+# Fonction de traduction
+# --------------------------------------------------
+def translate(text, src_lang, tgt_lang):
     if not text.strip():
-        return ""
-
-    src_lang, tgt_lang = LANG_CODES[direction]
-
-    # Tokenization avec la langue source explicitement définie
-    inputs = tokenizer(
-        text,
-        return_tensors="pt",
-        padding=True,
-        truncation=True,
-        src_lang=src_lang
-    ).to(device)
-
-    # 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
+        return "⚠️ Veuillez entrer un texte à traduire."
+
+    try:
+        tokenizer.src_lang = LANGUAGES[src_lang]
+        inputs = tokenizer(text, return_tensors="pt").to(device)
+
+        forced_bos_token_id = tokenizer.lang_code_to_id[LANGUAGES[tgt_lang]]
+        generated_tokens = model.generate(
+            **inputs,
+            forced_bos_token_id=forced_bos_token_id,
+            max_length=512,
+            num_beams=4,
+            early_stopping=True
+        )
+
+        translation = tokenizer.decode(generated_tokens[0], skip_special_tokens=True)
+        return translation
+
+    except Exception as e:
+        return f"❌ Erreur : {str(e)}"
+
+# --------------------------------------------------
+# Interface Gradio (version stylée pour Hugging Face)
+# --------------------------------------------------
+with gr.Blocks(title="🌍 NLLB-200 Traduction Multilingue") as demo:
+    gr.Markdown(
+        """
+        <div style="text-align: center;">
+            <h1>🌐 Traducteur IA Multilingue</h1>
+            <p style="font-size: 18px;">
+                Propulsé par <b>Meta NLLB-200 (3.3B)</b> — Traduction de plus de <b>200 langues</b>, 
+                dont <b>l’Ewe</b>, <b>le Fon</b>, et plusieurs langues africaines 🌍
+            </p>
+        </div>
+        """
     )
 
-    # 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)")
     with gr.Row():
-        inp = gr.Textbox(lines=6, label="Texte à traduire", placeholder="Entrez le texte ici...")
-        out = gr.Textbox(lines=6, label="Traduction", interactive=False)
+        src_lang = gr.Dropdown(
+            choices=list(LANGUAGES.keys()),
+            value="Français",
+            label="Langue source 🌍"
+        )
+        tgt_lang = gr.Dropdown(
+            choices=list(LANGUAGES.keys()),
+            value="Ewe",
+            label="Langue cible 🌍"
+        )
 
-    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"
-    )
+    with gr.Row():
+        text_input = gr.Textbox(
+            placeholder="Entre ton texte ici...",
+            lines=6,
+            label="Texte à traduire"
+        )
+        text_output = gr.Textbox(
+            placeholder="Résultat de la traduction...",
+            lines=6,
+            label="Traduction"
+        )
 
     translate_btn = gr.Button("🔁 Traduire")
-
-    translate_btn.click(
-        fn=translate,
-        inputs=[inp, direction, max_len],
-        outputs=[out]
+    translate_btn.click(translate, [text_input, src_lang, tgt_lang], text_output)
+
+    gr.Markdown(
+        """
+        ---
+        <div style="text-align: center; font-size: 14px;">
+            🧠 Propulsé par <a href="https://huggingface.co/facebook/nllb-200-3.3B" target="_blank">Meta NLLB-200 (3.3B)</a><br>
+            Créé avec ❤️ par <b>Joël Elisée Adzonya</b>
+        </div>
+        """
     )
 
 demo.launch()