Update app.py

app.py

@@ -1,9 +1,11 @@
 import gradio as gr
 from dataclasses import dataclass
 
-from pytorch_ie import AnnotationList, LabeledSpan, Pipeline, TextDocument, annotation_field
-from pytorch_ie.models import TransformerSpanClassificationModel
-from pytorch_ie.taskmodules import TransformerSpanClassificationTaskModule
+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
 
 
@@ -13,10 +15,8 @@ class ExampleDocument(TextDocument):
 
 
 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)
+ner_pipeline = AutoPipeline.from_pretrained(model_name_or_path, device=-1, num_workers=0)
 
 
 def predict(text):
@@ -44,5 +44,12 @@ def predict(text):
     return html
 
 
-iface = gr.Interface(fn=predict, inputs="textbox", outputs="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()