Update handler.py

handler.py

@@ -7,55 +7,64 @@ class EndpointHandler:
         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 process(self, inputs):
+    def process_single_text(self, text):
         try:
-            # Validation et préparation
-            texts = self._validate_input(inputs)
-            
-            # Configuration explicite des langues source et cible
+            # Configuration de la langue source
             self.tokenizer.src_lang = "ru_RU"
             
             # Tokenization
             inputs = self.tokenizer(
-                texts,
+                text,
                 return_tensors="pt",
                 padding=True,
                 truncation=True,
-                max_length=self.max_length
+                max_length=512  # Réduit pour plus d'efficacité
             ).to(self.device)
             
-            # Inférence avec langue cible explicite
+            # Génération
             with torch.no_grad():
-                outputs = self.model.generate(
+                generated_tokens = self.model.generate(
                     **inputs,
                     forced_bos_token_id=self.tokenizer.lang_code_to_id["en_XX"],
-                    max_length=self.max_length,
-                    num_beams=5,
+                    max_length=512,
+                    num_beams=4,
+                    length_penalty=1.0,
                     do_sample=False
                 )
             
-            # Post-traitement
-            translations = self.tokenizer.batch_decode(outputs, skip_special_tokens=True)
-            
-            # Retourne à la fois le texte brut et le JSON
-            text_output = translations[0] if translations else ""
-            json_output = [{"output": translation} for translation in translations]
-            
-            return text_output, json_output
+            # Décodage
+            translation = self.tokenizer.decode(generated_tokens[0], skip_special_tokens=True)
+            return translation
             
         except Exception as e:
-            return "", [{"error": str(e)}]
+            return f"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"])
+        try:
+            if "inputs" not in data:
+                raise ValueError("Request must contain 'inputs' field")
+            
+            inputs = data["inputs"]
+            
+            # Si l'entrée est une chaîne de caractères
+            if isinstance(inputs, str):
+                translation = self.process_single_text(inputs)
+                return translation  # Retourne directement la traduction comme texte
+            
+            # Si l'entrée est une liste
+            elif isinstance(inputs, list):
+                translations = []
+                for item in inputs:
+                    if isinstance(item, dict) and "input" in item:
+                        translation = self.process_single_text(item["input"])
+                        translations.append({"output": translation})
+                    else:
+                        translations.append({"error": "Invalid input format"})
+                return translations
+            
+            else:
+                raise ValueError("Invalid input format")
+                
+        except Exception as e:
+            return str(e)  # Retourne l'erreur comme texte