Update src/Translate.py

src/Translate.py

@@ -36,7 +36,7 @@ class Translators:
     def translationpipe(self):
         translation  = pipeline('translation', model = self.model_name)
         return translation(self.input_text)[0]['translation_text'], self.message
-    def mbartlarge(self):
+    def mbartlarge25(self):
         from transformers import MBartForConditionalGeneration, MBart50TokenizerFast, MBartTokenizer
         src_lang = f"{self.sl}_XX"
         tgt_lang = f"{self.tl}_{self.tl.upper()}"
@@ -60,6 +60,35 @@ class Translators:
         print(src_lang, tgt_lang, tokenizer.lang_code_to_id[tgt_lang])
         translation = tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0]
         return translation, self.message
+    def mbartlarge(self):
+        from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+        
+        model_name = "facebook/mbart-large-cc25"
+        
+        # load tokenizer and model
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+        
+        # tell tokenizer the source language
+        tokenizer.src_lang = "en_XX"
+        
+        # find the id for the target language and force it at generation
+        forced_bos_token_id = tokenizer.lang_code_to_id["ro_RO"]
+        
+        # create the pipeline (pass tokenizer and model explicitly)
+        pipe = pipeline("translation", model=model, tokenizer=tokenizer)
+        
+        # call the pipeline; generation kwargs are forwarded to model.generate
+        src_text = "This is a test sentence."
+        result = pipe(
+            src_text,
+            num_beams=4,
+            max_length=512,
+            forced_bos_token_id=forced_bos_token_id
+        )
+        
+        return result[0]["translation_text"], self.message
+
 
 def paraphraseTranslateMethod(requestValue: str, model: str):
     nltk.download('punkt')