NER / app.py
asahi417's picture
Update app.py
416b0e4
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()