Update handler.py

handler.py

@@ -16,35 +16,36 @@ class EndpointHandler:
         self.model, self.tokenizer = load_model(path)
         self.pipeline = TextGenerationPipeline(
             model=self.model,
-            tokenizer=self.tokenizer,
-            max_new_tokens=512,
-            temperature=0.7,
-            top_p=0.95,
-            repetition_penalty=1.15,
-            do_sample=True
+            tokenizer=self.tokenizer
         )
 
     def __call__(self, data):
-        inputs = data.pop("inputs", data)
-        parameters = data.pop("parameters", {})
-        
-        generation_kwargs = {
-            "max_new_tokens": 512,
-            "temperature": 0.7,
-            "top_p": 0.95,
-            "repetition_penalty": 1.15,
-            "do_sample": True
+        # Extraire le texte d'entrée
+        if isinstance(data, dict):
+            text = data.pop("inputs", "")
+        else:
+            text = data
+            
+        # Paramètres de génération
+        params = {
+            "max_new_tokens": data.get("max_new_tokens", 512),
+            "temperature": data.get("temperature", 0.7),
+            "top_p": data.get("top_p", 0.95),
+            "repetition_penalty": data.get("repetition_penalty", 1.15),
+            "do_sample": data.get("do_sample", True)
         }
-        generation_kwargs.update(parameters)
-        
-        if isinstance(inputs, str):
-            inputs = [inputs]
+
+        try:
+            # Générer le texte
+            result = self.pipeline(
+                text,
+                **params
+            )
             
-        outputs = self.pipeline(
-            inputs,
-            **generation_kwargs
-        )
-        
-        if len(outputs) == 1:
-            return {"generated_text": outputs[0]["generated_text"]}
-        return [{"generated_text": o["generated_text"]} for o in outputs]
+            # Formater la sortie
+            if isinstance(result, list):
+                return {"generated_text": result[0]["generated_text"]}
+            return {"generated_text": result["generated_text"]}
+            
+        except Exception as e:
+            return {"error": str(e)}