Upload handler.py with huggingface_hub

handler.py

@@ -0,0 +1,70 @@
+from typing import Dict, List, Any
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import PeftModel, PeftConfig
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        # 1. Load the adapter config from the local path (where the repo is cloned on the endpoint)
+        self.peft_config = PeftConfig.from_pretrained(path)
+        
+        # 2. Load the Base Model
+        # We use device_map="auto" to use the GPU available in the endpoint
+        # torch_dtype=torch.float16 is standard for inference on T4/A10G
+        self.base_model = AutoModelForCausalLM.from_pretrained(
+            self.peft_config.base_model_name_or_path,
+            return_dict=True,
+            torch_dtype=torch.float16,
+            device_map="auto",
+            trust_remote_code=True
+        )
+        
+        # 3. Load the Tokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            self.peft_config.base_model_name_or_path, 
+            trust_remote_code=True
+        )
+        
+        # 4. Load the Adapter (Fine-tuned weights)
+        self.model = PeftModel.from_pretrained(self.base_model, path)
+        self.model.eval()
+
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        """
+        Args:
+            data (:obj: `Dict[str, Any]`):
+                Input data payload. Expects a key 'inputs' containing the prompt text.
+                Optional parameters: 'temperature', 'max_new_tokens', 'top_p', etc.
+        """
+        # Get inputs
+        inputs = data.pop("inputs", data)
+        parameters = data.pop("parameters", {})
+        
+        # Default generation parameters
+        max_new_tokens = parameters.get("max_new_tokens", 512)
+        temperature = parameters.get("temperature", 0.7)
+        top_p = parameters.get("top_p", 0.9)
+        
+        # Handle list of inputs or single string
+        if isinstance(inputs, list):
+            inputs = inputs[0] # Simplification for single-turn
+
+        # Tokenize
+        input_ids = self.tokenizer(inputs, return_tensors="pt").input_ids.to(self.model.device)
+        
+        # Generate
+        with torch.no_grad():
+            output_ids = self.model.generate(
+                input_ids=input_ids,
+                max_new_tokens=max_new_tokens,
+                temperature=temperature,
+                top_p=top_p,
+                do_sample=True,
+                pad_token_id=self.tokenizer.eos_token_id
+            )
+        
+        # Decode
+        # We slice [input_ids.shape[1]:] to return ONLY the generated response, not the prompt
+        generated_text = self.tokenizer.decode(output_ids[0][input_ids.shape[1]:], skip_special_tokens=True)
+        
+        return [{"generated_text": generated_text}]