fusion-bench 0.2.9__py3-none-any.whl

fusion_bench/__init__.py

@@ -0,0 +1,20 @@
+# flake8: noqa: F401
+from . import (
+    constants,
+    dataset,
+    method,
+    metrics,
+    mixins,
+    modelpool,
+    models,
+    optim,
+    programs,
+    taskpool,
+    tasks,
+    utils,
+)
+from .method import BaseAlgorithm, BaseModelFusionAlgorithm
+from .modelpool import BaseModelPool
+from .models import separate_io
+from .taskpool import BaseTaskPool
+from .utils import parse_dtype, print_parameters, timeit_context

fusion_bench/__main__.py

@@ -0,0 +1,4 @@
+from fusion_bench.scripts.cli import main
+
+if __name__ == "__main__":
+    main()

fusion_bench/compat/method/__init__.py

@@ -0,0 +1,109 @@
+import warnings
+
+from omegaconf import DictConfig
+
+from .base_algorithm import ModelFusionAlgorithm
+
+
+class AlgorithmFactory:
+    """
+    Factory class to create and manage different model fusion algorithms.
+
+    This class provides methods to create algorithms based on a given configuration,
+    register new algorithms, and list available algorithms.
+    """
+
+    _aglorithms = {
+        # single task learning (fine-tuning)
+        "clip_finetune": ".classification.clip_finetune.ImageClassificationFineTuningForCLIP",
+        # analysis
+        # model merging methods
+        "clip_task_wise_adamerging": ".adamerging.clip_task_wise_adamerging.CLIPTaskWiseAdaMergingAlgorithm",
+        "clip_layer_wise_adamerging": ".adamerging.clip_layer_wise_adamerging.CLIPLayerWiseAdaMergingAlgorithm",
+        "singular_projection_merging": "fusion_bench.method.smile_upscaling.singular_projection_merging.SingularProjectionMergingAlgorithm",
+        "clip_layer_wise_adamerging_surgery": ".surgery.clip_layer_wise_adamerging_surgery.CLIPLayerWiseAdaMergingSurgeryAlgorithm",
+        # plug-and-play model merging methods
+        "clip_concrete_task_arithmetic": ".concrete_subspace.clip_concrete_task_arithmetic.ConcreteTaskArithmeticAlgorithmForCLIP",
+        "clip_concrete_task_wise_adamerging": ".concrete_subspace.clip_concrete_adamerging.ConcreteTaskWiseAdaMergingForCLIP",
+        "clip_concrete_layer_wise_adamerging": ".concrete_subspace.clip_concrete_adamerging.ConcreteLayerWiseAdaMergingForCLIP",
+        # model mixing methods
+        "clip_weight_ensembling_moe": ".we_moe.clip_we_moe.CLIPWeightEnsemblingMoEAlgorithm",
+        "sparse_clip_weight_ensembling_moe": "fusion_bench.method.SparseCLIPWeightEnsemblingMoEAlgorithm",
+        "smile_mistral_upscaling": ".smile_upscaling.smile_mistral_upscaling.SmileMistralUpscalingAlgorithm",
+        "rankone_moe": ".rankone_moe.clip_rankone_moe.CLIPRankOneMoEAlgorithm",
+    }
+
+    @staticmethod
+    def create_algorithm(method_config: DictConfig) -> ModelFusionAlgorithm:
+        """
+        Create an instance of a model fusion algorithm based on the provided configuration.
+
+        Args:
+            method_config (DictConfig): The configuration for the algorithm. Must contain a 'name' attribute that specifies the type of the algorithm.
+
+        Returns:
+            ModelFusionAlgorithm: An instance of the specified algorithm.
+
+        Raises:
+            ValueError: If 'name' attribute is not found in the configuration or does not match any known algorithm names.
+        """
+        warnings.warn(
+            "AlgorithmFactory.create_algorithm() is deprecated and will be removed in future versions. "
+            "Please implement new model fusion algorithm using `fusion_bench.method.BaseModelFusionAlgorithm` instead.",
+            DeprecationWarning,
+        )
+
+        from fusion_bench.utils import import_object
+
+        algorithm_name = method_config.name
+        if algorithm_name not in AlgorithmFactory._aglorithms:
+            raise ValueError(
+                f"Unknown algorithm: {algorithm_name}, available algorithms: {AlgorithmFactory._aglorithms.keys()}."
+                "You can register a new algorithm using `AlgorithmFactory.register_algorithm()` method."
+            )
+        algorithm_cls = AlgorithmFactory._aglorithms[algorithm_name]
+        if isinstance(algorithm_cls, str):
+            if algorithm_cls.startswith("."):
+                algorithm_cls = f"fusion_bench.method.{algorithm_cls[1:]}"
+            algorithm_cls = import_object(algorithm_cls)
+        return algorithm_cls(method_config)
+
+    @staticmethod
+    def register_algorithm(name: str, algorithm_cls):
+        """
+        Register a new algorithm with the factory.
+
+        Args:
+            name (str): The name of the algorithm.
+            algorithm_cls: The class of the algorithm to register.
+        """
+        AlgorithmFactory._aglorithms[name] = algorithm_cls
+
+    @classmethod
+    def available_algorithms(cls):
+        """
+        Get a list of available algorithms.
+
+        Returns:
+            list: A list of available algorithm names.
+        """
+        return list(cls._aglorithms.keys())
+
+
+def load_algorithm_from_config(method_config: DictConfig):
+    """
+    Loads an algorithm based on the provided configuration.
+
+    The function checks the 'name' attribute of the configuration and returns an instance of the corresponding algorithm.
+    If the 'name' attribute is not found or does not match any known algorithm names, a ValueError is raised.
+
+    Args:
+        method_config (DictConfig): The configuration for the algorithm. Must contain a 'name' attribute that specifies the type of the algorithm.
+
+    Returns:
+        An instance of the specified algorithm.
+
+    Raises:
+        ValueError: If 'name' attribute is not found in the configuration or does not match any known algorithm names.
+    """
+    return AlgorithmFactory.create_algorithm(method_config)

fusion_bench/compat/method/base_algorithm.py

@@ -0,0 +1,58 @@
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Optional
+
+from omegaconf import DictConfig
+
+if TYPE_CHECKING:
+    from fusion_bench.programs.base_program import BaseHydraProgram
+
+__all__ = ["ModelFusionAlgorithm"]
+
+
+class ModelFusionAlgorithm(ABC):
+    """
+    Abstract base class for model fusion algorithms (for v0.1.x versions, deprecated).
+    For implementing new method, use `fusion_bench.method.BaseModelFusionAlgorithm` instead.
+
+    This class provides a template for implementing model fusion algorithms.
+    Subclasses must implement the `run` method to define the fusion logic.
+
+    Attributes:
+        config (DictConfig): Configuration for the algorithm.
+    """
+
+    _program: "BaseHydraProgram" = None
+    """A reference to the program that is running the algorithm."""
+
+    def __init__(self, algorithm_config: Optional[DictConfig] = None):
+        """
+        Initialize the model fusion algorithm with the given configuration.
+
+        Args:
+            algorithm_config (Optional[DictConfig]): Configuration for the algorithm. Defaults to an empty configuration if not provided.
+                Get access to the configuration using `self.config`.
+        """
+        if algorithm_config is None:
+            algorithm_config = DictConfig({})
+        self.config = algorithm_config
+
+    @abstractmethod
+    def run(self, modelpool):
+        """
+        Fuse the models in the given model pool.
+
+        This method must be implemented by subclasses to define the fusion logic.
+        `modelpool` is an object responsible for managing the models to be fused and optional datasets to be used for fusion.
+
+        Args:
+            modelpool: The pool of models to fuse.
+
+        Returns:
+            The fused model.
+
+        Examples:
+            >>> algorithm = SimpleAverageAlgorithm()
+            >>> modelpool = ModelPool()
+            >>> merged_model = algorithm.fuse(modelpool)
+        """
+        pass

fusion_bench/compat/modelpool/AutoModelForSeq2SeqLM.py

@@ -0,0 +1,34 @@
+import logging
+
+from omegaconf import DictConfig
+from transformers import AutoModelForSeq2SeqLM
+
+from fusion_bench.utils import timeit_context
+
+from .base_pool import ModelPool
+
+log = logging.getLogger(__name__)
+
+
+class AutoModelForSeq2SeqLMPool(ModelPool):
+    def load_model(self, model_config: str | DictConfig):
+        """
+        Load a Seq2Seq language model based on the provided configuration.
+
+        The configuration options are:
+
+        - name: The name of the model.
+        - path: The path where the model is stored.
+
+        Args:
+            model_config (str | DictConfig): The configuration for the model. This can be either a string (name of the model) or a DictConfig object containing the model configuration.
+
+        Returns:
+            model: The loaded Seq2Seq language model.
+        """
+
+        if isinstance(model_config, str):
+            model_config = self.get_model_config(model_config)
+        with timeit_context(f"Loading model {model_config['name']}"):
+            model = AutoModelForSeq2SeqLM.from_pretrained(model_config["path"])
+            return model

fusion_bench/compat/modelpool/__init__.py

@@ -0,0 +1,116 @@
+# flake8: noqa F401
+import warnings
+
+from omegaconf import DictConfig
+
+from fusion_bench.modelpool.huggingface_gpt2_classification import (
+    GPT2ForSequenceClassificationPool,
+)
+from fusion_bench.modelpool.PeftModelForSeq2SeqLM import PeftModelForSeq2SeqLMPool
+
+from .AutoModelForSeq2SeqLM import AutoModelForSeq2SeqLMPool
+from .base_pool import DictModelPool, ListModelPool, ModelPool, to_modelpool
+from .huggingface_clip_vision import HuggingFaceClipVisionPool
+
+
+class ModelPoolFactory:
+    """
+    Factory class to create and manage different model pools.
+
+    This class provides methods to create model pools based on a given configuration,
+    register new model pools, and list available model pools.
+    """
+
+    _modelpool = {
+        "NYUv2ModelPool": "fusion_bench.modelpool.nyuv2_modelpool.NYUv2ModelPool",
+        "huggingface_clip_vision": HuggingFaceClipVisionPool,
+        "HF_GPT2ForSequenceClassification": GPT2ForSequenceClassificationPool,
+        "AutoModelPool": ".huggingface_automodel.AutoModelPool",
+        # CausualLM
+        "AutoModelForCausalLMPool": ".huggingface_llm.AutoModelForCausalLMPool",
+        "LLamaForCausalLMPool": ".huggingface_llm.LLamaForCausalLMPool",
+        "MistralForCausalLMPool": ".huggingface_llm.MistralForCausalLMPool",
+        # Seq2SeqLM
+        "AutoModelForSeq2SeqLMPool": AutoModelForSeq2SeqLMPool,
+        "PeftModelForSeq2SeqLMPool": PeftModelForSeq2SeqLMPool,
+    }
+
+    @staticmethod
+    def create_modelpool(modelpool_config: DictConfig) -> ModelPool:
+        """
+        Create an instance of a model pool based on the provided configuration.
+        This is for v0.1.x versions, deprecated.
+        For implementing new model pool, use `fusion_bench.modelpool.BaseModelPool` instead.
+
+        Args:
+            modelpool_config (DictConfig): The configuration for the model pool.
+            Must contain a 'type' attribute that specifies the type of the model pool.
+
+        Returns:
+            ModelPool: An instance of the specified model pool.
+
+        Raises:
+            ValueError: If 'type' attribute is not found in the configuration or does not match any known model pool types.
+        """
+        warnings.warn(
+            "ModelPoolFactory.create_modelpool() is deprecated and will be removed in future versions. "
+            "Please implement new model pool using `fusion_bench.modelpool.BaseModelPool` instead.",
+            DeprecationWarning,
+        )
+
+        from fusion_bench.utils import import_object
+
+        modelpool_type = modelpool_config.get("type")
+        if modelpool_type is None:
+            raise ValueError("Model pool type not specified")
+
+        if modelpool_type not in ModelPoolFactory._modelpool:
+            raise ValueError(
+                f"Unknown model pool: {modelpool_type}, available model pools: {ModelPoolFactory._modelpool.keys()}. You can register a new model pool using `ModelPoolFactory.register_modelpool()` method."
+            )
+        modelpool_cls = ModelPoolFactory._modelpool[modelpool_type]
+        if isinstance(modelpool_cls, str):
+            if modelpool_cls.startswith("."):
+                modelpool_cls = f"fusion_bench.compat.modelpool.{modelpool_cls[1:]}"
+            modelpool_cls = import_object(modelpool_cls)
+        return modelpool_cls(modelpool_config)
+
+    @staticmethod
+    def register_modelpool(name: str, modelpool_cls):
+        """
+        Register a new model pool with the factory.
+
+        Args:
+            name (str): The name of the model pool.
+            modelpool_cls: The class of the model pool to register.
+        """
+        ModelPoolFactory._modelpool[name] = modelpool_cls
+
+    @classmethod
+    def available_modelpools(cls):
+        """
+        Get a list of available model pools.
+
+        Returns:
+            list: A list of available model pool names.
+        """
+        return list(cls._modelpool.keys())
+
+
+def load_modelpool_from_config(modelpool_config: DictConfig):
+    """
+    Loads a model pool based on the provided configuration.
+
+    The function checks the 'type' attribute of the configuration and returns an instance of the corresponding model pool.
+    If the 'type' attribute is not found or does not match any known model pool types, a ValueError is raised.
+
+    Args:
+        modelpool_config (DictConfig): The configuration for the model pool. Must contain a 'type' attribute that specifies the type of the model pool.
+
+    Returns:
+        An instance of the specified model pool.
+
+    Raises:
+        ValueError: If 'type' attribute is not found in the configuration or does not match any known model pool types.
+    """
+    return ModelPoolFactory.create_modelpool(modelpool_config)

fusion_bench/compat/modelpool/base_pool.py

@@ -0,0 +1,328 @@
+import logging
+from abc import ABC
+from copy import deepcopy
+from typing import Dict, List, Optional, Union
+
+import torch
+from omegaconf import DictConfig
+from torch import nn
+
+from fusion_bench.modelpool import BaseModelPool
+from fusion_bench.utils import timeit_context
+
+__all__ = ["ModelPool", "DictModelPool", "ListModelPool", "to_modelpool"]
+
+log = logging.getLogger(__name__)
+
+
+class ModelPool(ABC):
+    """
+    This is the base class for all modelpools.
+    For verison v0.1.x, deprecated.
+    Please implemente new algorithms use `fusion_bench.modelpool.BaseModelPool`.
+    """
+
+    _model_names = None
+
+    def __init__(self, modelpool_config: Optional[DictConfig] = None):
+        """
+        Initialize the ModelPool with the given configuration.
+
+        Args:
+            modelpool_config (Optional[DictConfig]): The configuration for the model pool.
+        """
+        super().__init__()
+        self.config = modelpool_config
+
+        # check for duplicate model names
+        if self.config is not None and self.config.get("models", None) is not None:
+            model_names = [model["name"] for model in self.config["models"]]
+            assert len(model_names) == len(
+                set(model_names)
+            ), "Duplicate model names found in model pool"
+            self._model_names = model_names
+
+    def __len__(self):
+        """
+        Return the number of models in the model pool, exclude special models such as `_pretrained_`.
+
+        Returns:
+            int: The number of models in the model pool.
+        """
+        return len(self.model_names)
+
+    @property
+    def model_names(self) -> List[str]:
+        """
+        This property returns a list of model names from the configuration, excluding any names that start or end with an underscore.
+        To obtain all model names, including those starting or ending with an underscore, use the `_model_names` attribute.
+
+        Returns:
+            list: A list of model names.
+        """
+        names = [
+            name for name in self._model_names if name[0] != "_" and name[-1] != "_"
+        ]
+        return names
+
+    @property
+    def has_pretrained(self):
+        """
+        Check if the pretrained model is available in the model pool.
+
+        Returns:
+            bool: True if the pretrained model is available, False otherwise.
+        """
+        for model_config in self.config["models"]:
+            if model_config.get("name", None) == "_pretrained_":
+                return True
+        return False
+
+    def get_model_config(self, model_name: str):
+        """
+        Retrieves the configuration for a specific model from the model pool.
+
+        Args:
+            model_name (str): The name of the model for which to retrieve the configuration.
+
+        Returns:
+            dict: The configuration dictionary for the specified model.
+
+        Raises:
+            ValueError: If the specified model is not found in the model pool.
+        """
+        for model in self.config["models"]:
+            if model["name"] == model_name:
+                return model
+        raise ValueError(f"Model {model_name} not found in model pool")
+
+    def load_model(self, model_config: Union[str, DictConfig]) -> nn.Module:
+        """
+        The models are load lazily, so this method should be implemented to load the model from the model pool.
+
+        Load the model from the model pool.
+
+        Args:
+            model_config (Union[str, DictConfig]): The configuration dictionary for the model to load.
+
+        Returns:
+            Any: The loaded model.
+        """
+        raise NotImplementedError
+
+    def load_pretrained_or_first_model(self, *args, **kwargs):
+        """
+        Load the pretrained model if available, otherwise load the first model in the list.
+
+        This method checks if a pretrained model is available. If it is, it loads the pretrained model.
+        If not, it loads the first model from the list of model names.
+
+        Returns:
+            nn.Module: The loaded model.
+        """
+        if self.has_pretrained:
+            model = self.load_model("_pretrained_", *args, **kwargs)
+        else:
+            model = self.load_model(self.model_names[0], *args, **kwargs)
+        return model
+
+    def save_model(self, model: nn.Module, path: str):
+        """
+        Save the state dictionary of the model to the specified path.
+
+        Args:
+            model (nn.Module): The model whose state dictionary is to be saved.
+            path (str): The path where the state dictionary will be saved.
+        """
+        with timeit_context(f"Saving the state dict of model to {path}"):
+            torch.save(model.state_dict(), path)
+
+    def models(self):
+        """
+        Generator that yields models from the model pool.
+
+        Yields:
+            nn.Module: The next model in the model pool.
+        """
+        for model_name in self.model_names:
+            yield self.load_model(model_name)
+
+    def named_models(self):
+        """
+        Generator that yields model names and models from the model pool.
+
+        Yields:
+            tuple: A tuple containing the model name and the model.
+        """
+        for model_name in self.model_names:
+            yield model_name, self.load_model(model_name)
+
+    def get_train_dataset(self, model_name: str):
+        """
+        Get the training dataset for the model.
+
+        Args:
+            model_name (str): The name of the model for which to get the training dataset.
+
+        Returns:
+            Any: The training dataset for the model.
+        """
+        raise NotImplementedError
+
+    def get_test_dataset(self, model_name: str):
+        """
+        Get the testing dataset for the model.
+
+        Args:
+            model_name (str): The name of the model for which to get the testing dataset.
+
+        Returns:
+            Any: The testing dataset for the model.
+        """
+        raise NotImplementedError
+
+    def setup_taskpool(self, taskpool):
+        """
+        Setup the taskpool before evaluation.
+        Such as setting the fabric, processor, tokenizer, etc.
+
+        Args:
+            taskpool (Any): The taskpool to setup.
+        """
+        pass
+
+    def to_modellist(self) -> List[nn.Module]:
+        """
+        Convert the model pool to a list of models.
+
+        Returns:
+            list: A list of models.
+        """
+        return [self.load_model(m) for m in self.model_names]
+
+    def to_modeldict(self) -> Dict[str, nn.Module]:
+        """
+        Convert the model pool to a dictionary of models.
+
+        Returns:
+            dict: A dictionary of models.
+        """
+        return {m: self.load_model(m) for m in self.model_names}
+
+
+class ListModelPool(ModelPool):
+    """
+    ModelPool from a list of models.
+    """
+
+    def __init__(
+        self,
+        models: List[nn.Module],
+        has_pretraned: bool = False,
+    ):
+        """
+        Initialize the ListModelPool with the given list of models.
+
+        Args:
+            models (List[nn.Module]): The list of models.
+            has_pretraned (bool): Whether the first model in the list is pretrained.
+        """
+        modelpool_config = {}
+        modelpool_config["models"] = []
+        model_dict = {}
+        if has_pretraned:
+            model_dict["_pretrained_"] = models[0]
+            modelpool_config["models"].append({"name": "_pretrained_"})
+            for i, model in enumerate(models[1:]):
+                model_dict[f"model_{i}"] = model
+                modelpool_config["models"].append({"name": f"model_{i}"})
+        else:
+            for i, model in enumerate(models):
+                model_dict[f"model_{i}"] = model
+                modelpool_config["models"].append({"name": f"model_{i}"})
+
+        self.model_dict = model_dict
+        super().__init__(DictConfig(modelpool_config))
+
+    def load_model(self, model_config: str | DictConfig, copy=True) -> nn.Module:
+        """
+        Load the model from the model pool.
+
+        Args:
+            model_config (str | DictConfig): The model name or the configuration dictionary for the model to load.
+            copy (bool): Whether to return a copy of the model, defaults to `True`.
+
+        Returns:
+            nn.Module: The loaded model.
+        """
+        if isinstance(model_config, str):
+            model_config = self.get_model_config(model_config)
+        model_name = model_config["name"]
+        model = self.model_dict[model_name]
+        if copy:
+            model = deepcopy(model)
+        return model
+
+
+class DictModelPool(ModelPool):
+    """
+    ModelPool from a dictionary of models.
+    """
+
+    def __init__(self, model_dict: Dict[str, nn.Module]):
+        """
+        Initialize the DictModelPool with the given dictionary of models.
+
+        Args:
+            model_dict (Dict[str, nn.Module]): The dictionary of models.
+        """
+        modelpool_config = {}
+        modelpool_config["models"] = []
+        for model_name, model in model_dict.items():
+            modelpool_config["models"].append({"name": model_name})
+        self.model_dict = model_dict
+        super().__init__(DictConfig(modelpool_config))
+
+    def load_model(self, model_config: str | DictConfig, copy=True) -> nn.Module:
+        """
+        Load the model from the model pool.
+
+        Args:
+            model_config (str | DictConfig): The configuration dictionary for the model to load.
+            copy (bool): Whether to return a copy of the model.
+
+        Returns:
+            nn.Module: The loaded model.
+        """
+        if isinstance(model_config, str):
+            model_config = self.get_model_config(model_config)
+        model_name = model_config["name"]
+        model = self.model_dict[model_name]
+        if copy:
+            model = deepcopy(model)
+        return model
+
+
+def to_modelpool(obj: List[nn.Module], **kwargs):
+    """
+    Convert the given object to a model pool.
+
+    Args:
+        obj (List[nn.Module]): The object to convert to a model pool.
+
+    Returns:
+        ModelPool: The converted model pool.
+
+    Raises:
+        ValueError: If the object cannot be converted to a model pool.
+    """
+    if isinstance(obj, (ModelPool, BaseModelPool)):
+        return obj
+    elif isinstance(obj, (list, tuple)) and all(isinstance(m, nn.Module) for m in obj):
+        return ListModelPool(models=obj, **kwargs)
+    elif isinstance(obj, Dict) and all(isinstance(m, nn.Module) for m in obj.values()):
+        return DictModelPool(model_dict=obj, **kwargs)
+    elif isinstance(obj, nn.Module):
+        return ListModelPool(models=[obj], **kwargs)
+    else:
+        raise ValueError(f"Invalid modelpool object: {obj}")