Upload handler.py with huggingface_hub

handler.py

@@ -1,45 +1,82 @@
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from typing import Dict, List, Any
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
+
 class EndpointHandler:
-    def __init__(self):
-        self.tokenizer = None
-        self.model = None
-
-    def load_model(self):
-        """Load model and tokenizer"""
-        self.tokenizer = AutoTokenizer.from_pretrained(".")
-        self.model = AutoModelForCausalLM.from_pretrained(
-            ".",
-            torch_dtype=torch.float16,
-            device_map="auto",
-            use_safetensors=True
-        )
+    def __init__(self, path=""):
+        # Initialize model and tokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained(path)
+        self.model = AutoModelForCausalLM.from_pretrained(path)
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.model.to(self.device)
 
-    def __call__(self, data):
-        """Inference request handler"""
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        """
+        Args:
+            data: JSON input with structure:
+            {
+                "inputs": "your text prompt here",
+                "parameters": {
+                    "max_new_tokens": 50,
+                    "temperature": 0.7,
+                    "top_p": 0.9,
+                    "do_sample": true
+                }
+            }
+        """
+        # Get input text and parameters
         inputs = data.pop("inputs", data)
         parameters = data.pop("parameters", {})
-        
-        # Set default parameters if not provided
-        max_length = parameters.get("max_length", 100)
-        temperature = parameters.get("temperature", 0.7)
-        top_p = parameters.get("top_p", 0.9)
-        
-        # Tokenize inputs
-        input_ids = self.tokenizer(inputs, return_tensors="pt").input_ids.to(self.model.device)
-        
-        # Generate
+
+        # Default generation parameters
+        generation_config = {
+            "max_new_tokens": parameters.get("max_new_tokens", 50),
+            "temperature": parameters.get("temperature", 0.7),
+            "top_p": parameters.get("top_p", 0.9),
+            "do_sample": parameters.get("do_sample", True),
+            "pad_token_id": self.tokenizer.eos_token_id,
+            "num_return_sequences": parameters.get("num_return_sequences", 1)
+        }
+
+        # Tokenize
+        inputs = self.tokenizer(
+            inputs,
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=512
+        ).to(self.device)
+
+        # Generate text
         with torch.no_grad():
-            outputs = self.model.generate(
-                input_ids,
-                max_length=max_length,
-                temperature=temperature,
-                top_p=top_p,
-                pad_token_id=self.tokenizer.pad_token_id,
-                eos_token_id=self.tokenizer.eos_token_id,
+            generated_ids = self.model.generate(
+                inputs.input_ids,
+                attention_mask=inputs.attention_mask,
+                **generation_config
             )
-        
-        # Decode and return response
-        generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-        return {"generated_text": generated_text}
+
+        # Decode and return generated text
+        generated_texts = self.tokenizer.batch_decode(
+            generated_ids,
+            skip_special_tokens=True
+        )
+
+        return {
+            "generated_text": generated_texts[0],  # Return first generation if multiple
+            "all_generations": generated_texts  # All generations if num_return_sequences > 1
+        }
+
+    def preprocess(self, data):
+        """
+        Handle different input formats
+        """
+        if isinstance(data, str):
+            return {"inputs": data}
+        return data
+
+    def postprocess(self, data):
+        """
+        Clean up output if needed
+        """
+        return data