register all models @pie_core.Model (#503)

This PR registers all models `@pie_core.Model` instead of
`@PyTorchIEModel` to make them work with
`pie_core.AutoAnnotationPipeline.from_config` in the future.

Additional changes:
 - remove `AutoPyTorchIEModel` (just use `AutoModel` instead)
- re-export `AutoModel` from `pie_core`. NOTE: Previously
`AutoPyTorchIEModel` was exported as `AutoModel`! But if used correctly,
it should not make a difference.
 - register `PyTorchIEPipeline` at name `"pytorch-ie"`
- remove backwards compatibility: `Pipeline` and `AutoPipeline`(use
`PyTorchIEPipeline` in either case)
 - adjust the `README.md`

Note: It is not yet possible to load an "old" pipeline (e.g.
[pie/example-ner-spanclf-conll03](https://huggingface.co/pie/example-ner-spanclf-conll03))
via `AutoAnnotationPipeline` since its config does not has a
`pipeline_type` key. In the future, it should be possible to provide
that as argument in the following manner, but **this still needs a minor
adjustment in `pie_core.Auto.from_config` (TODO: reference issue/PR in
pie-core)**:
```
pipeline = AutoAnnotationPipeline.from_pretrained("pie/example-ner-spanclf-conll03", pipeline_type="pytorch-ie")
```

Author: ArneBinder

README.md

@@ -318,8 +318,8 @@ for details.
 ```python
 from dataclasses import dataclass
 
+from pytorch_ie import PyTorchIEPipeline
 from pytorch_ie.annotations import LabeledSpan
-from pytorch_ie.auto import AutoPipeline
 from pytorch_ie.core import AnnotationLayer, annotation_field
 from pytorch_ie.documents import TextDocument
 
@@ -334,7 +334,7 @@ document = ExampleDocument(
 )
 
 # see below for the long version
-ner_pipeline = AutoPipeline.from_pretrained("pie/example-ner-spanclf-conll03", device=-1, num_workers=0)
+ner_pipeline = PyTorchIEPipeline.from_pretrained("pie/example-ner-spanclf-conll03", device=-1, num_workers=0)
 
 ner_pipeline(document)
 
@@ -349,7 +349,7 @@ for entity in document.entities.predictions:
 
 <details>
 <summary>
-To create the same pipeline as above without `AutoPipeline`
+Under the hood, the following happens when calling `PyTorchIEPipeline.from_pretrained`
 </summary>
 
 ```python
@@ -364,23 +364,6 @@ ner_pipeline = PyTorchIEPipeline(model=ner_model, taskmodule=ner_taskmodule, dev
 
 </details>
 
-<details>
-<summary>
-Or, without `Auto` classes at all
-</summary>
-
-```python
-from pytorch_ie.pipeline import PyTorchIEPipeline
-from pytorch_ie.models import TransformerSpanClassificationModel
-from pytorch_ie.taskmodules import TransformerSpanClassificationTaskModule
-
-model_name_or_path = "pie/example-ner-spanclf-conll03"
-ner_taskmodule = TransformerSpanClassificationTaskModule.from_pretrained(model_name_or_path)
-ner_model = TransformerSpanClassificationModel.from_pretrained(model_name_or_path)
-ner_pipeline = PyTorchIEPipeline(model=ner_model, taskmodule=ner_taskmodule, device=-1, num_workers=0)
-```
-
-</details>
 <details>
 <summary>
 
@@ -391,8 +374,8 @@ ner_pipeline = PyTorchIEPipeline(model=ner_model, taskmodule=ner_taskmodule, dev
 ```python
 from dataclasses import dataclass
 
+from pytorch_ie import PyTorchIEPipeline
 from pytorch_ie.annotations import BinaryRelation, LabeledSpan
-from pytorch_ie.auto import AutoPipeline
 from pytorch_ie.core import AnnotationLayer, annotation_field
 from pytorch_ie.documents import TextDocument
 
@@ -407,7 +390,7 @@ 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)
+re_pipeline = PyTorchIEPipeline.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))

examples/predict/ner_span_classification.py

@@ -2,10 +2,10 @@
 
 from pie_core import AnnotationLayer, annotation_field
 
+from pytorch_ie import PyTorchIEPipeline
 from pytorch_ie.annotations import LabeledSpan
 from pytorch_ie.documents import TextDocument
 from pytorch_ie.models import TransformerSpanClassificationModel
-from pytorch_ie.pipeline import Pipeline
 from pytorch_ie.taskmodules import TransformerSpanClassificationTaskModule
 
 
@@ -19,7 +19,7 @@ def main():
     ner_taskmodule = TransformerSpanClassificationTaskModule.from_pretrained(model_name_or_path)
     ner_model = TransformerSpanClassificationModel.from_pretrained(model_name_or_path)
 
-    ner_pipeline = Pipeline(model=ner_model, taskmodule=ner_taskmodule, device=-1)
+    ner_pipeline = PyTorchIEPipeline(model=ner_model, taskmodule=ner_taskmodule, device=-1)
 
     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."

examples/predict/re_generative.py

@@ -2,10 +2,10 @@
 
 from pie_core import AnnotationLayer, annotation_field
 
+from pytorch_ie import PyTorchIEPipeline
 from pytorch_ie.annotations import BinaryRelation, LabeledSpan
 from pytorch_ie.documents import TextDocument
 from pytorch_ie.models import TransformerSeq2SeqModel
-from pytorch_ie.pipeline import Pipeline
 from pytorch_ie.taskmodules import TransformerSeq2SeqTaskModule
 
 
@@ -30,7 +30,7 @@ def main():
         model_name_or_path=model_name_or_path,
     )
 
-    pipeline = Pipeline(model=model, taskmodule=taskmodule, device=-1)
+    pipeline = PyTorchIEPipeline(model=model, taskmodule=taskmodule, device=-1)
 
     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."

examples/predict/re_text_classification.py

@@ -2,10 +2,10 @@
 
 from pie_core import AnnotationLayer, annotation_field
 
+from pytorch_ie import PyTorchIEPipeline
 from pytorch_ie.annotations import BinaryRelation, LabeledSpan
 from pytorch_ie.documents import TextDocument
 from pytorch_ie.models import TransformerTextClassificationModel
-from pytorch_ie.pipeline import Pipeline
 from pytorch_ie.taskmodules import TransformerRETextClassificationTaskModule
 
 
@@ -22,7 +22,7 @@ def main():
     )
     re_model = TransformerTextClassificationModel.from_pretrained(model_name_or_path)
 
-    re_pipeline = Pipeline(model=re_model, taskmodule=re_taskmodule, device=-1)
+    re_pipeline = PyTorchIEPipeline(model=re_model, taskmodule=re_taskmodule, device=-1)
 
     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."

src/pytorch_ie/__init__.py

@@ -1,10 +1,7 @@
 # flake8: noqa
 
-from pytorch_ie.auto import AutoModel, AutoPipeline, AutoTaskModule
+from pytorch_ie.auto import AutoModel, AutoTaskModule
 from pytorch_ie.core import *
 from pytorch_ie.datamodule import PieDataModule
 from pytorch_ie.dataset import IterableTaskEncodingDataset, TaskEncodingDataset
 from pytorch_ie.pipeline import PyTorchIEPipeline
-
-# kept for backward compatibility
-Pipeline = PyTorchIEPipeline

src/pytorch_ie/auto.py

@@ -1,6 +1,2 @@
 # kept for backward compatibility
-from pie_core import AutoTaskModule
-
-# kept for backward compatibility
-from pytorch_ie.model import AutoPyTorchIEModel as AutoModel
-from pytorch_ie.pipeline import PyTorchIEPipeline as AutoPipeline
+from pie_core import AutoModel, AutoTaskModule

src/pytorch_ie/model.py

@@ -38,10 +38,3 @@ def predict(self, inputs: Any, **kwargs) -> Any:
         outputs = self(inputs, **kwargs)
         decoded_outputs = self.decode(inputs=inputs, outputs=outputs)
         return decoded_outputs
-
-
-# TODO: remove this class when all models are registered with @Model.register()
-#   also see notes in PyTorchIEPipeline
-class AutoPyTorchIEModel(Model, Auto[PyTorchIEModel]):
-
-    BASE_CLASS = PyTorchIEModel

src/pytorch_ie/models/sequence_classification_with_pooler.py

@@ -14,14 +14,14 @@
 )
 
 import torch
+from pie_core import Model
 from torch import FloatTensor, LongTensor, nn
 from torch.nn import Parameter
 from torch.optim import AdamW
 from transformers import AutoConfig, AutoModel, PreTrainedModel, get_linear_schedule_with_warmup
 from transformers.modeling_outputs import SequenceClassifierOutput
 from typing_extensions import TypeAlias
 
-from pytorch_ie import PyTorchIEModel
 from pytorch_ie.models.interface import RequiresModelNameOrPath, RequiresNumClasses
 
 from .common import ModelWithBoilerplate
@@ -236,7 +236,7 @@ def configure_optimizers(self):
             return optimizer
 
 
-@PyTorchIEModel.register()
+@Model.register()
 class SequenceClassificationModelWithPooler(
     SequenceClassificationModelWithPoolerBase,
     RequiresNumClasses,
@@ -286,7 +286,7 @@ def decode(self, inputs: InputType, outputs: OutputType) -> TargetType:
         return {"labels": labels, "probabilities": probabilities}
 
 
-@PyTorchIEModel.register()
+@Model.register()
 class SequencePairSimilarityModelWithPooler(
     SequenceClassificationModelWithPoolerBase,
 ):

src/pytorch_ie/models/simple_generative.py

@@ -3,6 +3,7 @@
 from typing import Any, Dict, Optional, Tuple, Type, Union
 
 import torch
+from pie_core import Model
 from pie_core.utils.hydra import resolve_type
 from pytorch_lightning.utilities.types import OptimizerLRScheduler
 from torch import FloatTensor, LongTensor
@@ -11,8 +12,6 @@
 from transformers.modeling_outputs import Seq2SeqLMOutput
 from typing_extensions import TypeAlias
 
-from pytorch_ie import PyTorchIEModel
-
 from .common import ModelWithBoilerplate
 
 logger = logging.getLogger(__name__)
@@ -26,7 +25,7 @@
 StepOutputType: TypeAlias = FloatTensor
 
 
-@PyTorchIEModel.register()
+@Model.register()
 class SimpleGenerativeModel(
     ModelWithBoilerplate[InputType, OutputType, TargetType, StepOutputType],
 ):

src/pytorch_ie/models/simple_sequence_classification.py

@@ -2,6 +2,7 @@
 from typing import Iterator, MutableMapping, Optional, Tuple, Union
 
 import torch.nn
+from pie_core import Model
 from torch import FloatTensor, LongTensor
 from torch.nn import Parameter
 from torch.optim import AdamW
@@ -13,7 +14,6 @@
 from transformers.modeling_outputs import SequenceClassifierOutput
 from typing_extensions import TypeAlias
 
-from pytorch_ie import PyTorchIEModel
 from pytorch_ie.models.common import ModelWithBoilerplate
 from pytorch_ie.models.interface import RequiresModelNameOrPath, RequiresNumClasses
 
@@ -29,7 +29,7 @@
 logger = logging.getLogger(__name__)
 
 
-@PyTorchIEModel.register()
+@Model.register()
 class SimpleSequenceClassificationModel(
     ModelWithBoilerplate[InputType, OutputType, TargetType, StepOutputType],
     RequiresModelNameOrPath,