Update handler.py

handler.py

@@ -86,19 +86,22 @@ class EndpointHandler:
                 item = item["messages"]
 
             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(
+                rendered = self.tokenizer.apply_chat_template(
                     item,
-                    return_tensors="pt",
+                    tokenize=False,
                     add_generation_prompt=True,
                 )
+                enc = self.tokenizer(rendered, return_tensors="pt")
+                input_ids = enc["input_ids"]
+                attention_mask = enc.get("attention_mask", None)
             else:
-                if not isinstance(item, str):
-                    item = str(item)
-                enc = self.tokenizer(item, return_tensors="pt")
+                enc = self.tokenizer(str(item), return_tensors="pt")
                 input_ids = enc["input_ids"]
-
+                attention_mask = enc.get("attention_mask", None)
             input_ids = input_ids.to(self.model.device)
+            if attention_mask is not None:
+                attention_mask = attention_mask.to(self.model.device)
+
             input_len = input_ids.shape[-1]
 
             gen_ids = self.model.generate(