Create handler.py

handler.py

@@ -0,0 +1,57 @@
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextGenerationPipeline
+import torch
+
+def load_model(model_id):
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id,
+        device_map="auto",
+        torch_dtype=torch.float16,
+        load_in_4bit=True
+    )
+    return model, tokenizer
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        self.model, self.tokenizer = load_model(path)
+        self.pipeline = TextGenerationPipeline(
+            model=self.model,
+            tokenizer=self.tokenizer
+        )
+
+    def __call__(self, data):
+        # Extraire le texte d'entrée
+        if isinstance(data, dict):
+            text = data.get("inputs", "")
+        else:
+            text = data
+
+        # 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 la réponse
+            outputs = self.pipeline(
+                text,
+                **generation_kwargs
+            )
+
+            # Formater la sortie en tableau comme requis par l'API
+            if isinstance(outputs, list):
+                return [{"generated_text": output["generated_text"]} for output in outputs]
+            return [{"generated_text": outputs["generated_text"]}]
+
+        except Exception as e:
+            return [{"error": str(e)}]