import gradio as gr
from transformers import AutoTokenizer, AutoModelForTokenClassification,pipeline 

tokenizer = AutoTokenizer.from_pretrained("dbmdz/electra-large-discriminator-finetuned-conll03-english")

model = AutoModelForTokenClassification.from_pretrained("dbmdz/electra-large-discriminator-finetuned-conll03-english")



ner_pipeline = pipeline("ner",model=model, 
                tokenizer=tokenizer)

examples = [
    "where did Wandobire's laptop come from, was it africa or uganda?",
]

examples_2 = [
    "The Intern was oriented on ICT setup and Infrastructure of Soroti University, drafted workplan and started off the Internship. Simon was encouraged to take the Internship seriously as there was a lot to learn.",
]

examples_3 = [
    "Partially done, expected a better result based on Steven's experienced. More effort needed ...",
]




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


gr.Interface(ner_electra,
             gr.Textbox(placeholder="Enter sentence here..."), 
             gr.HighlightedText(),
             examples=[[examples],[examples_2],[examples_3],],
    title="Comparative Natural Entity Recognition Model by Brian Joram Wandobire",
    description="takes in a comment as an input and outputs the Entities",
             ).launch()