Update handle.py

handle.py

@@ -1,10 +1,10 @@
 from typing import Dict, Any
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
 class EndpointHandler:
     def __init__(self, model_dir: str, **kwargs):
-        # Load tokenizer and model
+        # 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,
@@ -13,10 +13,7 @@ class EndpointHandler:
             trust_remote_code=True
         )
 
-        # Chat template setup
-        self.tokenizer.chat_template = self.tokenizer.chat_template or "{% for message in messages %}{{ message['role'] }}: {{ message['content'] }}\n{% endfor %}assistant:"
-
-        # Generation config (can be tuned)
+        # Paramètres de génération
         self.generation_config = {
             "max_new_tokens": 512,
             "do_sample": True,
@@ -32,17 +29,27 @@ class EndpointHandler:
         if not inputs:
             return {"error": "No 'inputs' provided in request."}
 
-        # Wrap input in a basic chat message structure
+        # Message format type ChatML
         messages = [{"role": "user", "content": inputs}]
 
-        # Apply chat template
-        prompt = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-
+        # 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)
 
-        # Decode output
+        # 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}