fusion-bench 0.2.6__py3-none-any.whl → 0.2.8__py3-none-any.whl

fusion_bench/modelpool/clip_vision/modelpool.py

@@ -1,8 +1,11 @@
 import logging
 from copy import deepcopy
-from typing import Optional
+from typing import Optional, Union
 
+from datasets import load_dataset
 from omegaconf import DictConfig, open_dict
+from torch import nn
+from torch.utils.data import Dataset
 from transformers import CLIPModel, CLIPProcessor, CLIPVisionModel
 from typing_extensions import override
 
@@ -36,17 +39,29 @@ class CLIPVisionModelPool(BaseModelPool):
 
     def load_processor(self, *args, **kwargs) -> CLIPProcessor:
         assert self._processor is not None, "Processor is not defined in the config"
-        processor = instantiate(self._processor, *args, **kwargs)
+        if isinstance(self._processor, str):
+            log.info(f"Loading `transformers.CLIPProcessor`: {self._processor}")
+            processor = CLIPProcessor.from_pretrained(self._processor)
+        else:
+            processor = instantiate(self._processor, *args, **kwargs)
         return processor
 
     def load_clip_model(self, model_name: str, *args, **kwargs) -> CLIPModel:
         model_config = self._models[model_name]
-        assert isinstance(model_config, DictConfig), "Model config must be a DictConfig"
-        model_config = deepcopy(model_config)
-        with open_dict(model_config):
-            model_config._target_ = "transformers.CLIPModel.from_pretrained"
-        clip_model = instantiate(model_config, *args, **kwargs)
-        return clip_model
+
+        if isinstance(model_config, str):
+            log.info(f"Loading `transformers.CLIPModel`: {model_config}")
+            clip_model = CLIPModel.from_pretrained(model_config, *args, **kwargs)
+            return clip_model
+        else:
+            assert isinstance(
+                model_config, DictConfig
+            ), "Model config must be a DictConfig"
+            model_config = deepcopy(model_config)
+            with open_dict(model_config):
+                model_config._target_ = "transformers.CLIPModel.from_pretrained"
+            clip_model = instantiate(model_config, *args, **kwargs)
+            return clip_model
 
     @override
     def save_model(self, model: CLIPVisionModel, path: str):
@@ -59,3 +74,72 @@ class CLIPVisionModelPool(BaseModelPool):
         """
         with timeit_context(f'Saving clip vision model to "{path}"'):
             model.save_pretrained(path)
+
+    def load_model(
+        self, model_name_or_config: Union[str, DictConfig], *args, **kwargs
+    ) -> CLIPVisionModel:
+        """
+        This method is used to load a CLIPVisionModel from the model pool.
+
+        Example configuration could be:
+
+        ```yaml
+        models:
+            cifar10: tanganke/clip-vit-base-patch32_cifar10
+            sun397: tanganke/clip-vit-base-patch32_sun397
+            stanford-cars: tanganke/clip-vit-base-patch32_stanford-cars
+        ```
+
+        Args:
+            model_name_or_config (Union[str, DictConfig]): The name of the model or the model configuration.
+
+        Returns:
+            CLIPVisionModel: The loaded CLIPVisionModel.
+        """
+        if (
+            isinstance(model_name_or_config, str)
+            and model_name_or_config in self._models
+        ):
+            model = self._models[model_name_or_config]
+            if isinstance(model, str):
+                log.info(f"Loading `transformers.CLIPVisionModel`: {model}")
+                return CLIPVisionModel.from_pretrained(model, *args, **kwargs)
+            if isinstance(model, nn.Module):
+                log.info(f"Returning existing model: {model}")
+                return model
+
+        # If the model is not a string, we use the default load_model method
+        return super().load_model(model_name_or_config, *args, **kwargs)
+
+    def load_train_dataset(self, dataset_name: str, *args, **kwargs):
+        dataset_config = self._train_datasets[dataset_name]
+        if isinstance(dataset_config, str):
+            log.info(
+                f"Loading train dataset using `datasets.load_dataset`: {dataset_config}"
+            )
+            dataset = load_dataset(dataset_config, split="train")
+        else:
+            dataset = super().load_train_dataset(dataset_name, *args, **kwargs)
+        return dataset
+
+    def load_val_dataset(self, dataset_name: str, *args, **kwargs):
+        dataset_config = self._val_datasets[dataset_name]
+        if isinstance(dataset_config, str):
+            log.info(
+                f"Loading validation dataset using `datasets.load_dataset`: {dataset_config}"
+            )
+            dataset = load_dataset(dataset_config, split="validation")
+        else:
+            dataset = super().load_val_dataset(dataset_name, *args, **kwargs)
+        return dataset
+
+    def load_test_dataset(self, dataset_name: str, *args, **kwargs):
+        dataset_config = self._test_datasets[dataset_name]
+        if isinstance(dataset_config, str):
+            log.info(
+                f"Loading test dataset using `datasets.load_dataset`: {dataset_config}"
+            )
+            dataset = load_dataset(dataset_config, split="test")
+        else:
+            dataset = super().load_test_dataset(dataset_name, *args, **kwargs)
+        return dataset

fusion_bench/modelpool/seq_classification_lm/__init__.py

@@ -0,0 +1,2 @@
+from .reward_model import create_reward_model_from_pretrained
+from .seq_classification_lm import SeqenceClassificationModelPool

fusion_bench/modelpool/seq_classification_lm/reward_model.py

@@ -0,0 +1,15 @@
+from transformers import AutoModelForSequenceClassification
+
+
+def create_reward_model_from_pretrained(pretrained_model_name_or_path: str, **kwargs):
+    """
+    Create a reward model for reward modeling (RLHF).
+
+    Args:
+        pretrained_model_name_or_path (str): The name or path of the pretrained model.
+        **kwargs: Additional keyword arguments passed to the model class.
+    """
+    model = AutoModelForSequenceClassification.from_pretrained(
+        pretrained_model_name_or_path, num_labels=1, **kwargs
+    )
+    return model

fusion_bench/modelpool/seq_classification_lm/seq_classification_lm.py

@@ -0,0 +1,98 @@
+import logging
+import os
+from copy import deepcopy
+from typing import TYPE_CHECKING, Any, Optional, TypeAlias, Union, cast  # noqa: F401
+
+from omegaconf import DictConfig, flag_override
+from transformers import PreTrainedModel, PreTrainedTokenizer
+from typing_extensions import override
+
+from fusion_bench.modelpool import BaseModelPool
+from fusion_bench.utils import instantiate
+from fusion_bench.utils.dtype import parse_dtype
+
+if TYPE_CHECKING:
+    from transformers import LlamaForSequenceClassification
+
+log = logging.getLogger(__name__)
+
+
+class SeqenceClassificationModelPool(BaseModelPool):
+
+    def __init__(
+        self,
+        models,
+        *,
+        tokenizer: Optional[DictConfig],
+        model_kwargs: Optional[DictConfig] = None,
+        **kwargs,
+    ):
+        super().__init__(models, **kwargs)
+        # process `model_kwargs`
+        self._tokenizer = tokenizer
+        self._model_kwargs = model_kwargs
+        if self._model_kwargs is None:
+            self._model_kwargs = DictConfig({})
+        with flag_override(self._model_kwargs, "allow_objects", True):
+            if hasattr(self._model_kwargs, "torch_dtype"):
+                self._model_kwargs.torch_dtype = parse_dtype(
+                    self._model_kwargs.torch_dtype
+                )
+
+    @override
+    def load_model(
+        self,
+        model_name_or_config: str | DictConfig,
+        *args,
+        **kwargs,
+    ) -> Union[PreTrainedModel, "LlamaForSequenceClassification"]:
+        model_kwargs = deepcopy(self._model_kwargs)
+        model_kwargs.update(kwargs)
+        if isinstance(model_name_or_config, str):
+            log.info(f"Loading model: {model_name_or_config}", stacklevel=2)
+        return super().load_model(model_name_or_config, *args, **model_kwargs)
+
+    def load_tokenizer(self, *args, **kwargs) -> PreTrainedTokenizer:
+        assert self._tokenizer is not None, "Tokenizer is not defined in the config"
+        log.info("Loading tokenizer.", stacklevel=2)
+        tokenizer = instantiate(self._tokenizer, *args, **kwargs)
+        return tokenizer
+
+    @override
+    def save_model(
+        self,
+        model: PreTrainedModel,
+        path: str,
+        push_to_hub: bool = False,
+        model_dtype: Optional[str] = None,
+        save_tokenizer: bool = False,
+        tokenizer_kwargs=None,
+        **kwargs,
+    ):
+        """
+        Save the model to the specified path.
+
+        Args:
+            model (PreTrainedModel): The model to be saved.
+            path (str): The path where the model will be saved.
+            push_to_hub (bool, optional): Whether to push the model to the Hugging Face Hub. Defaults to False.
+            save_tokenizer (bool, optional): Whether to save the tokenizer along with the model. Defaults to False.
+            **kwargs: Additional keyword arguments passed to the `save_pretrained` method.
+        """
+        path = os.path.expanduser(path)
+        if save_tokenizer:
+            if tokenizer_kwargs is None:
+                tokenizer_kwargs = {}
+            # load the tokenizer
+            tokenizer = self.load_tokenizer(**tokenizer_kwargs)
+            tokenizer.save_pretrained(
+                path,
+                push_to_hub=push_to_hub,
+            )
+        if model_dtype is not None:
+            model.to(dtype=parse_dtype(model_dtype))
+        model.save_pretrained(
+            path,
+            push_to_hub=push_to_hub,
+            **kwargs,
+        )

fusion_bench/models/chat_templates/__init__.py

@@ -0,0 +1 @@
+from .load_tokenizer import chat_template_mapping, load_tokenizer_with_chat_template

fusion_bench/models/chat_templates/llama_3_Instruct.py

@@ -0,0 +1 @@
+CHAT_TEMPLATE = '{{- bos_token }}\n{%- if custom_tools is defined %}\n    {%- set tools = custom_tools %}\n{%- endif %}\n{%- if not tools_in_user_message is defined %}\n    {%- set tools_in_user_message = true %}\n{%- endif %}\n{%- if not date_string is defined %}\n    {%- if strftime_now is defined %}\n        {%- set date_string = strftime_now("%d %b %Y") %}\n    {%- else %}\n        {%- set date_string = "26 Jul 2024" %}\n    {%- endif %}\n{%- endif %}\n{%- if not tools is defined %}\n    {%- set tools = none %}\n{%- endif %}\n\n{#- This block extracts the system message, so we can slot it into the right place. #}\n{%- if messages[0][\'role\'] == \'system\' %}\n    {%- set system_message = messages[0][\'content\']|trim %}\n    {%- set messages = messages[1:] %}\n{%- else %}\n    {%- set system_message = "" %}\n{%- endif %}\n\n{#- System message #}\n{{- "<|start_header_id|>system<|end_header_id|>\\n\\n" }}\n{%- if tools is not none %}\n    {{- "Environment: ipython\\n" }}\n{%- endif %}\n{{- "Cutting Knowledge Date: December 2023\\n" }}\n{{- "Today Date: " + date_string + "\\n\\n" }}\n{%- if tools is not none and not tools_in_user_message %}\n    {{- "You have access to the following functions. To call a function, please respond with JSON for a function call." }}\n    {{- \'Respond in the format {"name": function name, "parameters": dictionary of argument name and its value}.\' }}\n    {{- "Do not use variables.\\n\\n" }}\n    {%- for t in tools %}\n        {{- t | tojson(indent=4) }}\n        {{- "\\n\\n" }}\n    {%- endfor %}\n{%- endif %}\n{{- system_message }}\n{{- "<|eot_id|>" }}\n\n{#- Custom tools are passed in a user message with some extra guidance #}\n{%- if tools_in_user_message and not tools is none %}\n    {#- Extract the first user message so we can plug it in here #}\n    {%- if messages | length != 0 %}\n        {%- set first_user_message = messages[0][\'content\']|trim %}\n        {%- set messages = messages[1:] %}\n    {%- else %}\n        {{- raise_exception("Cannot put tools in the first user message when there\'s no first user message!") }}\n{%- endif %}\n    {{- \'<|start_header_id|>user<|end_header_id|>\\n\\n\' -}}\n    {{- "Given the following functions, please respond with a JSON for a function call " }}\n    {{- "with its proper arguments that best answers the given prompt.\\n\\n" }}\n    {{- \'Respond in the format {"name": function name, "parameters": dictionary of argument name and its value}.\' }}\n    {{- "Do not use variables.\\n\\n" }}\n    {%- for t in tools %}\n        {{- t | tojson(indent=4) }}\n        {{- "\\n\\n" }}\n    {%- endfor %}\n    {{- first_user_message + "<|eot_id|>"}}\n{%- endif %}\n\n{%- for message in messages %}\n    {%- if not (message.role == \'ipython\' or message.role == \'tool\' or \'tool_calls\' in message) %}\n        {{- \'<|start_header_id|>\' + message[\'role\'] + \'<|end_header_id|>\\n\\n\'+ message[\'content\'] | trim + \'<|eot_id|>\' }}\n    {%- elif \'tool_calls\' in message %}\n        {%- if not message.tool_calls|length == 1 %}\n            {{- raise_exception("This model only supports single tool-calls at once!") }}\n        {%- endif %}\n        {%- set tool_call = message.tool_calls[0].function %}\n        {{- \'<|start_header_id|>assistant<|end_header_id|>\\n\\n\' -}}\n        {{- \'{"name": "\' + tool_call.name + \'", \' }}\n        {{- \'"parameters": \' }}\n        {{- tool_call.arguments | tojson }}\n        {{- "}" }}\n        {{- "<|eot_id|>" }}\n    {%- elif message.role == "tool" or message.role == "ipython" %}\n        {{- "<|start_header_id|>ipython<|end_header_id|>\\n\\n" }}\n        {%- if message.content is mapping or message.content is iterable %}\n            {{- message.content | tojson }}\n        {%- else %}\n            {{- message.content }}\n        {%- endif %}\n        {{- "<|eot_id|>" }}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- \'<|start_header_id|>assistant<|end_header_id|>\\n\\n\' }}\n{%- endif %}\n'

fusion_bench/models/chat_templates/load_tokenizer.py

@@ -0,0 +1,43 @@
+import logging
+
+from transformers import AutoTokenizer
+
+from .llama_3_Instruct import CHAT_TEMPLATE as LLAMA_3_INSTRUCT_CHAT_TEMPLATE
+
+chat_template_mapping = {"llama_3_instruct": LLAMA_3_INSTRUCT_CHAT_TEMPLATE}
+
+log = logging.getLogger(__name__)
+
+
+def load_tokenizer_with_chat_template(
+    pretrained_model_name_or_path: str,
+    model_family: str,
+    overwrite_chat_template: bool = True,
+    **kwargs,
+):
+    """
+    Load the tokenizer for Llama 3 model.
+
+    Args:
+        pretrained_model_name_or_path (str): The name or path of the pretrained model.
+        model_family (str): The model family.
+        **kwargs: Additional keyword arguments passed to the tokenizer class.
+    """
+    assert (
+        model_family in chat_template_mapping
+    ), f"Model family {model_family} not found. Available model families: {chat_template_mapping.keys()}"
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        pretrained_model_name_or_path,
+        **kwargs,
+    )
+
+    if tokenizer.chat_template is None:
+        tokenizer.chat_template = chat_template_mapping[model_family]
+    else:
+        if overwrite_chat_template:
+            log.warning("Overwriting the chat template with the default chat template.")
+            tokenizer.chat_template = chat_template_mapping[model_family]
+        else:
+            log.warning("Chat template already exists. Skipping overwriting.")
+    return tokenizer

fusion_bench/models/hf_clip.py

@@ -1,4 +1,5 @@
-from typing import Callable, Iterable, List  # noqa: F401
+import logging
+from typing import TYPE_CHECKING, Callable, Iterable, List  # noqa: F401
 
 import torch
 from torch import Tensor, nn
@@ -7,6 +8,11 @@ from transformers.models.clip.modeling_clip import BaseModelOutputWithPooling
 
 from fusion_bench.utils.devices import get_device
 
+if TYPE_CHECKING:
+    from fusion_bench.models.surgery.surgerymodelwrapper import SurgeryModelWrapper
+
+log = logging.getLogger(__name__)
+
 default_templates = [
     lambda c: f"a photo of a {c}",
 ]
@@ -33,6 +39,7 @@ class HFCLIPClassifier(nn.Module):
         self,
         clip_model: CLIPModel,
         processor: CLIPProcessor,
+        extra_module=None,
     ):
         """
         Initialize the HFCLIPClassifier.
@@ -56,6 +63,8 @@ class HFCLIPClassifier(nn.Module):
             persistent=False,
         )
 
+        self.extra_module = extra_module
+
     @property
     def text_model(self):
         """Get the text model component of CLIP."""
@@ -111,7 +120,13 @@ class HFCLIPClassifier(nn.Module):
 
         self.zeroshot_weights = zeroshot_weights
 
-    def forward(self, images, return_image_embeds=False, return_dict=False):
+    def forward(
+        self,
+        images: Tensor,
+        return_image_embeds=False,
+        return_dict=False,
+        task_name=None,
+    ):
         """
         Perform forward pass for zero-shot image classification.
 
@@ -120,6 +135,9 @@ class HFCLIPClassifier(nn.Module):
 
         Args:
             images (Tensor): Input images to classify.
+            return_image_embeds (bool): Whether to return the image embeddings.
+            return_dict (bool): Whether to return a dictionary with logits and image embeddings.
+            task_name (Optional[str]): The name of the task.
 
         Returns:
             Tensor: Classification logits for each input image.
@@ -131,16 +149,22 @@ class HFCLIPClassifier(nn.Module):
             raise ValueError("Must set classification task before forward pass")
         text_embeds = self.zeroshot_weights
 
-        image_embeds = self.vision_model(images)
-        if isinstance(image_embeds, Tensor):
-            pass
-        elif isinstance(image_embeds, BaseModelOutputWithPooling):
-            image_embeds = image_embeds[1]
-        image_embeds = self.clip_model.visual_projection(image_embeds)
-
+        image_embeds = self.get_image_features(images)
         # normalize embeddings
         image_embeds = image_embeds / image_embeds.norm(p=2, dim=-1, keepdim=True)
 
+        if (
+            hasattr(self.vision_model, "is_surgery_model")
+            and self.vision_model.is_surgery_model
+        ):
+            # Dealing with the surgery model, for more details, please refer to:
+            # (ICML 2024) Yang, et.al. Representation Surgery for Multi-Task Model Merging
+            # https://arxiv.org/abs/2402.02705
+            self.vision_model: "SurgeryModelWrapper" = self.vision_model
+            image_embeds, _, _ = self.vision_model.compute_surgery_features(
+                image_embeds, dataset_name=task_name
+            )
+
         # cosine similarity
         logit_scale = self.clip_model.logit_scale.exp()
         logits_per_text = torch.matmul(text_embeds, image_embeds.t()) * logit_scale
@@ -156,3 +180,20 @@ class HFCLIPClassifier(nn.Module):
                 return logits_per_image, image_embeds
             else:
                 return logits_per_image
+
+    def get_image_features(self, images: Tensor) -> Tensor:
+        """
+        Compute the image embeddings.
+
+        Returns:
+            image_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The image embeddings obtained by
+            applying the projection layer to the pooled output of [`CLIPVisionModel`].
+        """
+
+        image_embeds = self.vision_model(images)
+        if isinstance(image_embeds, Tensor):
+            pass
+        elif isinstance(image_embeds, BaseModelOutputWithPooling):
+            image_embeds = image_embeds[1]
+        image_embeds = self.clip_model.visual_projection(image_embeds)
+        return image_embeds

fusion_bench/models/surgery/__init__.py

@@ -0,0 +1 @@
+from .surgerymodelwrapper import SurgeryModelWrapper

fusion_bench/models/surgery/surgerymodelwrapper.py

@@ -0,0 +1,158 @@
+import math
+from typing import TYPE_CHECKING, Callable, Generic, List, Union
+
+import torch
+from torch import nn
+from transformers.models.clip.modeling_clip import (
+    CLIPVisionModel,
+    CLIPVisionTransformer,
+)
+
+from fusion_bench.utils.type import TorchModelType
+
+
+def regularize_name(name: str):
+    name = name.replace("-", "_")
+    name = name.replace(".", "_")
+    return name
+
+
+class SurgeryModelWrapper(torch.nn.Module, Generic[TorchModelType]):
+
+    is_surgery_model = True
+    """A flag to indicate that this is a surgery model."""
+
+    def __init__(
+        self,
+        model: TorchModelType,
+        test_datasets: List[str],
+        projection_dim: int = 512,
+        hidden_dim: int = 16,
+    ):
+        super(SurgeryModelWrapper, self).__init__()
+        self.model = model
+        self.model.requires_grad_(False)
+
+        self.test_datasets = test_datasets
+        self.non_linear_func = torch.nn.ReLU()
+
+        self.projection_dim = projection_dim
+        self.hidden_dim = hidden_dim
+
+        for dataset_name in test_datasets:
+            self.add_surgery_module(dataset_name)
+
+    def add_surgery_module(self, dataset_name: str):
+        """
+        Add a surgery module for a given dataset.
+
+        Args:
+            dataset_name (str): The name of the dataset.
+        """
+        dataset_name = regularize_name(dataset_name)
+
+        down_proj = torch.nn.Linear(self.projection_dim, self.hidden_dim, bias=False)
+        up_proj = torch.nn.Linear(self.hidden_dim, self.projection_dim, bias=False)
+
+        torch.nn.init.kaiming_uniform_(down_proj.weight, a=math.sqrt(5))
+        torch.nn.init.zeros_(up_proj.weight)
+
+        self.add_module(
+            "feature_mapping_to_head_down_proj_{}".format(dataset_name), down_proj
+        )
+        self.add_module(
+            "feature_mapping_to_head_up_proj_{}".format(dataset_name), up_proj
+        )
+
+    def collect_trainable_params(self):
+        trainable_params = []
+
+        # surgery parameter
+        for dataset_name in self.test_datasets:
+            dataset_name = regularize_name(dataset_name)
+            down_proj = getattr(
+                self, "feature_mapping_to_head_down_proj_{}".format(dataset_name)
+            )
+            up_proj = getattr(
+                self, "feature_mapping_to_head_up_proj_{}".format(dataset_name)
+            )
+            trainable_params.append(down_proj.weight)
+            trainable_params.append(up_proj.weight)
+        return trainable_params
+
+    def collect_surgery_module(self):
+        surgery_module = {}
+
+        # surgery parameter
+        for dataset_name in self.test_datasets:
+            dataset_name = regularize_name(dataset_name)
+            down_proj = getattr(
+                self, "feature_mapping_to_head_down_proj_{}".format(dataset_name)
+            )
+            up_proj = getattr(
+                self, "feature_mapping_to_head_up_proj_{}".format(dataset_name)
+            )
+            surgery_module[
+                "feature_mapping_to_head_down_proj_{}".format(dataset_name)
+            ] = down_proj
+            surgery_module[
+                "feature_mapping_to_head_up_proj_{}".format(dataset_name)
+            ] = up_proj
+
+        surgery_module["non_linear_func"] = self.non_linear_func
+
+        return surgery_module
+
+    def compute_surgery_features(
+        self,
+        compute_features_fn: Union[
+            torch.Tensor, Callable[[TorchModelType], torch.Tensor]
+        ],
+        dataset_name: str,
+    ):
+        """
+        Compute the surgery features.
+
+        Args:
+            compute_features_fn (Union[torch.Tensor, Callable[[nn.Module], torch.Tensor]]): A function that computes the features or a tensor that represents the features.
+            dataset_name (str): The name of the dataset.
+
+        Returns:
+            feature (torch.Tensor): The surgery features.
+            feature0 (torch.Tensor): The original features.
+            feature_sub (torch.Tensor): feature0 - feature.
+        """
+        dataset_name = regularize_name(dataset_name)
+
+        if isinstance(compute_features_fn, torch.Tensor):
+            feature = compute_features_fn
+        elif callable(compute_features_fn):
+            feature = compute_features_fn(self.model)
+        else:
+            raise ValueError(
+                "compute_features_fn must be a tensor or a callable, but got {}".format(
+                    type(compute_features_fn)
+                )
+            )
+
+        feature0 = feature
+
+        # feature bias
+        down_proj = getattr(
+            self, "feature_mapping_to_head_down_proj_{}".format(dataset_name)
+        )
+        up_proj = getattr(
+            self, "feature_mapping_to_head_up_proj_{}".format(dataset_name)
+        )
+        feature_sub = down_proj(feature)
+        feature_sub = self.non_linear_func(feature_sub)
+        feature_sub = up_proj(feature_sub)
+
+        # surgery feature
+        feature = feature0 - feature_sub
+
+        return feature, feature0, feature_sub
+
+    def forward(self, *args, **kwargs):
+        """The wrappered model should just forward like normal."""
+        return self.model(*args, **kwargs)

fusion_bench/models/utils.py

@@ -1,5 +1,6 @@
 from typing import List
 
+import torch
 from torch import nn
 
 
@@ -70,3 +71,10 @@ def find_layers_with_type(
         if isinstance(submodule, tuple(layer_types)):
             res[name] = submodule
     return res
+
+
+def disable_dropout(model: torch.nn.Module):
+    """Disable dropout in a model."""
+    for module in model.modules():
+        if isinstance(module, torch.nn.Dropout):
+            module.p = 0

fusion_bench/models/wrappers/layer_wise_fusion.py

@@ -1,13 +1,22 @@
 import functools
 import logging
 from copy import deepcopy
-from typing import Any, Callable, Dict, Iterator, List, Optional  # noqa: F401
+from typing import (  # noqa: F401
+    Any,
+    Callable,
+    Dict,
+    Generic,
+    Iterator,
+    List,
+    Optional,
+    TypeVar,
+)
 
 import torch
 from torch import Tensor, nn
 from torch.func import functional_call
 
-from fusion_bench.utils.type import StateDictType
+from fusion_bench.utils.type import StateDictType, TorchModelType
 
 __all__ = ["get_layer_wise_weights", "fuse_weights", "LayerWiseMergedModel"]
 
@@ -132,14 +141,14 @@ def fuse_weights(
     }
 
 
-class LayerWiseMergedModel(nn.Module):
+class LayerWiseMergedModel(nn.Module, Generic[TorchModelType]):
     _merged_state_dict: StateDictType = None
 
     def __init__(
         self,
         layer_wise_weight: Tensor,
-        pretrained_model: nn.Module,
-        finetuned_models: List[nn.Module],
+        pretrained_model: TorchModelType,
+        finetuned_models: List[TorchModelType],
         clamp_weights: bool = True,
         tie_weights: bool = False,
         strict: bool = True,

fusion_bench/models/wrappers/task_wise_fusion.py

@@ -16,13 +16,13 @@ outputs = merged_model(inputs)
 
 import functools
 import logging
-from typing import Any, Callable, Dict, Iterator, List, Optional  # noqa: F401
+from typing import Any, Callable, Dict, Generic, Iterator, List, Optional  # noqa: F401
 
 import torch
 from torch import Tensor, nn
 from torch.func import functional_call
 
-from fusion_bench.utils.type import StateDictType
+from fusion_bench.utils.type import StateDictType, TorchModelType
 
 log = logging.getLogger(__name__)
 
@@ -157,14 +157,14 @@ def fuse_weights(
     }
 
 
-class TaskWiseMergedModel(nn.Module):
+class TaskWiseMergedModel(nn.Module, Generic[TorchModelType]):
     _merged_state_dict: StateDictType = None
 
     def __init__(
         self,
         task_wise_weight: Tensor,
-        pretrained_model: nn.Module,
-        finetuned_models: List[nn.Module],
+        pretrained_model: TorchModelType,
+        finetuned_models: List[TorchModelType],
         clamp_weights: bool = True,
         tie_weights: bool = False,
         strict: bool = True,

fusion_bench/optim/__init__.py

@@ -0,0 +1,2 @@
+from . import exception, lr_scheduler
+from .mezo import MeZO