Update handle.py

handle.py

@@ -1,55 +1,29 @@
-from typing import Dict, Any
+from typing import Dict, List, Any
 from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
 
 class EndpointHandler:
-    def __init__(self, model_dir: str, **kwargs):
-        # Charger le tokenizer et le modèle
-        self.tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
-        self.model = AutoModelForCausalLM.from_pretrained(
-            model_dir,
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-            device_map="auto",
-            trust_remote_code=True
-        )
-
-        # Paramètres de génération
-        self.generation_config = {
-            "max_new_tokens": 512,
-            "do_sample": True,
-            "temperature": 0.7,
-            "top_p": 0.9,
-            "top_k": 50,
-            "pad_token_id": self.tokenizer.eos_token_id,
-            "eos_token_id": self.tokenizer.eos_token_id,
-        }
-
-    def __call__(self, data: Dict[str, Any]) -> Dict[str, str]:
-        inputs = data.get("inputs") or data.get("text")
-        if not inputs:
-            return {"error": "No 'inputs' provided in request."}
-
-        # Message format type ChatML
-        messages = [{"role": "user", "content": inputs}]
-
-        # Appliquer le template si possible
-        if hasattr(self.tokenizer, "apply_chat_template"):
-            prompt = self.tokenizer.apply_chat_template(
-                messages,
-                tokenize=False,
-                add_generation_prompt=True
-            )
-        else:
-            # Fallback simple si pas de template
-            prompt = "user: " + inputs + "\nassistant:"
-
-        # Tokeniser et générer
-        input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids.to(self.model.device)
-
-        with torch.no_grad():
-            output_ids = self.model.generate(input_ids, **self.generation_config)
-
-        # Décoder la sortie après le prompt
-        response = self.tokenizer.decode(output_ids[0][input_ids.shape[-1]:], skip_special_tokens=True)
-
-        return {"generated_text": response}
+    def __init__(self, path: str):
+        # Charger le modèle et le tokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained(path)
+        self.model = AutoModelForCausalLM.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
+        inputs = data.pop("inputs", data)
+
+        # Tokenize les entrées
+        input_ids = self.tokenizer.encode(inputs, return_tensors="pt")
+
+        # 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)
+
+        # Retourner le texte généré
+        return {"generated_text": generated_text}