fusion-bench 0.2.22__py3-none-any.whl → 0.2.23__py3-none-any.whl

fusion_bench/modelpool/causal_lm/causal_lm.py

@@ -33,6 +33,44 @@ log = logging.getLogger(__name__)
 
 @auto_register_config
 class CausalLMPool(BaseModelPool):
+    """A model pool for managing and loading causal language models.
+
+    This class provides a unified interface for loading and managing multiple
+    causal language models, typically used in model fusion and ensemble scenarios.
+    It supports both eager and lazy loading strategies, and handles model
+    configuration through YAML configs or direct instantiation.
+
+    The pool can manage models from Hugging Face Hub, local paths, or custom
+    configurations. It also provides tokenizer management and model saving
+    capabilities with optional Hugging Face Hub integration.
+
+    Args:
+        models: Dictionary or configuration specifying the models to be managed.
+            Can contain model names mapped to paths or detailed configurations.
+        tokenizer: Tokenizer configuration, either a string path/name or
+            a DictConfig with detailed tokenizer settings.
+        model_kwargs: Additional keyword arguments passed to model loading.
+            Common options include torch_dtype, device_map, etc.
+        enable_lazy_loading: Whether to use lazy loading for models. When True,
+            models are loaded as LazyStateDict objects instead of actual models,
+            which can save memory for large model collections.
+        **kwargs: Additional arguments passed to the parent BaseModelPool.
+
+    Example:
+        ```python
+        >>> pool = CausalLMPool(
+        ...     models={
+        ...         "model_a": "microsoft/DialoGPT-medium",
+        ...         "model_b": "/path/to/local/model"
+        ...     },
+        ...     tokenizer="microsoft/DialoGPT-medium",
+        ...     model_kwargs={"torch_dtype": "bfloat16"}
+        ... )
+        >>> model = pool.load_model("model_a")
+        >>> tokenizer = pool.load_tokenizer()
+        ```
+    """
+
     def __init__(
         self,
         models,
@@ -47,6 +85,20 @@ class CausalLMPool(BaseModelPool):
             self.model_kwargs = DictConfig({})
 
     def get_model_path(self, model_name: str):
+        """Extract the model path from the model configuration.
+
+        Args:
+            model_name: The name of the model as defined in the models configuration.
+
+        Returns:
+            str: The path or identifier for the model. For string configurations,
+                returns the string directly. For dict configurations, extracts
+                the 'pretrained_model_name_or_path' field.
+
+        Raises:
+            RuntimeError: If the model configuration is invalid or the model
+                name is not found in the configuration.
+        """
         model_name_or_config = self._models[model_name]
         if isinstance(model_name_or_config, str):
             return model_name_or_config
@@ -56,6 +108,16 @@ class CausalLMPool(BaseModelPool):
             raise RuntimeError("Invalid model configuration")
 
     def get_model_kwargs(self):
+        """Get processed model keyword arguments for model loading.
+
+        Converts the stored `model_kwargs` from DictConfig to a regular dictionary
+        and processes special arguments like torch_dtype for proper model loading.
+
+        Returns:
+            dict: Processed keyword arguments ready to be passed to model
+                loading functions. The torch_dtype field, if present, is
+                converted from string to the appropriate torch dtype object.
+        """
         model_kwargs = (
             OmegaConf.to_container(self.model_kwargs, resolve=True)
             if isinstance(self.model_kwargs, DictConfig)
@@ -71,31 +133,52 @@ class CausalLMPool(BaseModelPool):
         model_name_or_config: str | DictConfig,
         *args,
         **kwargs,
-    ) -> PreTrainedModel:
-        """
-        Example of YAML config:
+    ) -> Union[PreTrainedModel, LazyStateDict]:
+        """Load a causal language model from the model pool.
 
-        ```yaml
-        models:
-          _pretrained_: path_to_pretrained_model # if a plain string, it will be passed to AutoModelForCausalLM.from_pretrained
-          model_a: path_to_model_a
-          model_b: path_to_model_b
-        ```
+        This method supports multiple loading strategies:
+        1. Loading by model name from the configured model pool
+        2. Loading from a direct configuration dictionary
+        3. Lazy loading using LazyStateDict for memory efficiency
 
-        or equivalently,
-
-        ```yaml
-        models:
-          _pretrained_:
-            _target_: transformers.AutoModelForCausalLM # any callable that returns a model
-            pretrained_model_name_or_path: path_to_pretrained_model
-          model_a:
-            _target_: transformers.AutoModelForCausalLM
-            pretrained_model_name_or_path: path_to_model_a
-          model_b:
-            _target_: transformers.AutoModelForCausalLM
-            pretrained_model_name_or_path: path_to_model_b
-        ```
+        The method automatically handles different model configuration formats
+        and applies the appropriate loading strategy based on the enable_lazy_loading flag.
+
+        Args:
+            model_name_or_config: Either a string model name that exists in the
+                model pool configuration, or a DictConfig/dict containing the
+                model configuration directly.
+            *args: Additional positional arguments passed to the model constructor.
+            **kwargs: Additional keyword arguments passed to the model constructor.
+                These will be merged with the pool's model_kwargs.
+
+        Returns:
+            Union[PreTrainedModel, LazyStateDict]: The loaded model. Returns a
+                PreTrainedModel for normal loading or a LazyStateDict for lazy loading.
+
+        Raises:
+            RuntimeError: If the model configuration is invalid.
+            KeyError: If the model name is not found in the model pool.
+
+        Example YAML configurations:
+            Simple string configuration:
+            ```yaml
+            models:
+              _pretrained_: path_to_pretrained_model
+              model_a: path_to_model_a
+              model_b: path_to_model_b
+            ```
+
+            Detailed configuration:
+            ```yaml
+            models:
+              _pretrained_:
+                _target_: transformers.AutoModelForCausalLM
+                pretrained_model_name_or_path: path_to_pretrained_model
+              model_a:
+                _target_: transformers.AutoModelForCausalLM
+                pretrained_model_name_or_path: path_to_model_a
+            ```
         """
         model_kwargs = self.get_model_kwargs()
         model_kwargs.update(kwargs)
@@ -139,23 +222,36 @@ class CausalLMPool(BaseModelPool):
         return model
 
     def load_tokenizer(self, *args, **kwargs) -> PreTrainedTokenizer:
-        """
-        Example of YAML config:
-
-        ```yaml
-        tokenizer: google/gemma-2-2b-it # if a plain string, it will be passed to AutoTokenizer.from_pretrained
-        ```
+        """Load the tokenizer associated with this model pool.
 
-        or equivalently,
+        Loads a tokenizer based on the tokenizer configuration provided during
+        pool initialization. Supports both simple string paths and detailed
+        configuration dictionaries.
 
-        ```yaml
-        tokenizer:
-          _target_: transformers.AutoTokenizer # any callable that returns a tokenizer
-          pretrained_model_name_or_path: google/gemma-2-2b-it
-        ```
+        Args:
+            *args: Additional positional arguments passed to the tokenizer constructor.
+            **kwargs: Additional keyword arguments passed to the tokenizer constructor.
 
         Returns:
-            PreTrainedTokenizer: The tokenizer.
+            PreTrainedTokenizer: The loaded tokenizer instance.
+
+        Raises:
+            AssertionError: If no tokenizer is defined in the configuration.
+
+        Example YAML configurations:
+            Simple string configuration:
+            ```yaml
+            tokenizer: google/gemma-2-2b-it
+            ```
+
+            Detailed configuration:
+            ```yaml
+            tokenizer:
+              _target_: transformers.AutoTokenizer
+              pretrained_model_name_or_path: google/gemma-2-2b-it
+              use_fast: true
+              padding_side: left
+            ```
         """
         assert self.tokenizer is not None, "Tokenizer is not defined in the config"
         log.info("Loading tokenizer.", stacklevel=2)
@@ -177,15 +273,49 @@ class CausalLMPool(BaseModelPool):
         tokenizer: Optional[PreTrainedTokenizer] = None,
         **kwargs,
     ):
-        """
-        Save the model to the specified path.
+        """Save a model to the specified path with optional tokenizer and Hub upload.
+
+        This method provides comprehensive model saving capabilities including
+        optional tokenizer saving, dtype conversion, and Hugging Face Hub upload.
+        The model is saved in the standard Hugging Face format.
 
         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.
+            model: The PreTrainedModel instance to be saved.
+            path: The local path where the model will be saved. Supports tilde
+                expansion for home directory paths.
+            push_to_hub: Whether to push the saved model to the Hugging Face Hub.
+                Requires proper authentication and repository permissions.
+            model_dtype: Optional string specifying the target dtype for the model
+                before saving (e.g., "float16", "bfloat16"). The model will be
+                converted to this dtype before saving.
+            save_tokenizer: Whether to save the tokenizer alongside the model.
+                If True, the tokenizer will be loaded using the pool's tokenizer
+                configuration and saved to the same path.
+            tokenizer_kwargs: Additional keyword arguments for tokenizer loading
+                when save_tokenizer is True.
+            tokenizer: Optional pre-loaded tokenizer instance. If provided, this
+                tokenizer will be saved regardless of the save_tokenizer flag.
+            **kwargs: Additional keyword arguments passed to the model's
+                save_pretrained method.
+
+        Side Effects:
+            - Creates model files in the specified directory
+            - Optionally creates tokenizer files in the same directory
+            - May convert the model to a different dtype
+            - May upload files to Hugging Face Hub
+
+        Example:
+            ```python
+            >>> pool = CausalLMPool(models=..., tokenizer=...)
+            >>> model = pool.load_model("my_model")
+            >>> pool.save_model(
+            ...     model,
+            ...     "/path/to/save",
+            ...     save_tokenizer=True,
+            ...     model_dtype="float16",
+            ...     push_to_hub=True
+            ... )
+            ```
         """
         path = os.path.expanduser(path)
         # NOTE: if tokenizer is provided, it will be saved regardless of `save_tokenizer`
@@ -209,9 +339,55 @@ class CausalLMPool(BaseModelPool):
 
 
 class CausalLMBackbonePool(CausalLMPool):
+    """A specialized model pool that loads only the transformer backbone layers.
+
+    This class extends CausalLMPool to provide access to just the transformer
+    layers (backbone) of causal language models, excluding the language modeling
+    head and embeddings. This is useful for model fusion scenarios where only
+    the core transformer layers are needed.
+
+    The class automatically extracts the `model.layers` component from loaded
+    AutoModelForCausalLM instances, providing direct access to the transformer
+    blocks. Lazy loading is not supported for this pool type.
+
+    Note:
+        This pool automatically disables lazy loading as it needs to access
+        the internal structure of the model to extract the backbone layers.
+
+    Example:
+        ```python
+        >>> backbone_pool = CausalLMBackbonePool(
+        ...     models={"model_a": "microsoft/DialoGPT-medium"},
+        ...     tokenizer="microsoft/DialoGPT-medium"
+        ... )
+        >>> layers = backbone_pool.load_model("model_a")  # Returns nn.ModuleList of transformer layers
+        ```
+    """
+
     def load_model(
         self, model_name_or_config: str | DictConfig, *args, **kwargs
     ) -> Module:
+        """Load only the transformer backbone layers from a causal language model.
+
+        This method loads a complete causal language model and then extracts
+        only the transformer layers (backbone), discarding the embedding layers
+        and language modeling head. This is useful for model fusion scenarios
+        where only the core transformer computation is needed.
+
+        Args:
+            model_name_or_config: Either a string model name from the pool
+                configuration or a DictConfig with model loading parameters.
+            *args: Additional positional arguments passed to the parent load_model method.
+            **kwargs: Additional keyword arguments passed to the parent load_model method.
+
+        Returns:
+            Module: The transformer layers (typically a nn.ModuleList) containing
+                the core transformer blocks without embeddings or output heads.
+
+        Note:
+            Lazy loading is automatically disabled for this method as it needs
+            to access the internal model structure to extract the layers.
+        """
         if self.enable_lazy_loading:
             log.warning(
                 "CausalLMBackbonePool does not support lazy loading. "
@@ -231,6 +407,49 @@ def load_peft_causal_lm(
     is_trainable: bool = True,
     merge_and_unload: bool = False,
 ):
+    """Load a causal language model with PEFT (Parameter-Efficient Fine-Tuning) adapters.
+
+    This function loads a base causal language model and applies PEFT adapters
+    (such as LoRA, AdaLoRA, or other parameter-efficient fine-tuning methods)
+    to create a fine-tuned model. It supports both keeping the adapters separate
+    or merging them into the base model.
+
+    Args:
+        base_model_path: Path or identifier for the base causal language model.
+            Can be a Hugging Face model name or local path.
+        peft_model_path: Path to the PEFT adapter configuration and weights.
+            This should contain the adapter_config.json and adapter weights.
+        torch_dtype: The torch data type to use for the model. Common options
+            include "float16", "bfloat16", "float32". Defaults to "bfloat16".
+        is_trainable: Whether the loaded PEFT model should be trainable.
+            Set to False for inference-only usage to save memory.
+        merge_and_unload: Whether to merge the PEFT adapters into the base model
+            and unload the adapter weights. When True, returns a standard
+            PreTrainedModel instead of a PeftModel.
+
+    Returns:
+        Union[PeftModel, PreTrainedModel]: The loaded model with PEFT adapters.
+            Returns a PeftModel if merge_and_unload is False, or a PreTrainedModel
+            if the adapters are merged and unloaded.
+
+    Example:
+        ```python
+        >>> # Load model with adapters for training
+        >>> model = load_peft_causal_lm(
+        ...     "microsoft/DialoGPT-medium",
+        ...     "/path/to/lora/adapters",
+        ...     is_trainable=True
+        ... )
+
+        >>> # Load and merge adapters for inference
+        >>> merged_model = load_peft_causal_lm(
+        ...     "microsoft/DialoGPT-medium",
+        ...     "/path/to/lora/adapters",
+        ...     merge_and_unload=True,
+        ...     is_trainable=False
+        ... )
+        ```
+    """
     base_model = AutoModelForCausalLM.from_pretrained(
         base_model_path, torch_dtype=torch_dtype
     )

fusion_bench/modelpool/seq2seq_lm/modelpool.py

@@ -18,6 +18,48 @@ def load_lora_model(
     is_trainable: bool = True,
     merge_and_unload: bool = True,
 ):
+    """Load a sequence-to-sequence model with LoRA (Low-Rank Adaptation) fine-tuning.
+
+    This function loads a base sequence-to-sequence language model and applies
+    LoRA adapters for parameter-efficient fine-tuning. LoRA allows for efficient
+    adaptation of large models by adding trainable low-rank matrices to the
+    existing weights without modifying the original parameters.
+
+    Args:
+        base_model_path: Path or identifier for the base sequence-to-sequence model.
+            Can be a Hugging Face model name (e.g., "t5-base") or local path.
+        peft_model_path: Path to the directory containing LoRA adapter weights
+            and configuration. Should include adapter_config.json and adapter weights.
+        is_trainable: Whether the loaded model should be trainable. Set to False
+            for inference-only usage to save memory and computation.
+        merge_and_unload: Whether to merge the LoRA weights into the base model
+            and unload the adapter. When True, returns a standard model instead
+            of a PeftModel, which can be more efficient for inference.
+
+    Returns:
+        Union[PeftModel, AutoModelForSeq2SeqLM]: The loaded model with LoRA
+            adapters. Returns a PeftModel if merge_and_unload is False, or
+            a standard AutoModelForSeq2SeqLM if adapters are merged.
+
+    Example:
+        ```python
+        >>> # Load model with separate adapters for training
+        >>> model = load_lora_model(
+        ...     "t5-base",
+        ...     "/path/to/lora/adapters",
+        ...     is_trainable=True,
+        ...     merge_and_unload=False
+        ... )
+
+        >>> # Load and merge adapters for efficient inference
+        >>> merged_model = load_lora_model(
+        ...     "t5-base",
+        ...     "/path/to/lora/adapters",
+        ...     is_trainable=False,
+        ...     merge_and_unload=True
+        ... )
+        ```
+    """
     base_model = AutoModelForSeq2SeqLM.from_pretrained(base_model_path)
     model = PeftModel.from_pretrained(
         base_model,
@@ -30,6 +72,46 @@ def load_lora_model(
 
 
 class Seq2SeqLMPool(BaseModelPool):
+    """A model pool specialized for sequence-to-sequence language models.
+
+    This model pool provides management and loading capabilities for sequence-to-sequence
+    (seq2seq) language models such as T5, BART, and mT5. It extends the base model pool
+    functionality with seq2seq-specific features including tokenizer management and
+    model configuration handling.
+
+    Seq2seq models are particularly useful for tasks that require generating output
+    sequences from input sequences, such as translation, summarization, question
+    answering, and text generation. This pool streamlines the process of loading
+    and configuring multiple seq2seq models for fusion and ensemble scenarios.
+
+    Key Features:
+        - Specialized loading for AutoModelForSeq2SeqLM models
+        - Integrated tokenizer management
+        - Support for model-specific keyword arguments
+        - Automatic dtype parsing and configuration
+        - Compatible with PEFT (Parameter-Efficient Fine-Tuning) adapters
+
+    Attributes:
+        _tokenizer: Configuration for the tokenizer associated with the models
+        _model_kwargs: Default keyword arguments applied to all model loading operations
+
+    Example:
+        ```python
+        pool = Seq2SeqLMPool(
+            models={
+                "t5_base": "t5-base",
+                "t5_large": "t5-large",
+                "custom_model": "/path/to/local/model"
+            },
+            tokenizer={"_target_": "transformers.T5Tokenizer",
+                      "pretrained_model_name_or_path": "t5-base"},
+            model_kwargs={"torch_dtype": "float16", "device_map": "auto"}
+        )
+        model = pool.load_model("t5_base")
+        tokenizer = pool.load_tokenizer()
+        ```
+    """
+
     _config_mapping = BaseModelPool._config_mapping | {
         "_tokenizer": "tokenizer",
         "_model_kwargs": "model_kwargs",
@@ -43,6 +125,35 @@ class Seq2SeqLMPool(BaseModelPool):
         model_kwargs: Optional[DictConfig] = None,
         **kwargs,
     ):
+        """Initialize the sequence-to-sequence language model pool.
+
+        Sets up the model pool with configurations for models, tokenizer, and
+        default model loading parameters. Automatically processes model kwargs
+        to handle special configurations like torch_dtype parsing.
+
+        Args:
+            models: Configuration dictionary specifying the seq2seq models to manage.
+                Keys are model names, values can be model paths/names or detailed configs.
+            tokenizer: Configuration for the tokenizer to use with the models.
+                Can be a simple path/name or detailed configuration with _target_.
+            model_kwargs: Default keyword arguments applied to all model loading
+                operations. Common options include torch_dtype, device_map, etc.
+                The torch_dtype field is automatically parsed from string to dtype.
+            **kwargs: Additional arguments passed to the parent BaseModelPool.
+
+        Example:
+            ```python
+            pool = Seq2SeqLMPool(
+                models={
+                    "base": "t5-base",
+                    "large": {"_target_": "transformers.AutoModelForSeq2SeqLM",
+                             "pretrained_model_name_or_path": "t5-large"}
+                },
+                tokenizer="t5-base",
+                model_kwargs={"torch_dtype": "bfloat16"}
+            )
+            ```
+        """
         super().__init__(models, **kwargs)
         self._tokenizer = tokenizer
         self._model_kwargs = model_kwargs
@@ -55,11 +166,46 @@ class Seq2SeqLMPool(BaseModelPool):
                 )
 
     def load_model(self, model_name_or_config: str | DictConfig, *args, **kwargs):
+        """Load a sequence-to-sequence language model from the pool.
+
+        Loads a seq2seq model using the parent class loading mechanism while
+        automatically applying the pool's default model kwargs. The method
+        merges the pool's model_kwargs with any additional kwargs provided,
+        giving priority to the explicitly provided kwargs.
+
+        Args:
+            model_name_or_config: Either a string model name from the pool
+                configuration or a DictConfig containing model loading parameters.
+            *args: Additional positional arguments passed to the parent load_model method.
+            **kwargs: Additional keyword arguments that override the pool's default
+                model_kwargs. Common options include device, torch_dtype, etc.
+
+        Returns:
+            AutoModelForSeq2SeqLM: The loaded sequence-to-sequence language model.
+        """
         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):
+        """Load the tokenizer associated with the sequence-to-sequence models.
+
+        Loads a tokenizer based on the tokenizer configuration provided during
+        pool initialization. The tokenizer should be compatible with the seq2seq
+        models in the pool and is typically used for preprocessing input text
+        and postprocessing generated output.
+
+        Args:
+            *args: Additional positional arguments passed to the tokenizer constructor.
+            **kwargs: Additional keyword arguments passed to the tokenizer constructor.
+
+        Returns:
+            PreTrainedTokenizer: The loaded tokenizer instance compatible with
+                the seq2seq models in this pool.
+
+        Raises:
+            AssertionError: If no tokenizer configuration is provided.
+        """
         assert self._tokenizer is not None, "Tokenizer is not defined in the config"
         tokenizer = isinstance(self._tokenizer, *args, **kwargs)
         return tokenizer

fusion_bench/models/hf_utils.py

@@ -42,12 +42,14 @@ def load_model_card_template(basename: str) -> str:
         FileNotFoundError: If the template file is not found in any of the search locations.
     """
     if os.path.exists(basename):
-        return open(basename).read()
+        with open(basename, "r") as f:
+            return f.read()
 
     for template_dir in MODEL_CARD_TEMPLATE_DIRS:
         template_path = os.path.join(template_dir, basename)
         if os.path.exists(template_path):
-            return open(template_path).read()
+            with open(template_path, "r") as f:
+                return f.read()
 
     raise FileNotFoundError(f"Model card template '{basename}' not found.")
 
@@ -141,6 +143,9 @@ def save_pretrained_with_remote_code(
 
 def create_default_model_card(
     models: list[str],
+    *,
+    title: str = "Deep Model Fusion",
+    tags: list[str] = ["fusion-bench", "merge"],
     description=None,
     algorithm_config: DictConfig = None,
     modelpool_config: DictConfig = None,
@@ -151,8 +156,8 @@ def create_default_model_card(
     card = template.render(
         models=models,
         library_name="transformers",
-        tags=["fusion-bench", "merge"],
-        title="Deep Model Fusion",
+        title=title,
+        tags=tags,
         description=description,
         algorithm_config_str=try_to_yaml(algorithm_config),
         modelpool_config_str=try_to_yaml(modelpool_config),

fusion_bench/models/linearized/vision_model.py

@@ -45,21 +45,21 @@ def linearize_lora_model_(model):
 
 
 def load_fft_vision_model_hf(
-    model_name: str, return_vison_model=True
+    model_name: str, return_vision_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.
+        return_vision_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:
+    if return_vision_model:
         return CLIPVisionModel.from_pretrained(model_name).vision_model
     else:
         return model
@@ -69,7 +69,7 @@ def load_lora_vision_model_hf(
     base_model_name: str,
     peft_name: str,
     merge_and_unload: bool = False,
-    return_vison_model=True,
+    return_vision_model=True,
 ) -> PeftModel:
     """
     Load a LoRA (Low-Rank Adaptation) vision model from Hugging Face.
@@ -80,7 +80,7 @@ def load_lora_vision_model_hf(
         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.
+        return_vision_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.
@@ -97,7 +97,7 @@ def load_lora_vision_model_hf(
         vision_model = peft_model
 
     # Return the vision model
-    if return_vison_model:
+    if return_vision_model:
         return vision_model
     else:
         model.vision_model = vision_model

fusion_bench/models/modeling_smile_mistral/__init__.py

@@ -1,6 +1,7 @@
 from . import register
 from .configuration_smile_mistral import SmileMistralConfig
 from .modeling_smile_mistral import (
+    SmileMistralDecoderLayer,
     SmileMistralForCausalLM,
     SmileMistralModel,
 )