Update src/Translate.py

src/Translate.py

@@ -36,6 +36,18 @@ 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):
+        from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
+        src_lang = f"{self.sl}_XX"
+        tgt_lang = f"{self.sl}_{self.sl.upper()}"
+        # Load model and tokenizer
+        model = MBartForConditionalGeneration.from_pretrained(self.model_name)
+        tokenizer = MBart50TokenizerFast.from_pretrained(self.model_name, src_lang=src_lang)
+        # Tokenize and translate
+        inputs = tokenizer(self.input_text, return_tensors="pt")
+        translated_tokens = model.generate(**inputs, forced_bos_token_id=tokenizer.lang_code_to_id[tgt_lang])
+        translation = tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0]
+        return translation, self.message
 
 def paraphraseTranslateMethod(requestValue: str, model: str):
     nltk.download('punkt')