Update handler.py

handler.py

@@ -22,23 +22,10 @@ def _is_messages(x: Any) -> bool:
 class EndpointHandler:
     """
     Hugging Face Inference Endpoints custom handler.
-
-    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
+    Expects:
+      - request body is a dict
+      - always contains `inputs`
+      - may contain `parameters` for generation
     """
 
     def __init__(self, model_dir: str):
@@ -86,7 +73,6 @@ 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))
@@ -101,21 +87,11 @@ class EndpointHandler:
 
             if _is_messages(item):
                 # Chat template path exists in repo; tokenizer.apply_chat_template will use it if configured
-                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"]
+                input_ids = self.tokenizer.apply_chat_template(
+                    item,
+                    return_tensors="pt",
+                    add_generation_prompt=True,
+                )
             else:
                 if not isinstance(item, str):
                     item = str(item)
@@ -138,12 +114,9 @@ class EndpointHandler:
                 eos_token_id=self.tokenizer.eos_token_id,
             )
 
-            # 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)
+            # Only return newly generated tokens
+            new_tokens = gen_ids[0, input_len:]
+            text = self.tokenizer.decode(new_tokens, skip_special_tokens=True)
             return {"generated_text": text}
 
         # Batch support