Create handler.py

handler.py

@@ -0,0 +1,56 @@
+import torch
+from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
+
+class EndpointHandler:
+    def __init__(self, model_dir):
+        self.tokenizer = MBart50TokenizerFast.from_pretrained(model_dir)
+        self.model = MBartForConditionalGeneration.from_pretrained(model_dir)
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.model.to(self.device)
+        self.max_length = 1024
+        
+    def _validate_input(self, inputs):
+        if isinstance(inputs, str):
+            return [inputs]
+        elif isinstance(inputs, list) and all(isinstance(item, dict) and "input" in item for item in inputs):
+            return [item["input"] for item in inputs]
+        raise ValueError("Input must be a string or a list of dictionaries with 'input' key")
+
+    def _prepare_input(self, text):
+        return f"ru_RU {text}"
+
+    def process(self, inputs):
+        try:
+            # Validation et préparation
+            texts = self._validate_input(inputs)
+            prepared_texts = [self._prepare_input(text) for text in texts]
+            
+            # Tokenization
+            inputs = self.tokenizer(
+                prepared_texts,
+                return_tensors="pt",
+                padding=True,
+                truncation=True,
+                max_length=self.max_length
+            ).to(self.device)
+            
+            # Inférence
+            with torch.no_grad():
+                outputs = self.model.generate(
+                    **inputs,
+                    max_length=self.max_length,
+                    num_beams=5,
+                    do_sample=False
+                )
+            
+            # Post-traitement
+            translations = self.tokenizer.batch_decode(outputs, skip_special_tokens=True)
+            return [{"output": translation} for translation in translations]
+            
+        except Exception as e:
+            return [{"error": str(e)}]
+
+    def __call__(self, data):
+        if not isinstance(data, dict) or "inputs" not in data:
+            return [{"error": "Request must contain 'inputs' field"}]
+        return self.process(data["inputs"])