Update app.py

Annotate query

app.py

@@ -73,12 +73,29 @@ def parse_query(query: str, labels: Union[str, list], threshold: float = 0.3, ne
                 entities.append(entity)
         else:
             entities.append(entity)
-
+    
     payload = {"query": query, "entities": entities}
     print(f"{datetime.now()} :: parse_query :: {json.dumps(payload)}\n")
 
     return payload
 
+def annotate_query(query: str, labels: Union[str, list], threshold: float = 0.3, nested_ner: bool = False, model_name: str = None) -> Dict[str, Union[str, list]]:
+    payload = parse_query(query, labels, threshold, nested_ner, model_name)
+
+    return {
+        "text": query,
+        "entities": [
+            {
+                "entity": entity["label"],
+                "word": entity["text"],
+                "start": entity["start"],
+                "end": entity["end"],
+                "score": entity["score"],
+            }
+            for entity in payload["entities"]
+        ],
+    }
+
 
 # Initialize model here.
 print("Initializing models...")
@@ -132,27 +149,29 @@ with gr.Blocks(title="GLiNER-query-parser") as demo:
             scale=0,
         )
 
-    output = gr.JSON(label="Extracted entities")
+    # output = gr.JSON(label="Extracted entities")
+
+    output = gr.HighlightedText(label="Annotated entities")
     submit_btn = gr.Button("Submit")
 
     # Submitting
     query.submit(
-        fn=parse_query, inputs=[query, entities, threshold, is_nested, model_name], outputs=output
+        fn=annotate_query, inputs=[query, entities, threshold, is_nested, model_name], outputs=output
     )
     entities.submit(
-        fn=parse_query, inputs=[query, entities, threshold, is_nested, model_name], outputs=output
+        fn=annotate_query, inputs=[query, entities, threshold, is_nested, model_name], outputs=output
     )
     threshold.release(
-        fn=parse_query, inputs=[query, entities, threshold, is_nested, model_name], outputs=output
+        fn=annotate_query, inputs=[query, entities, threshold, is_nested, model_name], outputs=output
     )
     submit_btn.click(
-        fn=parse_query, inputs=[query, entities, threshold, is_nested, model_name], outputs=output
+        fn=annotate_query, inputs=[query, entities, threshold, is_nested, model_name], outputs=output
     )
     is_nested.change(
-        fn=parse_query, inputs=[query, entities, threshold, is_nested, model_name], outputs=output
+        fn=annotate_query, inputs=[query, entities, threshold, is_nested, model_name], outputs=output
     )
     model_name.change(
-        fn=parse_query, inputs=[query, entities, threshold, is_nested, model_name], outputs=output
+        fn=annotate_query, inputs=[query, entities, threshold, is_nested, model_name], outputs=output
     )
 
 demo.queue(default_concurrency_limit=5)