Rename wrap_model() -> init() (speechbrain#29)

calpt · calpt · commit efaa650e4b4d · 2023-08-25T16:00:50.000+02:00
diff --git a/README.md b/README.md
@@ -27,11 +27,11 @@ This branch disentangles `adapter-transformers` from HF Transformers and adds Tr
     ```
 - Built-in HF model classes can be adapted for usage with adapters via a wrapper method, e.g.:
     ```python
+    import adapters
     from transformers import BertModel
-    from adapters import wrap_model
 
     model = BertModel.from_pretrained("bert-base-uncased")
-    model = wrap_model(model)
+    adapters.init(model)
     ```
 
 ### Model support
diff --git a/examples/pytorch/translation/run_translation.py b/examples/pytorch/translation/run_translation.py
@@ -28,10 +28,10 @@
 import numpy as np
 from datasets import load_dataset
 
+import adapters
 import evaluate
 import transformers
 from adapters import AdapterArguments, Seq2SeqAdapterTrainer, setup_adapter_training
-from adapters.wrappers.model import wrap_model
 from transformers import (
     AutoConfig,
     AutoModelForSeq2SeqLM,
@@ -389,7 +389,7 @@ def main():
         revision=model_args.model_revision,
         use_auth_token=True if model_args.use_auth_token else None,
     )
-    model = wrap_model(model)
+    adapters.init(model)
 
     # We resize the embeddings only when necessary to avoid index errors. If you are creating a model from scratch
     # on a small vocab and want a smaller embedding size, remove this test.
diff --git a/src/adapters/__init__.py b/src/adapters/__init__.py
@@ -146,9 +146,9 @@
         "list_adapters",
     ],
     "wrappers": [
+        "init",
         "load_model",
         "wrap_config",
-        "wrap_model",
     ],
 }
 
@@ -240,7 +240,7 @@
         get_adapter_info,
         list_adapters,
     )
-    from .wrappers import load_model, wrap_config, wrap_model
+    from .wrappers import init, load_model, wrap_config
 
 else:
     import sys
diff --git a/src/adapters/models/albert/adapter_model.py b/src/adapters/models/albert/adapter_model.py
@@ -17,7 +17,7 @@
     TaggingHead,
 )
 from ...model_mixin import EmbeddingAdaptersWrapperMixin
-from ...wrappers import wrap_model
+from ...wrappers import init
 
 
 @add_start_docstrings(
@@ -28,7 +28,8 @@ class AlbertAdapterModel(EmbeddingAdaptersWrapperMixin, ModelWithFlexibleHeadsAd
     def __init__(self, config):
         super().__init__(config)
 
-        self.albert = wrap_model(AlbertModel(config))
+        self.albert = AlbertModel(config)
+        init(self.albert)
 
         self._init_head_modules()
 
diff --git a/src/adapters/models/bart/adapter_model.py b/src/adapters/models/bart/adapter_model.py
@@ -21,7 +21,7 @@
     Seq2SeqLMHead,
 )
 from ...model_mixin import EmbeddingAdaptersWrapperMixin
-from ...wrappers import wrap_model
+from ...wrappers import init
 
 
 @add_start_docstrings(
@@ -32,7 +32,8 @@ class BartAdapterModel(EmbeddingAdaptersWrapperMixin, ModelWithFlexibleHeadsAdap
 
     def __init__(self, config: BartConfig, **kwargs):
         super().__init__(config, **kwargs)
-        self.model = wrap_model(BartModel(config))
+        self.model = BartModel(config)
+        init(self.model)
 
         self._init_head_modules()
 
diff --git a/src/adapters/models/beit/adapter_model.py b/src/adapters/models/beit/adapter_model.py
@@ -12,7 +12,7 @@
 
 from ...context import AdapterSetup
 from ...heads import ImageClassificationHead, ModelWithFlexibleHeadsAdaptersMixin
-from ...wrappers import wrap_model
+from ...wrappers import init
 
 
 @add_start_docstrings(
@@ -23,7 +23,8 @@ class BeitAdapterModel(ModelWithFlexibleHeadsAdaptersMixin, BeitPreTrainedModel)
     def __init__(self, config):
         super().__init__(config)
 
-        self.beit = wrap_model(BeitModel(config))
+        self.beit = BeitModel(config)
+        init(self.beit)
 
         self._init_head_modules()
 
diff --git a/src/adapters/models/bert/adapter_model.py b/src/adapters/models/bert/adapter_model.py
@@ -21,7 +21,7 @@
     TaggingHead,
 )
 from ...model_mixin import EmbeddingAdaptersWrapperMixin
-from ...wrappers import wrap_model
+from ...wrappers import init
 
 
 @add_start_docstrings(
@@ -32,7 +32,8 @@ class BertAdapterModel(EmbeddingAdaptersWrapperMixin, ModelWithFlexibleHeadsAdap
     def __init__(self, config):
         super().__init__(config)
 
-        self.bert = wrap_model(BertModel(config))
+        self.bert = BertModel(config)
+        init(self.bert)
 
         self._init_head_modules()
 
diff --git a/src/adapters/models/bert_generation/adapter_model.py b/src/adapters/models/bert_generation/adapter_model.py
@@ -9,7 +9,7 @@
 from ...context import AdapterSetup
 from ...heads import BertStyleMaskedLMHead, CausalLMHead, ModelWithFlexibleHeadsAdaptersMixin
 from ...model_mixin import EmbeddingAdaptersWrapperMixin
-from ...wrappers import wrap_model
+from ...wrappers import init
 
 
 @add_start_docstrings(
@@ -24,7 +24,8 @@ class BertGenerationAdapterModel(
     def __init__(self, config):
         super().__init__(config)
 
-        self.bert = wrap_model(BertGenerationEncoder(config))
+        self.bert = BertGenerationEncoder(config)
+        init(self.bert)
 
         self._init_head_modules()
 
diff --git a/src/adapters/models/clip/adapter_model.py b/src/adapters/models/clip/adapter_model.py
@@ -13,15 +13,16 @@
 from ...context import AdapterSetup
 from ...heads import ModelWithFlexibleHeadsAdaptersMixin
 from ...model_mixin import EmbeddingAdaptersWrapperMixin
-from ...wrappers import wrap_model
+from ...wrappers import init
 
 
 @add_start_docstrings(CLIP_START_DOCSTRING)
 class CLIPAdapterModel(EmbeddingAdaptersWrapperMixin, ModelWithFlexibleHeadsAdaptersMixin, CLIPPreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
 
-        self.clip = wrap_model(CLIPModel(config))
+        self.clip = CLIPModel(config)
+        init(self.clip)
 
         self._init_head_modules()
 
diff --git a/src/adapters/models/deberta/adapter_model.py b/src/adapters/models/deberta/adapter_model.py
@@ -11,7 +11,7 @@
     TaggingHead,
 )
 from ...model_mixin import EmbeddingAdaptersWrapperMixin
-from ...wrappers import wrap_model
+from ...wrappers import init
 
 
 @add_start_docstrings(
@@ -23,7 +23,8 @@ class DebertaAdapterModel(EmbeddingAdaptersWrapperMixin, ModelWithFlexibleHeadsA
     def __init__(self, config):
         super().__init__(config)
 
-        self.deberta = wrap_model(DebertaModel(config))
+        self.deberta = DebertaModel(config)
+        init(self.deberta)
 
         self._init_head_modules()
 
diff --git a/src/adapters/models/deberta_v2/adapter_model.py b/src/adapters/models/deberta_v2/adapter_model.py
@@ -12,7 +12,7 @@
     TaggingHead,
 )
 from ...model_mixin import EmbeddingAdaptersWrapperMixin
-from ...wrappers import wrap_model
+from ...wrappers import init
 
 
 @add_start_docstrings(
@@ -26,7 +26,8 @@ class DebertaV2AdapterModel(
     def __init__(self, config):
         super().__init__(config)
 
-        self.deberta = wrap_model(DebertaV2Model(config))
+        self.deberta = DebertaV2Model(config)
+        init(self.deberta)
 
         self._init_head_modules()
 
diff --git a/src/adapters/models/distilbert/adapter_model.py b/src/adapters/models/distilbert/adapter_model.py
@@ -20,7 +20,7 @@
     TaggingHead,
 )
 from ...model_mixin import EmbeddingAdaptersWrapperMixin
-from ...wrappers import wrap_model
+from ...wrappers import init
 
 
 @add_start_docstrings(
@@ -32,7 +32,8 @@ class DistilBertAdapterModel(
 ):
     def __init__(self, config):
         super().__init__(config)
-        self.distilbert = wrap_model(DistilBertModel(config))
+        self.distilbert = DistilBertModel(config)
+        init(self.distilbert)
 
         self._init_head_modules()
 
diff --git a/src/adapters/models/gpt2/adapter_model.py b/src/adapters/models/gpt2/adapter_model.py
@@ -16,7 +16,7 @@
     TaggingHead,
 )
 from ...model_mixin import EmbeddingAdaptersWrapperMixin
-from ...wrappers import wrap_model
+from ...wrappers import init
 
 
 logger = logging.getLogger(__name__)
@@ -36,7 +36,8 @@
 class GPT2AdapterModel(EmbeddingAdaptersWrapperMixin, ModelWithFlexibleHeadsAdaptersMixin, GPT2PreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
-        self.transformer = wrap_model(GPT2Model(config))
+        self.transformer = GPT2Model(config)
+        init(self.transformer)
 
         self._init_head_modules()
 
diff --git a/src/adapters/models/gptj/adapter_model.py b/src/adapters/models/gptj/adapter_model.py
@@ -15,7 +15,7 @@
     TaggingHead,
 )
 from ...model_mixin import EmbeddingAdaptersWrapperMixin
-from ...wrappers import wrap_model
+from ...wrappers import init
 
 
 logger = logging.getLogger(__name__)
@@ -35,7 +35,8 @@
 class GPTJAdapterModel(EmbeddingAdaptersWrapperMixin, ModelWithFlexibleHeadsAdaptersMixin, GPTJPreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
-        self.transformer = wrap_model(GPTJModel(config))
+        self.transformer = GPTJModel(config)
+        init(self.transformer)
 
         self._init_head_modules()
 
diff --git a/src/adapters/models/mbart/adapter_model.py b/src/adapters/models/mbart/adapter_model.py
@@ -21,7 +21,7 @@
     Seq2SeqLMHead,
 )
 from ...model_mixin import EmbeddingAdaptersWrapperMixin
-from ...wrappers import wrap_model
+from ...wrappers import init
 
 
 @add_start_docstrings(
@@ -32,7 +32,8 @@ class MBartAdapterModel(EmbeddingAdaptersWrapperMixin, ModelWithFlexibleHeadsAda
 
     def __init__(self, config: MBartConfig, **kwargs):
         super().__init__(config, **kwargs)
-        self.model = wrap_model(MBartModel(config))
+        self.model = MBartModel(config)
+        init(self.model)
 
         self._init_head_modules()
 
diff --git a/src/adapters/models/roberta/adapter_model.py b/src/adapters/models/roberta/adapter_model.py
@@ -21,7 +21,7 @@
     TaggingHead,
 )
 from ...model_mixin import EmbeddingAdaptersWrapperMixin
-from ...wrappers import wrap_model
+from ...wrappers import init
 
 
 @add_start_docstrings(
@@ -32,7 +32,8 @@ class RobertaAdapterModel(EmbeddingAdaptersWrapperMixin, ModelWithFlexibleHeadsA
     def __init__(self, config):
         super().__init__(config)
 
-        self.roberta = wrap_model(RobertaModel(config))
+        self.roberta = RobertaModel(config)
+        init(self.roberta)
 
         self._init_head_modules()
 
diff --git a/src/adapters/models/t5/adapter_model.py b/src/adapters/models/t5/adapter_model.py
@@ -7,7 +7,7 @@
 
 from ...heads import ModelWithFlexibleHeadsAdaptersMixin, QuestionAnsweringHead, Seq2SeqLMHead
 from ...model_mixin import EmbeddingAdaptersWrapperMixin
-from ...wrappers import wrap_model
+from ...wrappers import init
 
 
 logger = logging.getLogger(__name__)
@@ -24,7 +24,8 @@ class T5AdapterModel(EmbeddingAdaptersWrapperMixin, ModelWithFlexibleHeadsAdapte
     def __init__(self, config):
         super().__init__(config)
 
-        self.transformer = wrap_model(T5Model(config))
+        self.transformer = T5Model(config)
+        init(self.transformer)
 
         self._init_head_modules()
 
diff --git a/src/adapters/models/vit/adapter_model.py b/src/adapters/models/vit/adapter_model.py
@@ -12,7 +12,7 @@
 
 from ...context import AdapterSetup
 from ...heads import ImageClassificationHead, ModelWithFlexibleHeadsAdaptersMixin
-from ...wrappers import wrap_model
+from ...wrappers import init
 
 
 @add_start_docstrings(
@@ -23,7 +23,8 @@ class ViTAdapterModel(ModelWithFlexibleHeadsAdaptersMixin, ViTPreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
 
-        self.vit = wrap_model(ViTModel(config))
+        self.vit = ViTModel(config)
+        init(self.vit)
 
         self._init_head_modules()
 
diff --git a/src/adapters/models/xlm_roberta/adapter_model.py b/src/adapters/models/xlm_roberta/adapter_model.py
@@ -21,7 +21,7 @@
     TaggingHead,
 )
 from ...model_mixin import EmbeddingAdaptersWrapperMixin
-from ...wrappers import wrap_model
+from ...wrappers import init
 
 
 @add_start_docstrings(
@@ -34,7 +34,8 @@ class XLMRobertaAdapterModel(
     def __init__(self, config):
         super().__init__(config)
 
-        self.roberta = wrap_model(XLMRobertaModel(config))
+        self.roberta = XLMRobertaModel(config)
+        init(self.roberta)
 
         self._init_head_modules()
 
diff --git a/src/adapters/wrappers/__init__.py b/src/adapters/wrappers/__init__.py
@@ -21,12 +21,12 @@
 from transformers.utils import _LazyModule
 
 
-_import_structure = {"configuration": ["wrap_config"], "model": ["wrap_model", "load_model"]}
+_import_structure = {"configuration": ["wrap_config"], "model": ["init", "load_model"]}
 
 
 if TYPE_CHECKING:
     from .configuration import wrap_config
-    from .model import load_model, wrap_model
+    from .model import init, load_model
 
 else:
     import sys
diff --git a/src/adapters/wrappers/model.py b/src/adapters/wrappers/model.py
diff --git a/tests_adapters/composition/test_adapter_composition.py b/tests_adapters/composition/test_adapter_composition.py
diff --git a/tests_adapters/methods/base.py b/tests_adapters/methods/base.py
diff --git a/tests_adapters/methods/test_adapter_common.py b/tests_adapters/methods/test_adapter_common.py
diff --git a/tests_adapters/test_adapter.py b/tests_adapters/test_adapter.py
diff --git a/tests_adapters/test_adapter_conversion.py b/tests_adapters/test_adapter_conversion.py
diff --git a/tests_adapters/test_adapter_heads.py b/tests_adapters/test_adapter_heads.py
diff --git a/tests_adapters/test_adapter_hub.py b/tests_adapters/test_adapter_hub.py
diff --git a/tests_adapters/test_adapter_save_id2label.py b/tests_adapters/test_adapter_save_id2label.py
diff --git a/tests_adapters/test_adapter_trainer.py b/tests_adapters/test_adapter_trainer.py
