Create handler.py

handler.py

@@ -0,0 +1,24 @@
+from typing import Dict, List, Any
+from transformers import pipeline, AutoTokenizer
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        # Load the optimized model
+        tokenizer = AutoTokenizer.from_pretrained(path)
+        # Create inference pipeline for text classification
+        self.pipeline = pipeline("text-classification", model=path, tokenizer=tokenizer)
+
+    def __call__(self, data: str) -> List[List[Dict[str, float]]]:
+        """
+        Args:
+            data (str): A raw string input for inference.
+        Returns:
+            A list containing the prediction results:
+            A list of one list, e.g., [[{"label": "LABEL", "score": 0.99}]]
+        """
+        # Pass the data as `text` directly
+        inputs = data.pop("inputs", data)
+        prediction = self.pipeline(inputs)
+
+        # Return the prediction result
+        return prediction