Update goai_helpers/goai_traduction.py

goai_helpers/goai_traduction.py

@@ -21,50 +21,30 @@ def goai_traduction(text, src_lang, tgt_lang):
         model_id = "ArissBandoss/mos2fr-3B-1200"
     else:
         model_id = "ArissBandoss/nllb-200-distilled-600M-finetuned-fr-to-mos-V4"
-    
-    tokenizer = AutoTokenizer.from_pretrained(model_id, token=auth_token, truncation=True, max_length=512)
-    model = AutoModelForSeq2SeqLM.from_pretrained(model_id, token=auth_token).to(device)
-    print(model.lm_head.weight.shape)  # doit être [vocab_size, hidden_size]
-    print(model.model.shared.weight.shape)  # idem
-
-    tgt_lang_id = tokenizer.convert_tokens_to_ids(tgt_lang)
 
-    generation_config = GenerationConfig(
-    max_new_tokens=1024,
-    early_stopping=False,
-    decoder_start_token_id=tokenizer.convert_tokens_to_ids(tgt_lang),
-    forced_bos_token_id=tokenizer.convert_tokens_to_ids(tgt_lang),
-    eos_token_id=tokenizer.eos_token_id,
-    )
+    tokenizer = AutoTokenizer.from_pretrained(model_id, token=auth_token)
+    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)
-    print(inputs)
-    
-    # Utilisation de convert_tokens_to_ids au lieu de lang_code_to_id
+
+    # Ajout du code de langue cible
     tgt_lang_id = tokenizer.convert_tokens_to_ids(tgt_lang)
-    
-    # Génération avec paramètres améliorés
+
+    # Génération contrôlée
     outputs = model.generate(
-    **inputs,
-    generation_config=generation_config
+        **inputs,
+        forced_bos_token_id=tgt_lang_id,
+        eos_token_id=tokenizer.eos_token_id,  # S’assurer que le modèle peut s’arrêter
+        max_length=512,  # Teste avec 256 puis augmente progressivement
+        do_sample=False, 
+        early_stopping=True
     )
 
-
-    print("Token IDs:", outputs)
-    print("Tokens:", [tokenizer.decode([tok]) for tok in outputs[0]])
-
-    
-    # Décodage de la sortie
     translation = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
     print("ici translation")
     print(translation)
     return translation
 
-
 def real_time_traduction(input_text, src_lang, tgt_lang):
     return goai_traduction(input_text, src_lang, tgt_lang)