Update goai_helpers/goai_traduction.py

goai_helpers/goai_traduction.py

@@ -25,6 +25,11 @@ def goai_traduction(text, src_lang, tgt_lang):
     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)
 
+    tgt_lang_id = tokenizer.convert_tokens_to_ids(tgt_lang)
+
+    model.config.forced_bos_token_id = tgt_lang_id
+    model.config.decoder_start_token_id = tgt_lang_id
+
     print("decoder_start_token_id:", model.config.decoder_start_token_id)
     print("forced_bos_token_id:", model.config.forced_bos_token_id)
     print("eos_token_id:", model.config.eos_token_id)
@@ -45,10 +50,7 @@ def goai_traduction(text, src_lang, tgt_lang):
         **inputs,
         forced_bos_token_id=tgt_lang_id,
         eos_token_id=tokenizer.eos_token_id,
-        max_length=512,
-           num_beams=4,
-    do_sample=False,
-    no_repeat_ngram_size=3
+        max_length=512
     )
 
     print("Token IDs:", outputs)