import from pie_core directly

src/pie_modules/annotations.py

@@ -1,5 +1,7 @@
 import dataclasses
-from typing import Optional, Tuple
+from typing import Optional
+
+from pie_core import Annotation
 
 # re-export all annotations from pytorch_ie to have a single entry point
 from pytorch_ie.annotations import (
@@ -15,7 +17,6 @@
     Span,
     _post_init_single_label,
 )
-from pytorch_ie.core import Annotation
 
 
 @dataclasses.dataclass(eq=True, frozen=True)

src/pie_modules/document/processing/merge_multi_spans.py

@@ -2,7 +2,7 @@
 
 from typing import Generic, TypeVar, get_args
 
-from pytorch_ie.core import Annotation, Document
+from pie_core import Annotation, Document
 
 from pie_modules.annotations import MultiSpan, Span
 

src/pie_modules/document/processing/merge_spans_via_relation.py

@@ -1,7 +1,7 @@
 import logging
 from typing import Optional, Sequence, Set, Tuple, TypeVar, Union
 
-from pytorch_ie.core import AnnotationLayer, Document
+from pie_core import AnnotationLayer, Document
 
 from pie_modules.annotations import BinaryRelation, LabeledMultiSpan, LabeledSpan
 from pie_modules.utils import resolve_type

src/pie_modules/document/processing/regex_partitioner.py

@@ -6,7 +6,7 @@
 import statistics
 from typing import Any, Callable, Iterable, Iterator, Match, TypeVar
 
-from pytorch_ie.core import EnterDatasetMixin, ExitDatasetMixin
+from pie_core import EnterDatasetMixin, ExitDatasetMixin
 
 from pie_modules.annotations import LabeledSpan
 from pie_modules.documents import TextBasedDocument

src/pie_modules/document/processing/relation_argument_sorter.py

@@ -3,7 +3,7 @@
 import logging
 from typing import TypeVar
 
-from pytorch_ie.core import Annotation, AnnotationLayer, Document
+from pie_core import Annotation, AnnotationLayer, Document
 
 from pie_modules.annotations import BinaryRelation, LabeledMultiSpan, LabeledSpan
 

src/pie_modules/document/processing/text_span_trimmer.py

@@ -3,7 +3,7 @@
 import logging
 from typing import TypeVar
 
-from pytorch_ie.core import AnnotationLayer, Document
+from pie_core import AnnotationLayer, Document
 
 from pie_modules.annotations import LabeledMultiSpan, LabeledSpan, Span
 

src/pie_modules/document/processing/tokenization.py

@@ -16,7 +16,7 @@
     Union,
 )
 
-from pytorch_ie.core import Annotation
+from pie_core import Annotation
 from transformers import PreTrainedTokenizer
 
 from pie_modules.annotations import MultiSpan, Span

src/pie_modules/documents.py

@@ -1,6 +1,6 @@
 import dataclasses
 
-from pytorch_ie.core import AnnotationLayer, annotation_field
+from pie_core import AnnotationLayer, annotation_field
 
 # re-export all documents from pytorch_ie to have a single entry point
 from pytorch_ie.documents import (

src/pie_modules/metrics/confusion_matrix.py

@@ -3,8 +3,8 @@
 from typing import Callable, Dict, Optional, Tuple, Union
 
 import pandas as pd
-from pytorch_ie.core import Annotation, Document, DocumentMetric
-from pytorch_ie.utils.hydra import resolve_target
+from pie_core import Annotation, Document, DocumentMetric
+from pie_core.utils.hydra import resolve_target
 
 logger = logging.getLogger(__name__)
 

src/pie_modules/metrics/f1.py

@@ -4,8 +4,8 @@
 from typing import Callable, Collection, Dict, Hashable, Optional, Tuple, Union
 
 import pandas as pd
-from pytorch_ie.core import Annotation, Document, DocumentMetric
-from pytorch_ie.utils.hydra import resolve_target
+from pie_core import Annotation, Document, DocumentMetric
+from pie_core.utils.hydra import resolve_target
 
 logger = logging.getLogger(__name__)
 

src/pie_modules/metrics/relation_argument_distance_collector.py

@@ -1,8 +1,8 @@
 from collections import defaultdict
 from typing import Any, Dict, List, Optional, Type, Union
 
-from pytorch_ie.core import Document, DocumentStatistic
-from pytorch_ie.utils.hydra import resolve_target
+from pie_core import Document, DocumentStatistic
+from pie_core.utils.hydra import resolve_target
 from transformers import AutoTokenizer, PreTrainedTokenizer
 
 from pie_modules.annotations import BinaryRelation, NaryRelation, Span

src/pie_modules/metrics/span_coverage_collector.py

@@ -1,7 +1,7 @@
 import logging
 from typing import Any, Dict, List, Optional, Set, Type, Union
 
-from pytorch_ie.core import Document, DocumentStatistic
+from pie_core import Document, DocumentStatistic
 from transformers import AutoTokenizer, PreTrainedTokenizer
 
 from pie_modules.annotations import LabeledMultiSpan, Span

src/pie_modules/metrics/span_length_collector.py

@@ -2,7 +2,7 @@
 from collections import defaultdict
 from typing import Any, Callable, Dict, List, Optional, Type, Union
 
-from pytorch_ie.core import Document, DocumentStatistic
+from pie_core import Document, DocumentStatistic
 from transformers import AutoTokenizer, PreTrainedTokenizer
 
 from pie_modules.annotations import Span

src/pie_modules/metrics/squad_f1.py

@@ -6,7 +6,7 @@
 from typing import Dict, List
 
 import pandas as pd
-from pytorch_ie.core import DocumentMetric
+from pie_core import DocumentMetric
 
 from pie_modules.documents import TextDocumentWithQuestionsAndExtractiveAnswers
 

src/pie_modules/metrics/statistics.py

@@ -2,7 +2,7 @@
 from collections import defaultdict
 from typing import Any, Callable, Dict, List, Optional, Type, Union
 
-from pytorch_ie.core import Document, DocumentStatistic
+from pie_core import Document, DocumentStatistic
 from transformers import AutoTokenizer, PreTrainedTokenizer
 
 from pie_modules.documents import TextBasedDocument

src/pie_modules/models/common/has_taskmodule.py

@@ -1,7 +1,6 @@
 from typing import Any, Dict, Optional
 
-from pytorch_ie.auto import AutoTaskModule
-from pytorch_ie.core import TaskModule
+from pie_core import AutoTaskModule, TaskModule
 
 from pie_modules.models.interface import RequiresTaskmoduleConfig
 

src/pie_modules/taskmodules/common/interfaces.py

@@ -1,7 +1,7 @@
 import abc
 from typing import Any, Dict, Generic, List, Optional, Tuple, TypeVar
 
-from pytorch_ie.core import Annotation
+from pie_core import Annotation
 
 # Annotation Encoding type: encoding for a single annotation
 AE = TypeVar("AE")

src/pie_modules/taskmodules/common/mixins.py

@@ -7,7 +7,7 @@
 import pandas as pd
 import torch
 import torch.nn.functional as F
-from pytorch_ie.core import Annotation
+from pie_core import Annotation
 from torch import Tensor
 
 logger = logging.getLogger(__name__)

src/pie_modules/taskmodules/common/taskmodule_with_document_converter.py

@@ -1,12 +1,13 @@
 from abc import ABC, abstractmethod
 from typing import Generic, Iterable, Iterator, Optional, Sequence, Type, TypeVar, Union
 
-from pytorch_ie.core import Document, TaskModule
-from pytorch_ie.core.taskmodule import (
+from pie_core import (
+    Document,
     IterableTaskEncodingDataset,
     TaskEncoding,
     TaskEncodingDataset,
     TaskEncodingSequence,
+    TaskModule,
 )
 from typing_extensions import TypeAlias
 

src/pie_modules/taskmodules/cross_text_binary_coref.py

@@ -15,7 +15,7 @@
 )
 
 import torch
-from pytorch_ie.core import Annotation, TaskEncoding, TaskModule
+from pie_core import Annotation, TaskEncoding, TaskModule
 from pytorch_ie.utils.window import get_window_around_slice
 from torchmetrics import MetricCollection
 from torchmetrics.classification import (

src/pie_modules/taskmodules/extractive_question_answering.py

@@ -4,7 +4,7 @@
 
 import numpy as np
 import torch
-from pytorch_ie.core import Annotation, AnnotationLayer, TaskEncoding, TaskModule
+from pie_core import Annotation, AnnotationLayer, TaskEncoding, TaskModule
 from tokenizers import Encoding
 from transformers import AutoTokenizer, BatchEncoding, PreTrainedTokenizer
 from transformers.modeling_outputs import QuestionAnsweringModelOutput

src/pie_modules/taskmodules/labeled_span_extraction_by_token_classification.py

@@ -24,7 +24,7 @@
 )
 
 import torch
-from pytorch_ie.core import AnnotationLayer, TaskEncoding, TaskModule
+from pie_core import AnnotationLayer, TaskEncoding, TaskModule
 from tokenizers import Encoding
 from torchmetrics import F1Score, Metric, MetricCollection, Precision, Recall
 from transformers import AutoTokenizer

src/pie_modules/taskmodules/metrics/precision_recall_and_f1_for_labeled_annotations.py

@@ -3,7 +3,7 @@
 from typing import Any, Collection, Dict, Iterable, Optional, Union
 
 import torch
-from pytorch_ie.core import Annotation
+from pie_core import Annotation
 from torch import LongTensor
 
 from pie_modules.utils import flatten_dict

src/pie_modules/taskmodules/pointer_network_for_end2end_re.py

@@ -18,18 +18,18 @@
 )
 
 import torch
-from pytorch_ie.core import (
+from pie_core import (
     Annotation,
     AnnotationLayer,
     Document,
+    TaskEncoding,
     TaskModule,
 )
-from pytorch_ie.core.taskmodule import (
+from pie_core.taskmodule import (
     InputEncoding,
     ModelBatchOutput,
     TargetEncoding,
     TaskBatchEncoding,
-    TaskEncoding,
 )
 from torchmetrics import Metric
 from transformers import AutoTokenizer, LogitsProcessorList, PreTrainedTokenizer

src/pie_modules/taskmodules/re_span_pair_classification.py

@@ -27,7 +27,7 @@
 
 import pandas as pd
 import torch
-from pytorch_ie.core import (
+from pie_core import (
     Annotation,
     AnnotationLayer,
     Document,

src/pie_modules/taskmodules/re_text_classification_with_indices.py

@@ -27,7 +27,7 @@
 
 import numpy as np
 import torch
-from pytorch_ie.core import (
+from pie_core import (
     Annotation,
     AnnotationLayer,
     Document,

src/pie_modules/taskmodules/text_to_text.py

@@ -15,18 +15,18 @@
 )
 
 import torch
-from pytorch_ie.core import (
+from pie_core import (
     Annotation,
     AnnotationLayer,
     Document,
+    TaskEncoding,
     TaskModule,
 )
-from pytorch_ie.core.taskmodule import (
+from pie_core.taskmodule import (
     InputEncoding,
     ModelBatchOutput,
     TargetEncoding,
     TaskBatchEncoding,
-    TaskEncoding,
 )
 from torchmetrics import Metric
 from transformers import AutoTokenizer, PreTrainedTokenizer

src/pie_modules/utils/hydra.py

@@ -1,7 +1,7 @@
 from typing import Optional, Type, TypeVar, Union
 
-from pytorch_ie.core import Document
-from pytorch_ie.utils.hydra import resolve_target
+from pie_core import Document
+from pie_core.utils.hydra import resolve_target
 
 T = TypeVar("T", bound=Document)
 T_super = TypeVar("T_super", bound=Document)

tests/conftest.py

@@ -3,7 +3,7 @@
 
 import pkg_resources
 import pytest
-from pytorch_ie.core import AnnotationLayer, annotation_field
+from pie_core import AnnotationLayer, annotation_field
 
 from pie_modules.annotations import BinaryRelation, LabeledSpan, Span
 from pie_modules.documents import TextBasedDocument

tests/document/processing/test_regex_partitioner.py

@@ -4,7 +4,7 @@
 from typing import Tuple
 
 import pytest
-from pytorch_ie.core import AnnotationLayer, annotation_field
+from pie_core import AnnotationLayer, annotation_field
 
 from pie_modules.annotations import LabeledSpan
 from pie_modules.document.processing import RegexPartitioner

tests/document/processing/test_relation_argument_sorter.py

@@ -2,7 +2,7 @@
 import logging
 
 import pytest
-from pytorch_ie.core import Annotation, AnnotationLayer, annotation_field
+from pie_core import Annotation, AnnotationLayer, annotation_field
 
 from pie_modules.annotations import (
     BinaryRelation,

tests/document/processing/test_text_span_trimmer.py

@@ -1,7 +1,7 @@
 import dataclasses
 
 import pytest
-from pytorch_ie.core import AnnotationLayer, annotation_field
+from pie_core import AnnotationLayer, annotation_field
 
 from pie_modules.annotations import BinaryRelation, LabeledMultiSpan, LabeledSpan
 from pie_modules.document.processing import TextSpanTrimmer

tests/document/processing/test_tokenization.py

@@ -3,7 +3,7 @@
 from typing import Dict
 
 import pytest
-from pytorch_ie.core import Annotation, AnnotationLayer, Document, annotation_field
+from pie_core import Annotation, AnnotationLayer, Document, annotation_field
 from transformers import AutoTokenizer, PreTrainedTokenizer
 
 from pie_modules.annotations import (

tests/metrics/test_confusion_matrix.py

@@ -4,7 +4,7 @@
 from typing import Dict
 
 import pytest
-from pytorch_ie.core import Annotation, AnnotationLayer, annotation_field
+from pie_core import Annotation, AnnotationLayer, annotation_field
 
 from pie_modules.annotations import LabeledSpan
 from pie_modules.documents import TextBasedDocument

tests/metrics/test_f1.py

@@ -1,7 +1,7 @@
 from dataclasses import dataclass
 
 import pytest
-from pytorch_ie.core import AnnotationLayer, annotation_field
+from pie_core import AnnotationLayer, annotation_field
 
 from pie_modules.annotations import LabeledSpan
 from pie_modules.documents import TextBasedDocument

tests/metrics/test_relation_argument_distance_collector.py

@@ -1,7 +1,7 @@
 import dataclasses
 
 import pytest
-from pytorch_ie.core import Annotation, AnnotationLayer, Document, annotation_field
+from pie_core import Annotation, AnnotationLayer, Document, annotation_field
 
 from pie_modules.annotations import BinaryRelation, LabeledSpan, NaryRelation
 from pie_modules.documents import TextBasedDocument, TokenBasedDocument

tests/metrics/test_span_coverage_collector.py

@@ -1,7 +1,7 @@
 import dataclasses
 
 import pytest
-from pytorch_ie.core import Annotation, AnnotationLayer, Document, annotation_field
+from pie_core import Annotation, AnnotationLayer, Document, annotation_field
 
 from pie_modules.annotations import LabeledMultiSpan, LabeledSpan
 from pie_modules.documents import TextBasedDocument, TokenBasedDocument

tests/metrics/test_span_length_collector.py

@@ -1,7 +1,7 @@
 import dataclasses
 
 import pytest
-from pytorch_ie.core import AnnotationLayer, Document, annotation_field
+from pie_core import AnnotationLayer, Document, annotation_field
 
 from pie_modules.annotations import Label, LabeledSpan
 from pie_modules.documents import TextBasedDocument, TokenBasedDocument

tests/taskmodules/common/test_mixins.py

@@ -3,7 +3,7 @@
 from typing import List
 
 import torch
-from pytorch_ie import Annotation
+from pie_core import Annotation
 
 from pie_modules.taskmodules.common import BatchableMixin
 from pie_modules.taskmodules.common.mixins import RelationStatisticsMixin