|
import re |
|
import gradio as gr |
|
from dataclasses import dataclass |
|
from prettytable import PrettyTable |
|
|
|
from pytorch_ie import AnnotationList, BinaryRelation, Span, LabeledSpan, Pipeline, TextDocument, annotation_field |
|
from pytorch_ie.models import TransformerSpanClassificationModel, TransformerTextClassificationModel |
|
from pytorch_ie.taskmodules import TransformerSpanClassificationTaskModule, TransformerRETextClassificationTaskModule |
|
|
|
from typing import List |
|
|
|
|
|
@dataclass |
|
class ExampleDocument(TextDocument): |
|
entities: AnnotationList[LabeledSpan] = annotation_field(target="text") |
|
relations: AnnotationList[BinaryRelation] = annotation_field(target="entities") |
|
|
|
|
|
model_name_or_path = "pie/example-ner-spanclf-conll03" |
|
ner_taskmodule = TransformerSpanClassificationTaskModule.from_pretrained(model_name_or_path) |
|
ner_model = TransformerSpanClassificationModel.from_pretrained(model_name_or_path) |
|
|
|
ner_pipeline = Pipeline(model=ner_model, taskmodule=ner_taskmodule, device=-1, num_workers=0) |
|
|
|
model_name_or_path = "pie/example-re-textclf-tacred" |
|
re_taskmodule = TransformerRETextClassificationTaskModule.from_pretrained(model_name_or_path) |
|
re_model = TransformerTextClassificationModel.from_pretrained(model_name_or_path) |
|
|
|
re_pipeline = Pipeline(model=re_model, taskmodule=re_taskmodule, device=-1, num_workers=0) |
|
|
|
|
|
def predict(text): |
|
document = ExampleDocument(text) |
|
|
|
ner_pipeline(document, predict_field="entities") |
|
|
|
for entity in document.entities.predictions: |
|
document.entities.append(entity) |
|
|
|
re_pipeline(document, predict_field="relations") |
|
|
|
t = PrettyTable() |
|
t.field_names = ["head", "tail", "relation"] |
|
t.align = "l" |
|
for relation in document.relations.predictions: |
|
t.add_row([str(relation.head), str(relation.tail), relation.label]) |
|
|
|
html = t.get_html_string(format=True) |
|
html = ( |
|
"<div style='max-width:100%; max-height:360px; overflow:auto'>" |
|
+ html |
|
+ "</div>" |
|
) |
|
|
|
return html |
|
|
|
|
|
iface = gr.Interface( |
|
fn=predict, |
|
inputs="textbox", |
|
outputs="html", |
|
) |
|
iface.launch() |
|
|