Update app.py

app.py

@@ -26,6 +26,7 @@ def ner(text: str):
     # Predict entities using the fine-tuned GLiNER model
     entities = model.predict_entities(text, labels, flat_ner=True, threshold=threshold)
 
+    # Filter for "Textile Machinery" entities
     textile_entities = [
         {
             "entity": ent["label"],
@@ -38,15 +39,14 @@ def ner(text: str):
         if ent["label"] == "Textile Machinery"
     ]
 
-    highlighted_text = text
-    for ent in sorted(textile_entities, key=lambda x: x['start'], reverse=True):
-        highlighted_text = (
-            highlighted_text[:ent['start']] +
-            f"<span style='background-color: yellow; font-weight: bold;'>{highlighted_text[ent['start']:ent['end']]}</span>" +
-            highlighted_text[ent['end']:]
-        )
+    # Prepare entities for color-coded display using gr.HighlightedText in the required dictionary format
+    highlights = [{"start": ent["start"], "end": ent["end"], "entity": ent["entity"]} for ent in textile_entities]
 
-    return highlighted_text, textile_entities
+    # Return two outputs: one for the highlighted text and one for the entities in JSON format
+    return {
+        "text": text,
+        "entities": highlights
+    }, textile_entities
 
 # Gradio Interface
 with gr.Blocks(title="Textile Machinery NER Demo") as demo: