Update handler.py

handler.py

@@ -22,10 +22,23 @@ def _is_messages(x: Any) -> bool:
 class EndpointHandler:
     """
     Hugging Face Inference Endpoints custom handler.
-    Expects:
-      - request body is a dict
-      - always contains `inputs`
-      - may contain `parameters` for generation
+
+    Supports both text and chat formats:
+
+    Text format:
+      {"inputs": "Hello, how are you?"}
+
+    Chat format (recommended):
+      {"inputs": [{"role": "user", "content": "Hello!"}]}
+      or
+      {"inputs": {"messages": [{"role": "user", "content": "Hello!"}]}}
+
+    Parameters:
+      - max_new_tokens (default: 256): Max tokens to generate
+      - temperature (default: 0.7): Sampling temperature
+      - top_p (default: 0.95): Nucleus sampling
+      - repetition_penalty (default: 1.0): Penalize repetitions
+      - return_full_text (default: False): If True, return full conversation; if False, only new tokens
     """
 
     def __init__(self, model_dir: str):
@@ -73,6 +86,7 @@ class EndpointHandler:
         top_p = float(params.get("top_p", 0.95))
         top_k = int(params.get("top_k", 0))
         repetition_penalty = float(params.get("repetition_penalty", 1.0))
+        return_full_text = bool(params.get("return_full_text", False))
 
         do_sample = bool(params.get("do_sample", temperature > 0))
         num_beams = int(params.get("num_beams", 1))
@@ -87,11 +101,21 @@ class EndpointHandler:
 
             if _is_messages(item):
                 # Chat template path exists in repo; tokenizer.apply_chat_template will use it if configured
-                input_ids = self.tokenizer.apply_chat_template(
-                    item,
-                    return_tensors="pt",
-                    add_generation_prompt=True,
-                )
+                try:
+                    # Use tokenize=False to get the formatted string first
+                    prompt = self.tokenizer.apply_chat_template(
+                        item,
+                        tokenize=False,
+                        add_generation_prompt=True,
+                    )
+                    # Then tokenize it separately to avoid unpacking issues
+                    enc = self.tokenizer(prompt, return_tensors="pt")
+                    input_ids = enc["input_ids"]
+                except Exception:
+                    # Fallback: if chat template fails, use the last user message
+                    last_user_msg = next((m["content"] for m in reversed(item) if m.get("role") == "user"), "")
+                    enc = self.tokenizer(last_user_msg, return_tensors="pt")
+                    input_ids = enc["input_ids"]
             else:
                 if not isinstance(item, str):
                     item = str(item)
@@ -114,9 +138,12 @@ class EndpointHandler:
                 eos_token_id=self.tokenizer.eos_token_id,
             )
 
-            # Only return newly generated tokens
-            new_tokens = gen_ids[0, input_len:]
-            text = self.tokenizer.decode(new_tokens, skip_special_tokens=True)
+            # Return newly generated tokens by default, or full text if requested
+            if return_full_text:
+                text = self.tokenizer.decode(gen_ids[0], skip_special_tokens=True)
+            else:
+                new_tokens = gen_ids[0, input_len:]
+                text = self.tokenizer.decode(new_tokens, skip_special_tokens=True)
             return {"generated_text": text}
 
         # Batch support