Upload handler.py

handler.py

@@ -1,20 +1,18 @@
 from typing import Dict, List, Any
 from transformers import pipeline
 
-# Load the model and tokenizer at startup
-classifier = pipeline("text-classification")
+class EndpointHandler:
+    def __init__(self, path=""):
+        # Load the model and tokenizer at startup
+        self.classifier = pipeline("text-classification", model=path if path else None)
 
-def preprocess(inputs: Dict[str, Any]) -> List[str]:
-    # Hugging Face Inference API sends {"inputs": ...}
-    if isinstance(inputs["inputs"], list):
-        return inputs["inputs"]
-    return [inputs["inputs"]]
+    def preprocess(self, inputs: Dict[str, Any]) -> List[str]:
+        # Hugging Face Inference API sends {"inputs": ...}
+        if isinstance(inputs["inputs"], list):
+            return inputs["inputs"]
+        return [inputs["inputs"]]
 
-def inference(inputs: Dict[str, Any]) -> List[Dict[str, Any]]:
-    texts = preprocess(inputs)
-    results = classifier(texts)
-    return results
-
-# The entrypoint for the Inference API
-def handle(inputs: Dict[str, Any], context: Dict[str, Any] = None) -> Any:
-    return inference(inputs)
+    def __call__(self, inputs: Dict[str, Any]) -> List[Dict[str, Any]]:
+        texts = self.preprocess(inputs)
+        results = self.classifier(texts)
+        return results