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

fusion_bench/modelpool/causal_lm/causal_lm.py

@@ -8,7 +8,7 @@ from copy import deepcopy
 from typing import Any, Dict, Optional, TypeAlias, Union, cast  # noqa: F401
 
 import peft
-from omegaconf import DictConfig, flag_override
+from omegaconf import DictConfig, OmegaConf, flag_override
 from torch import nn
 from torch.nn.modules import Module
 from transformers import (
@@ -19,45 +19,86 @@ from transformers import (
 )
 from typing_extensions import override
 
-from fusion_bench.modelpool import BaseModelPool
-from fusion_bench.utils import instantiate
-from fusion_bench.utils.dtype import parse_dtype
+from fusion_bench import (
+    BaseModelPool,
+    auto_register_config,
+    import_object,
+    instantiate,
+    parse_dtype,
+)
 from fusion_bench.utils.lazy_state_dict import LazyStateDict
-from fusion_bench.utils.packages import import_object
 
 log = logging.getLogger(__name__)
 
 
+@auto_register_config
 class CausalLMPool(BaseModelPool):
-    _config_mapping = BaseModelPool._config_mapping | {
-        "_tokenizer": "tokenizer",
-        "_model_kwargs": "model_kwargs",
-        "load_lazy": "load_lazy",
-    }
+    """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,
         *,
-        tokenizer: Optional[DictConfig],
+        tokenizer: Optional[DictConfig | str],
         model_kwargs: Optional[DictConfig] = None,
-        load_lazy: bool = False,
+        enable_lazy_loading: bool = False,
         **kwargs,
     ):
         super().__init__(models, **kwargs)
-        # process `model_kwargs`
-        self._tokenizer = tokenizer
-        self._model_kwargs = model_kwargs
-        if self._model_kwargs is None:
-            self._model_kwargs = DictConfig({})
-        with flag_override(self._model_kwargs, "allow_objects", True):
-            if hasattr(self._model_kwargs, "torch_dtype"):
-                self._model_kwargs.torch_dtype = parse_dtype(
-                    self._model_kwargs.torch_dtype
-                )
-        self.load_lazy = load_lazy
+        if model_kwargs is None:
+            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
@@ -66,39 +107,80 @@ class CausalLMPool(BaseModelPool):
         else:
             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)
+            else self.model_kwargs
+        )
+        if "torch_dtype" in model_kwargs:
+            model_kwargs["torch_dtype"] = parse_dtype(model_kwargs["torch_dtype"])
+        return model_kwargs
+
     @override
     def load_model(
         self,
         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 = deepcopy(self._model_kwargs)
+        model_kwargs = self.get_model_kwargs()
         model_kwargs.update(kwargs)
 
         if isinstance(model_name_or_config, str):
@@ -108,7 +190,7 @@ class CausalLMPool(BaseModelPool):
                 model_config = self._models[model_name_or_config]
                 if isinstance(model_config, str):
                     # model_config is a string
-                    if not self.load_lazy:
+                    if not self.enable_lazy_loading:
                         model = AutoModelForCausalLM.from_pretrained(
                             model_config,
                             *args,
@@ -126,7 +208,7 @@ class CausalLMPool(BaseModelPool):
         elif isinstance(model_name_or_config, (DictConfig, Dict)):
             model_config = model_name_or_config
 
-        if not self.load_lazy:
+        if not self.enable_lazy_loading:
             model = instantiate(model_config, *args, **model_kwargs)
         else:
             meta_module_class = model_config.pop("_target_")
@@ -140,30 +222,43 @@ class CausalLMPool(BaseModelPool):
         return model
 
     def load_tokenizer(self, *args, **kwargs) -> PreTrainedTokenizer:
-        """
-        Example of YAML config:
+        """Load the tokenizer associated with this model pool.
 
-        ```yaml
-        tokenizer: google/gemma-2-2b-it # if a plain string, it will be passed to AutoTokenizer.from_pretrained
-        ```
+        Loads a tokenizer based on the tokenizer configuration provided during
+        pool initialization. Supports both simple string paths and detailed
+        configuration dictionaries.
 
-        or equivalently,
-
-        ```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"
+        assert self.tokenizer is not None, "Tokenizer is not defined in the config"
         log.info("Loading tokenizer.", stacklevel=2)
-        if isinstance(self._tokenizer, str):
-            tokenizer = AutoTokenizer.from_pretrained(self._tokenizer, *args, **kwargs)
+        if isinstance(self.tokenizer, str):
+            tokenizer = AutoTokenizer.from_pretrained(self.tokenizer, *args, **kwargs)
         else:
-            tokenizer = instantiate(self._tokenizer, *args, **kwargs)
+            tokenizer = instantiate(self.tokenizer, *args, **kwargs)
         return tokenizer
 
     @override
@@ -178,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`
@@ -210,15 +339,61 @@ 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:
-        if self.load_lazy:
+        """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. "
                 "Falling back to normal loading."
             )
-            self.load_lazy = False
+            self.enable_lazy_loading = False
         model: AutoModelForCausalLM = super().load_model(
             model_name_or_config, *args, **kwargs
         )
@@ -232,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/__init__.py

@@ -2,4 +2,9 @@
 from fusion_bench.utils import LazyStateDict
 
 from . import separate_io, utils
+from .hf_utils import (
+    create_default_model_card,
+    load_model_card_template,
+    save_pretrained_with_remote_code,
+)
 from .parameter_dict import ParameterDictModel