app.py

# Necessary imports
import warnings

warnings.filterwarnings("ignore")

import gradio as gr
from ner import pii_ner


# Examples to display in the interface
examples = [
    [
        "Hi support, I can't log in! My account username is 'johndoe88'. Every time I try, it says 'invalid credentials'. Please reset my password. You can reach me at (555) 123-4567 or johnd@example.com",
        "email, phone_number, user_name",
        0.3,
        False,
        False,
    ],
    [
        "Patient John Doe, MRN 123456, diagnosed with diabetes.",
        "name, medical record number, condition",
        0.3,
        False,
        False,
    ],
    [
        "Client Jane Doe vs. Corporation ABC, case #2024-CV-001",
        "name, organization, case number",
        0.3,
        False,
        False,
    ],
]


# Launch the gradio UI
with gr.Blocks(title="GLiNER-PII") as demo:
    gr.Markdown(
        """
        # GLiNER-PII
        GLiNER-PII is a successor to the Gretel GLiNER PII/PHI models. Built on the GLiNER bi-large base, it detects and classifies a broad range of Personally Identifiable Information (PII) and Protected Health Information (PHI) in structured and unstructured text. It is non-generative and produces span-level entity annotations with confidence scores across 55+ categories. This model was developed by NVIDIA.
        
        ## Links
        * Model: https://huggingface.co/nvidia/gliner-pii
        * Training dataset: https://huggingface.co/datasets/nvidia/nemotron-pii
        * GLiNER library: https://pypi.org/project/gliner/
        """
    )

    # Text input
    input_text = gr.TextArea(label="Text input", placeholder="Enter your text here")

    # Labels, threshold, nested NER
    with gr.Row() as row:
        labels = gr.Textbox(
            label="Labels",
            placeholder="Enter your PII/PHI labels here (comma separated)",
            scale=2,
        )
        threshold = gr.Slider(
            0,
            1,
            value=0.5,
            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,
        )
        multi_label = gr.Checkbox(
            value=False,
            label="Multi-label",
            info="Allow for multi-label?",
            scale=0,
        )

    # Output
    output = gr.HighlightedText(label="Predicted PII Entities")

    # Submit button
    submit_btn = gr.Button("Submit")
    examples = gr.Examples(
        examples,
        fn=pii_ner,
        inputs=[input_text, labels, threshold, nested_ner, multi_label],
        outputs=output,
        cache_examples=True,
        cache_mode="lazy",
    )

    # Events when submitting
    input_text.submit(
        fn=pii_ner,
        inputs=[input_text, labels, threshold, nested_ner, multi_label],
        outputs=output,
    )
    labels.submit(
        fn=pii_ner,
        inputs=[input_text, labels, threshold, nested_ner, multi_label],
        outputs=output,
    )
    threshold.release(
        fn=pii_ner,
        inputs=[input_text, labels, threshold, nested_ner, multi_label],
        outputs=output,
    )
    submit_btn.click(
        fn=pii_ner,
        inputs=[input_text, labels, threshold, nested_ner, multi_label],
        outputs=output,
    )
    nested_ner.change(
        fn=pii_ner,
        inputs=[input_text, labels, threshold, nested_ner, multi_label],
        outputs=output,
    )
    multi_label.change(
        fn=pii_ner,
        inputs=[input_text, labels, threshold, nested_ner, multi_label],
        outputs=output,
    )

demo.launch(debug=False, theme=gr.themes.Soft())