Update goai_helpers/goai_traduction.py

goai_helpers/goai_traduction.py

@@ -13,30 +13,39 @@ login(token=auth_token)
 @spaces.GPU
 def goai_traduction(text, src_lang, tgt_lang):
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
+    
     if src_lang == "fra_Latn" and tgt_lang == "mos_Latn":
         model_id = "ArissBandoss/nllb-200-distilled-600M-finetuned-fr-to-mos-V4"
-        
     elif src_lang == "mos_Latn" and tgt_lang == "fra_Latn":
-        model_id = "ArissBandoss/mos2fr-5B-800"
-
+        model_id = "ArissBandoss/mos2fr-3B"
     else:
         model_id = "ArissBandoss/nllb-200-distilled-600M-finetuned-fr-to-mos-V4"
-
+    
     tokenizer = AutoTokenizer.from_pretrained(model_id, token=auth_token)
-    model     = AutoModelForSeq2SeqLM.from_pretrained(model_id, token=auth_token)
-
-        
-        
-    trans_pipe = pipeline("translation", 
-                          model=model, tokenizer=tokenizer, 
-                          src_lang=src_lang, tgt_lang=tgt_lang, 
-                          max_length=max_length,
-                           max_new_tokens=512,
-                          device=device
-                         )
-    
-    return trans_pipe(text)[0]["translation_text"]
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_id, token=auth_token).to(device)
+    
+    # Ajout du code de langue source
+    tokenizer.src_lang = src_lang
+    
+    # Tokenisation du texte d'entrée
+    inputs = tokenizer(text, return_tensors="pt").to(device)
+    
+    # Utilisation de convert_tokens_to_ids au lieu de lang_code_to_id
+    tgt_lang_id = tokenizer.convert_tokens_to_ids(tgt_lang)
+    
+    # Génération avec paramètres améliorés
+    outputs = model.generate(
+        **inputs,
+        forced_bos_token_id=tgt_lang_id,  
+        max_new_tokens=1024,  
+        num_beams=5, 
+        early_stopping=True
+    )
+    
+    # Décodage de la sortie
+    translation = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
+    
+    return translation
 
 
 def real_time_traduction(input_text, src_lang, tgt_lang):