Update handler.py

handler.py

@@ -12,30 +12,22 @@ class EndpointHandler:
             "additional_special_tokens": ["[QUERY]", "[LABEL_NAME]", "[LABEL_DESCRIPTION]"]
         })
         self.model = AutoModel.from_pretrained(path).to(self.device)
-        
+
         head_path = os.path.join(path, "classifier_head.json")
         with open(head_path, "r") as f:
             head = json.load(f)
-        
+
         self.classifier = torch.nn.Linear(self.model.config.hidden_size, 1).to(self.device)
         self.classifier.weight.data = torch.tensor(head["scorer_weight"]).to(self.device)
         self.classifier.bias.data = torch.tensor(head["scorer_bias"]).to(self.device)
-        
+
         self.model.eval()
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
-        """
-        Expected input format:
-        {
-            "query": "how to sharpen kitchen knives",
-            "candidates": [
-                {"label": "Tool-Specific", "description": "..."},
-                {"label": "Local Intent", "description": "..."}
-            ]
-        }
-        """
-        query = data["query"]
-        candidates = data["candidates"]
+        payload = data.get("inputs", data)
+
+        query = payload["query"]
+        candidates = payload["candidates"]
         results = []
 
         with torch.no_grad():