Update app.py

app.py

@@ -15,14 +15,14 @@ dataset_iter = iter(
 # Load the model
 model = GLiNER.from_pretrained("max-long/textile_machines_3_oct", trust_remote_code=True)
 
-def ner(text: str, labels: str, threshold: float, nested_ner: bool):
-    # Split and clean labels
-    labels = [label.strip() for label in labels.split(",")]
+def ner(text: str):
+    labels = ["Textile Machinery"]  # Capitalized label for entity matching
+    threshold = 0.5
     
     # Predict entities using the fine-tuned GLiNER model
-    entities = model.predict_entities(text, labels, flat_ner=not nested_ner, threshold=threshold)
+    entities = model.predict_entities(text, labels, flat_ner=True, threshold=threshold)
     
-    # Filter for "textile machinery" entities
+    # Filter for "Textile Machinery" entities
     textile_entities = [
         {
             "entity": ent["label"],
@@ -32,7 +32,7 @@ def ner(text: str, labels: str, threshold: float, nested_ner: bool):
             "score": ent.get("score", 0),
         }
         for ent in entities
-        if ent["label"].lower() == "textile machinery"
+        if ent["label"] == "Textile Machinery"  # Exact match with capitalization
     ]
     
     # Prepare data for HighlightedText
@@ -50,7 +50,7 @@ with gr.Blocks(title="Textile Machinery NER Demo") as demo:
     gr.Markdown(
         """
         # Textile Machinery Entity Recognition Demo
-        This demo selects a random text snippet from the British Library's books dataset and identifies "textile machinery" entities using a fine-tuned GLiNER model.
+        This demo selects a random text snippet from a subset of the British Library's books dataset and identifies "Textile Machinery" entities using a fine-tuned GLiNER model.
         """
     )
     
@@ -62,34 +62,11 @@ with gr.Blocks(title="Textile Machinery NER Demo") as demo:
         lines=5
     )
     
-    with gr.Row():
-        labels = gr.Textbox(
-            value="textile machinery",
-            label="Labels",
-            placeholder="Enter your labels here (comma separated)",
-            scale=2,
-        )
-        threshold = gr.Slider(
-            0,
-            1,
-            value=0.3,
-            step=0.01,
-            label="Threshold",
-            info="Lower the threshold to increase how many entities get predicted.",
-            scale=1,
-        )
-        nested_ner = gr.Checkbox(
-            value=False,
-            label="Nested NER",
-            info="Allow for nested NER?",
-            scale=0,
-        )
-    
     # Define output components
     output_highlighted = gr.HTML(label="Predicted Entities")
     output_entities = gr.JSON(label="Entities")
     
-    submit_btn = gr.Button("Analyze Random Snippet")
+    submit_btn = gr.Button("Find Textile Machinery!")
     refresh_btn = gr.Button("Get New Snippet")
 
     def get_new_snippet():
@@ -110,7 +87,7 @@ with gr.Blocks(title="Textile Machinery NER Demo") as demo:
     # Connect submit button
     submit_btn.click(
         fn=ner,
-        inputs=[input_text, labels, threshold, nested_ner],
+        inputs=[input_text],
         outputs=[output_highlighted, output_entities]
     )