maj distil et pipe

app.py

@@ -1,25 +1,20 @@
 import torch
-from transformers import pipeline
 import gradio as gr
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 
 # --------------------------------------------------
-# Chargement du pipeline NLLB
+# Chargement du modèle NLLB
 # --------------------------------------------------
 MODEL_NAME = "facebook/nllb-200-distilled-1.3B"
 
-device = 0 if torch.cuda.is_available() else -1
-print(f"🚀 Chargement du modèle {MODEL_NAME} sur {'GPU' if device == 0 else 'CPU'}...")
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"🚀 Chargement du modèle {MODEL_NAME} sur {device}...")
 
-translator = pipeline(
-    "translation",
-    model=MODEL_NAME,
-    device=device,
-    src_lang="fra_Latn",
-    tgt_lang="ewe_Latn"
-)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME).to(device)
 
 # --------------------------------------------------
-# Dictionnaire de langues (tu peux en ajouter)
+# Dictionnaire de langues
 # --------------------------------------------------
 LANGUAGES = {
     "Français": "fra_Latn",
@@ -41,13 +36,23 @@ def translate(text, src_lang, tgt_lang):
         return "⚠️ Veuillez entrer un texte à traduire."
 
     try:
-        result = translator(
-            text,
-            src_lang=LANGUAGES[src_lang],
-            tgt_lang=LANGUAGES[tgt_lang],
+        # Configuration des langues
+        src_code = LANGUAGES[src_lang]
+        tgt_code = LANGUAGES[tgt_lang]
+        
+        # Préparation de l'entrée
+        tokenizer.src_lang = src_code
+        inputs = tokenizer(text, return_tensors="pt").to(device)
+        
+        # Génération
+        translated_tokens = model.generate(
+            **inputs,
+            forced_bos_token_id=tokenizer.convert_tokens_to_ids(tgt_code),
             max_length=512
         )
-        return result[0]["translation_text"]
+        
+        # Décodage
+        return tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0]
 
     except Exception as e:
         return f"❌ Erreur : {str(e)}"