From 2156973b0599b504d600c94098f36c00babc2be8 Mon Sep 17 00:00:00 2001
From: Arne Binder <arne.binder@dfki.de>
Date: Fri, 2 Feb 2024 15:06:38 +0100
Subject: [PATCH 1/2] re-export annotations from pytorch-ie

---
 src/pie_modules/annotations.py | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/src/pie_modules/annotations.py b/src/pie_modules/annotations.py
index 5268f3c57..5b6946c86 100644
--- a/src/pie_modules/annotations.py
+++ b/src/pie_modules/annotations.py
@@ -1,7 +1,19 @@
 import dataclasses
 from typing import Optional
 
-from pytorch_ie.annotations import Span
+# re-export all annotations from pytorch_ie to have a single entry point
+from pytorch_ie.annotations import (
+    BinaryRelation,
+    Label,
+    LabeledMultiSpan,
+    LabeledSpan,
+    MultiLabel,
+    MultiLabeledBinaryRelation,
+    MultiLabeledMultiSpan,
+    MultiLabeledSpan,
+    NaryRelation,
+    Span,
+)
 from pytorch_ie.core import Annotation
 
 

From 9f2f392b766f85ea8db765ada729d66f10878423 Mon Sep 17 00:00:00 2001
From: Arne Binder <arne.binder@dfki.de>
Date: Fri, 2 Feb 2024 15:06:58 +0100
Subject: [PATCH 2/2] re-export documents from pytorch-ie

---
 src/pie_modules/documents.py | 26 ++++++++++++++++++++++----
 1 file changed, 22 insertions(+), 4 deletions(-)

diff --git a/src/pie_modules/documents.py b/src/pie_modules/documents.py
index 26ae4c549..ead62e2ef 100644
--- a/src/pie_modules/documents.py
+++ b/src/pie_modules/documents.py
@@ -1,14 +1,32 @@
 import dataclasses
 
-from pytorch_ie import AnnotationLayer
-from pytorch_ie.annotations import BinaryRelation, LabeledSpan
-from pytorch_ie.core import AnnotationList, annotation_field
-from pytorch_ie.documents import TextBasedDocument, TokenBasedDocument
+from pytorch_ie.core import AnnotationLayer, AnnotationList, annotation_field
+
+# re-export all documents from pytorch_ie to have a single entry point
+from pytorch_ie.documents import (
+    TextBasedDocument,
+    TextDocumentWithLabel,
+    TextDocumentWithLabeledPartitions,
+    TextDocumentWithLabeledSpans,
+    TextDocumentWithLabeledSpansAndBinaryRelations,
+    TextDocumentWithLabeledSpansAndLabeledPartitions,
+    TextDocumentWithLabeledSpansAndSentences,
+    TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions,
+    TextDocumentWithMultiLabel,
+    TextDocumentWithSentences,
+    TextDocumentWithSpans,
+    TextDocumentWithSpansAndBinaryRelations,
+    TextDocumentWithSpansAndLabeledPartitions,
+    TextDocumentWithSpansBinaryRelationsAndLabeledPartitions,
+    TokenBasedDocument,
+)
 
 from pie_modules.annotations import (
     AbstractiveSummary,
+    BinaryRelation,
     ExtractiveAnswer,
     GenerativeAnswer,
+    LabeledSpan,
     Question,
 )