Update handler.py

handler.py

@@ -22,30 +22,36 @@ class EndpointHandler:
     def __call__(self, data):
         # Extraire le texte d'entrée
         if isinstance(data, dict):
-            text = data.pop("inputs", "")
+            text = data.get("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)
+
+        # Paramètres de génération par défaut
+        generation_kwargs = {
+            "max_new_tokens": 512,
+            "temperature": 0.7,
+            "top_p": 0.95,
+            "repetition_penalty": 1.15,
+            "do_sample": True,
+            "pad_token_id": self.tokenizer.pad_token_id,
+            "eos_token_id": self.tokenizer.eos_token_id,
         }
 
+        # Mettre à jour avec les paramètres de la requête si fournis
+        if isinstance(data, dict) and "parameters" in data:
+            generation_kwargs.update(data["parameters"])
+
         try:
-            # Générer le texte
-            result = self.pipeline(
+            # Générer la réponse
+            outputs = self.pipeline(
                 text,
-                **params
+                **generation_kwargs
             )
-            
+
             # Formater la sortie
-            if isinstance(result, list):
-                return {"generated_text": result[0]["generated_text"]}
-            return {"generated_text": result["generated_text"]}
-            
+            if isinstance(outputs, list):
+                return {"generated_text": outputs[0]["generated_text"]}
+            return {"generated_text": outputs["generated_text"]}
+
         except Exception as e:
             return {"error": str(e)}