ArneBinder's picture
replace example input
8b928dc
import re
import gradio as gr
from dataclasses import dataclass
from prettytable import PrettyTable
from pytorch_ie.annotations import LabeledSpan, BinaryRelation
from pytorch_ie.auto import AutoPipeline
from pytorch_ie.core import AnnotationList, annotation_field
from pytorch_ie.documents import TextBasedDocument
from pytorch_ie.taskmodules import *
from pytorch_ie.models import *
from typing import List
@dataclass
class ExampleDocument(TextBasedDocument):
entities: AnnotationList[LabeledSpan] = annotation_field(target="text")
relations: AnnotationList[BinaryRelation] = annotation_field(target="entities")
ner_model_name_or_path = "pie/example-ner-spanclf-conll03"
re_model_name_or_path = "pie/example-re-textclf-tacred"
ner_pipeline = AutoPipeline.from_pretrained(ner_model_name_or_path, device=-1, num_workers=0)
re_pipeline = AutoPipeline.from_pretrained(re_model_name_or_path, device=-1, num_workers=0, taskmodule_kwargs=dict(create_relation_candidates=True))
def predict(text):
document = ExampleDocument(text)
# execute NER pipeline
ner_pipeline(document)
# show predicted entities and promote them from predictions to ground-truth annotations
print(f"detected entities:")
for entity in document.entities.predictions:
print(f"'{entity}', label={entity.label}, score={entity.score}")
document.entities.append(entity.copy())
# execute RE pipeline
re_pipeline(document)
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=gr.inputs.Textbox(
lines=5,
default="“Making a super tasty alt-chicken wing is only half of it,” said Po Bronson, general partner at SOSV and managing director of IndieBio.",
),
outputs="html",
)
iface.launch()