import gradio as gr
from transformers import pipeline

# Model names (keeping it programmatic)
model_names = [
    "dslim/bert-base-NER",
    "dslim/bert-base-NER-uncased",
    "dslim/bert-large-NER",
    "dslim/distilbert-NER",
]

example_sent = (
    "Nim Chimpsky was a chimpanzee at Columbia University named after Noam Chomsky."
)

# Programmatically build the model info dict
model_info = {
    model_name: {
        "link": f"https://huggingface.co/{model_name}",
        "usage": f"""from transformers import pipeline
ner = pipeline("ner", model="{model_name}", grouped_entities=True)
result = ner("{example_sent}")
print(result)""",
    }
    for model_name in model_names
}

# Load models into a dictionary programmatically for the analyze function
models = {
    model_name: pipeline("ner", model=model_name, grouped_entities=True)
    for model_name in model_names
}


# Function to display model info (link and usage code)
def display_model_info(model_name):
    info = model_info[model_name]
    usage_code = info["usage"]
    link_button = f'[Open model page for {model_name} ]({info["link"]})'
    return usage_code, link_button


# Function to run NER on input text
def analyze_text(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 (NER) with BERT Models")

    # Dropdown for model selection
    model_selector = gr.Dropdown(
        choices=list(model_info.keys()),
        value=list(model_info.keys())[0],
        label="Select Model",
    )

    # Textbox for input text
    text_input = gr.Textbox(
        label="Enter Text",
        lines=5,
        value=example_sent,
    )
    analyze_button = gr.Button("Run NER Model")
    output = gr.HighlightedText(label="NER Result", combine_adjacent=True)

    # Outputs: usage code, model page link, and analyze button
    code_output = gr.Code(label="Use this model", visible=True)
    link_output = gr.Markdown(
        f"[Open model page for {model_selector} ]({model_selector})"
    )
    # Button for analyzing the input text
    analyze_button.click(
        analyze_text, inputs=[text_input, model_selector], outputs=output
    )

    # Trigger the code output and model link when model is changed
    model_selector.change(
        display_model_info, inputs=[model_selector], outputs=[code_output, link_output]
    )

    # Call the display_model_info function on load to set initial values
    demo.load(
        fn=display_model_info,
        inputs=[model_selector],
        outputs=[code_output, link_output],
    )

demo.launch()