Upload handler.py

handler.py

@@ -101,18 +101,21 @@ class EndpointHandler:
 
             if _is_messages(item):
                 # Chat template path exists in repo; tokenizer.apply_chat_template will use it if configured
-                result = self.tokenizer.apply_chat_template(
-                    item,
-                    return_tensors="pt",
-                    add_generation_prompt=True,
-                )
-                # Handle both single tensor and tuple/dict returns
-                if isinstance(result, dict):
-                    input_ids = result["input_ids"]
-                elif isinstance(result, tuple):
-                    input_ids = result[0]
-                else:
-                    input_ids = result
+                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)