update handler.py to support chat completions

handler.py

@@ -10,7 +10,10 @@ class EndpointHandler:
         self.peft_config = PeftConfig.from_pretrained(path)
 
         # Load tokenizer from base model
-        self.tokenizer = AutoTokenizer.from_pretrained(self.peft_config.base_model_name_or_path, use_fast=False)
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            self.peft_config.base_model_name_or_path,
+            use_fast=False
+        )
 
         # Load base model
         base_model = AutoModelForCausalLM.from_pretrained(
@@ -24,20 +27,29 @@ class EndpointHandler:
         self.model.eval()
 
     def __call__(self, data: Dict[str, str]) -> Dict[str, str]:
-        user_input = data.get("inputs", "")
-        if not user_input:
-            return {"error": "No input provided."}
-
-        # Format input as chat message using chat template
-        messages = [
-            {"role": "user", "content": user_input}
-        ]
-        prompt = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        # Handle both chat-style and plain input
+        if "messages" in data:
+            messages = data["messages"]
+            prompt = self.tokenizer.apply_chat_template(
+                messages,
+                tokenize=False,
+                add_generation_prompt=True
+            )
+        else:
+            user_input = data.get("inputs", "")
+            if not user_input:
+                return {"error": "No input provided."}
+            messages = [{"role": "user", "content": user_input}]
+            prompt = self.tokenizer.apply_chat_template(
+                messages,
+                tokenize=False,
+                add_generation_prompt=True
+            )
 
-        # Tokenize
+        # Tokenize input
         inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
 
-        # Generate
+        # Generate output
         with torch.no_grad():
             output_ids = self.model.generate(
                 **inputs,
@@ -48,6 +60,7 @@ class EndpointHandler:
                 pad_token_id=self.tokenizer.eos_token_id
             )
 
+        # Decode and return result
         output_text = self.tokenizer.decode(output_ids[0], skip_special_tokens=True)
-
         return {"generated_text": output_text}
+