update app

app.py

@@ -25,6 +25,23 @@ def print_nicely(entities):
     return "\n".join(entity_details)
 
 
+# Helper function to flatten entities and prepare them for HighlightedText
+def prepare_entities_for_highlight(text, results):
+    entities = []
+    for category, entity_list in results.items():
+        for entity in entity_list:
+            # Appending each entity's word, start and end for highlighting, including the entity label and score
+            entities.append(
+                (
+                    entity["start"],
+                    entity["end"],
+                    f"{entity['entity']} ({entity['score']:.2f}%)",
+                )
+            )
+
+    return {"text": text, "entities": entities}
+
+
 # Function to process the sentence and extract entities
 def extract_entities(sentence):
     results = ner_pipeline(sentence)
@@ -36,7 +53,7 @@ def extract_entities(sentence):
 
     print(results)
     print(entity_results)
-    return results
+    return prepare_entities_for_highlight(sentence, results)
 
 
 # Create Gradio interface