|
import gradio as gr |
|
from dataclasses import dataclass |
|
|
|
from pytorch_ie.annotations import LabeledSpan |
|
from pytorch_ie.auto import AutoPipeline |
|
from pytorch_ie.core import AnnotationList, annotation_field |
|
from pytorch_ie.documents import TextDocument |
|
|
|
from spacy import displacy |
|
|
|
|
|
@dataclass |
|
class ExampleDocument(TextDocument): |
|
entities: AnnotationList[LabeledSpan] = annotation_field(target="text") |
|
|
|
|
|
model_name_or_path = "pie/example-ner-spanclf-conll03" |
|
|
|
ner_pipeline = AutoPipeline.from_pretrained(model_name_or_path, device=-1, num_workers=0) |
|
|
|
|
|
|
|
def predict(text): |
|
document = ExampleDocument(text) |
|
|
|
ner_pipeline(document) |
|
|
|
doc = { |
|
"text": document.text, |
|
"ents": [{ |
|
"start": entity.start, |
|
"end": entity.end, |
|
"label": entity.label |
|
} for entity in sorted(document.entities.predictions, key=lambda e: e.start)], |
|
"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 |
|
|
|
|
|
iface = gr.Interface( |
|
fn=predict, |
|
inputs=gr.inputs.Textbox( |
|
lines=5, |
|
default="There is still some uncertainty that Musk - also chief executive of electric car maker Tesla and rocket company SpaceX - will pull off his planned buyout.", |
|
), |
|
outputs="html", |
|
) |
|
iface.launch() |
|
|
