Update handler.py

handler.py

@@ -19,17 +19,22 @@ class EndpointHandler:
         # Get inputs and labels
         inputs = data.get("inputs", "")
         labels = ["party", "document title"]
-        print('labels',labels)
         # Predict entities using GLiNER
         entities = self.model.predict_entities(inputs, labels)
 
-        # Format the results to match the expected output structure
-        formatted_results = []
-        for entity in entities:
-            formatted_entity = {
-               entity["label"]: entity["text"],
-            }
-            print(formatted_entity)
-            formatted_results.append(formatted_entity)
+        # Initialize a dictionary to store organized entities
+            organized_entities = {label: {"labels": [], "scores": []} for label in labels}
 
-        return formatted_results
+            for entity in entities:
+                label = entity['label']
+                text = entity['text']
+                score = entity['score']
+
+                # Append text and score to the corresponding label
+                organized_entities[label]["labels"].append(text)
+                organized_entities[label]["scores"].append(score)
+
+            # Store organized entities in document metadata
+            doc.meta["entities"] = organized_entities
+
+        return {"documents": documents}