diff --git a/src/adapters/heads/model_mixin.py b/src/adapters/heads/model_mixin.py index 89947732c..24f27602b 100644 --- a/src/adapters/heads/model_mixin.py +++ b/src/adapters/heads/model_mixin.py @@ -53,8 +53,8 @@ class ModelWithFlexibleHeadsAdaptersMixin(ModelWithHeadsAdaptersMixin): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self._convert_to_flex_head = True - if not hasattr(self.config, "custom_heads"): - self.config.custom_heads = {} + if not hasattr(self, "custom_heads"): + self.custom_heads = {} self._active_heads = [] def head_type(head_type_str: str): @@ -176,7 +176,7 @@ def add_prediction_head_from_config( head_class = MODEL_HEAD_MAP[head_type] head = head_class(self, head_name, **config) self.add_prediction_head(head, overwrite_ok=overwrite_ok, set_active=set_active) - elif head_type in self.config.custom_heads: + elif head_type in self.custom_heads: # we have to re-add the head type for custom heads self.add_custom_head(head_type, head_name, overwrite_ok=overwrite_ok, **config) else: @@ -193,7 +193,7 @@ def get_prediction_heads_config(self): return heads def register_custom_head(self, identifier, head): - self.config.custom_heads[identifier] = head + self.custom_heads[identifier] = head @property def active_head(self) -> Union[str, List[str]]: @@ -253,8 +253,8 @@ def set_active_adapters( ) def add_custom_head(self, head_type, head_name, overwrite_ok=False, set_active=True, **kwargs): - if head_type in self.config.custom_heads: - head = self.config.custom_heads[head_type](self, head_name, **kwargs) + if head_type in self.custom_heads: + head = self.custom_heads[head_type](self, head_name, **kwargs) # When a build-in head is added as a custom head it does not have the head_type property if not hasattr(head.config, "head_type"): head.config["head_type"] = head_type diff --git a/tests/test_adapter_custom_head.py b/tests/test_adapter_custom_head.py index b68662bfc..8e29636d0 100644 --- a/tests/test_adapter_custom_head.py +++ b/tests/test_adapter_custom_head.py @@ -48,9 +48,10 @@ def test_add_custom_head(self): def test_save_load_custom_head(self): model_name = "bert-base-uncased" model_config = AutoConfig.from_pretrained(model_name) - model_config.custom_heads = {"tag": CustomHead} model1 = AutoAdapterModel.from_pretrained(model_name, config=model_config) model2 = AutoAdapterModel.from_pretrained(model_name, config=model_config) + model1.custom_heads = {"tag": CustomHead} + model2.custom_heads = {"tag": CustomHead} config = {"num_labels": 3, "layers": 2, "activation_function": "tanh"} model1.add_custom_head(head_type="tag", head_name="custom_head", **config)