reoarganize pytorch_ie.core (#166)

* move taskmodule to pytorch_ie.core * move document to pytorch_ie.core * move Annotation to pytorch_ie.core.document * move Metadata to pytorch_ie.core.taskmodule * move TextDocument to documents * rename registerable.py to registrable.py * rename pytorch_ie.py to model.py * export TaskEncoding and Auto classes * do not export any entry from annotations and documents on package level Co-authored-by: Arne Binder <[email protected]>
ArneBinder · May 5, 2022 · 5e452b0 · 5e452b0
1 parent 138046f
commit 5e452b0
Show file tree

Hide file tree

Showing 52 changed files with 175 additions and 172 deletions.
diff --git a/datasets/conll2002/conll2002.py b/datasets/conll2002/conll2002.py
@@ -2,7 +2,9 @@
 
 import datasets
 import pytorch_ie.data.builder
-from pytorch_ie import AnnotationList, LabeledSpan, TextDocument, annotation_field
+from pytorch_ie import AnnotationList, annotation_field
+from pytorch_ie.annotations import LabeledSpan
+from pytorch_ie.documents import TextDocument
 from pytorch_ie.utils.span import tokens_and_tags_to_text_and_labeled_spans
 
 

diff --git a/datasets/conll2003/conll2003.py b/datasets/conll2003/conll2003.py
@@ -2,7 +2,9 @@
 
 import datasets
 import pytorch_ie.data.builder
-from pytorch_ie import AnnotationList, LabeledSpan, TextDocument, annotation_field
+from pytorch_ie import AnnotationList, annotation_field
+from pytorch_ie.annotations import LabeledSpan
+from pytorch_ie.documents import TextDocument
 from pytorch_ie.utils.span import tokens_and_tags_to_text_and_labeled_spans
 
 

diff --git a/datasets/conllpp/conllpp.py b/datasets/conllpp/conllpp.py
@@ -2,7 +2,9 @@
 
 import datasets
 import pytorch_ie.data.builder
-from pytorch_ie import AnnotationList, LabeledSpan, TextDocument, annotation_field
+from pytorch_ie import AnnotationList, annotation_field
+from pytorch_ie.annotations import LabeledSpan
+from pytorch_ie.documents import TextDocument
 from pytorch_ie.utils.span import tokens_and_tags_to_text_and_labeled_spans
 
 

diff --git a/datasets/german_legal_entity_recognition/german_legal_entity_recognition.py b/datasets/german_legal_entity_recognition/german_legal_entity_recognition.py
@@ -2,7 +2,9 @@
 
 import datasets
 import pytorch_ie.data.builder
-from pytorch_ie import AnnotationList, LabeledSpan, TextDocument, annotation_field
+from pytorch_ie import AnnotationList, annotation_field
+from pytorch_ie.annotations import LabeledSpan
+from pytorch_ie.documents import TextDocument
 from pytorch_ie.utils.span import tokens_and_tags_to_text_and_labeled_spans
 
 _VERSION = "1.0.0"

diff --git a/datasets/germaner/germaner.py b/datasets/germaner/germaner.py
@@ -2,7 +2,9 @@
 
 import datasets
 import pytorch_ie.data.builder
-from pytorch_ie import AnnotationList, LabeledSpan, TextDocument, annotation_field
+from pytorch_ie import AnnotationList, annotation_field
+from pytorch_ie.annotations import LabeledSpan
+from pytorch_ie.documents import TextDocument
 from pytorch_ie.utils.span import tokens_and_tags_to_text_and_labeled_spans
 
 

diff --git a/datasets/germeval_14/germeval_14.py b/datasets/germeval_14/germeval_14.py
@@ -2,7 +2,9 @@
 
 import datasets
 import pytorch_ie.data.builder
-from pytorch_ie import AnnotationList, LabeledSpan, TextDocument, annotation_field
+from pytorch_ie import AnnotationList, annotation_field
+from pytorch_ie.annotations import LabeledSpan
+from pytorch_ie.documents import TextDocument
 from pytorch_ie.utils.span import tokens_and_tags_to_text_and_labeled_spans
 
 

diff --git a/datasets/ncbi_disease/ncbi_disease.py b/datasets/ncbi_disease/ncbi_disease.py
@@ -2,7 +2,9 @@
 
 import datasets
 import pytorch_ie.data.builder
-from pytorch_ie import AnnotationList, LabeledSpan, TextDocument, annotation_field
+from pytorch_ie import AnnotationList, annotation_field
+from pytorch_ie.annotations import LabeledSpan
+from pytorch_ie.documents import TextDocument
 from pytorch_ie.utils.span import tokens_and_tags_to_text_and_labeled_spans
 
 

diff --git a/datasets/wikiann/wikiann.py b/datasets/wikiann/wikiann.py
@@ -2,7 +2,9 @@
 
 import datasets
 import pytorch_ie.data.builder
-from pytorch_ie import AnnotationList, LabeledSpan, TextDocument, annotation_field
+from pytorch_ie import AnnotationList, annotation_field
+from pytorch_ie.annotations import LabeledSpan
+from pytorch_ie.documents import TextDocument
 from pytorch_ie.utils.span import tokens_and_tags_to_text_and_labeled_spans
 
 _VERSION = "1.1.0"

diff --git a/datasets/wnut_17/wnut_17.py b/datasets/wnut_17/wnut_17.py
@@ -2,7 +2,9 @@
 
 import datasets
 import pytorch_ie.data.builder
-from pytorch_ie import AnnotationList, LabeledSpan, TextDocument, annotation_field
+from pytorch_ie import AnnotationList, annotation_field
+from pytorch_ie.annotations import LabeledSpan
+from pytorch_ie.documents import TextDocument
 from pytorch_ie.utils.span import tokens_and_tags_to_text_and_labeled_spans
 
 

diff --git a/examples/predict/ner_span_classification.py b/examples/predict/ner_span_classification.py
@@ -1,6 +1,8 @@
 from dataclasses import dataclass
 
-from pytorch_ie import AnnotationList, LabeledSpan, Pipeline, TextDocument, annotation_field
+from pytorch_ie import AnnotationList, Pipeline, annotation_field
+from pytorch_ie.annotations import LabeledSpan
+from pytorch_ie.documents import TextDocument
 from pytorch_ie.models import TransformerSpanClassificationModel
 from pytorch_ie.taskmodules import TransformerSpanClassificationTaskModule
 

diff --git a/examples/predict/re_generative.py b/examples/predict/re_generative.py
@@ -1,13 +1,8 @@
 from dataclasses import dataclass
 
-from pytorch_ie import (
-    AnnotationList,
-    BinaryRelation,
-    LabeledSpan,
-    Pipeline,
-    TextDocument,
-    annotation_field,
-)
+from pytorch_ie import AnnotationList, Pipeline, annotation_field
+from pytorch_ie.annotations import BinaryRelation, LabeledSpan
+from pytorch_ie.documents import TextDocument
 from pytorch_ie.models import TransformerSeq2SeqModel
 from pytorch_ie.taskmodules import TransformerSeq2SeqTaskModule
 

diff --git a/examples/predict/re_text_classification.py b/examples/predict/re_text_classification.py
@@ -1,13 +1,8 @@
 from dataclasses import dataclass
 
-from pytorch_ie import (
-    AnnotationList,
-    BinaryRelation,
-    LabeledSpan,
-    Pipeline,
-    TextDocument,
-    annotation_field,
-)
+from pytorch_ie import AnnotationList, Pipeline, annotation_field
+from pytorch_ie.annotations import BinaryRelation, LabeledSpan
+from pytorch_ie.documents import TextDocument
 from pytorch_ie.models import TransformerTextClassificationModel
 from pytorch_ie.taskmodules import TransformerRETextClassificationTaskModule
 

diff --git a/src/pytorch_ie/__init__.py b/src/pytorch_ie/__init__.py
@@ -1,16 +1,6 @@
 # flake8: noqa
 
-from .annotations import (
-    BinaryRelation,
-    Label,
-    LabeledMultiSpan,
-    LabeledSpan,
-    MultiLabel,
-    MultiLabeledBinaryRelation,
-    MultiLabeledMultiSpan,
-    MultiLabeledSpan,
-    Span,
-)
-from .data import *
-from .document import AnnotationList, Document, TextDocument, annotation_field
-from .pipeline import Pipeline
+from pytorch_ie.auto import AutoModel, AutoPipeline, AutoTaskModule
+from pytorch_ie.core import *
+from pytorch_ie.data import *
+from pytorch_ie.pipeline import Pipeline
diff --git a/src/pytorch_ie/annotations.py b/src/pytorch_ie/annotations.py
@@ -1,7 +1,7 @@
-from dataclasses import asdict, dataclass, field
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Tuple
 
-from pytorch_ie.document import AnnotationList
+from pytorch_ie.core.document import Annotation
 
 
 def _validate_single_label(self):
@@ -29,36 +29,6 @@ def _validate_multi_label(self):
         )
 
 
-@dataclass(eq=True, frozen=True)
-class Annotation:
-    _target: Optional[Union[AnnotationList, str]] = field(
-        default=None, init=False, repr=False, hash=False
-    )
-
-    def set_target(self, value: Union[AnnotationList, str, None]):
-        object.__setattr__(self, "_target", value)
-
-    @property
-    def target(self) -> Optional[Union[AnnotationList, str]]:
-        return self._target
-
-    def asdict(self) -> Dict[str, Any]:
-        dct = asdict(self)
-        dct["_id"] = hash(self)
-        del dct["_target"]
-        return dct
-
-    @classmethod
-    def fromdict(
-        cls,
-        dct: Dict[str, Any],
-        annotation_store: Optional[Dict[int, Tuple[str, "Annotation"]]] = None,
-    ):
-        tmp_dct = dict(dct)
-        tmp_dct.pop("_id", None)
-        return cls(**tmp_dct)
-
-
 @dataclass(eq=True, frozen=True)
 class Label(Annotation):
     label: str

diff --git a/src/pytorch_ie/auto.py b/src/pytorch_ie/auto.py
@@ -5,10 +5,9 @@
 from huggingface_hub.constants import PYTORCH_WEIGHTS_NAME
 from huggingface_hub.file_download import hf_hub_download
 
-from pytorch_ie import Pipeline
-from pytorch_ie.core import PyTorchIEModel
+from pytorch_ie.core import PyTorchIEModel, TaskModule
 from pytorch_ie.core.hf_hub_mixin import PyTorchIEModelHubMixin, PyTorchIETaskmoduleModelHubMixin
-from pytorch_ie.taskmodules import TaskModule
+from pytorch_ie.pipeline import Pipeline
 
 
 class AutoTaskModule(PyTorchIETaskmoduleModelHubMixin):

diff --git a/src/pytorch_ie/core/__init__.py b/src/pytorch_ie/core/__init__.py
@@ -1 +1,3 @@
-from .pytorch_ie import PyTorchIEModel
+from .document import Annotation, AnnotationList, Document, annotation_field
+from .model import PyTorchIEModel
+from .taskmodule import TaskEncoding, TaskModule
diff --git a/src/pytorch_ie/document.py → src/pytorch_ie/core/document.py b/src/pytorch_ie/document.py → src/pytorch_ie/core/document.py
@@ -1,10 +1,7 @@
 import dataclasses
 import typing
 from collections.abc import Mapping, Sequence
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, TypeVar, Union, overload
-
-if TYPE_CHECKING:
-    from pytorch_ie.annotations import Annotation
+from typing import Any, Dict, List, Optional, Set, Tuple, TypeVar, Union, overload
 
 
 def _depth_first_search(lst: List[str], visited: Set[str], graph: Dict[str, List[str]], node: str):
@@ -24,6 +21,36 @@ def annotation_field(target: Optional[str] = None):
     return dataclasses.field(metadata=dict(target=target), init=False, repr=False)
 
 
+@dataclasses.dataclass(eq=True, frozen=True)
+class Annotation:
+    _target: Optional[Union["AnnotationList", str]] = dataclasses.field(
+        default=None, init=False, repr=False, hash=False
+    )
+
+    def set_target(self, value: Union["AnnotationList", str, None]):
+        object.__setattr__(self, "_target", value)
+
+    @property
+    def target(self) -> Optional[Union["AnnotationList", str]]:
+        return self._target
+
+    def asdict(self) -> Dict[str, Any]:
+        dct = dataclasses.asdict(self)
+        dct["_id"] = hash(self)
+        del dct["_target"]
+        return dct
+
+    @classmethod
+    def fromdict(
+        cls,
+        dct: Dict[str, Any],
+        annotation_store: Optional[Dict[int, Tuple[str, "Annotation"]]] = None,
+    ):
+        tmp_dct = dict(dct)
+        tmp_dct.pop("_id", None)
+        return cls(**tmp_dct)
+
+
 T = TypeVar("T", covariant=False, bound="Annotation")
 
 
@@ -200,11 +227,3 @@ def fromdict(cls, dct):
             getattr(doc, field_name).append(annotation)
 
         return doc
-
-
-@dataclasses.dataclass
-class TextDocument(Document):
-    text: str
-    id: Optional[str] = None
-    metadata: Dict[str, Any] = dataclasses.field(default_factory=dict)
-    _root_annotation: str = dataclasses.field(default="text", init=False, repr=False)
diff --git a/src/pytorch_ie/core/pytorch_ie.py → src/pytorch_ie/core/model.py b/src/pytorch_ie/core/pytorch_ie.py → src/pytorch_ie/core/model.py
@@ -3,7 +3,7 @@
 from pytorch_lightning import LightningModule
 
 from pytorch_ie.core.hf_hub_mixin import PyTorchIEModelHubMixin
-from pytorch_ie.core.registerable import Registrable
+from pytorch_ie.core.registrable import Registrable
 
 
 class PyTorchIEModel(LightningModule, Registrable, PyTorchIEModelHubMixin):

diff --git a/src/pytorch_ie/core/registerable.py → src/pytorch_ie/core/registrable.py b/src/pytorch_ie/core/registerable.py → src/pytorch_ie/core/registrable.py
diff --git a/src/pytorch_ie/taskmodules/taskmodule.py → src/pytorch_ie/core/taskmodule.py b/src/pytorch_ie/taskmodules/taskmodule.py → src/pytorch_ie/core/taskmodule.py
@@ -16,11 +16,10 @@
     overload,
 )
 
-from pytorch_ie import Dataset, Document
-from pytorch_ie.annotations import Annotation
+from pytorch_ie.core.document import Annotation, Document
 from pytorch_ie.core.hf_hub_mixin import PyTorchIETaskmoduleModelHubMixin
-from pytorch_ie.core.registerable import Registrable
-from pytorch_ie.data import Metadata
+from pytorch_ie.core.registrable import Registrable
+from pytorch_ie.data import Dataset
 
 """
 workflow:
@@ -48,6 +47,9 @@ class InplaceNotSupportedException(Exception):
     pass
 
 
+Metadata = Dict[str, Any]
+
+
 class TaskEncoding(Generic[DocumentType, InputEncoding, TargetEncoding]):
     def __init__(
         self,

diff --git a/src/pytorch_ie/data/__init__.py b/src/pytorch_ie/data/__init__.py
@@ -1,13 +1,8 @@
-from typing import Any, Dict
-
 from .builder import GeneratorBasedBuilder
 from .dataset import Dataset
 from .dataset_formatter import DocumentFormatter
 
-Metadata = Dict[str, Any]
-
 __all__ = [
-    "Metadata",
     "GeneratorBasedBuilder",
     "Dataset",
     "DocumentFormatter",

diff --git a/src/pytorch_ie/data/builder.py b/src/pytorch_ie/data/builder.py
@@ -4,8 +4,8 @@
 from datasets.load import load_dataset_builder
 
 import datasets
+from pytorch_ie.core.document import Document
 from pytorch_ie.data.dataset import Dataset, decorate_convert_to_dict_of_lists
-from pytorch_ie.document import Document
 
 
 class GeneratorBasedBuilder(datasets.builder.GeneratorBasedBuilder):

diff --git a/src/pytorch_ie/data/datamodules/datamodule.py b/src/pytorch_ie/data/datamodules/datamodule.py
@@ -1,11 +1,10 @@
-from typing import Any, Dict, Generic, List, Optional, Sequence
+from typing import Any, Dict, Generic, Optional, Sequence
 
 from pytorch_lightning import LightningDataModule
 from torch.utils.data import DataLoader
 from torch.utils.data.dataset import Dataset
 
-from pytorch_ie.data.datasets import PIEDatasetDict
-from pytorch_ie.taskmodules.taskmodule import (
+from pytorch_ie.core.taskmodule import (
     DocumentType,
     InputEncoding,
     TargetEncoding,

diff --git a/src/pytorch_ie/data/dataset.py b/src/pytorch_ie/data/dataset.py
@@ -1,12 +1,12 @@
 from functools import wraps
-from typing import TYPE_CHECKING, Callable, List, Optional, Type, Union
+from typing import Callable, List, Optional, Type, Union
 
 import pandas as pd
 from datasets.formatting import _register_formatter
 
 import datasets
+from pytorch_ie.core.document import Document
 from pytorch_ie.data.dataset_formatter import DocumentFormatter
-from pytorch_ie.document import Document
 
 _register_formatter(DocumentFormatter, "document")
 

diff --git a/src/pytorch_ie/data/dataset_formatter.py b/src/pytorch_ie/data/dataset_formatter.py
@@ -3,7 +3,7 @@
 import pyarrow as pa
 from datasets.formatting.formatting import Formatter
 
-from pytorch_ie.document import Document
+from pytorch_ie.core.document import Document
 
 
 class DocumentFormatter(Formatter[Document, list, List[Document]]):

diff --git a/src/pytorch_ie/data/datasets/brat.py b/src/pytorch_ie/data/datasets/brat.py
@@ -6,8 +6,8 @@
 from typing import Any, Dict, Iterable, List, Optional, Tuple
 
 from datasets import DatasetDict, load_dataset
-from pytorch_ie import BinaryRelation, Document, LabeledMultiSpan, LabeledSpan
-from pytorch_ie.annotations import Annotation
+from pytorch_ie import Document
+from pytorch_ie.annotations import Annotation, BinaryRelation, LabeledMultiSpan, LabeledSpan
 from pytorch_ie.data.datasets import HF_DATASETS_ROOT
 
 DEFAULT_HEAD_ARGUMENT_NAME: str = "Arg1"