Update app.py

app.py

@@ -6,13 +6,13 @@ 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 TextDocument
+from pytorch_ie.documents import TextBasedDocument
 
 from typing import List
 
 
 @dataclass
-class ExampleDocument(TextDocument):
+class ExampleDocument(TextBasedDocument):
     entities: AnnotationList[LabeledSpan] = annotation_field(target="text")
     relations: AnnotationList[BinaryRelation] = annotation_field(target="entities")
 
@@ -21,20 +21,22 @@ 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)
+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)
 
-    print(f"list detected entities:")
-    while len(document.entities.predictions) > 0:
-        entity = document.entities.predictions.pop(0)
-        print(f"entity detected: {entity}")
-        document.entities.append(entity)
+    # show predicted entities and promote them from predictions to ground-truth annotations   
+    print(f"detected entities:\n")
+    for entity in document.entities.predictions:
+        print(f"{entity}")
+        document.entities.append(entity.copy())
 
+    # execute RE pipeline
     re_pipeline(document)
 
     t = PrettyTable()