Update app.py

app.py

@@ -3,15 +3,17 @@ from gliner import GLiNER
 import gradio as gr
 from datasets import load_dataset
 
-# Load the subset dataset from Hugging Face Hub
-subset_dataset = load_dataset("TheBritishLibrary/blbooks", split="train", streaming=True, trust_remote_code=True)
+# Load the large dataset with streaming
+dataset_iter = iter(load_dataset("TheBritishLibrary/blbooks", split="train", streaming=True, trust_remote_code=True))
 
 # Load the GLiNER model
 model = GLiNER.from_pretrained("max-long/textile_machines_3_oct", trust_remote_code=True)
 
-# Define the NER function
 def ner(text: str, labels: str, threshold: float, nested_ner: bool):
+    # Split and clean labels
     labels = [label.strip() for label in labels.split(",")]
+    
+    # Predict entities using GLiNER
     entities = model.predict_entities(text, labels, flat_ner=not nested_ner, threshold=threshold)
     
     # Filter for "textile machinery" entities
@@ -27,18 +29,17 @@ def ner(text: str, labels: str, threshold: float, nested_ner: bool):
         if ent["label"].lower() == "textile machinery"
     ]
     
-    # Highlight entities with HTML
+    # 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'>{highlighted_text[ent['start']:ent['end']]}</span>" +
+            f"<span style='background-color: yellow; font-weight: bold;'>{highlighted_text[ent['start']:ent['end']]}</span>" +
             highlighted_text[ent['end']:]
         )
     
-    return gr.HTML(highlighted_text), textile_entities
+    return highlighted_text, textile_entities
 
-# Build Gradio interface
 with gr.Blocks(title="Textile Machinery NER Demo") as demo:
     gr.Markdown(
         """
@@ -79,6 +80,7 @@ Textile Machine 2 => textile machinery
             """
         )
     
+    # Display a random example
     input_text = gr.Textbox(
         value="Amelia Earhart flew her single engine Lockheed Vega 5B across the Atlantic to Paris.",
         label="Text input",
@@ -109,28 +111,32 @@ Textile Machine 2 => textile machinery
             scale=0,
         )
     
-    output = gr.HighlightedText(label="Predicted Entities")
+    # Define output components
+    output_highlighted = gr.HTML(label="Predicted Entities")
+    output_entities = gr.JSON(label="Entities")
     
     submit_btn = gr.Button("Analyze Random Snippet")
     refresh_btn = gr.Button("Get New Snippet")
     
     # Function to fetch a new random snippet
     def get_new_snippet():
-        # WARNING: Streaming datasets may have performance implications
         try:
-            sample = next(iter(subset_dataset))['text']
+            sample = next(dataset_iter)['text']
             return sample
         except StopIteration:
             return "No more snippets available."
     
+    # Connect refresh button
     refresh_btn.click(fn=get_new_snippet, outputs=input_text)
     
+    # Connect submit button
     submit_btn.click(
         fn=ner,
         inputs=[input_text, labels, threshold, nested_ner],
-        outputs=[output, gr.JSON(label="Entities")]
+        outputs=[output_highlighted, output_entities]
     )
     
+    # Define examples
     examples = [
         [
             "However, both models lack other frequent DM symptoms including the fibre-type dependent atrophy, myotonia, cataract and male-infertility.",
@@ -138,13 +144,25 @@ Textile Machine 2 => textile machinery
             0.3,
             False,
         ],
+        [
+            "Amelia Earhart flew her single engine Lockheed Vega 5B across the Atlantic to Paris.",
+            "textile machinery",
+            0.3,
+            False,
+        ],
+        [
+            "The Shore Line route of the CNS & M until 1955 served, from south to north, the Illinois communities of Chicago, Evanston, Wilmette...",
+            "textile machinery",
+            0.3,
+            False,
+        ],
         # Add more examples as needed
     ]
     
     gr.Examples(
         examples=examples,
         inputs=[input_text, labels, threshold, nested_ner],
-        outputs=[output, gr.JSON(label="Entities")],
+        outputs=[output_highlighted, output_entities],
         fn=ner,
         label="Examples",
         cache_examples=True,