Upload handler.py

handler.py

@@ -0,0 +1,30 @@
+from transformers import MarianMTModel, MarianTokenizer
+from typing import Any, List, Dict
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        # Load the model and tokenizer
+        self.model = MarianMTModel.from_pretrained(path)
+        self.tokenizer = MarianTokenizer.from_pretrained(path)
+
+    def __call__(self, data: Any) -> List[Dict[str, str]]:
+        """
+        Args:
+            data (dict): The request payload with an "inputs" key containing the text to translate.
+        Returns:
+            List[Dict]: A list containing the translated text.
+        """
+        # Get the input text from the request
+        text = data.get("inputs", "")
+
+        # Tokenize the input text
+        inputs = self.tokenizer(text, return_tensors="pt", padding=True)
+
+        # Perform the translation
+        translated = self.model.generate(**inputs)
+
+        # Decode the translated text
+        translated_text = self.tokenizer.decode(translated[0], skip_special_tokens=True)
+
+        # Return the translated text as a response
+        return [{"translation_text": translated_text}]