Create handler.py

handler.py

@@ -0,0 +1,98 @@
+import json
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+import torch
+from typing import Dict, List, Any
+
+class EndpointHandler():
+    def __init__(self, path=""):
+        """
+        Initializes the model and tokenizer.
+        Args:
+            path (str): Path to the directory containing the model files.
+        """
+        # Load model and tokenizer from the path provided by Inference Endpoints
+        self.tokenizer = AutoTokenizer.from_pretrained(path)
+        self.model = AutoModelForSequenceClassification.from_pretrained(path)
+        
+        # Determine device - Inference Endpoints will handle GPU allocation if available
+        if torch.cuda.is_available():
+            self.device = torch.device("cuda")
+        else:
+            self.device = torch.device("cpu")
+        self.model.to(self.device)
+        self.model.eval() # Set model to evaluation mode
+        print("Model and tokenizer loaded successfully.")
+
+
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        """
+        Processes inference requests.
+        Args:
+            data (Dict[str, Any]): A dictionary containing the input data. 
+                                   Expected key: "inputs" (string or list of strings).
+                                   Optional key: "parameters" (dictionary).
+        Returns:
+            List[Dict[str, Any]]: A list of dictionaries, where each dictionary 
+                                  contains the prediction results for one input string 
+                                  (e.g., [{"label": "AI", "score": 0.98}]).
+        """
+        # Get inputs from the payload
+        inputs = data.pop("inputs", None)
+        parameters = data.pop("parameters", {}) # Optional parameters
+
+        if inputs is None:
+            raise ValueError("Missing 'inputs' key in request data")
+
+        # Preprocessing: Tokenize the input text
+        # Handle both single string and list of strings inputs
+        # Padding=True and truncation=True are good defaults for batch processing
+        tokenized_inputs = self.tokenizer(
+            inputs, 
+            return_tensors="pt", 
+            padding=True, 
+            truncation=True,
+            max_length=self.tokenizer.model_max_length # Use model's max sequence length
+        ).to(self.device)
+
+        # Inference: Run prediction
+        with torch.no_grad(): # Disable gradient calculations for inference
+            outputs = self.model(**tokenized_inputs)
+
+        # Postprocessing: Convert logits to probabilities and get labels
+        logits = outputs.logits
+        probabilities = torch.softmax(logits, dim=-1)
+        
+        results = []
+        # Iterate through each input in the batch
+        for i in range(probabilities.shape[0]):
+            scores = probabilities[i].tolist() # Get probabilities for the i-th input
+            predictions = []
+            for j, score in enumerate(scores):
+                # Map the class index (j) to the actual label string
+                label = self.model.config.id2label[j]
+                predictions.append({"label": label, "score": score})
+            
+            # Sort predictions by score descending if needed, or just return all
+            # predictions.sort(key=lambda x: x["score"], reverse=True) 
+            results.append(predictions) # Append all label scores for this input
+
+        # If the original input was a single string, return just the first result list
+        if isinstance(inputs, str):
+             # However, the standard API often expects a list even for single inputs
+             # So we return results which is already a list containing one list of predictions
+             # Let's adjust to return a flat list of predictions if input was single string,
+             # matching common pipeline output. But check what your consumer expects.
+             # For now, return the list of lists structure for consistency.
+             pass # Keep results as list of lists: [[{'label': '...', 'score': ...}]]
+        
+        # If you want to return ONLY the top prediction per input string:
+        # top_results = []
+        # for i in range(probabilities.shape[0]):
+        #    top_prob, top_idx = torch.max(probabilities[i], dim=0)
+        #    label = self.model.config.id2label[top_idx.item()]
+        #    score = top_prob.item()
+        #    top_results.append({"label": label, "score": score})
+        # return top_results
+
+        # Return all labels and scores per input
+        return results