import gradio as gr
from transformers import pipeline

# Load the NER models
# Load the NER models
models = {
    "dslim/bert-base-NER": pipeline(
        "ner", model="dslim/bert-base-NER", grouped_entities=True
    ),
    "dslim/bert-base-NER-uncased": pipeline(
        "ner", model="dslim/bert-base-NER-uncased", grouped_entities=True
    ),
    "dslim/bert-large-NER": pipeline(
        "ner", model="dslim/bert-large-NER", grouped_entities=True
    ),
    "dslim/distilbert-NER": pipeline(
        "ner", model="dslim/distilbert-NER", grouped_entities=True
    ),
}


def process(text, model_name):
    ner = models[model_name]
    ner_results = ner(text)
    highlighted_text = []
    last_idx = 0
    for entity in ner_results:
        start = entity["start"]
        end = entity["end"]
        label = entity["entity_group"]
        # Add non-entity text
        if start > last_idx:
            highlighted_text.append((text[last_idx:start], None))
        # Add entity text
        highlighted_text.append((text[start:end], label))
        last_idx = end
    # Add any remaining text after the last entity
    if last_idx < len(text):
        highlighted_text.append((text[last_idx:], None))
    return highlighted_text


with gr.Blocks() as demo:
    gr.Markdown("# Named Entity Recognition with BERT Models")
    with gr.Row():
        model_selector = gr.Dropdown(
            choices=list(models.keys()),
            value=list(models.keys())[0],
            label="Select Model",
        )
    text_input = gr.Textbox(
        label="Enter Text",
        lines=5,
        value="Hugging Face Inc. is a company based in New York City. Its headquarters are in DUMBO, therefore very close to the Manhattan Bridge.",
    )
    output = gr.HighlightedText(label="Named Entities")
    analyze_button = gr.Button("Analyze")
    analyze_button.click(process, inputs=[text_input, model_selector], outputs=output)

demo.launch()