From 9d8335951d9f997c12cde879d93690d0b5f26a69 Mon Sep 17 00:00:00 2001 From: ArneBinder Date: Thu, 5 May 2022 19:13:04 +0200 Subject: [PATCH] dont export core and fix readme examples (#173) * dont export core and fix readme examples * move PIEDatasetDict from data.datasets.__init__ to data.__init__ and rename to DatasetDict Co-authored-by: Arne Binder --- README.rst | 24 ++++++++++--------- datasets/conll2002/conll2002.py | 2 +- datasets/conll2003/conll2003.py | 2 +- datasets/conllpp/conllpp.py | 2 +- .../german_legal_entity_recognition.py | 2 +- datasets/germaner/germaner.py | 2 +- datasets/germeval_14/germeval_14.py | 2 +- datasets/ncbi_disease/ncbi_disease.py | 2 +- datasets/wikiann/wikiann.py | 2 +- datasets/wnut_17/wnut_17.py | 2 +- examples/predict/ner_span_classification.py | 3 ++- examples/predict/re_generative.py | 3 ++- examples/predict/re_text_classification.py | 3 ++- src/pytorch_ie/__init__.py | 3 ++- src/pytorch_ie/data/__init__.py | 7 ++++++ src/pytorch_ie/data/datasets/__init__.py | 6 ----- src/pytorch_ie/documents.py | 2 +- src/pytorch_ie/models/transformer_seq2seq.py | 2 +- .../models/transformer_span_classification.py | 2 +- .../models/transformer_text_classification.py | 2 +- .../transformer_token_classification.py | 2 +- .../transformer_re_text_classification.py | 2 +- .../taskmodules/transformer_seq2seq.py | 2 +- .../transformer_span_classification.py | 2 +- .../transformer_text_classification.py | 2 +- .../transformer_token_classification.py | 2 +- tests/conftest.py | 3 ++- tests/data/datasets/test_brat.py | 3 ++- .../pipeline/test_ner_span_classification.py | 3 ++- tests/pipeline/test_re_generative.py | 3 ++- tests/pipeline/test_re_text_classification.py | 3 ++- .../test_transformer_span_classification.py | 2 +- .../test_transformer_token_classification.py | 2 +- tests/test_annotations.py | 2 +- tests/test_auto.py | 10 ++------ tests/test_document.py | 2 +- tests/test_pipeline.py | 2 +- tests/train/test_training.py | 2 +- 38 files changed, 65 insertions(+), 59 deletions(-) diff --git a/README.rst b/README.rst index 47b08e4b..bf581e30 100644 --- a/README.rst +++ b/README.rst @@ -60,21 +60,22 @@ Span-classification-based Named Entity Recognition from dataclasses import dataclass - from pytorch_ie import AnnotationList, LabeledSpan, TextDocument, annotation_field + from pytorch_ie.annotations import LabeledSpan from pytorch_ie.auto import AutoPipeline - + from pytorch_ie.core import AnnotationList, annotation_field + from pytorch_ie.documents import TextDocument @dataclass class ExampleDocument(TextDocument): entities: AnnotationList[LabeledSpan] = annotation_field(target="text") - # see below for the long version - ner_pipeline = AutoPipeline.from_pretrained("pie/example-ner-spanclf-conll03", device=-1, num_workers=0) - document = ExampleDocument( "“Making a super tasty alt-chicken wing is only half of it,” said Po Bronson, general partner at SOSV and managing director of IndieBio." ) + # see below for the long version + ner_pipeline = AutoPipeline.from_pretrained("pie/example-ner-spanclf-conll03", device=-1, num_workers=0) + ner_pipeline(document, predict_field="entities") for entity in document.entities.predictions: @@ -89,8 +90,8 @@ To create the same pipeline as above without `AutoPipeline`: .. code:: python - from pytorch_ie import Pipeline from pytorch_ie.auto import AutoTaskModule, AutoModel + from pytorch_ie.pipeline import Pipeline model_name_or_path = "pie/example-ner-spanclf-conll03" ner_taskmodule = AutoTaskModule.from_pretrained(model_name_or_path) @@ -101,7 +102,7 @@ Or, without `Auto` classes at all: .. code:: python - from pytorch_ie import Pipeline + from pytorch_ie.pipeline import Pipeline from pytorch_ie.models import TransformerSpanClassificationModel from pytorch_ie.taskmodules import TransformerSpanClassificationTaskModule @@ -118,8 +119,10 @@ Text-classification-based Relation Extraction from dataclasses import dataclass - from pytorch_ie import AnnotationList, BinaryRelation, LabeledSpan, TextDocument, annotation_field + from pytorch_ie.annotations import BinaryRelation, LabeledSpan from pytorch_ie.auto import AutoPipeline + from pytorch_ie.core import AnnotationList, annotation_field + from pytorch_ie.documents import TextDocument @dataclass @@ -127,13 +130,12 @@ Text-classification-based Relation Extraction entities: AnnotationList[LabeledSpan] = annotation_field(target="text") relations: AnnotationList[BinaryRelation] = annotation_field(target="entities") - - re_pipeline = AutoPipeline.from_pretrained("pie/example-re-textclf-tacred", device=-1, num_workers=0) - document = ExampleDocument( "“Making a super tasty alt-chicken wing is only half of it,” said Po Bronson, general partner at SOSV and managing director of IndieBio." ) + re_pipeline = AutoPipeline.from_pretrained("pie/example-re-textclf-tacred", device=-1, num_workers=0) + for start, end, label in [(65, 75, "PER"), (96, 100, "ORG"), (126, 134, "ORG")]: document.entities.append(LabeledSpan(start=start, end=end, label=label)) diff --git a/datasets/conll2002/conll2002.py b/datasets/conll2002/conll2002.py index 511e8452..5e81da14 100644 --- a/datasets/conll2002/conll2002.py +++ b/datasets/conll2002/conll2002.py @@ -2,8 +2,8 @@ import datasets import pytorch_ie.data.builder -from pytorch_ie import AnnotationList, annotation_field from pytorch_ie.annotations import LabeledSpan +from pytorch_ie.core import AnnotationList, annotation_field 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 index 00b637d3..0c2354a9 100644 --- a/datasets/conll2003/conll2003.py +++ b/datasets/conll2003/conll2003.py @@ -2,8 +2,8 @@ import datasets import pytorch_ie.data.builder -from pytorch_ie import AnnotationList, annotation_field from pytorch_ie.annotations import LabeledSpan +from pytorch_ie.core import AnnotationList, annotation_field 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 index 3b33d7cb..a67995f1 100644 --- a/datasets/conllpp/conllpp.py +++ b/datasets/conllpp/conllpp.py @@ -2,8 +2,8 @@ import datasets import pytorch_ie.data.builder -from pytorch_ie import AnnotationList, annotation_field from pytorch_ie.annotations import LabeledSpan +from pytorch_ie.core import AnnotationList, annotation_field 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 index f138bdd0..7ba5c2d3 100644 --- a/datasets/german_legal_entity_recognition/german_legal_entity_recognition.py +++ b/datasets/german_legal_entity_recognition/german_legal_entity_recognition.py @@ -2,8 +2,8 @@ import datasets import pytorch_ie.data.builder -from pytorch_ie import AnnotationList, annotation_field from pytorch_ie.annotations import LabeledSpan +from pytorch_ie.core import AnnotationList, annotation_field from pytorch_ie.documents import TextDocument from pytorch_ie.utils.span import tokens_and_tags_to_text_and_labeled_spans diff --git a/datasets/germaner/germaner.py b/datasets/germaner/germaner.py index 01cacc31..e1cade6e 100644 --- a/datasets/germaner/germaner.py +++ b/datasets/germaner/germaner.py @@ -2,8 +2,8 @@ import datasets import pytorch_ie.data.builder -from pytorch_ie import AnnotationList, annotation_field from pytorch_ie.annotations import LabeledSpan +from pytorch_ie.core import AnnotationList, annotation_field 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 index 548c5884..259725ef 100644 --- a/datasets/germeval_14/germeval_14.py +++ b/datasets/germeval_14/germeval_14.py @@ -2,8 +2,8 @@ import datasets import pytorch_ie.data.builder -from pytorch_ie import AnnotationList, annotation_field from pytorch_ie.annotations import LabeledSpan +from pytorch_ie.core import AnnotationList, annotation_field 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 index 84446deb..b5a7bc03 100644 --- a/datasets/ncbi_disease/ncbi_disease.py +++ b/datasets/ncbi_disease/ncbi_disease.py @@ -2,8 +2,8 @@ import datasets import pytorch_ie.data.builder -from pytorch_ie import AnnotationList, annotation_field from pytorch_ie.annotations import LabeledSpan +from pytorch_ie.core import AnnotationList, annotation_field 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 index dce5d348..b9c437fd 100644 --- a/datasets/wikiann/wikiann.py +++ b/datasets/wikiann/wikiann.py @@ -2,8 +2,8 @@ import datasets import pytorch_ie.data.builder -from pytorch_ie import AnnotationList, annotation_field from pytorch_ie.annotations import LabeledSpan +from pytorch_ie.core import AnnotationList, annotation_field from pytorch_ie.documents import TextDocument from pytorch_ie.utils.span import tokens_and_tags_to_text_and_labeled_spans diff --git a/datasets/wnut_17/wnut_17.py b/datasets/wnut_17/wnut_17.py index cb42a563..6a499959 100644 --- a/datasets/wnut_17/wnut_17.py +++ b/datasets/wnut_17/wnut_17.py @@ -2,8 +2,8 @@ import datasets import pytorch_ie.data.builder -from pytorch_ie import AnnotationList, annotation_field from pytorch_ie.annotations import LabeledSpan +from pytorch_ie.core import AnnotationList, annotation_field 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 index 7c8965fb..9437d969 100644 --- a/examples/predict/ner_span_classification.py +++ b/examples/predict/ner_span_classification.py @@ -1,9 +1,10 @@ from dataclasses import dataclass -from pytorch_ie import AnnotationList, Pipeline, annotation_field from pytorch_ie.annotations import LabeledSpan +from pytorch_ie.core import AnnotationList, annotation_field from pytorch_ie.documents import TextDocument from pytorch_ie.models import TransformerSpanClassificationModel +from pytorch_ie.pipeline import Pipeline from pytorch_ie.taskmodules import TransformerSpanClassificationTaskModule diff --git a/examples/predict/re_generative.py b/examples/predict/re_generative.py index 9361df84..1f5d76a5 100644 --- a/examples/predict/re_generative.py +++ b/examples/predict/re_generative.py @@ -1,9 +1,10 @@ from dataclasses import dataclass -from pytorch_ie import AnnotationList, Pipeline, annotation_field from pytorch_ie.annotations import BinaryRelation, LabeledSpan +from pytorch_ie.core import AnnotationList, annotation_field from pytorch_ie.documents import TextDocument from pytorch_ie.models import TransformerSeq2SeqModel +from pytorch_ie.pipeline import Pipeline from pytorch_ie.taskmodules import TransformerSeq2SeqTaskModule diff --git a/examples/predict/re_text_classification.py b/examples/predict/re_text_classification.py index e398293d..218e70de 100644 --- a/examples/predict/re_text_classification.py +++ b/examples/predict/re_text_classification.py @@ -1,9 +1,10 @@ from dataclasses import dataclass -from pytorch_ie import AnnotationList, Pipeline, annotation_field from pytorch_ie.annotations import BinaryRelation, LabeledSpan +from pytorch_ie.core import AnnotationList, annotation_field from pytorch_ie.documents import TextDocument from pytorch_ie.models import TransformerTextClassificationModel +from pytorch_ie.pipeline import Pipeline from pytorch_ie.taskmodules import TransformerRETextClassificationTaskModule diff --git a/src/pytorch_ie/__init__.py b/src/pytorch_ie/__init__.py index 494fd451..29badd60 100644 --- a/src/pytorch_ie/__init__.py +++ b/src/pytorch_ie/__init__.py @@ -1,6 +1,7 @@ # flake8: noqa from pytorch_ie.auto import AutoModel, AutoPipeline, AutoTaskModule -from pytorch_ie.core import * from pytorch_ie.data import * +from pytorch_ie.models import * from pytorch_ie.pipeline import Pipeline +from pytorch_ie.taskmodules import * diff --git a/src/pytorch_ie/data/__init__.py b/src/pytorch_ie/data/__init__.py index cc1545e8..269be2db 100644 --- a/src/pytorch_ie/data/__init__.py +++ b/src/pytorch_ie/data/__init__.py @@ -1,9 +1,16 @@ +from typing import Dict, Union + +from datasets import Split + from .builder import GeneratorBasedBuilder from .dataset import Dataset from .dataset_formatter import DocumentFormatter +DatasetDict = Dict[Union[str, Split], Dataset] + __all__ = [ "GeneratorBasedBuilder", "Dataset", + "DatasetDict", "DocumentFormatter", ] diff --git a/src/pytorch_ie/data/datasets/__init__.py b/src/pytorch_ie/data/datasets/__init__.py index 702174f0..ffbe5d00 100644 --- a/src/pytorch_ie/data/datasets/__init__.py +++ b/src/pytorch_ie/data/datasets/__init__.py @@ -1,9 +1,3 @@ import pathlib -from typing import Dict, List, Union - -from datasets import Split -from pytorch_ie import Document HF_DATASETS_ROOT = pathlib.Path(__file__).parent / "hf_datasets" - -PIEDatasetDict = Dict[Union[str, Split], List[Document]] diff --git a/src/pytorch_ie/documents.py b/src/pytorch_ie/documents.py index dd41c36b..1e6b2ea8 100644 --- a/src/pytorch_ie/documents.py +++ b/src/pytorch_ie/documents.py @@ -1,7 +1,7 @@ import dataclasses from typing import Any, Dict, Optional -from pytorch_ie import Document +from pytorch_ie.core import Document @dataclasses.dataclass diff --git a/src/pytorch_ie/models/transformer_seq2seq.py b/src/pytorch_ie/models/transformer_seq2seq.py index 9ea40db9..ed54f7cd 100644 --- a/src/pytorch_ie/models/transformer_seq2seq.py +++ b/src/pytorch_ie/models/transformer_seq2seq.py @@ -4,7 +4,7 @@ from transformers import AutoModelForSeq2SeqLM, BatchEncoding from transformers.modeling_outputs import Seq2SeqLMOutput -from pytorch_ie import PyTorchIEModel +from pytorch_ie.core import PyTorchIEModel from pytorch_ie.core.taskmodule import Metadata from pytorch_ie.documents import TextDocument diff --git a/src/pytorch_ie/models/transformer_span_classification.py b/src/pytorch_ie/models/transformer_span_classification.py index 1808991b..3d5b3d84 100644 --- a/src/pytorch_ie/models/transformer_span_classification.py +++ b/src/pytorch_ie/models/transformer_span_classification.py @@ -11,7 +11,7 @@ get_linear_schedule_with_warmup, ) -from pytorch_ie import PyTorchIEModel +from pytorch_ie.core import PyTorchIEModel from pytorch_ie.models.modules.mlp import MLP TransformerSpanClassificationModelBatchEncoding = BatchEncoding diff --git a/src/pytorch_ie/models/transformer_text_classification.py b/src/pytorch_ie/models/transformer_text_classification.py index 1f1dcf3f..1352c356 100644 --- a/src/pytorch_ie/models/transformer_text_classification.py +++ b/src/pytorch_ie/models/transformer_text_classification.py @@ -4,7 +4,7 @@ from torch import Tensor, nn from transformers import AdamW, AutoConfig, AutoModel, get_linear_schedule_with_warmup -from pytorch_ie import PyTorchIEModel +from pytorch_ie.core import PyTorchIEModel TransformerTextClassificationModelBatchEncoding = MutableMapping[str, Any] TransformerTextClassificationModelBatchOutput = Dict[str, Any] diff --git a/src/pytorch_ie/models/transformer_token_classification.py b/src/pytorch_ie/models/transformer_token_classification.py index 73978da0..7a54db90 100644 --- a/src/pytorch_ie/models/transformer_token_classification.py +++ b/src/pytorch_ie/models/transformer_token_classification.py @@ -5,7 +5,7 @@ from torch import Tensor, nn from transformers import AutoConfig, AutoModelForTokenClassification, BatchEncoding -from pytorch_ie import PyTorchIEModel +from pytorch_ie.core import PyTorchIEModel TransformerTokenClassificationModelBatchEncoding = BatchEncoding TransformerTokenClassificationModelBatchOutput = Dict[str, Any] diff --git a/src/pytorch_ie/taskmodules/transformer_re_text_classification.py b/src/pytorch_ie/taskmodules/transformer_re_text_classification.py index aeeed259..3d9e8fbf 100644 --- a/src/pytorch_ie/taskmodules/transformer_re_text_classification.py +++ b/src/pytorch_ie/taskmodules/transformer_re_text_classification.py @@ -7,8 +7,8 @@ from transformers.file_utils import PaddingStrategy from transformers.tokenization_utils_base import BatchEncoding, TruncationStrategy -from pytorch_ie import TaskEncoding, TaskModule from pytorch_ie.annotations import BinaryRelation, LabeledSpan, MultiLabeledBinaryRelation, Span +from pytorch_ie.core import TaskEncoding, TaskModule from pytorch_ie.documents import TextDocument from pytorch_ie.models import ( TransformerTextClassificationModelBatchOutput, diff --git a/src/pytorch_ie/taskmodules/transformer_seq2seq.py b/src/pytorch_ie/taskmodules/transformer_seq2seq.py index ff00232c..e8bda111 100644 --- a/src/pytorch_ie/taskmodules/transformer_seq2seq.py +++ b/src/pytorch_ie/taskmodules/transformer_seq2seq.py @@ -6,8 +6,8 @@ from transformers.file_utils import PaddingStrategy from transformers.tokenization_utils_base import TruncationStrategy -from pytorch_ie import Annotation, TaskEncoding, TaskModule from pytorch_ie.annotations import BinaryRelation, LabeledSpan, Span +from pytorch_ie.core import Annotation, TaskEncoding, TaskModule from pytorch_ie.documents import TextDocument from pytorch_ie.models import ( TransformerSeq2SeqModelBatchOutput, diff --git a/src/pytorch_ie/taskmodules/transformer_span_classification.py b/src/pytorch_ie/taskmodules/transformer_span_classification.py index 5774589d..b4ca546e 100644 --- a/src/pytorch_ie/taskmodules/transformer_span_classification.py +++ b/src/pytorch_ie/taskmodules/transformer_span_classification.py @@ -8,8 +8,8 @@ from transformers.file_utils import PaddingStrategy from transformers.tokenization_utils_base import BatchEncoding, TruncationStrategy -from pytorch_ie import TaskEncoding, TaskModule from pytorch_ie.annotations import LabeledSpan, MultiLabeledSpan, Span +from pytorch_ie.core import TaskEncoding, TaskModule from pytorch_ie.documents import TextDocument from pytorch_ie.models.transformer_span_classification import ( TransformerSpanClassificationModelBatchOutput, diff --git a/src/pytorch_ie/taskmodules/transformer_text_classification.py b/src/pytorch_ie/taskmodules/transformer_text_classification.py index 8f848a7e..6b23940b 100644 --- a/src/pytorch_ie/taskmodules/transformer_text_classification.py +++ b/src/pytorch_ie/taskmodules/transformer_text_classification.py @@ -17,8 +17,8 @@ from transformers.file_utils import PaddingStrategy from transformers.tokenization_utils_base import TruncationStrategy -from pytorch_ie import TaskEncoding, TaskModule from pytorch_ie.annotations import Label, MultiLabel +from pytorch_ie.core import TaskEncoding, TaskModule from pytorch_ie.documents import TextDocument from pytorch_ie.models.transformer_text_classification import ( TransformerTextClassificationModelBatchOutput, diff --git a/src/pytorch_ie/taskmodules/transformer_token_classification.py b/src/pytorch_ie/taskmodules/transformer_token_classification.py index bf0add0d..93c26993 100644 --- a/src/pytorch_ie/taskmodules/transformer_token_classification.py +++ b/src/pytorch_ie/taskmodules/transformer_token_classification.py @@ -8,8 +8,8 @@ from transformers.file_utils import PaddingStrategy from transformers.tokenization_utils_base import BatchEncoding, TruncationStrategy -from pytorch_ie import TaskEncoding, TaskModule from pytorch_ie.annotations import LabeledSpan, Span +from pytorch_ie.core import TaskEncoding, TaskModule from pytorch_ie.documents import TextDocument from pytorch_ie.models.transformer_token_classification import ( TransformerTokenClassificationModelBatchOutput, diff --git a/tests/conftest.py b/tests/conftest.py index 6a5a79fe..9b41afc9 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -3,8 +3,9 @@ import pytest import datasets -from pytorch_ie import AnnotationList, Dataset, annotation_field from pytorch_ie.annotations import BinaryRelation, LabeledSpan, Span +from pytorch_ie.core import AnnotationList, annotation_field +from pytorch_ie.data import Dataset from pytorch_ie.documents import TextDocument from tests import FIXTURES_ROOT diff --git a/tests/data/datasets/test_brat.py b/tests/data/datasets/test_brat.py index 882dbe96..af75836d 100644 --- a/tests/data/datasets/test_brat.py +++ b/tests/data/datasets/test_brat.py @@ -1,5 +1,5 @@ # type: ignore - +""" import os import pytest @@ -29,6 +29,7 @@ "T2\tperson 25 37\tJenny Durkan\n", "R1\tmayor_of head:T2 tail:T1\n", ] +""" # def get_doc1(with_ids: bool = False, **kwargs) -> TextDocument: diff --git a/tests/pipeline/test_ner_span_classification.py b/tests/pipeline/test_ner_span_classification.py index b68f08f2..74e44bf7 100644 --- a/tests/pipeline/test_ner_span_classification.py +++ b/tests/pipeline/test_ner_span_classification.py @@ -2,10 +2,11 @@ import pytest -from pytorch_ie import AnnotationList, Pipeline, annotation_field from pytorch_ie.annotations import LabeledSpan +from pytorch_ie.core import AnnotationList, annotation_field from pytorch_ie.documents import TextDocument from pytorch_ie.models import TransformerSpanClassificationModel +from pytorch_ie.pipeline import Pipeline from pytorch_ie.taskmodules import TransformerSpanClassificationTaskModule diff --git a/tests/pipeline/test_re_generative.py b/tests/pipeline/test_re_generative.py index a0492161..b57d8088 100644 --- a/tests/pipeline/test_re_generative.py +++ b/tests/pipeline/test_re_generative.py @@ -2,10 +2,11 @@ import pytest -from pytorch_ie import AnnotationList, Pipeline, annotation_field from pytorch_ie.annotations import BinaryRelation, LabeledSpan +from pytorch_ie.core import AnnotationList, annotation_field from pytorch_ie.documents import TextDocument from pytorch_ie.models import TransformerSeq2SeqModel +from pytorch_ie.pipeline import Pipeline from pytorch_ie.taskmodules import TransformerSeq2SeqTaskModule diff --git a/tests/pipeline/test_re_text_classification.py b/tests/pipeline/test_re_text_classification.py index 994da4c8..50c8ce6c 100644 --- a/tests/pipeline/test_re_text_classification.py +++ b/tests/pipeline/test_re_text_classification.py @@ -3,10 +3,11 @@ import pytest -from pytorch_ie import AnnotationList, Pipeline, annotation_field from pytorch_ie.annotations import BinaryRelation, LabeledSpan +from pytorch_ie.core import AnnotationList, annotation_field from pytorch_ie.documents import TextDocument from pytorch_ie.models import TransformerTextClassificationModel +from pytorch_ie.pipeline import Pipeline from pytorch_ie.taskmodules import TransformerRETextClassificationTaskModule diff --git a/tests/taskmodules/test_transformer_span_classification.py b/tests/taskmodules/test_transformer_span_classification.py index 6952118e..00e04c46 100644 --- a/tests/taskmodules/test_transformer_span_classification.py +++ b/tests/taskmodules/test_transformer_span_classification.py @@ -2,7 +2,7 @@ import pytest import torch -from pytorch_ie import TaskModule +from pytorch_ie.core import TaskModule from pytorch_ie.taskmodules.transformer_span_classification import ( TransformerSpanClassificationTaskModule, ) diff --git a/tests/taskmodules/test_transformer_token_classification.py b/tests/taskmodules/test_transformer_token_classification.py index 0fcae901..62486e0b 100644 --- a/tests/taskmodules/test_transformer_token_classification.py +++ b/tests/taskmodules/test_transformer_token_classification.py @@ -5,7 +5,7 @@ # import torch # from numpy.testing import assert_almost_equal -# from pytorch_ie import LabeledSpan +# from pytorch_ie.annotations import LabeledSpan # from pytorch_ie.taskmodules import TransformerTokenClassificationTaskModule # from tests.fixtures.document import ( # DOC1_ENTITY_BERLIN, diff --git a/tests/test_annotations.py b/tests/test_annotations.py index 527fb5f9..082ad5d3 100644 --- a/tests/test_annotations.py +++ b/tests/test_annotations.py @@ -3,7 +3,6 @@ import pytest -from pytorch_ie import AnnotationList, annotation_field from pytorch_ie.annotations import ( BinaryRelation, Label, @@ -15,6 +14,7 @@ MultiLabeledSpan, Span, ) +from pytorch_ie.core import AnnotationList, annotation_field from pytorch_ie.documents import TextDocument diff --git a/tests/test_auto.py b/tests/test_auto.py index 69c8cca1..8a038bb1 100644 --- a/tests/test_auto.py +++ b/tests/test_auto.py @@ -2,15 +2,9 @@ import pytest -from pytorch_ie import ( - AnnotationList, - AutoModel, - AutoPipeline, - AutoTaskModule, - TaskModule, - annotation_field, -) from pytorch_ie.annotations import LabeledSpan +from pytorch_ie.auto import AutoModel, AutoPipeline, AutoTaskModule +from pytorch_ie.core import AnnotationList, TaskModule, annotation_field from pytorch_ie.documents import TextDocument from pytorch_ie.models import TransformerSpanClassificationModel from pytorch_ie.taskmodules import TransformerSpanClassificationTaskModule diff --git a/tests/test_document.py b/tests/test_document.py index 2c9eab50..011163fb 100644 --- a/tests/test_document.py +++ b/tests/test_document.py @@ -3,8 +3,8 @@ import pytest -from pytorch_ie import AnnotationList, annotation_field from pytorch_ie.annotations import BinaryRelation, LabeledSpan, Span +from pytorch_ie.core import AnnotationList, annotation_field from pytorch_ie.documents import TextDocument diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index d3ed0a07..8ca845fc 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -6,9 +6,9 @@ from transformers.modeling_outputs import BaseModelOutputWithPooling import pytorch_ie.models.modules.mlp -from pytorch_ie import Pipeline from pytorch_ie.core.taskmodule import InplaceNotSupportedException from pytorch_ie.models.transformer_span_classification import TransformerSpanClassificationModel +from pytorch_ie.pipeline import Pipeline from pytorch_ie.taskmodules.transformer_span_classification import ( TransformerSpanClassificationTaskModule, ) diff --git a/tests/train/test_training.py b/tests/train/test_training.py index 616eb50d..df34e39a 100644 --- a/tests/train/test_training.py +++ b/tests/train/test_training.py @@ -4,7 +4,7 @@ import pytorch_lightning as pl from torch.utils.data import DataLoader -from pytorch_ie import Document, PyTorchIEModel, TaskModule +from pytorch_ie.core import Document, PyTorchIEModel, TaskModule from pytorch_ie.models import TransformerTokenClassificationModel from pytorch_ie.taskmodules import TransformerTokenClassificationTaskModule