fusion-bench 0.2.9__py3-none-any.whl

fusion_bench/modelpool/seq2seq_lm/modelpool.py

@@ -0,0 +1,65 @@
+import logging
+from copy import deepcopy
+from typing import Optional
+
+from omegaconf import DictConfig, flag_override
+from peft import PeftModel
+from transformers import AutoModelForSeq2SeqLM
+
+from fusion_bench.modelpool import BaseModelPool
+from fusion_bench.utils import parse_dtype
+
+log = logging.getLogger(__name__)
+
+
+def load_lora_model(
+    base_model_path: str,
+    peft_model_path: str,
+    is_trainable: bool = True,
+    merge_and_unload: bool = True,
+):
+    base_model = AutoModelForSeq2SeqLM.from_pretrained(base_model_path)
+    model = PeftModel.from_pretrained(
+        base_model,
+        peft_model_path,
+        is_trainable=is_trainable,
+    )
+    if merge_and_unload:
+        model = model.merge_and_unload()
+    return model
+
+
+class Seq2SeqLMPool(BaseModelPool):
+    _config_mapping = BaseModelPool._config_mapping | {
+        "_tokenizer": "tokenizer",
+        "_model_kwargs": "model_kwargs",
+    }
+
+    def __init__(
+        self,
+        models: DictConfig,
+        *,
+        tokenizer: Optional[DictConfig],
+        model_kwargs: Optional[DictConfig] = None,
+        **kwargs,
+    ):
+        super().__init__(models, **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
+                )
+
+    def load_model(self, model_name_or_config: str | DictConfig, *args, **kwargs):
+        model_kwargs = deepcopy(self._model_kwargs)
+        model_kwargs.update(kwargs)
+        return super().load_model(model_name_or_config, *args, **model_kwargs)
+
+    def load_tokenizer(self, *args, **kwargs):
+        assert self._tokenizer is not None, "Tokenizer is not defined in the config"
+        tokenizer = isinstance(self._tokenizer, *args, **kwargs)
+        return tokenizer

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/__init__.py

@@ -0,0 +1,3 @@
+# flake8: noqa F401
+from . import separate_io, utils
+from .parameter_dict import ParameterDictModel

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

@@ -0,0 +1,199 @@
+import logging
+from typing import TYPE_CHECKING, Callable, Iterable, List  # noqa: F401
+
+import torch
+from torch import Tensor, nn
+from transformers import CLIPModel, CLIPProcessor
+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}",
+]
+
+
+class HFCLIPClassifier(nn.Module):
+    """
+    A classifier based on the CLIP (Contrastive Language-Image Pre-training) model.
+
+    This class wraps a CLIP model and provides functionality for image classification
+    using zero-shot learning. It allows setting a classification task with custom
+    class names and text templates.
+
+    Attributes:
+        clip_model (CLIPModel): The underlying CLIP model.
+        processor (CLIPProcessor): The CLIP processor for preparing inputs.
+        zeroshot_weights (Tensor): Computed text embeddings for zero-shot classification.
+        classnames (List[str]): List of class names for the current classification task.
+        templates (List[Callable[[str], str]]): List of template functions for generating text prompts.
+
+    """
+
+    def __init__(
+        self,
+        clip_model: CLIPModel,
+        processor: CLIPProcessor,
+        extra_module=None,
+    ):
+        """
+        Initialize the HFCLIPClassifier.
+
+        Args:
+            clip_model (CLIPModel): The CLIP model to use for classification.
+            processor (CLIPProcessor): The CLIP processor for preparing inputs.
+        """
+        super().__init__()
+        # we only fine-tune the vision model
+        clip_model.visual_projection.requires_grad_(False)
+        clip_model.text_model.requires_grad_(False)
+        clip_model.text_projection.requires_grad_(False)
+        clip_model.logit_scale.requires_grad_(False)
+
+        self.clip_model = clip_model
+        self.processor = processor
+        self.register_buffer(
+            "zeroshot_weights",
+            None,
+            persistent=False,
+        )
+
+        self.extra_module = extra_module
+
+    @property
+    def text_model(self):
+        """Get the text model component of CLIP."""
+        return self.clip_model.text_model
+
+    @property
+    def vision_model(self):
+        """Get the vision model component of CLIP."""
+        return self.clip_model.vision_model
+
+    def set_classification_task(
+        self,
+        classnames: List[str],
+        templates: List[Callable[[str], str]] = default_templates,
+    ):
+        """
+        Set up the zero-shot classification task.
+
+        This method computes text embeddings for the given class names using the
+        provided templates. These embeddings are then used for classification.
+
+        Args:
+            classnames (List[str]): List of class names for the classification task.
+            templates (List[Callable[[str], str]], optional): List of template functions
+                for generating text prompts. Defaults to `default_templates`, i.e.
+                ["a photo of a {classname}"].
+        """
+        processor = self.processor
+
+        self.classnames = classnames
+        self.templates = templates
+
+        with torch.no_grad():
+            zeroshot_weights = []
+            for classname in classnames:
+                text = [template(classname) for template in templates]
+                inputs = processor(text=text, return_tensors="pt", padding=True)
+                inputs = {
+                    k: v.to(get_device(self.text_model)) for k, v in inputs.items()
+                }
+                embeddings = self.text_model(**inputs)[1]
+                embeddings = self.clip_model.text_projection(embeddings)
+
+                # normalize embeddings
+                embeddings = embeddings / embeddings.norm(p=2, dim=-1, keepdim=True)
+
+                embeddings = embeddings.mean(dim=0)
+                embeddings = embeddings / embeddings.norm(p=2, dim=-1, keepdim=True)
+
+                zeroshot_weights.append(embeddings)
+
+            zeroshot_weights = torch.stack(zeroshot_weights, dim=0)
+
+        self.zeroshot_weights = zeroshot_weights
+
+    def forward(
+        self,
+        images: Tensor,
+        return_image_embeds=False,
+        return_dict=False,
+        task_name=None,
+    ):
+        """
+        Perform forward pass for zero-shot image classification.
+
+        This method computes image embeddings for the input images and calculates
+        the similarity with the pre-computed text embeddings to produce classification logits.
+
+        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.
+
+        Raises:
+            ValueError: If the classification task hasn't been set using set_classification_task.
+        """
+        if self.zeroshot_weights is None:
+            raise ValueError("Must set classification task before forward pass")
+        text_embeds = self.zeroshot_weights
+
+        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
+        logits_per_image = logits_per_text.t()
+
+        if return_dict:
+            ret = {"logits": logits_per_image}
+            if return_image_embeds:
+                ret.update({"image_embeds": image_embeds})
+            return ret
+        else:
+            if return_image_embeds:
+                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/linearized/linearized_model_utils.py

@@ -0,0 +1,91 @@
+import logging
+from collections import OrderedDict
+from copy import deepcopy
+from typing import Optional
+
+import torch.nn as nn
+from torch.func import functional_call, jvp
+
+log = logging.getLogger(__name__)
+
+
+def dict_params_to_tuple(dict_params: dict):
+    return tuple(v for k, v in dict_params.items())
+
+
+class LinearizedModelWraper(nn.Module):
+    def __init__(self, model: nn.Module, init_model: Optional[nn.Module] = None):
+        """
+        Initializes a linearized model.
+
+        Args:
+            model (nn.Module): The underlying PyTorch model to be linearized.
+            init_model (nn.Module): The initial PyTorch model used to compute the linearization parameters (default: None).
+        """
+        super().__init__()
+        self.model = model
+        if init_model is None:
+            init_model = model
+        assert not hasattr(self, "params0")
+        params0 = deepcopy([(k, v.detach()) for k, v in init_model.named_parameters()])
+        self.params0_keys = [k for k, v in params0]
+        self.params0_values = nn.ParameterList([v for k, v in params0])
+        for p in self.params0_values:
+            p.requires_grad_(False)
+
+    def tuple_params_to_dict(self, tuple_params):
+        """
+        Converts a tuple of parameters to a dictionary with keys corresponding to the parameter names.
+
+        Args:
+            tuple_params (Tuple[Tensor, ...]): A tuple of parameters.
+
+        Returns:
+            Dict[str, Tensor]: A dictionary with keys corresponding to the parameter names and values corresponding to the
+            parameter values.
+        """
+        assert len(tuple_params) == len(self.params0_keys)
+        state_dict = {}
+        for k, p in zip(self.params0_keys, tuple_params):
+            state_dict[k] = p
+        return state_dict
+
+    def forward(self, *args, **kwargs):
+        """
+        Computes the linearized model output using a first-order Taylor decomposition.
+
+        Args:
+            *args: Positional arguments to be passed to the model.
+            **kwargs: Keyword arguments to be passed to the model.
+
+        Returns:
+            torch.Tensor: The output of the linearized model, computed using a first-order Taylor decomposition.
+        """
+        params0 = tuple(self.params0_values)
+        params = dict_params_to_tuple(OrderedDict(self.model.named_parameters()))
+        dparams = tuple(p - p0 for p, p0 in zip(params, params0))
+        out, dp = jvp(
+            lambda *param: functional_call(
+                self.model, self.tuple_params_to_dict(param), args, kwargs
+            ),
+            params0,
+            dparams,
+        )
+        return out + dp
+
+    @staticmethod
+    def unload_linearized_modules_(module: nn.Module):
+        """
+        Unloads the linearized module and returns the original module.
+
+        Args:
+            module (nn.Module): The linearized module to be unloaded.
+
+        Returns:
+            nn.Module: The original module.
+        """
+        for name, model in module.named_children():
+            if isinstance(model, LinearizedModelWraper):
+                setattr(module, name, model.model)
+            else:
+                LinearizedModelWraper.unload_linearized_modules_(model)

fusion_bench/models/linearized/vision_model.py

@@ -0,0 +1,122 @@
+import os
+from typing import Tuple, Union
+
+from huggingface_hub import hf_hub_download
+from peft import LoraConfig, PeftModel, get_peft_model
+from peft.tuners.lora import LoraLayer
+from safetensors.torch import load_file
+from transformers import CLIPVisionModel
+from transformers.models.clip.modeling_clip import CLIPVisionTransformer
+
+from .linearized_model_utils import LinearizedModelWraper
+
+
+def get_file_path(peft_name, filename):
+    if os.path.isdir(peft_name):
+        # If peft_name is a local directory path
+        return os.path.join(peft_name, filename)
+    else:
+        # If peft_name is a Hugging Face model name
+        return hf_hub_download(peft_name, filename)
+
+
+def _get_submodules(model, key) -> Tuple:
+    """
+    Retrieves the parent module, target module, and target module name for a given key in a PyTorch model.
+    """
+    parent = model.get_submodule(".".join(key.split(".")[:-1]))
+    target_name = key.split(".")[-1]
+    target = model.get_submodule(key)
+    return parent, target, target_name
+
+
+def linearize_lora_model_(model):
+    """
+    Linearizes the LoraLayer modules in a PyTorch model according to the PETA paper.
+    """
+    for key, module in model.named_modules():
+        # if isinstance(module, LoraLayer) and isinstance(module, nn.Linear):
+        if isinstance(module, LoraLayer):
+            # print("L-LoRA MODULE : ", module)
+            parent, target, target_name = _get_submodules(model, key)
+            setattr(parent, target_name, LinearizedModelWraper(target))
+            # print("Linearized Lora Layer")
+    return model
+
+
+def load_fft_vision_model_hf(
+    model_name: str, return_vison_model=True
+) -> Union[CLIPVisionTransformer, CLIPVisionModel]:
+    """
+    Load a CLIP vision model from Hugging Face.
+
+    Args:
+        model_name (str): The name of the CLIP vision model to load from Hugging Face.
+        return_vison_model (bool, optional): If False, the full CLIPVisionModel is returned. If True, only the vision model (`CLIPVisionTransformer`) is returned. Defaults to True.
+
+    Returns:
+        Union[CLIPVisionTransformer, CLIPVisionModel]: The vision model.
+    """
+    model = CLIPVisionModel.from_pretrained(model_name)
+
+    if return_vison_model:
+        return CLIPVisionModel.from_pretrained(model_name).vision_model
+    else:
+        return model
+
+
+def load_lora_vision_model_hf(
+    base_model_name: str,
+    peft_name: str,
+    merge_and_unload: bool = False,
+    return_vison_model=True,
+):
+    """
+    Load a LoRA (Low-Rank Adaptation) vision model from Hugging Face.
+
+    This function loads a vision model and applies a LoRA adaptation to it. The model can be optionally merged and unloaded.
+
+    Parameters:
+        base_model_name (str): The name of the base vision model to load from Hugging Face.
+        peft_name (str): The name of the LoRA adaptation to apply to the base model.
+        merge_and_unload (bool, optional): If True, the LoRA adaptation is merged into the base model and the LoRA layers are removed. Defaults to False.
+        return_vison_model (bool, optional): If False, the full CLIPVisionModel is returned. If True, only the vision model (`CLIPVisionTransformer`) is returned. Defaults to True.
+
+    Returns:
+        PeftModel: The adapted vision model, optionally merged and unloaded.
+    """
+    model = CLIPVisionModel.from_pretrained(base_model_name)
+
+    # Load the Peft model
+    # note that we apply lora on type `CLIPVisionTransformer` instead of `CLIPVisionModel`
+    vision_model = model.vision_model
+    peft_model = PeftModel.from_pretrained(vision_model, peft_name, is_trainable=True)
+    if merge_and_unload:
+        vision_model = peft_model.merge_and_unload()
+    else:
+        vision_model = peft_model
+
+    # Return the vision model
+    if return_vison_model:
+        return vision_model
+    else:
+        model.vision_model = vision_model
+        return model
+
+
+def load_l_lora_vision_model_hf(base_model_name: str, peft_name: str):
+    """
+    Load a linearized L-LoRA model from a base model and a Peft model (HuggingFace).
+    """
+    base_model = CLIPVisionModel.from_pretrained(base_model_name).vision_model
+    peft_config = LoraConfig.from_pretrained(peft_name)
+    peft_config.inference_mode = False  # This is important, make the model trainable
+    model = get_peft_model(base_model, peft_config)
+    linearize_lora_model_(model)
+    for filename in ["linearized_adapter_model.safetensors"]:
+        path = get_file_path(peft_name, filename)
+        state_dict = load_file(path)
+        for name, param in state_dict.items():
+            model.get_parameter(name).data = param
+
+    return model

fusion_bench/models/llama/__init__.py

@@ -0,0 +1,16 @@
+# Copyright 2024 the LlamaFactory team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .model_utils.misc import find_all_linear_modules
+from .tokenizer_loader import load_config, load_tokenizer