Update app.py

app.py

@@ -3,9 +3,10 @@ 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 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 TextDocument
 
 from typing import List
 
@@ -16,17 +17,11 @@ class ExampleDocument(TextDocument):
     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_model_name_or_path = "pie/example-ner-spanclf-conll03"
+re_model_name_or_path = "pie/example-re-textclf-tacred"
 
-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)
+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)
 
 
 def predict(text):
@@ -57,7 +52,10 @@ def predict(text):
 
 iface = gr.Interface(
     fn=predict,
-    inputs="textbox",
+    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()