Update goai_traduction.py

goai_traduction.py

@@ -2,6 +2,7 @@ from transformers import pipeline
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 
 max_length = 512
+device = 0 if torch.cuda.is_available() else "cpu"
 model_id = "ArissBandoss/nllb-200-distilled-600M-finetuned-fr-to-mos-V1"
 
 tokenizer = AutoTokenizer.from_pretrained(model_id)
@@ -13,6 +14,7 @@ def goai_traduction(text, src_lang, tgt_lang):
                         model=model, tokenizer=tokenizer, 
                         src_lang=src_lang, tgt_lang=tgt_lang, 
                         max_length=max_length
+                        device=device
                        )
     
   return trans_pipe(text)[0]["translation_text"]