from transformers import pipeline
import gradio as gr


# List of NER models
models = ["dslim/bert-base-NER", "dslim/bert-base-NER-uncased", "dslim/bert-large-NER"]


def ner(text, model_choice):
    ner_pipeline = pipeline("ner", model=model_choice)
    output = ner_pipeline(text)
    return {"text": text, "entities": output}


examples = [
    "Does Chicago have any stores and does Joe live here?",
]

demo = gr.Interface(
    fn=ner,
    inputs=[
        gr.Textbox(placeholder="Enter sentence here..."),
        gr.Dropdown(choices=models, label="Choose NER Model"),
    ],
    outputs=gr.HighlightedText(),
    examples=examples,
)

demo.launch()