fusion-bench 0.2.28__py3-none-any.whl → 0.2.30__py3-none-any.whl

fusion_bench/constants/__init__.py

@@ -5,4 +5,8 @@ from .paths import *
 from .runtime import RuntimeConstants
 
 # fusionbench version
-FUSION_BENCH_VERSION = importlib.metadata.version("fusion-bench")
+try:
+    FUSION_BENCH_VERSION = importlib.metadata.version("fusion-bench")
+except importlib.metadata.PackageNotFoundError:
+    # Fallback when package is not installed (e.g., during development)
+    FUSION_BENCH_VERSION = "0.0.0.dev"

fusion_bench/constants/runtime.py

@@ -1,3 +1,29 @@
+"""
+Runtime Constants Module.
+
+This module provides a thread-safe singleton class for managing runtime configuration
+and constants across the Fusion Bench framework. It centralizes access to runtime
+settings like cache directories, debug flags, and logging preferences.
+
+Example:
+    ```python
+    from fusion_bench.constants.runtime import RuntimeConstants
+
+    # Get the singleton instance
+    runtime = RuntimeConstants()
+
+    # Configure cache directory
+    runtime.cache_dir = "/custom/cache/path"
+
+    # Enable debug mode
+    runtime.debug = True
+
+    # Control function call logging
+    runtime.print_function_call = True
+    ```
+"""
+
+import os
 import threading
 from pathlib import Path
 from typing import Optional, Union
@@ -5,18 +31,46 @@ from typing import Optional, Union
 
 class RuntimeConstants:
     """
-    This class holds constants related to the runtime environment of the Fusion Bench framework.
-    It includes default values for cache directories and other runtime configurations.
+    Thread-safe singleton for managing runtime configuration and constants.
+
+    This class provides centralized access to runtime settings that affect the
+    behavior of the entire Fusion Bench framework. It ensures consistent
+    configuration across all modules and supports thread-safe access in
+    multi-threaded environments.
+
+    Attributes:
+        debug: Global debug flag for enabling verbose logging and debugging features.
+
+    Example:
+        ```python
+        runtime = RuntimeConstants()
+
+        # Configure caching
+        runtime.cache_dir = Path.home() / ".cache" / "fusion_bench"
+
+        # Enable debugging
+        runtime.debug = True
+        runtime.print_function_call = True
+        ```
 
-    Implemented as a thread-safe singleton to ensure consistent runtime configuration
-    across the entire application.
+    Note:
+        This class implements the singleton pattern with thread-safe initialization.
+        Multiple calls to the constructor will return the same instance.
     """
 
     _instance: Optional["RuntimeConstants"] = None
     _lock = threading.Lock()
 
     def __new__(cls) -> "RuntimeConstants":
-        """Create a new instance using singleton pattern with thread safety."""
+        """
+        Create or return the singleton instance with thread safety.
+
+        Uses double-check locking pattern to ensure thread-safe singleton creation
+        while minimizing synchronization overhead.
+
+        Returns:
+            The singleton RuntimeConstants instance.
+        """
         with cls._lock:
             # Double-check locking pattern
             if cls._instance is None:
@@ -25,33 +79,83 @@ class RuntimeConstants:
             return cls._instance
 
     def __init__(self):
-        """Initialize the singleton instance only once."""
+        """
+        Initialize the singleton instance only once.
+
+        Subsequent calls to __init__ are no-ops to maintain singleton behavior.
+        """
         if not self._initialized:
-            # Add your runtime constants here
+            # Initialize default values
             self._initialized = True
 
     debug = False
+    """Global debug flag for enabling verbose logging and debugging features."""
 
     @property
     def cache_dir(self) -> Path:
+        """
+        Get the default cache directory for models and datasets.
+
+        Returns:
+            Path object pointing to the cache directory.
+
+        Example:
+            ```python
+            runtime = RuntimeConstants()
+            print(f"Cache directory: {runtime.cache_dir}")
+            ```
+        """
         from fusion_bench.utils.cache_utils import DEFAULT_CACHE_DIR
 
         return DEFAULT_CACHE_DIR
 
     @cache_dir.setter
     def cache_dir(self, path: Union[str, Path]) -> None:
+        """
+        Set the default cache directory for models and datasets.
+
+        Args:
+            path: New cache directory path as string or Path object.
+
+        Example:
+            ```python
+            runtime = RuntimeConstants()
+            runtime.cache_dir = "/data/fusion_bench_cache"
+            ```
+        """
         from fusion_bench.utils.cache_utils import set_default_cache_dir
 
         set_default_cache_dir(path)
 
     @property
     def print_function_call(self) -> bool:
+        """
+        Get whether function calls are printed during instantiation.
+
+        Returns:
+            True if function call printing is enabled, False otherwise.
+        """
         from fusion_bench.utils.instantiate_utils import PRINT_FUNCTION_CALL
 
         return PRINT_FUNCTION_CALL
 
     @print_function_call.setter
     def print_function_call(self, enable: bool) -> None:
+        """
+        Set whether to print function calls during instantiation.
+
+        Useful for debugging to see which functions are being called
+        when instantiating objects from configuration.
+
+        Args:
+            enable: True to enable printing, False to disable.
+
+        Example:
+            ```python
+            runtime = RuntimeConstants()
+            runtime.print_function_call = True  # Enable verbose logging
+            ```
+        """
         from fusion_bench.utils.instantiate_utils import set_print_function_call
 
         set_print_function_call(enable)

fusion_bench/dataset/gsm8k.py

@@ -13,8 +13,12 @@ def load_gsm8k_question_label_data(
 
     An example in the dataset:
 
-    {'question': 'Natalia sold clips to 48 of her friends in April, and then she sold half as many clips in May. How many clips did Natalia sell altogether in April and May?',
-     'answer': 'Natalia sold 48/2 = <<48/2=24>>24 clips in May.\nNatalia sold 48+24 = <<48+24=72>>72 clips altogether in April and May.\n#### 72'}
+    ```python
+    {
+        'question': 'Natalia sold clips to 48 of her friends in April, and then she sold half as many clips in May. How many clips did Natalia sell altogether in April and May?',
+        'answer': 'Natalia sold 48/2 = <<48/2=24>>24 clips in May.\nNatalia sold 48+24 = <<48+24=72>>72 clips altogether in April and May.\n#### 72'
+    }
+    ```
 
     Args:
         dataset_name (Literal["train", "test", "train_socratic", "test_socratic"]): The name of the dataset to load.

fusion_bench/dataset/image_corruption/make_corruption.py

@@ -1,4 +1,31 @@
 # -*- coding: utf-8 -*-
+"""
+Image Corruption Module for Robustness Testing.
+
+This module provides various image corruption functions to test model robustness.
+It implements common corruptions such as noise, blur, compression artifacts, and
+weather effects. These corruptions are commonly used in benchmark datasets like
+ImageNet-C and CIFAR-10-C.
+
+The corruptions can be applied at different severity levels (1-5), where higher
+levels indicate stronger corruption effects.
+
+Example:
+    ```python
+    from PIL import Image
+    from fusion_bench.dataset.image_corruption.make_corruption import gaussian_noise, motion_blur
+
+    # Load an image
+    img = Image.open("example.jpg")
+
+    # Apply gaussian noise at severity level 3
+    corrupted_img = gaussian_noise(img, severity=3)
+
+    # Apply motion blur at severity level 2
+    blurred_img = motion_blur(img, severity=2)
+    ```
+"""
+
 import logging
 
 logger = logging.getLogger(__name__)
@@ -37,11 +64,39 @@ warnings.simplefilter("ignore", UserWarning)
 
 # /////////////// Distortions ///////////////
 class MotionImage(WandImage):
+    """
+    Extended WandImage class with motion blur capability.
+
+    This class wraps ImageMagick's motion blur functionality through the Wand library.
+    """
+
     def motion_blur(self, radius=0.0, sigma=0.0, angle=0.0):
+        """
+        Apply motion blur effect to the image.
+
+        Args:
+            radius: The radius of the Gaussian, in pixels, not counting the center pixel.
+            sigma: The standard deviation of the Gaussian, in pixels.
+            angle: Apply the effect along this angle in degrees.
+        """
         wandlibrary.MagickMotionBlurImage(self.wand, radius, sigma, angle)
 
 
 def gaussian_noise(x, severity=1):
+    """
+    Apply Gaussian noise corruption to an image.
+
+    Adds random Gaussian noise to the image, simulating sensor noise or
+    environmental interference.
+
+    Args:
+        x: Input image as PIL Image or numpy array. If numpy array, should be in
+           range [0, 255].
+        severity: Corruption severity level from 1 (mild) to 5 (severe).
+
+    Returns:
+        numpy.ndarray: Corrupted image as numpy array in range [0, 255].
+    """
     c = [0.04, 0.06, 0.08, 0.09, 0.10][severity - 1]
 
     x = np.array(x) / 255.0
@@ -49,6 +104,20 @@ def gaussian_noise(x, severity=1):
 
 
 def impulse_noise(x, severity=1):
+    """
+    Apply impulse (salt-and-pepper) noise corruption to an image.
+
+    Randomly replaces pixels with either maximum or minimum intensity values,
+    simulating transmission errors or faulty pixels.
+
+    Args:
+        x: Input image as PIL Image or numpy array. If numpy array, should be in
+           range [0, 255].
+        severity: Corruption severity level from 1 (mild) to 5 (severe).
+
+    Returns:
+        numpy.ndarray: Corrupted image as numpy array in range [0, 255].
+    """
     c = [0.01, 0.02, 0.03, 0.05, 0.07][severity - 1]
 
     x = sk.util.random_noise(np.array(x) / 255.0, mode="s&p", amount=c)
@@ -56,6 +125,21 @@ def impulse_noise(x, severity=1):
 
 
 def motion_blur(x, severity=1):
+    """
+    Apply motion blur corruption to an image.
+
+    Simulates camera shake or object motion during image capture by applying
+    directional blur at a random angle.
+
+    Args:
+        x: Input PIL Image.
+        severity: Corruption severity level from 1 (mild) to 5 (severe).
+           Higher severity increases blur radius and sigma.
+
+    Returns:
+        numpy.ndarray: Corrupted image as numpy array in range [0, 255].
+           Returns RGB image regardless of input format.
+    """
     c = [(6, 1), (6, 1.5), (6, 2), (8, 2), (9, 2.5)][severity - 1]
 
     output = BytesIO()
@@ -73,6 +157,21 @@ def motion_blur(x, severity=1):
 
 
 def spatter(x, severity=1):
+    """
+    Apply spatter corruption to an image.
+
+    Simulates liquid splatter effects (water or mud) on the image, creating
+    realistic occlusions similar to raindrops or dirt on a camera lens.
+
+    Args:
+        x: Input image as PIL Image or numpy array. If numpy array, should be in
+           range [0, 255].
+        severity: Corruption severity level from 1 (mild) to 5 (severe).
+           Levels 1-3 simulate water splatter, levels 4-5 simulate mud splatter.
+
+    Returns:
+        numpy.ndarray: Corrupted image as numpy array in range [0, 255].
+    """
     c = [
         (0.62, 0.1, 0.7, 0.7, 0.5, 0),
         (0.65, 0.1, 0.8, 0.7, 0.5, 0),
@@ -140,6 +239,21 @@ def spatter(x, severity=1):
 
 
 def contrast(x, severity=1):
+    """
+    Apply contrast reduction corruption to an image.
+
+    Reduces image contrast by blending pixels toward their mean values,
+    simulating poor lighting conditions or low-quality image sensors.
+
+    Args:
+        x: Input image as PIL Image or numpy array. If numpy array, should be in
+           range [0, 255].
+        severity: Corruption severity level from 1 (mild) to 5 (severe).
+           Higher severity results in lower contrast.
+
+    Returns:
+        numpy.ndarray: Corrupted image as numpy array in range [0, 255].
+    """
     c = [0.75, 0.5, 0.4, 0.3, 0.15][severity - 1]
 
     x = np.array(x) / 255.0
@@ -148,6 +262,20 @@ def contrast(x, severity=1):
 
 
 def jpeg_compression(x, severity=1):
+    """
+    Apply JPEG compression artifacts to an image.
+
+    Simulates compression artifacts from lossy JPEG encoding at various
+    quality levels, commonly seen in heavily compressed images.
+
+    Args:
+        x: Input PIL Image.
+        severity: Corruption severity level from 1 (mild) to 5 (severe).
+           Lower severity uses higher JPEG quality (less compression).
+
+    Returns:
+        PIL.Image: Corrupted image as PIL Image.
+    """
     c = [80, 65, 58, 50, 40][severity - 1]
 
     output = BytesIO()
@@ -158,6 +286,23 @@ def jpeg_compression(x, severity=1):
 
 
 def pixelate(x, severity=1):
+    """
+    Apply pixelation corruption to an image.
+
+    Reduces image resolution by downsampling and then upsampling,
+    creating a blocky, pixelated appearance.
+
+    Args:
+        x: Input PIL Image with size (32, 32).
+        severity: Corruption severity level from 1 (mild) to 5 (severe).
+           Higher severity results in more pixelation.
+
+    Returns:
+        PIL.Image: Corrupted image as PIL Image with original size (32, 32).
+
+    Note:
+        This function is designed for 32x32 images (e.g., CIFAR-10).
+    """
     c = [0.95, 0.9, 0.85, 0.75, 0.65][severity - 1]
 
     x = x.resize((int(32 * c), int(32 * c)), PILImage.BOX)
@@ -170,6 +315,29 @@ def pixelate(x, severity=1):
 
 
 distortion_methods = collections.OrderedDict()
+"""
+OrderedDict mapping corruption names to their corresponding functions.
+
+Available corruptions:
+    - "Gaussian Noise": Additive Gaussian noise
+    - "Impulse Noise": Salt-and-pepper noise
+    - "Motion Blur": Directional motion blur
+    - "Contrast": Reduced contrast
+    - "Pixelate": Resolution reduction
+    - "JPEG": JPEG compression artifacts
+    - "Spatter": Water or mud splatter effects
+
+Example:
+    ```python
+    from PIL import Image
+    from fusion_bench.dataset.image_corruption.make_corruption import distortion_methods
+
+    img = Image.open("example.jpg")
+    for name, corruption_fn in distortion_methods.items():
+        corrupted = corruption_fn(img, severity=3)
+        # Process corrupted image
+    ```
+"""
 distortion_methods["Gaussian Noise"] = gaussian_noise
 distortion_methods["Impulse Noise"] = impulse_noise
 distortion_methods["Motion Blur"] = motion_blur

fusion_bench/method/__init__.py

@@ -144,7 +144,15 @@ _extra_objects = {
 
 if TYPE_CHECKING:
     from .ada_svd import AdaSVDMergingForCLIPVisionModel
-    from .adamerging import *
+    from .adamerging import (
+        CLIPLayerWiseAdaMergingAlgorithm,
+        CLIPTaskWiseAdaMergingAlgorithm,
+        FlanT5LayerWiseAdaMergingAlgorithm,
+        GPT2LayerWiseAdaMergingAlgorithm,
+        LayerWiseAdaMergingForLlamaSFT,
+        ResNetLayerWiseAdamerging,
+        ResNetTaskWiseAdamerging,
+    )
     from .analysis import TaskVectorCosSimilarity, TaskVectorViolinPlot
     from .base_algorithm import BaseAlgorithm, BaseModelFusionAlgorithm
     from .bitdelta import BitDeltaAlgorithm
@@ -167,6 +175,7 @@ if TYPE_CHECKING:
     from .dawe import DataAdaptiveWeightEnsemblingForCLIP
     from .depth_upscaling import DepthUpscalingAlgorithm, DepthUpscalingForLlama
     from .doge_ta import DOGE_TA_Algorithm
+    from .dop import ContinualDOPForCLIP
     from .dummy import DummyAlgorithm
     from .ensemble import (
         MaxModelPredictorAlgorithm,
@@ -215,7 +224,6 @@ if TYPE_CHECKING:
     from .model_recombination import ModelRecombinationAlgorithm
     from .model_stock import ModelStock
     from .opcm import OPCMForCLIP
-    from .dop import ContinualDOPForCLIP
     from .pruning import (
         MagnitudeDiffPruningAlgorithm,
         MagnitudePruningForLlama,

fusion_bench/method/base_algorithm.py

@@ -40,6 +40,7 @@ from typing import Optional  # noqa: F401
 
 from fusion_bench.mixins import BaseYAMLSerializable
 from fusion_bench.modelpool import BaseModelPool
+from fusion_bench.utils.misc import DeprecationWarningMeta
 
 __all__ = ["BaseAlgorithm", "BaseModelFusionAlgorithm"]
 
@@ -202,27 +203,36 @@ class BaseAlgorithm(BaseYAMLSerializable):
         pass
 
 
-BaseModelFusionAlgorithm = BaseAlgorithm
-"""
-Alias for BaseAlgorithm class.
+# Create a deprecated wrapper class that inherits from BaseAlgorithm
+class BaseModelFusionAlgorithm(BaseAlgorithm, metaclass=DeprecationWarningMeta):
+    """
+    Alias for BaseAlgorithm class.
 
-This alias is provided for backward compatibility and semantic clarity.
-Some users may prefer the more explicit name 'BaseModelFusionAlgorithm'
-to emphasize that this class is specifically designed for model fusion
-tasks, while others may prefer the shorter 'BaseAlgorithm' name.
+    .. deprecated::
+        BaseModelFusionAlgorithm is deprecated and will be removed in a future version.
+        Use :class:`BaseAlgorithm` instead.
 
-Both names refer to the exact same class and can be used interchangeably.
+    This alias was provided for backward compatibility and semantic clarity.
+    Both names refer to the same base class and can be used interchangeably,
+    but BaseAlgorithm is now the preferred name for all implementations.
 
-Examples:
-    Using the original name:
-    >>> class MyAlgorithm(BaseAlgorithm):
-    ...     def run(self, modelpool): pass
+    Examples:
+        Preferred (using BaseAlgorithm):
 
-    Using the alias:
-    >>> class MyAlgorithm(BaseModelFusionAlgorithm):
-    ...     def run(self, modelpool): pass
+        >>> class MyAlgorithm(BaseAlgorithm):
+        ...     def run(self, modelpool): pass
 
-Note:
-    The alias is maintained for compatibility but BaseAlgorithm is the
-    preferred name for new implementations.
-"""
+        Deprecated (using BaseModelFusionAlgorithm):
+
+        >>> class MyAlgorithm(BaseModelFusionAlgorithm):  # Will trigger deprecation warning
+        ...     def run(self, modelpool): pass
+
+    Note:
+        New implementations should use :class:`BaseAlgorithm` exclusively.
+        The BaseModelFusionAlgorithm alias will be removed in a future release.
+
+    Warning:
+        Using BaseModelFusionAlgorithm will trigger a DeprecationWarning.
+    """
+
+    pass

fusion_bench/method/classification/image_classification_finetune.py

@@ -15,7 +15,7 @@ from lightning_utilities.core.rank_zero import rank_zero_only
 from lit_learn.lit_modules import ERM_LitModule
 from omegaconf import DictConfig
 from torch import nn
-from torch.utils.data import DataLoader
+from torch.utils.data import DataLoader, random_split
 from torchmetrics.classification import Accuracy
 
 from fusion_bench import (
@@ -29,7 +29,6 @@ from fusion_bench import (
 from fusion_bench.dataset import CLIPDataset
 from fusion_bench.modelpool import ResNetForImageClassificationPool
 from fusion_bench.tasks.clip_classification import get_num_classes
-from torch.utils.data import random_split
 
 log = get_rankzero_logger(__name__)
 

fusion_bench/method/gossip/clip_task_wise_gossip.py

@@ -13,41 +13,13 @@ from fusion_bench.modelpool import CLIPVisionModelPool
 from fusion_bench.models.hf_clip import HFCLIPClassifier
 from fusion_bench.tasks.clip_classification import get_classnames_and_templates
 from fusion_bench.utils import timeit_context
+from fusion_bench.utils.data import InfiniteDataLoader
 
 from .task_wise_gossip import TaskWiseGossipAlgorithm
 
 log = logging.getLogger(__name__)
 
 
-class InfiniteDataLoader:
-    """
-    A wrapper class for DataLoader to create an infinite data loader.
-    This is useful in case we are only interested in the number of steps and not the number of epochs.
-
-    This class wraps a DataLoader and provides an iterator that resets
-    when the end of the dataset is reached, creating an infinite loop.
-
-    Attributes:
-        data_loader (DataLoader): The DataLoader to wrap.
-        data_iter (iterator): An iterator over the DataLoader.
-    """
-
-    def __init__(self, data_loader):
-        self.data_loader = data_loader
-        self.data_iter = iter(data_loader)
-
-    def __iter__(self):
-        return self
-
-    def __next__(self):
-        try:
-            data = next(self.data_iter)
-        except StopIteration:
-            self.data_iter = iter(self.data_loader)  # Reset the data loader
-            data = next(self.data_iter)
-        return data
-
-
 class CLIPTaskWiseGossipAlgorithm(TaskWiseGossipAlgorithm):
     """
     A class for task-wise adaptive merging of CLIP models.

fusion_bench/metrics/model_kinship/__init__.py

@@ -0,0 +1,2 @@
+# Exploring Model Kinship for Merging LLMs
+# The implementation of this module is borrowed from: https://github.com/zjunlp/ModelKinship/

fusion_bench/metrics/model_kinship/calculate.py

@@ -0,0 +1,77 @@
+import logging
+from typing import List
+
+import numpy
+import torch
+
+from .utility import Metric
+
+
+def cosine_similarity(a, b):
+    similarity = numpy.sqrt(numpy.dot(a, b) ** 2 / (numpy.dot(a, a) * numpy.dot(b, b)))
+    return similarity
+
+
+def calculate_model_kinship(
+    delta1: numpy.ndarray, delta2: numpy.ndarray, metrics: List[str]
+) -> dict:
+    """
+    Calculate model kinship using specified metrics.
+
+    Args:
+        delta1: Delta parameters for first model
+        delta2: Delta parameters for second model
+        metrics: List of metrics to calculate
+
+    Returns:
+        dict: Dictionary of metric names and their calculated values
+    """
+    results = {}
+    for metric in metrics:
+        try:
+            if metric not in Metric.list():
+                raise ValueError(f"Unsupported metric: {metric}")
+            results[metric] = calculate_metric(delta1, delta2, metric)
+        except Exception as e:
+            results[metric] = f"Error calculating {metric}: {str(e)}"
+    return results
+
+
+def calculate_metric(
+    d_vector_1: torch.Tensor, d_vector_2: torch.Tensor, metric: str
+) -> str:
+    """
+    Calculate the specified metric between two delta vectors.
+
+    Args:
+        d_vector_1 (torch.Tensor): Delta parameters for model 1.
+        d_vector_2 (torch.Tensor): Delta parameters for model 2.
+        metric (str): The metric to calculate ('pcc', 'ed', 'cs').
+
+    Returns:
+        str: A formatted string with the result of the chosen metric.
+    """
+    logging.info(f"Starting calculation of {metric.upper()} metric...")
+
+    # Pearson Correlation Coefficient (PCC)
+    if metric == "pcc":
+        # Stack the two vectors and calculate the Pearson correlation coefficient
+        stack = torch.stack((d_vector_1, d_vector_2), dim=0)
+        pcc = torch.corrcoef(stack)[0, 1].item()
+        return f"Model Kinship based on Pearson Correlation Coefficient: {pcc}"
+
+    # Euclidean Distance (ED)
+    elif metric == "ed":
+        # Compute the Euclidean distance between the vectors
+        distance = torch.dist(d_vector_1, d_vector_2).item()
+        return f"Model Kinship based on Euclidean Distance: {distance}"
+
+    # Cosine Similarity (CS)
+    elif metric == "cs":
+        # Compute cosine similarity
+        cs = cosine_similarity(d_vector_1, d_vector_2)
+        return f"Model Kinship based on Cosine Similarity: {cs}"
+
+    # If metric is not recognized
+    else:
+        return "Invalid metric specified."