Update goai_helpers/goai_traduction.py

goai_helpers/goai_traduction.py

@@ -26,21 +26,25 @@ def goai_traduction(text, src_lang, tgt_lang):
     tokenizer = AutoTokenizer.from_pretrained(model_id, token=auth_token)
     model     = AutoModelForSeq2SeqLM.from_pretrained(model_id, token=auth_token)
 
-    if model_id == "ArissBandoss/mos2fr-3B":
-        try:
-            peft_config = PeftConfig.from_pretrained("ArissBandoss/mos2fr-3B")
-            base_model = AutoModelForSeq2SeqLM.from_pretrained(peft_config.base_model_name_or_path)
-            model = PeftModel.from_pretrained(base_model, "ArissBandoss/mos2fr-3B")
-        except ValueError:
-            from huggingface_hub import hf_hub_download
-            adapter_config_path = hf_hub_download(
-                repo_id="ArissBarndoss/mos2fr-3B",
-                filename="adapter_config.json"
-            )
+    if model_id == "ArissBarndoss/mos2fr-3B":
+        peft_config = PeftConfig.from_pretrained("ArissBandoss/mos2fr-3B")
+        base_model = AutoModelForSeq2SeqLM.from_pretrained(peft_config.base_model_name_or_path)
+        model = PeftModel.from_pretrained(base_model, "ArissBandoss/mos2fr-3B")
         
-            peft_config = PeftConfig.from_json_file(adapter_config_path)
-            base_model = AutoModelForSeq2SeqLM.from_pretrained(peft_config.base_model_name_or_path)
-            model = PeftModel.from_pretrained(base_model, "ArissBandoss/mos2fr-3B")
+        # Instead of using the pipeline, do direct generation
+        tokenizer = AutoTokenizer.from_pretrained(peft_config.base_model_name_or_path)
+        def translate(text, src_lang, tgt_lang, max_length=512):
+            inputs = tokenizer(text, return_tensors="pt", max_length=max_length, truncation=True)
+            inputs = {k: v.to(device) for k, v in inputs.items()}
+            
+            generation_kwargs = {}
+            if src_lang and tgt_lang:
+                generation_kwargs["forced_bos_token_id"] = tokenizer.lang_code_to_id[tgt_lang]
+            
+            outputs = model.generate(**inputs, max_length=max_length, **generation_kwargs)
+            return tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
+        translation_text = translate(text, src_lang, tgt_lang)
+        return translation_text