Update handler.py

handler.py

@@ -1,29 +1,22 @@
 from typing import Dict, List, Any
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-class EndpointHandler:
-    def __init__(self, path: str):
-        # Charger le modèle et le tokenizer
-        self.tokenizer = AutoTokenizer.from_pretrained(path)
+class EndpointHandler():
+    def __init__(self, path=""):
+        # Load the model and tokenizer from the specified path
         self.model = AutoModelForCausalLM.from_pretrained(path)
+        self.tokenizer = AutoTokenizer.from_pretrained(path)
 
-    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
-        """
-        Cette méthode est appelée à chaque requête.
-        :param data: un dictionnaire contenant les données d'entrée.
-        :return: un dictionnaire contenant la prédiction.
-        """
-        # Extraire les entrées du dictionnaire de données
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        # Extract input text from the request
         inputs = data.pop("inputs", data)
 
-        # Tokenize les entrées
+        # Tokenize input and generate text
         input_ids = self.tokenizer.encode(inputs, return_tensors="pt")
+        output_ids = self.model.generate(input_ids)
 
-        # Générer du texte
-        output_ids = self.model.generate(input_ids, max_length=100)
-
-        # Décoder les IDs de sortie en texte
-        generated_text = self.tokenizer.decode(output_ids[0], skip_special_tokens=True)
+        # Decode the generated output
+        output_text = self.tokenizer.decode(output_ids[0], skip_special_tokens=True)
 
-        # Retourner le texte généré
-        return {"generated_text": generated_text}
+        # Return the generated text
+        return [{"generated_text": output_text}]