Update app.py

app.py

@@ -3,7 +3,7 @@ from gliner import GLiNER
 import gradio as gr
 from datasets import load_dataset
 
-# Load the BL dataset with streaming
+# Load the BL dataset
 dataset_iter = iter(
     load_dataset(
         "max-long/bl_books_textile_filter",
@@ -13,10 +13,10 @@ dataset_iter = iter(
 )
 
 # Load the model
-model = GLiNER.from_pretrained("max-long/textile_machines_3_oct", trust_remote_code=True)
+model = GLiNER.from_pretrained("max-long/textile_machines_ner_5_oct", trust_remote_code=True)
 
 def ner(text: str):
-    labels = ["Textile Machinery"]  # Capitalized label for entity matching
+    labels = ["Textile Machinery"] 
     threshold = 0.5
     
     # Predict entities using the fine-tuned GLiNER model
@@ -35,16 +35,13 @@ def ner(text: str):
         if ent["label"] == "Textile Machinery"  # Exact match with capitalization
     ]
     
-    # Prepare data for HighlightedText
-    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
+    highlights = [(ent["start"], ent["end"], ent["entity"]) for ent in textile_entities]
     
-    return highlighted_text, textile_entities
+    return {
+        "text": text,
+        "entities": highlights
+    }
 
 with gr.Blocks(title="Textile Machinery NER Demo") as demo:
     gr.Markdown(
@@ -63,7 +60,7 @@ with gr.Blocks(title="Textile Machinery NER Demo") as demo:
     )
     
     # Define output components
-    output_highlighted = gr.HTML(label="Predicted Entities")
+    output_highlighted = gr.HighlightedText(label="Predicted Entities")
     output_entities = gr.JSON(label="Entities")
     
     submit_btn = gr.Button("Find Textile Machinery!")