fix: use documents and annotations directly from pie_documents (#499)

that was missed in #476

This also moves the helper method `_config_to_str` to tests package root
and fixes an import bug in an *unused* test fixture in
tests/taskmodules/test_simple_transformer_text_classification.py.

Author: ArneBinder

src/pytorch_ie/metrics/statistics/token_count_collector.py

@@ -2,10 +2,9 @@
 from typing import Any, Dict, Optional, Type, Union
 
 from pie_core import Document, DocumentStatistic
+from pie_documents.documents import TextBasedDocument
 from transformers import AutoTokenizer, PreTrainedTokenizer
 
-from pytorch_ie.documents import TextBasedDocument
-
 logger = logging.getLogger(__name__)
 
 

src/pytorch_ie/taskmodules/simple_transformer_text_classification.py

@@ -13,13 +13,13 @@
 import numpy as np
 import torch
 from pie_core import TaskEncoding, TaskModule
+from pie_documents.annotations import Label
+from pie_documents.documents import TextDocumentWithLabel
 from transformers import AutoTokenizer
 from transformers.file_utils import PaddingStrategy
 from transformers.tokenization_utils_base import TruncationStrategy
 from typing_extensions import TypeAlias
 
-from pytorch_ie.annotations import Label
-from pytorch_ie.documents import TextDocumentWithLabel
 from pytorch_ie.models.transformer_text_classification import ModelOutputType, ModelStepInputType
 
 logger = logging.getLogger(__name__)

src/pytorch_ie/taskmodules/transformer_re_text_classification.py

@@ -25,6 +25,11 @@
 import numpy as np
 import torch
 from pie_core import AnnotationLayer, Document, TaskEncoding, TaskModule
+from pie_documents.documents import (
+    TextDocument,
+    TextDocumentWithLabeledSpansAndBinaryRelations,
+    TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions,
+)
 from transformers import AutoTokenizer
 from transformers.file_utils import PaddingStrategy
 from transformers.tokenization_utils_base import TruncationStrategy
@@ -37,11 +42,6 @@
     NaryRelation,
     Span,
 )
-from pytorch_ie.documents import (
-    TextDocument,
-    TextDocumentWithLabeledSpansAndBinaryRelations,
-    TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions,
-)
 from pytorch_ie.models.transformer_text_classification import ModelOutputType, ModelStepInputType
 from pytorch_ie.taskmodules.interface import ChangesTokenizerVocabSize
 from pytorch_ie.utils.span import get_token_slice, is_contained_in

src/pytorch_ie/taskmodules/transformer_seq2seq.py

@@ -12,13 +12,13 @@
 from typing import Any, Dict, Iterator, List, Optional, Sequence, Tuple, Type, Union
 
 from pie_core import Annotation, TaskEncoding, TaskModule
+from pie_documents.annotations import BinaryRelation, LabeledSpan
+from pie_documents.documents import TextDocument, TextDocumentWithLabeledSpansAndBinaryRelations
 from transformers import AutoTokenizer
 from transformers.file_utils import PaddingStrategy
 from transformers.tokenization_utils_base import TruncationStrategy
 from typing_extensions import TypeAlias
 
-from pytorch_ie.annotations import BinaryRelation, LabeledSpan
-from pytorch_ie.documents import TextDocument, TextDocumentWithLabeledSpansAndBinaryRelations
 from pytorch_ie.models.transformer_seq2seq import ModelOutputType, ModelStepInputType
 
 InputEncodingType: TypeAlias = Dict[str, Sequence[int]]

src/pytorch_ie/taskmodules/transformer_span_classification.py

@@ -14,18 +14,18 @@
 import torch
 import torch.nn.functional as F
 from pie_core import TaskEncoding, TaskModule
-from transformers import AutoTokenizer
-from transformers.file_utils import PaddingStrategy
-from transformers.tokenization_utils_base import BatchEncoding, TruncationStrategy
-from typing_extensions import TypeAlias
-
-from pytorch_ie.annotations import LabeledSpan, MultiLabeledSpan, Span
-from pytorch_ie.documents import (
+from pie_documents.annotations import LabeledSpan, MultiLabeledSpan, Span
+from pie_documents.documents import (
     TextDocument,
     TextDocumentWithLabeledSpans,
     TextDocumentWithLabeledSpansAndLabeledPartitions,
     TextDocumentWithLabeledSpansAndSentences,
 )
+from transformers import AutoTokenizer
+from transformers.file_utils import PaddingStrategy
+from transformers.tokenization_utils_base import BatchEncoding, TruncationStrategy
+from typing_extensions import TypeAlias
+
 from pytorch_ie.models.transformer_span_classification import ModelOutputType, ModelStepInputType
 
 InputEncodingType: TypeAlias = BatchEncoding

src/pytorch_ie/taskmodules/transformer_text_classification.py

@@ -25,13 +25,13 @@
 import numpy as np
 import torch
 from pie_core import TaskEncoding, TaskModule
+from pie_documents.annotations import Label, MultiLabel
+from pie_documents.documents import TextDocument, TextDocumentWithLabel, TextDocumentWithMultiLabel
 from transformers import AutoTokenizer
 from transformers.file_utils import PaddingStrategy
 from transformers.tokenization_utils_base import TruncationStrategy
 from typing_extensions import TypeAlias
 
-from pytorch_ie.annotations import Label, MultiLabel
-from pytorch_ie.documents import TextDocument, TextDocumentWithLabel, TextDocumentWithMultiLabel
 from pytorch_ie.models.transformer_text_classification import ModelOutputType, ModelStepInputType
 
 logger = logging.getLogger(__name__)

src/pytorch_ie/taskmodules/transformer_token_classification.py

@@ -14,17 +14,17 @@
 import torch
 import torch.nn.functional as F
 from pie_core import TaskEncoding, TaskModule
+from pie_documents.annotations import LabeledSpan, Span
+from pie_documents.documents import (
+    TextDocument,
+    TextDocumentWithLabeledSpans,
+    TextDocumentWithLabeledSpansAndLabeledPartitions,
+)
 from transformers import AutoTokenizer
 from transformers.file_utils import PaddingStrategy
 from transformers.tokenization_utils_base import BatchEncoding, TruncationStrategy
 from typing_extensions import TypeAlias
 
-from pytorch_ie.annotations import LabeledSpan, Span
-from pytorch_ie.documents import (
-    TextDocument,
-    TextDocumentWithLabeledSpans,
-    TextDocumentWithLabeledSpansAndLabeledPartitions,
-)
 from pytorch_ie.models.transformer_token_classification import ModelOutputType, ModelStepInputType
 from pytorch_ie.utils.span import (
     bio_tags_to_spans,

src/pytorch_ie/utils/document.py

@@ -9,11 +9,9 @@
 from pie_core.document import BaseAnnotationList
 from pie_documents.annotations import Span
 from pie_documents.document.processing import text_based_document_to_token_based
-from pie_documents.documents import TextBasedDocument, TokenBasedDocument
+from pie_documents.documents import TextBasedDocument, TokenBasedDocument, WithMetadata
 from transformers import PreTrainedTokenizer
 
-from pytorch_ie.documents import WithMetadata
-
 logger = logging.getLogger(__name__)
 
 

tests/__init__.py

@@ -1,4 +1,10 @@
 from pathlib import Path
+from typing import Any, Dict
 
 TESTS_ROOT = Path(__file__).parent
 FIXTURES_ROOT = TESTS_ROOT / "fixtures"
+
+
+def _config_to_str(cfg: Dict[str, Any]) -> str:
+    result = "-".join([f"{k}={cfg[k]}" for k in sorted(cfg)])
+    return result

tests/conftest.py

@@ -3,9 +3,9 @@
 
 import pytest
 from pie_core import AnnotationLayer, annotation_field
+from pie_documents.annotations import BinaryRelation, LabeledSpan, Span
+from pie_documents.documents import TextDocument
 
-from pytorch_ie.annotations import BinaryRelation, LabeledSpan, Span
-from pytorch_ie.documents import TextDocument
 from tests import FIXTURES_ROOT