fusion-bench 0.2.19__py3-none-any.whl → 0.2.21__py3-none-any.whl

fusion_bench/mixins/serialization.py

@@ -1,20 +1,148 @@
+import inspect
 import logging
+from copy import deepcopy
+from functools import wraps
+from inspect import Parameter, _ParameterKind
 from pathlib import Path
 from typing import Dict, Optional, Union
 
-from omegaconf import OmegaConf
+from omegaconf import DictConfig, OmegaConf
 
+from fusion_bench.constants import FUSION_BENCH_VERSION
 from fusion_bench.utils import import_object, instantiate
+from fusion_bench.utils.instantiate_utils import set_print_function_call
 
 log = logging.getLogger(__name__)
 
+__all__ = [
+    "YAMLSerializationMixin",
+    "auto_register_config",
+    "BaseYAMLSerializable",
+]
+
+
+def auto_register_config(cls):
+    """
+    Decorator to automatically register __init__ parameters in _config_mapping.
+
+    This decorator enhances classes that inherit from YAMLSerializationMixin by
+    automatically mapping constructor parameters to configuration keys and
+    dynamically setting instance attributes based on provided arguments.
+
+    The decorator performs the following operations:
+    1. Inspects the class's __init__ method signature
+    2. Automatically populates the _config_mapping dictionary with parameter names
+    3. Wraps the __init__ method to handle both positional and keyword arguments
+    4. Sets instance attributes for all constructor parameters
+    5. Applies default values when parameters are not provided
+
+    Args:
+        cls (YAMLSerializationMixin): The class to be decorated. Must inherit from
+            YAMLSerializationMixin to ensure proper serialization capabilities.
+
+    Returns:
+        YAMLSerializationMixin: The decorated class with enhanced auto-registration
+            functionality and modified __init__ behavior.
+
+    Behavior:
+        - **Parameter Registration**: All non-variadic parameters (excluding *args, **kwargs)
+          from the __init__ method are automatically added to _config_mapping
+        - **Positional Arguments**: Handled in order and mapped to corresponding parameter names
+        - **Keyword Arguments**: Processed after positional arguments, overriding any conflicts
+        - **Default Values**: Applied when parameters are not provided via arguments
+        - **Attribute Setting**: All parameters become instance attributes accessible via dot notation
+
+    Example:
+        ```python
+        @auto_register_config
+        class MyAlgorithm(BaseYAMLSerializable):
+            def __init__(self, learning_rate: float = 0.001, batch_size: int = 32, model_name: str = "default"):
+                super().__init__()
+
+        # All instantiation methods work automatically:
+        algo1 = MyAlgorithm(0.01, 64)  # positional args
+        algo2 = MyAlgorithm(learning_rate=0.01, model_name="bert")  # keyword args
+        algo3 = MyAlgorithm(0.01, batch_size=128, model_name="gpt")  # mixed args
+
+        # Attributes are automatically set and can be serialized:
+        print(algo1.learning_rate)  # 0.01
+        print(algo1.batch_size)     # 64
+        print(algo1.model_name)     # "default" (from default value)
+
+        config = algo1.config
+        # DictConfig({'_target_': 'MyAlgorithm', 'learning_rate': 0.01, 'batch_size': 64, 'model_name': 'default'})
+        ```
+
+    Note:
+        - The decorator wraps the original __init__ method while preserving its signature for IDE support
+        - Parameters with *args or **kwargs signatures are ignored during registration
+        - The attributes are auto-registered, then the original __init__ method is called,
+        - Type hints, method name, and other metadata are preserved using functools.wraps
+        - This decorator is designed to work seamlessly with the YAML serialization system
+
+    Raises:
+        AttributeError: If the class does not have the required _config_mapping attribute
+            infrastructure (should inherit from YAMLSerializationMixin)
+    """
+    original_init = cls.__init__
+    sig = inspect.signature(original_init)
+
+    # Auto-register parameters in _config_mapping
+    if not "_config_mapping" in cls.__dict__:
+        cls._config_mapping = deepcopy(getattr(cls, "_config_mapping", {}))
+    for param_name in list(sig.parameters.keys())[1:]:  # Skip 'self'
+        if sig.parameters[param_name].kind not in [
+            _ParameterKind.VAR_POSITIONAL,
+            _ParameterKind.VAR_KEYWORD,
+        ]:
+            cls._config_mapping[param_name] = param_name
+
+    def __init__(self, *args, **kwargs):
+        # auto-register the attributes based on the signature
+        sig = inspect.signature(original_init)
+        param_names = list(sig.parameters.keys())[1:]  # Skip 'self'
+
+        # Handle positional arguments
+        for i, arg_value in enumerate(args):
+            if i < len(param_names):
+                param_name = param_names[i]
+                if sig.parameters[param_name].kind not in [
+                    _ParameterKind.VAR_POSITIONAL,
+                    _ParameterKind.VAR_KEYWORD,
+                ]:
+                    setattr(self, param_name, arg_value)
+
+        # Handle keyword arguments and defaults
+        for param_name in param_names:
+            if sig.parameters[param_name].kind not in [
+                _ParameterKind.VAR_POSITIONAL,
+                _ParameterKind.VAR_KEYWORD,
+            ]:
+                # Skip if already set by positional argument
+                param_index = param_names.index(param_name)
+                if param_index >= 0 and param_index < len(args):
+                    continue
+
+                if param_name in kwargs:
+                    setattr(self, param_name, kwargs[param_name])
+                else:
+                    # Set default value if available and attribute doesn't exist
+                    default_value = sig.parameters[param_name].default
+                    if default_value is not Parameter.empty:
+                        setattr(self, param_name, default_value)
+
+        # Call the original __init__
+        result = original_init(self, *args, **kwargs)
+        return result
+
+    # Replace the original __init__ method while preserving its signature
+    cls.__init__ = __init__
+    return cls
+
 
 class YAMLSerializationMixin:
-    _recursive_: bool = False
     _config_key: Optional[str] = None
-    _config_mapping: Dict[str, str] = {
-        "_recursive_": "_recursive_",
-    }
+    _config_mapping: Dict[str, str] = {}
     R"""
     `_config_mapping` is a dictionary mapping the attribute names of the class to the config option names. This is used to convert the class to a DictConfig.
 
@@ -47,46 +175,50 @@ class YAMLSerializationMixin:
     By default, the `_target_` key is set to the class name as `type(self).__name__`.
     """
 
-    def __init__(
-        self,
-        _recursive_: bool = False,
-        **kwargs,
-    ) -> None:
-        self._recursive_ = _recursive_
+    def __init__(self, **kwargs) -> None:
         for key, value in kwargs.items():
             log.warning(f"Unused argument: {key}={value}")
 
     @property
-    def config(self):
+    def config(self) -> DictConfig:
         R"""
         Returns the configuration of the model pool as a DictConfig.
 
-        This property calls the `to_config` method to convert the model pool
-        instance into a dictionary configuration, which can be used for
-        serialization or other purposes.
+        This property converts the model pool instance into a dictionary
+        configuration, which can be used for serialization or other purposes.
 
         Example:
-            >>> model = SomeModelFusionAlgorithm(hyper_param_1=1, hyper_param_2=2)
-            >>> config = model.config
-            >>> print(config)
-            DictConfig({'_target_': 'SomeModelFusionAlgorithm', 'hyper_param_1': 1, 'hyper_param_2': 2})
+
+        ```python
+        model = SomeModelFusionAlgorithm(hyper_param_1=1, hyper_param_2=2)
+        config = model.config
+        print(config)
+        # DictConfig({'_target_': 'SomeModelFusionAlgorithm', 'hyper_param_1': 1, 'hyper_param_2': 2})
+        ```
 
         This is useful for serializing the object to a YAML file or for debugging.
 
         Returns:
             DictConfig: The configuration of the model pool.
         """
-        return self.to_config()
+        config = {"_target_": f"{type(self).__module__}.{type(self).__qualname__}"}
+        for attr, key in self._config_mapping.items():
+            if hasattr(self, attr):
+                config[key] = getattr(self, attr)
 
-    def to_yaml(self, path: Union[str, Path]):
+        try:
+            return OmegaConf.create(config)
+        except Exception as e:
+            return OmegaConf.create(config, flags={"allow_objects": True})
+
+    def to_yaml(self, path: Union[str, Path], resolve: bool = True):
         """
         Save the model pool to a YAML file.
 
         Args:
             path (Union[str, Path]): The path to save the model pool to.
         """
-        config = self.to_config()
-        OmegaConf.save(config, path, resolve=True)
+        OmegaConf.save(self.config, path, resolve=resolve)
 
     @classmethod
     def from_yaml(cls, path: Union[str, Path]):
@@ -108,41 +240,126 @@ class YAMLSerializationMixin:
                 f"The class {target_cls.__name__} is not the same as the class {cls.__name__}. "
                 f"Instantiating the class {target_cls.__name__} instead."
             )
-        return instantiate(
-            config,
-            _recursive_=(
-                cls._recursive_
-                if config.get("_recursive_") is None
-                else config.get("_recursive_")
-            ),
-        )
+        with set_print_function_call(False):
+            return instantiate(config)
 
-    def to_config(self):
+    def register_parameter_to_config(
+        self,
+        attr_name: str,
+        param_name: str,
+        value,
+    ):
         """
-        Convert the model pool to a DictConfig.
+        Set an attribute value and register its config mapping.
 
-        Returns:
-            Dict: The model pool as a DictConfig.
+        This method allows dynamic setting of object attributes while simultaneously
+        updating the configuration mapping that defines how the attribute should
+        be serialized in the configuration output.
+
+        Args:
+            attr_name (str): The name of the attribute to set on this object.
+            arg_name (str): The corresponding parameter name to use in the config
+                serialization. This is how the attribute will appear in YAML output.
+            value: The value to assign to the attribute.
+
+        Example:
+            ```python
+            model = BaseYAMLSerializable()
+            model.set_option("learning_rate", "lr", 0.001)
+
+            # This sets model.learning_rate = 0.001
+            # and maps it to "lr" in the config output
+            config = model.config
+            # config will contain: {"lr": 0.001, ...}
+            ```
         """
-        config = {"_target_": type(self).__name__}
-        for attr, key in self._config_mapping.items():
-            if hasattr(self, attr):
-                config[key] = getattr(self, attr)
-        return OmegaConf.create(config)
+        setattr(self, attr_name, value)
+        self._config_mapping[attr_name] = param_name
+
+
+@auto_register_config
+class BaseYAMLSerializable(YAMLSerializationMixin):
+    """
+    A base class for YAML-serializable classes with enhanced metadata support.
+
+    This class extends `YAMLSerializationMixin` to provide additional metadata
+    fields commonly used in FusionBench classes, including usage information
+    and version tracking. It serves as a foundation for all serializable
+    model components in the framework.
+
+    The class automatically handles serialization of usage and version metadata
+    alongside the standard configuration parameters, making it easier to track
+    model provenance and intended usage patterns.
 
+    Attributes:
+        _usage_ (Optional[str]): Description of the model's intended usage or purpose.
+        _version_ (Optional[str]): Version information for the model or configuration.
 
-class BaseYAMLSerializableModel(YAMLSerializationMixin):
-    _config_mapping = YAMLSerializationMixin._config_mapping | {
-        "_usage_": "_usage_",
-        "_version_": "_version_",
-    }
+    Example:
+        ```python
+        class MyAlgorithm(BaseYAMLSerializable):
+            _config_mapping = BaseYAMLSerializable._config_mapping | {
+                "model_name": "model_name",
+                "num_layers": "num_layers",
+            }
+
+            def __init__(self, _usage_: str = None, _version_: str = None):
+                super().__init__(_usage_=_usage_, _version_=_version_)
+
+        # Usage with metadata
+        model = MyAlgorithm(
+            _usage_="Text classification fine-tuning",
+            _version_="1.0.0"
+        )
+
+        # Serialization includes metadata
+        config = model.config
+        # DictConfig({
+        #     '_target_': 'MyModel',
+        #     '_usage_': 'Text classification fine-tuning',
+        #     '_version_': '1.0.0'
+        # })
+        ```
+
+    Note:
+        The underscore prefix in `_usage_` and `_version_` follows the convention
+        for metadata fields that are not core model parameters but provide
+        important contextual information for model management and tracking.
+    """
 
     def __init__(
         self,
+        _recursive_: bool = False,
         _usage_: Optional[str] = None,
-        _version_: Optional[str] = None,
+        _version_: Optional[str] = FUSION_BENCH_VERSION,
         **kwargs,
     ):
+        """
+        Initialize a base YAML-serializable model with metadata support.
+
+        Args:
+            _usage_ (Optional[str], optional): Description of the model's intended
+                usage or purpose. This can include information about the training
+                domain, expected input types, or specific use cases. Defaults to None.
+            _version_ (Optional[str], optional): Version information for the model
+                or configuration. Can be used to track model iterations, dataset
+                versions, or compatibility information. Defaults to None.
+            **kwargs: Additional keyword arguments passed to the parent class.
+                Unused arguments will trigger warnings via the parent's initialization.
+
+        Example:
+            ```python
+            model = BaseYAMLSerializable(
+                _usage_="Image classification on CIFAR-10",
+                _version_="2.1.0"
+            )
+            ```
+        """
         super().__init__(**kwargs)
-        self._usage_ = _usage_
-        self._version_ = _version_
+        if _version_ != FUSION_BENCH_VERSION:
+            log.warning(
+                f"Current fusion-bench version is {FUSION_BENCH_VERSION}, but the serialized version is {_version_}. "
+                "Attempting to use current version."
+            )
+            # override _version_ with current fusion-bench version
+            self._version_ = FUSION_BENCH_VERSION

fusion_bench/modelpool/__init__.py

@@ -17,7 +17,7 @@ _import_structure = {
         "HuggingFaceGPT2ClassificationPool",
         "GPT2ForSequenceClassificationPool",
     ],
-    "seq_classification_lm": ["SeqenceClassificationModelPool"],
+    "seq_classification_lm": ["SequenceClassificationModelPool"],
 }
 
 
@@ -34,7 +34,7 @@ if TYPE_CHECKING:
     from .openclip_vision import OpenCLIPVisionModelPool
     from .PeftModelForSeq2SeqLM import PeftModelForSeq2SeqLMPool
     from .seq2seq_lm import Seq2SeqLMPool
-    from .seq_classification_lm import SeqenceClassificationModelPool
+    from .seq_classification_lm import SequenceClassificationModelPool
 
 else:
     sys.modules[__name__] = LazyImporter(

fusion_bench/modelpool/base_pool.py

@@ -1,13 +1,13 @@
 import logging
 from copy import deepcopy
-from typing import Dict, List, Optional, Union
+from typing import Dict, Generator, List, Optional, Tuple, Union
 
 import torch
 from omegaconf import DictConfig
 from torch import nn
 from torch.utils.data import Dataset
 
-from fusion_bench.mixins import BaseYAMLSerializableModel, HydraConfigMixin
+from fusion_bench.mixins import BaseYAMLSerializable, HydraConfigMixin
 from fusion_bench.utils import instantiate, timeit_context
 
 __all__ = ["BaseModelPool"]
@@ -15,7 +15,10 @@ __all__ = ["BaseModelPool"]
 log = logging.getLogger(__name__)
 
 
-class BaseModelPool(BaseYAMLSerializableModel, HydraConfigMixin):
+class BaseModelPool(
+    HydraConfigMixin,
+    BaseYAMLSerializable,
+):
     """
     A class for managing and interacting with a pool of models along with their associated datasets or other specifications. For example, a model pool may contain multiple models, each with its own training, validation, and testing datasets. As for the specifications, a vision model pool may contain image preprocessor, and a language model pool may contain a tokenizer.
 
@@ -31,7 +34,7 @@ class BaseModelPool(BaseYAMLSerializableModel, HydraConfigMixin):
     _program = None
     _config_key = "modelpool"
     _models: Union[DictConfig, Dict[str, nn.Module]]
-    _config_mapping = BaseYAMLSerializableModel._config_mapping | {
+    _config_mapping = BaseYAMLSerializable._config_mapping | {
         "_models": "models",
         "_train_datasets": "train_datasets",
         "_val_datasets": "val_datasets",
@@ -56,7 +59,7 @@ class BaseModelPool(BaseYAMLSerializableModel, HydraConfigMixin):
         super().__init__(**kwargs)
 
     @property
-    def has_pretrained(self):
+    def has_pretrained(self) -> bool:
         """
         Check if the model pool contains a pretrained model.
 
@@ -125,7 +128,7 @@ class BaseModelPool(BaseYAMLSerializableModel, HydraConfigMixin):
         return len(self.model_names)
 
     @staticmethod
-    def is_special_model(model_name: str):
+    def is_special_model(model_name: str) -> bool:
         """
         Determine if a model is special based on its name.
 
@@ -152,6 +155,23 @@ class BaseModelPool(BaseYAMLSerializableModel, HydraConfigMixin):
             model_config = deepcopy(model_config)
         return model_config
 
+    def get_model_path(self, model_name: str) -> str:
+        """
+        Get the path for the specified model.
+
+        Args:
+            model_name (str): The name of the model.
+
+        Returns:
+            str: The path for the specified model.
+        """
+        if isinstance(self._models[model_name], str):
+            return self._models[model_name]
+        else:
+            raise ValueError(
+                "Model path is not a string. Try to override this method in derived modelpool class."
+            )
+
     def load_model(
         self, model_name_or_config: Union[str, DictConfig], *args, **kwargs
     ) -> nn.Module:
@@ -159,7 +179,7 @@ class BaseModelPool(BaseYAMLSerializableModel, HydraConfigMixin):
         Load a model from the pool based on the provided configuration.
 
         Args:
-            model (Union[str, DictConfig]): The model name or configuration.
+            model_name_or_config (Union[str, DictConfig]): The model name or configuration.
 
         Returns:
             nn.Module: The instantiated model.
@@ -201,11 +221,11 @@ class BaseModelPool(BaseYAMLSerializableModel, HydraConfigMixin):
             model = self.load_model(self.model_names[0], *args, **kwargs)
         return model
 
-    def models(self):
+    def models(self) -> Generator[nn.Module, None, None]:
         for model_name in self.model_names:
             yield self.load_model(model_name)
 
-    def named_models(self):
+    def named_models(self) -> Generator[Tuple[str, nn.Module], None, None]:
         for model_name in self.model_names:
             yield model_name, self.load_model(model_name)
 

fusion_bench/modelpool/causal_lm/causal_lm.py

@@ -57,6 +57,15 @@ class CausalLMPool(BaseModelPool):
                 )
         self.load_lazy = load_lazy
 
+    def get_model_path(self, model_name: str):
+        model_name_or_config = self._models[model_name]
+        if isinstance(model_name_or_config, str):
+            return model_name_or_config
+        elif isinstance(model_name_or_config, (DictConfig, dict)):
+            return model_name_or_config.get("pretrained_model_name_or_path")
+        else:
+            raise RuntimeError("Invalid model configuration")
+
     @override
     def load_model(
         self,

fusion_bench/modelpool/clip_vision/modelpool.py

@@ -1,6 +1,6 @@
 import logging
 from copy import deepcopy
-from typing import Optional, Union
+from typing import Literal, Optional, Union
 
 from datasets import load_dataset
 from lightning.fabric.utilities import rank_zero_only
@@ -11,6 +11,7 @@ from transformers import CLIPModel, CLIPProcessor, CLIPVisionModel
 from typing_extensions import override
 
 from fusion_bench.utils import instantiate, timeit_context
+from fusion_bench.utils.modelscope import resolve_repo_path
 
 from ..base_pool import BaseModelPool
 
@@ -25,25 +26,32 @@ class CLIPVisionModelPool(BaseModelPool):
     the specifics of the CLIP Vision models provided by the Hugging Face Transformers library.
     """
 
-    _config_mapping = BaseModelPool._config_mapping | {"_processor": "processor"}
+    _config_mapping = BaseModelPool._config_mapping | {
+        "_processor": "processor",
+        "_platform": "hf",
+    }
 
     def __init__(
         self,
         models: DictConfig,
         *,
         processor: Optional[DictConfig] = None,
+        platform: Literal["hf", "huggingface", "modelscope"] = "hf",
         **kwargs,
     ):
         super().__init__(models, **kwargs)
-
         self._processor = processor
+        self._platform = platform
 
     def load_processor(self, *args, **kwargs) -> CLIPProcessor:
         assert self._processor is not None, "Processor is not defined in the config"
         if isinstance(self._processor, str):
             if rank_zero_only.rank == 0:
                 log.info(f"Loading `transformers.CLIPProcessor`: {self._processor}")
-            processor = CLIPProcessor.from_pretrained(self._processor)
+            repo_path = resolve_repo_path(
+                repo_id=self._processor, repo_type="model", platform=self._platform
+            )
+            processor = CLIPProcessor.from_pretrained(repo_path, *args, **kwargs)
         else:
             processor = instantiate(self._processor, *args, **kwargs)
         return processor
@@ -54,7 +62,10 @@ class CLIPVisionModelPool(BaseModelPool):
         if isinstance(model_config, str):
             if rank_zero_only.rank == 0:
                 log.info(f"Loading `transformers.CLIPModel`: {model_config}")
-            clip_model = CLIPModel.from_pretrained(model_config, *args, **kwargs)
+            repo_path = resolve_repo_path(
+                repo_id=model_config, repo_type="model", platform=self._platform
+            )
+            clip_model = CLIPModel.from_pretrained(repo_path, *args, **kwargs)
             return clip_model
         else:
             assert isinstance(
@@ -107,14 +118,17 @@ class CLIPVisionModelPool(BaseModelPool):
             if isinstance(model, str):
                 if rank_zero_only.rank == 0:
                     log.info(f"Loading `transformers.CLIPVisionModel`: {model}")
-                return CLIPVisionModel.from_pretrained(model, *args, **kwargs)
+                repo_path = resolve_repo_path(
+                    model, repo_type="model", platform=self._platform
+                )
+                return CLIPVisionModel.from_pretrained(repo_path, *args, **kwargs)
             if isinstance(model, nn.Module):
                 if rank_zero_only.rank == 0:
                     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)
+        else:
+            # 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]
@@ -123,7 +137,7 @@ class CLIPVisionModelPool(BaseModelPool):
                 log.info(
                     f"Loading train dataset using `datasets.load_dataset`: {dataset_config}"
                 )
-            dataset = load_dataset(dataset_config, split="train")
+            dataset = self._load_dataset(dataset_config, split="train")
         else:
             dataset = super().load_train_dataset(dataset_name, *args, **kwargs)
         return dataset
@@ -135,7 +149,7 @@ class CLIPVisionModelPool(BaseModelPool):
                 log.info(
                     f"Loading validation dataset using `datasets.load_dataset`: {dataset_config}"
                 )
-            dataset = load_dataset(dataset_config, split="validation")
+            dataset = self._load_dataset(dataset_config, split="validation")
         else:
             dataset = super().load_val_dataset(dataset_name, *args, **kwargs)
         return dataset
@@ -147,7 +161,24 @@ class CLIPVisionModelPool(BaseModelPool):
                 log.info(
                     f"Loading test dataset using `datasets.load_dataset`: {dataset_config}"
                 )
-            dataset = load_dataset(dataset_config, split="test")
+            dataset = self._load_dataset(dataset_config, split="test")
         else:
             dataset = super().load_test_dataset(dataset_name, *args, **kwargs)
         return dataset
+
+    def _load_dataset(self, name: str, split: str):
+        """
+        Load a dataset by its name and split.
+
+        Args:
+            dataset_name (str): The name of the dataset.
+            split (str): The split of the dataset to load (e.g., "train", "validation", "test").
+
+        Returns:
+            Dataset: The loaded dataset.
+        """
+        datset_dir = resolve_repo_path(
+            name, repo_type="dataset", platform=self._platform
+        )
+        dataset = load_dataset(datset_dir, split=split)
+        return dataset

fusion_bench/modelpool/seq_classification_lm/__init__.py

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

fusion_bench/modelpool/seq_classification_lm/seq_classification_lm.py

@@ -17,7 +17,7 @@ if TYPE_CHECKING:
 log = logging.getLogger(__name__)
 
 
-class SeqenceClassificationModelPool(BaseModelPool):
+class SequenceClassificationModelPool(BaseModelPool):
 
     def __init__(
         self,

fusion_bench/models/__init__.py

@@ -1,4 +1,5 @@
 # flake8: noqa F401
+from fusion_bench.utils import LazyStateDict
+
 from . import separate_io, utils
 from .parameter_dict import ParameterDictModel
-from fusion_bench.utils import LazyStateDict

fusion_bench/models/expert_sparsity/mixtral/__init__.py

@@ -10,6 +10,6 @@ Reference:
 """
 
 from .wrapper import (
-    PrunableMixtralSparseMoeBlockWrapper,
     DynamicSkippingMixtralSparseMoeBlockWrapper,
+    PrunableMixtralSparseMoeBlockWrapper,
 )