Upload 2 files

hgf.py

@@ -0,0 +1,23 @@
+from typing import Dict, List, Any
+
+from transformers import MBart50TokenizerFast, MBartForConditionalGeneration
+
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        self.tokenizer = MBart50TokenizerFast.from_pretrained("facebook/mbart-large-50-many-to-many-mmt")
+        self.model = MBartForConditionalGeneration.from_pretrained(path)
+
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        inputs = data.pop("inputs", data)
+        source_lang = data.pop("src_lang", None)
+        target_lang = data.pop("target_lang", None)
+
+        self.tokenizer.src_lang = source_lang
+
+        encoded = self.tokenizer(inputs, return_tensors='pt')
+        generated_tokens = self.model.generate(
+            **encoded, forced_bos_token_id=self.tokenizer.lang_code_to_id[target_lang]
+        )
+        d = self.tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
+        return d

requirements.txt

@@ -0,0 +1 @@
+transformers