BEiT model refactoring (speechbrain#9)

src/adapter_transformers/mixins/__init__.py

@@ -1,5 +1,6 @@
 from .albert import AlbertModelAdaptersMixin
 from .bart import BartDecoderAdaptersMixin, BartEncoderAdaptersMixin, BartModelAdaptersMixin
+from .beit import BeitIntermediateAdaptersMixin, BeitModelAdaptersMixin, BeitOutputAdaptersMixin
 from .bert import BertLayerAdaptersMixin, BertModelAdaptersMixin
 from .distilbert import DistilBertModelAdaptersMixin, DistilBertTransformerAdaptersMixin
 from .t5 import T5BlockAdaptersMixin, T5ModelAdaptersMixin, T5ModelAdaptersWithHeadsMixin
@@ -11,6 +12,9 @@
     "BartEncoder": BartEncoderAdaptersMixin,
     "BartDecoder": BartDecoderAdaptersMixin,
     "BartModel": BartModelAdaptersMixin,
+    "BeitIntermediate": BeitIntermediateAdaptersMixin,
+    "BeitOutput": BeitOutputAdaptersMixin,
+    "BeitModel": BeitModelAdaptersMixin,
     "BertLayer": BertLayerAdaptersMixin,
     "BertModel": BertModelAdaptersMixin,
     "Transformer": DistilBertTransformerAdaptersMixin,

src/adapter_transformers/mixins/beit.py

@@ -1,33 +1,54 @@
-import logging
 from typing import Iterable, Tuple
 
 import torch.nn as nn
 
 from ..layer import AdapterLayer
-from ..model_mixin import ModelAdaptersMixin, ModelWithHeadsAdaptersMixin
+from ..lora import Linear as LoRALinear
+from ..model_mixin import ModelBaseAdaptersMixin
+from ..prefix_tuning import PrefixTuningShim
 
 
-logger = logging.getLogger(__name__)
+class BeitSelfAttentionAdaptersMixin:
+    def init_adapters(self, config):
+        self.location_key = "self"
+
+        # Wrap layers for LoRA
+        self.query = LoRALinear.wrap(self.query, "selfattn", config, attn_key="q")
+        self.key = LoRALinear.wrap(self.key, "selfattn", config, attn_key="k")
+        self.value = LoRALinear.wrap(self.value, "selfattn", config, attn_key="v")
+
+        self.prefix_tuning = PrefixTuningShim(self.location_key + "_prefix" if self.location_key else None, config)
+
+
+class BeitIntermediateAdaptersMixin:
+    def init_adapters(self, config):
+        # Wrap layers for LoRA
+        self.dense = LoRALinear.wrap(self.dense, "intermediate", config)
+
+
+class BeitOutputAdaptersMixin:
+    def init_adapters(self, config):
+        # Wrap layers for LoRA
+        self.dense = LoRALinear.wrap(self.dense, "output", config)
 
 
 class BeitLayerAdaptersMixin:
     """Adds adapters to the BeitLayer module."""
 
-    def _init_adapter_modules(self):
-        self.attention_adapters = AdapterLayer("mh_adapter", self.config)
-        self.attention_adapters._init_adapter_modules()
-
-        self.output_adapters = AdapterLayer("output_adapter", self.config)
-        self.output_adapters._init_adapter_modules()
+    def init_adapters(self, config):
+        self.attention_adapters = AdapterLayer("mh_adapter")
+        self.output_adapters = AdapterLayer("output_adapter")
 
 
-class BeitModelAdaptersMixin(ModelAdaptersMixin):
+class BeitModelAdaptersMixin(ModelBaseAdaptersMixin):
     """Adds adapters to the BeitModel module."""
 
+    def init_adapters(self, config):
+        super().init_adapters(config)
+
     def iter_layers(self) -> Iterable[Tuple[int, nn.Module]]:
         for i, layer in enumerate(self.encoder.layer):
             yield i, layer
 
-
-class BeitModelWithHeadsAdaptersMixin(ModelWithHeadsAdaptersMixin):
-    pass
+    def set_input_embeddings(self, value):
+        self.embeddings.patch_embeddings = value

src/adapter_transformers/models/beit/__init__.py

@@ -23,27 +23,11 @@
 
 _import_structure = {
     "adapter_model": ["BeitAdapterModel"],
-    "modeling_bert": [
-        "BEIT_PRETRAINED_MODEL_ARCHIVE_LIST",
-        "BeitForImageClassification",
-        "BeitForMaskedImageModeling",
-        "BeitForSemanticSegmentation",
-        "BeitModel",
-        "BeitPreTrainedModel",
-    ],
 }
 
 
 if TYPE_CHECKING:
     from .adapter_model import BeitAdapterModel
-    from .modeling_beit import (
-        BEIT_PRETRAINED_MODEL_ARCHIVE_LIST,
-        BeitForImageClassification,
-        BeitForMaskedImageModeling,
-        BeitForSemanticSegmentation,
-        BeitModel,
-        BeitPreTrainedModel,
-    )
 
 else:
     import sys

src/adapter_transformers/models/beit/adapter_model.py

@@ -2,11 +2,17 @@
 
 import torch
 
+from transformers.models.beit.modeling_beit import (
+    BEIT_INPUTS_DOCSTRING,
+    BEIT_START_DOCSTRING,
+    BeitModel,
+    BeitPreTrainedModel,
+)
 from transformers.utils import add_start_docstrings, add_start_docstrings_to_model_forward
 
 from ...context import AdapterSetup
 from ...heads import ImageClassificationHead, ModelWithFlexibleHeadsAdaptersMixin
-from .modeling_beit import BEIT_INPUTS_DOCSTRING, BEIT_START_DOCSTRING, BeitModel, BeitPreTrainedModel
+from ...wrappers import wrap_model
 
 
 @add_start_docstrings(
@@ -17,7 +23,7 @@ class BeitAdapterModel(ModelWithFlexibleHeadsAdaptersMixin, BeitPreTrainedModel)
     def __init__(self, config):
         super().__init__(config)
 
-        self.beit = BeitModel(config)
+        self.beit = wrap_model(BeitModel(config))
 
         self._init_head_modules()