Update app.py

app.py

@@ -12,34 +12,39 @@ dataset_iter = load_dataset(
 ).shuffle(seed=42)  # Shuffle added
 
 # Load the model
-model = GLiNER.from_pretrained("urchade/gliner_multi-v2.1", trust_remote_code=True)
+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):
     # Convert user-provided labels (comma-separated string) into a list
     labels_list = [label.strip() for label in labels.split(",")]
     
+    # Truncate the text to avoid length exceeding model limits (e.g., 384 tokens)
+    max_length = 384
+    truncated_text = text[:max_length]
+    
     # Predict entities using the GLiNER model
-    entities = model.predict_entities(text, labels_list, flat_ner=not nested_ner, threshold=threshold)
+    entities = model.predict_entities(truncated_text, labels_list, flat_ner=not nested_ner, threshold=threshold)
     
     # Prepare entities for color-coded display using gr.HighlightedText
     highlights = [(ent["start"], ent["end"], ent["label"]) for ent in entities]
     
+    # Return both the highlighted text and the raw entities in JSON format
     return {
-        "text": text,
+        "text": truncated_text,
         "entities": highlights
-    }
+    }, entities  # Return both outputs: the first for HighlightedText, the second for JSON
 
 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 the GLiNER NER model. You can specify the entities you want to find, and results will be displayed in a color-coded format.
+        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, and results will be displayed in a color-coded format.
         """
     )
     
     # Display a random example
     input_text = gr.Textbox(
-        value="Click 'Get New Snippet' to load a random sample from the British Library Dataset",
+        value="The machine is fed by means of an endless apron, the wool entering at the smaller end...",
         label="Text input",
         placeholder="Enter your text here",
         lines=5
@@ -47,7 +52,7 @@ with gr.Blocks(title="General NER with Color-Coded Output") as demo:
 
     with gr.Row() as row:
         labels = gr.Textbox(
-            value="Places, People",  # Default example labels
+            value="Machine, Wool",  # Default example labels
             label="Labels",
             placeholder="Enter your labels here (comma separated)",
             scale=2,