Update app.py

app.py

@@ -12,37 +12,34 @@ dataset_iter = load_dataset(
 ).shuffle(seed=42)  # Shuffle added
 
 # Load the model
-model = GLiNER.from_pretrained("max-long/textile_machines_3_oct", trust_remote_code=True)
+model = GLiNER.from_pretrained("urchade/gliner_multi-v2.1", trust_remote_code=True)
 
-def ner(text: str, labels: str, threshold: float):
+def ner(text: str, labels: str, threshold: float, nested_ner: bool):
     # Convert user-provided labels (comma-separated string) into a list
     labels_list = [label.strip() for label in labels.split(",")]
     
-    # Predict entities using the fine-tuned GLiNER model
-    entities = model.predict_entities(text, labels_list, flat_ner=True, threshold=threshold)
+    # Predict entities using the GLiNER model
+    entities = model.predict_entities(text, labels_list, flat_ner=not nested_ner, threshold=threshold)
     
-    # Prepare data for HighlightedText
-    highlighted_text = text
-    for ent in sorted(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["label"]) for ent in entities]
     
-    return highlighted_text, entities
+    return {
+        "text": text,
+        "entities": highlights
+    }
 
-with gr.Blocks(title="General NER Demo") as demo:
+with gr.Blocks(title="General NER with Color-Coded Output") as demo:
     gr.Markdown(
         """
         # General Entity Recognition Demo
-        This demo selects a random text snippet from a subset of the British Library's books dataset and identifies entities using a fine-tuned GLiNER model. You can specify the entities you want to find.
+        This demo selects a random text snippet from a subset of the British Library's books dataset and identifies entities using the GLiNER NER model. You can specify the entities you want to find, and results will be displayed in a color-coded format.
         """
     )
     
     # Display a random example
     input_text = gr.Textbox(
-        value="The machine is fed by means of an endless apron, the wool entering at the smaller end...",
+        value="Click 'Get New Snippet' to load a random sample from the British Library Dataset",
         label="Text input",
         placeholder="Enter your text here",
         lines=5
@@ -50,7 +47,7 @@ with gr.Blocks(title="General NER Demo") as demo:
 
     with gr.Row() as row:
         labels = gr.Textbox(
-            value="Machine, Wool",  # Default example labels
+            value="Places, People",  # Default example labels
             label="Labels",
             placeholder="Enter your labels here (comma separated)",
             scale=2,
@@ -64,9 +61,14 @@ with gr.Blocks(title="General NER Demo") as demo:
             info="Lower the threshold to increase how many entities get predicted.",
             scale=1,
         )
+        nested_ner = gr.Checkbox(
+            value=False,
+            label="Nested NER",
+            info="Enable Nested NER?",
+        )
 
-    # Define output components
-    output_highlighted = gr.HTML(label="Predicted Entities")
+    # Define output components using HighlightedText for color-coded display
+    output_highlighted = gr.HighlightedText(label="Predicted Entities")
     output_entities = gr.JSON(label="Entities")
     
     submit_btn = gr.Button("Find Entities!")
@@ -84,8 +86,8 @@ with gr.Blocks(title="General NER Demo") as demo:
     
     # Connect submit button
     submit_btn.click(
-        fn=lambda text, labels, threshold: ner(text, labels, threshold),
-        inputs=[input_text, labels, threshold],
+        fn=ner,
+        inputs=[input_text, labels, threshold, nested_ner],
         outputs=[output_highlighted, output_entities]
     )