natayadev/pdf

1import re2import gradio as gr3from dataclasses import dataclass4from prettytable import PrettyTable5 6from pytorch_ie import AnnotationList, BinaryRelation, Span, LabeledSpan, Pipeline, TextDocument, annotation_field7from pytorch_ie.models import TransformerSpanClassificationModel, TransformerTextClassificationModel8from pytorch_ie.taskmodules import TransformerSpanClassificationTaskModule, TransformerRETextClassificationTaskModule9 10from typing import List11 12 13@dataclass14class ExampleDocument(TextDocument):15    entities: AnnotationList[LabeledSpan] = annotation_field(target="text")16    relations: AnnotationList[BinaryRelation] = annotation_field(target="entities")17 18 19model_name_or_path = "pie/example-ner-spanclf-conll03"20ner_taskmodule = TransformerSpanClassificationTaskModule.from_pretrained(model_name_or_path)21ner_model = TransformerSpanClassificationModel.from_pretrained(model_name_or_path)22 23ner_pipeline = Pipeline(model=ner_model, taskmodule=ner_taskmodule, device=-1, num_workers=0)24 25model_name_or_path = "pie/example-re-textclf-tacred"26re_taskmodule = TransformerRETextClassificationTaskModule.from_pretrained(model_name_or_path)27re_model = TransformerTextClassificationModel.from_pretrained(model_name_or_path)28 29re_pipeline = Pipeline(model=re_model, taskmodule=re_taskmodule, device=-1, num_workers=0)30 31 32def predict(text):33    document = ExampleDocument(text)34 35    ner_pipeline(document, predict_field="entities")36 37    for entity in document.entities.predictions:38        document.entities.append(entity)39 40    re_pipeline(document, predict_field="relations")41 42    t = PrettyTable()43    t.field_names = ["head", "tail", "relation"]44    t.align = "l"45    for relation in document.relations.predictions:46        t.add_row([str(relation.head), str(relation.tail), relation.label])47 48    html = t.get_html_string(format=True)49    html = (50        "<div style='max-width:100%; max-height:360px; overflow:auto'>"51        + html52        + "</div>"53    )54    55    return html56 57 58iface = gr.Interface(59    fn=predict,60    inputs="textbox",61    outputs="html",62)63iface.launch()