import gradio as gr
from tner import TransformersNER
from spacy import displacy

model = TransformersNER("tner/roberta-large-ontonotes5")
#model = TransformersNER("tner/bertweet-large-tweetner7-all")

examples = [
    "Jacob Collier is a Grammy awarded artist from England.",
    "When Sebastian Thrun started working on self-driving cars at Google in 2007 , few people outside of the company took him seriously.",
    "But Google is starting from behind. The company made a late push into hardware, and Apple’s Siri, available on iPhones, and Amazon’s Alexa software, which runs on its Echo and Dot devices, have clear leads in consumer adoption."
]


def predict(text):
    output = model.predict([text])
    tokens = output['input'][0]

    def retain_char_position(p):
        if p == 0:
            return 0
        return len(' '.join(tokens[:p])) + 1

    doc = {
        "text": text,
        "ents": [{
            "start": retain_char_position(entity['position'][0]),
            "end": retain_char_position(entity['position'][-1]) + len(entity['entity'][-1]),
            "label": entity['type']
        } for entity in output['entity_prediction'][0]],
        "title": None
    }

    html = displacy.render(doc, style="ent", page=True, manual=True, minify=True)
    html = (
        "<div style='max-width:100%; max-height:360px; overflow:auto'>"
        + html
        + "</div>"
    )
    
    return html


demo = gr.Interface(
    fn=predict,
    inputs=gr.inputs.Textbox(
        lines=5,
        placeholder="Input sentence...",
    ),
    outputs="html",
    examples=examples
)
demo.launch()