Update handler.py

handler.py

@@ -1,70 +1,86 @@
-from typing import Dict, Any, List
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
-
-class EndpointHandler:
-    def __init__(self, path=""):
-        """Initialize the model and tokenizer.
-        
-        Args:
-            path (str): Path to the model directory. Defaults to empty string.
-        """
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        self.model = AutoModelForCausalLM.from_pretrained(
-            path or "merged",
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-            device_map="auto"
-        )
-        self.tokenizer = AutoTokenizer.from_pretrained(path or "merged")
-
-    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
-        """Handle inference requests.
-        
-        Args:
-            data (Dict[str, Any]): The input data containing:
-                - instruction (str): The instruction for the model
-                - input (str, optional): Additional input text
-                - max_new_tokens (int, optional): Maximum number of tokens to generate. Defaults to 512
-                - temperature (float, optional): Sampling temperature. Defaults to 0.7
-        
-        Returns:
-            Dict[str, Any]: The model's response containing:
-                - response (str): The generated text
-        """
-        # Extract parameters from the request
-        instruction = data.get("instruction", "")
-        input_text = data.get("input", "")
-        max_new_tokens = data.get("max_new_tokens", 512)
-        temperature = data.get("temperature", 0.7)
-
-        # Create prompt
-        prompt = f"""Below is an instruction that describes a task. Write a response that appropriately completes the request.
-
-### Instruction:
-{instruction}"""
-
-        if input_text:
-            prompt += f"""
-
-### Input:
-{input_text}"""
-
-        prompt += """
-
-### Response:"""
-
-        # Generate response
-        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
-        outputs = self.model.generate(
-            **inputs,
-            max_new_tokens=max_new_tokens,
-            temperature=temperature,
-            do_sample=True,
-            pad_token_id=self.tokenizer.eos_token_id
-        )
-        
-        # Decode and extract response
-        full_response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-        response = full_response.split("### Response:")[-1].strip()
-
+from typing import Dict, Any, List
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        """Initialize the model and tokenizer.
+        
+        Args:
+            path (str): Path to the model directory. Defaults to empty string.
+        """
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.model = AutoModelForCausalLM.from_pretrained(
+            path or "merged",
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+            device_map="auto"
+        )
+        self.tokenizer = AutoTokenizer.from_pretrained(path or "merged")
+
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        """Handle inference requests.
+        
+        Args:
+            data (Dict[str, Any]): The input data. Can be in two formats:
+                1. Standard Hugging Face format:
+                    {
+                        "inputs": str,
+                        "parameters": Dict[str, Any]
+                    }
+                2. Custom format:
+                    {
+                        "instruction": str,
+                        "input": str (optional),
+                        "max_new_tokens": int (optional),
+                        "temperature": float (optional)
+                    }
+        
+        Returns:
+            Dict[str, Any]: The model's response containing:
+                - response (str): The generated text
+        """
+        # Handle standard Hugging Face format
+        if "inputs" in data:
+            instruction = data["inputs"]
+            parameters = data.get("parameters", {})
+            input_text = ""
+            max_new_tokens = parameters.get("max_new_tokens", 512)
+            temperature = parameters.get("temperature", 0.7)
+        # Handle custom format
+        else:
+            instruction = data.get("instruction", "")
+            input_text = data.get("input", "")
+            max_new_tokens = data.get("max_new_tokens", 512)
+            temperature = data.get("temperature", 0.7)
+
+        # Create prompt
+        prompt = f"""Below is an instruction that describes a task. Write a response that appropriately completes the request.
+
+### Instruction:
+{instruction}"""
+
+        if input_text:
+            prompt += f"""
+
+### Input:
+{input_text}"""
+
+        prompt += """
+
+### Response:"""
+
+        # Generate response
+        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
+        outputs = self.model.generate(
+            **inputs,
+            max_new_tokens=max_new_tokens,
+            temperature=temperature,
+            do_sample=True,
+            pad_token_id=self.tokenizer.eos_token_id
+        )
+        
+        # Decode and extract response
+        full_response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        response = full_response.split("### Response:")[-1].strip()
+
         return {"response": response}