Update handler.py

handler.py

@@ -1,25 +1,27 @@
+from typing import Dict, Any, List
 from transformers import AutoTokenizer
 from gliner import GLiNER
-from huggingface_inference_toolkit.base import BaseHandler
 
-class EndpointHandler(BaseHandler):
-    def __init__(self, path=""):
-        self.model = GLiNER.from_pretrained(path)
+
+class EndpointHandler:
+    def __init__(self, path: str = ""):
+        """Initialize the model and tokenizer"""
+        self.model = GLiNER.from_pretrained(path if path else "urchade/gliner_multi-v2.1")
         self.tokenizer = AutoTokenizer.from_pretrained("microsoft/deberta-v3-large")
         self.initialized = True
 
-    def __call__(self, data):
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, List[Dict[str, Any]]]:
         """
         Args:
-            data: Dictionary with:
-                - text (str): Input text
+            data (Dict[str, Any]): Dictionary containing:
+                - inputs/text (str): Input text
                 - labels (str): Comma-separated labels
                 - threshold (float, optional): Confidence threshold
                 - nested_ner (bool, optional): Enable nested NER
         Returns:
-            Dictionary with predicted entities
+            Dict[str, List[Dict[str, Any]]]: Dictionary with predicted entities
         """
-        # Get inputs
+        # Get inputs - handle both "inputs" and "text" keys for flexibility
         text = data.pop("inputs", data.get("text", ""))
         labels = data.get("labels", "").split(",")
         threshold = float(data.get("threshold", 0.3))