Update handler.py

handler.py

@@ -1,6 +1,5 @@
 from typing import Dict, List, Any
 from gliner import GLiNER
-import os
 
 class EndpointHandler:
     def __init__(self, path=""):
@@ -19,7 +18,7 @@ class EndpointHandler:
         """
         # Get inputs and labels
         inputs = data.get("inputs", "")
-        labels = data.get("labels", [])
+        labels = ["party", "document title"]
 
         # Predict entities using GLiNER
         entities = self.model.predict_entities(inputs, labels)
@@ -30,7 +29,7 @@ class EndpointHandler:
             formatted_entity = {
                 "word": entity["text"],
                 "entity_group": entity["label"],  # Assuming entity["label"] contains the label
-                "score": entity.get("score", 1.0)  # Assuming a default score of 1.0 if not provided
+                "score": entity.get("score", 1.0)  # Default to 1.0 if score not provided
             }
             formatted_results.append(formatted_entity)