update handler

handler.py

@@ -48,34 +48,68 @@ class EndpointHandler:
         """
         Process inference request
         data args:
-            inputs (:obj:`str`): The input text
+            inputs (:obj:`str` or :obj:`Dict`): The input text or messages
             parameters (:obj:`Dict`, optional): Parameters for generation
         """
         inputs = data.pop("inputs", data)
         parameters = data.pop("parameters", {})
         
-        # Format input using chat template for Llama
+        # Handle different input formats that BrickGPT sends
+        if isinstance(inputs, dict) and "messages" in inputs:
+            # BrickGPT format: {"messages": [{"role": "system", ...}, {"role": "user", ...}]}
+            messages = inputs["messages"]
+        elif isinstance(inputs, list):
+            # Direct messages array: [{"role": "system", ...}, {"role": "user", ...}]
+            messages = inputs
+        elif isinstance(inputs, str):
+            # Plain string input - create default messages
+            messages = [
+                {"role": "system", "content": "You are a helpful assistant."},
+                {"role": "user", "content": inputs}
+            ]
+        else:
+            # Fallback
+            messages = [{"role": "user", "content": str(inputs)}]
+        
+        # Format input using chat template
         if hasattr(self.tokenizer, 'chat_template') and self.tokenizer.chat_template:
-            # Use chat template if available
-            messages = [{"role": "user", "content": inputs}]
-            formatted_input = self.tokenizer.apply_chat_template(
-                messages, 
-                tokenize=False, 
-                add_generation_prompt=True
-            )
+            # Check if this is a continuation (has assistant message)
+            has_assistant = any(msg.get("role") == "assistant" for msg in messages)
+            
+            if has_assistant:
+                # For continuation, use continue_final_message=True
+                formatted_input = self.tokenizer.apply_chat_template(
+                    messages, 
+                    tokenize=False, 
+                    continue_final_message=True
+                )
+            else:
+                # For new generation, add generation prompt
+                formatted_input = self.tokenizer.apply_chat_template(
+                    messages, 
+                    tokenize=False, 
+                    add_generation_prompt=True
+                )
         else:
-            # Fallback to direct input
-            formatted_input = inputs
+            # Fallback to direct input (last user message)
+            user_messages = [msg["content"] for msg in messages if msg.get("role") == "user"]
+            formatted_input = user_messages[-1] if user_messages else str(inputs)
         
         # Default parameters optimized for BrickGPT
         generation_params = {
             "max_new_tokens": parameters.get("max_new_tokens", 512),
             "temperature": parameters.get("temperature", 0.6),
-            "top_p": parameters.get("top_p", 0.9),
-            "do_sample": True,
+            "top_k": parameters.get("top_k", 20),
+            "top_p": parameters.get("top_p", 1.0),
+            "do_sample": parameters.get("do_sample", True),
             "pad_token_id": self.tokenizer.pad_token_id,
         }
         
+        # Add stop tokens if provided
+        stop_tokens = parameters.get("stop", [])
+        if stop_tokens:
+            generation_params["stop_strings"] = stop_tokens
+        
         # Tokenize input
         input_ids = self.tokenizer(formatted_input, return_tensors="pt").input_ids.to(self.model.device)