Update app.py

app.py

@@ -5,7 +5,7 @@ from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 # --------------------------------------------------
 # Chargement du modèle NLLB
 # --------------------------------------------------
-MODEL_NAME = "facebook/nllb-200-3.3B"
+MODEL_NAME = "facebook/nllb-200-distilled-1.3B"
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 print(f"🚀 Chargement du modèle {MODEL_NAME} sur {device}...")
@@ -37,23 +37,31 @@ def translate(text, src_lang, tgt_lang):
 
     try:
         # Configuration des langues
-        src_code = LANGUAGES[src_lang]
-        tgt_code = LANGUAGES[tgt_lang]
+        src_code = LANGUAGES.get(src_lang, "fra_Latn")
+        tgt_code = LANGUAGES.get(tgt_lang, "ewe_Latn")
         
-        # Préparation de l'entrée
+        # Indispensable pour NLLB : définir la langue source dans le tokenizer
         tokenizer.src_lang = src_code
-        inputs = tokenizer(text, return_tensors="pt").to(device)
+        
+        # Tokenization
+        inputs = tokenizer(text, return_tensors="pt", padding=True).to(device)
         
         # Génération
-        translated_tokens = model.generate(
-            **inputs,
-            forced_bos_token_id=tokenizer.convert_tokens_to_ids(tgt_code),
-            max_length=512
-        )
+        with torch.no_grad():
+            translated_tokens = model.generate(
+                **inputs,
+                forced_bos_token_id=tokenizer.lang_code_to_id[tgt_code],
+                max_length=512
+            )
         
         # Décodage
-        return tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0]
-
+        result = tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0]
+        
+        # Nettoyage si le modèle renvoie du texte vide ou des espaces
+        if not result.strip():
+            return "⚠️ Le modèle n'a pas pu générer de traduction. Essayez une phrase plus simple."
+            
+        return result
     except Exception as e:
         return f"❌ Erreur : {str(e)}"