Add handler.py for Inference Endpoints

handler.py

@@ -1,27 +1,14 @@
-"""
-Custom Handler for DeepSeek-R1-Cybersecurity-8B-Merged
-HuggingFace Inference Endpoints
-"""
-
 from typing import Dict, Any
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-
 class EndpointHandler:
     def __init__(self, path: str = ""):
-        """Initialize the model and tokenizer."""
-        # Load tokenizer
-        self.tokenizer = AutoTokenizer.from_pretrained(
-            path,
-            trust_remote_code=True
-        )
-        
-        # Set pad token if not set
+        """Initialize model and tokenizer."""
+        self.tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True)
         if self.tokenizer.pad_token is None:
             self.tokenizer.pad_token = self.tokenizer.eos_token
         
-        # Load model
         self.model = AutoModelForCausalLM.from_pretrained(
             path,
             torch_dtype=torch.bfloat16,
@@ -29,43 +16,18 @@ class EndpointHandler:
             trust_remote_code=True
         )
         self.model.eval()
-        
-        # Get device
         self.device = next(self.model.parameters()).device
-        print(f"Model loaded on device: {self.device}")
+        print(f"✅ Model loaded on {self.device}")
 
     def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
-        """Process inference request."""
-        # Extract inputs
+        """Handle inference request."""
         inputs = data.get("inputs", data.get("input", ""))
-        parameters = data.get("parameters", {})
-        
-        # Handle both string and list inputs
-        if isinstance(inputs, str):
-            prompts = [inputs]
-        else:
-            prompts = inputs
-        
-        # Default generation parameters
-        generation_config = {
-            "max_new_tokens": parameters.get("max_new_tokens", 256),
-            "temperature": parameters.get("temperature", 0.7),
-            "top_p": parameters.get("top_p", 0.9),
-            "top_k": parameters.get("top_k", 50),
-            "do_sample": parameters.get("do_sample", True),
-            "repetition_penalty": parameters.get("repetition_penalty", 1.1),
-            "pad_token_id": self.tokenizer.pad_token_id,
-            "eos_token_id": self.tokenizer.eos_token_id,
-        }
-        
-        # Remove None values
-        generation_config = {k: v for k, v in generation_config.items() if v is not None}
+        params = data.get("parameters", {})
         
         # Tokenize
         encoded = self.tokenizer(
-            prompts,
+            inputs,
             return_tensors="pt",
-            padding=True,
             truncation=True,
             max_length=2048
         ).to(self.device)
@@ -74,20 +36,17 @@ class EndpointHandler:
         with torch.no_grad():
             outputs = self.model.generate(
                 **encoded,
-                **generation_config
+                max_new_tokens=params.get("max_new_tokens", 256),
+                temperature=params.get("temperature", 0.7),
+                top_p=params.get("top_p", 0.9),
+                do_sample=params.get("do_sample", True),
+                repetition_penalty=params.get("repetition_penalty", 1.1),
+                pad_token_id=self.tokenizer.pad_token_id,
+                eos_token_id=self.tokenizer.eos_token_id,
             )
         
-        # Decode
-        generated_texts = []
-        for i, output in enumerate(outputs):
-            # Remove the input tokens from the output
-            input_length = encoded["input_ids"][i].shape[0]
-            generated_tokens = output[input_length:]
-            text = self.tokenizer.decode(generated_tokens, skip_special_tokens=True)
-            generated_texts.append(text)
+        # Decode (remove input tokens)
+        generated = outputs[0][encoded["input_ids"].shape[1]:]
+        text = self.tokenizer.decode(generated, skip_special_tokens=True)
         
-        # Return single string if single input, else list
-        if isinstance(inputs, str):
-            return {"generated_text": generated_texts[0]}
-        else:
-            return {"generated_text": generated_texts}
+        return {"generated_text": text}