from transformers import pipeline
import gradio as gr

model_checkpoint = "MuntasirHossain/bert-finetuned-ner"

model = pipeline("token-classification", model=model_checkpoint, aggregation_strategy="simple")


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

description = "This AI model is trained to identify and classify named entities in unstructured text." 
title = "Named Entity Recognition"
theme = "grass"
examples=[["Mount Everest is Earth's highest mountain, located in the Mahalangur Himal sub-range of the Himalayas. The China-Nepal border runs across it."]]

gr.Interface(fn=predict,
    inputs="textbox",
    outputs="text",
    title=title,
    theme = theme,
    description=description,
    examples=examples,
).launch()