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

fusion_bench/models/hf_utils.py

@@ -0,0 +1,182 @@
+"""
+This module contains utilities for working with Hugging Face models.
+"""
+
+import inspect
+import os
+import shutil
+from typing import Optional, cast
+
+from omegaconf import OmegaConf
+from transformers.modeling_utils import PreTrainedModel
+
+from fusion_bench import BaseAlgorithm, BaseModelPool
+from fusion_bench.utils.pylogger import getRankZeroLogger
+
+log = getRankZeroLogger(__name__)
+
+__all__ = [
+    "save_pretrained_with_remote_code",
+    "generate_readme_head",
+    "generate_readme_body",
+    "generate_complete_readme",
+]
+
+
+def save_pretrained_with_remote_code(
+    model: PreTrainedModel,
+    auto_map: dict[str, object],
+    save_directory,
+    **kwargs,
+):
+    """
+    Saves a model with custom code to a directory.
+
+    This function facilitates saving a Hugging Face `PreTrainedModel` along with its
+    associated custom code. It inspects the objects provided in the `auto_map`,
+    copies their source files to the `save_directory`, and generates an `__init__.py`
+    to make them importable. It also updates the model's configuration with an
+    `auto_map` attribute, which allows `AutoModel.from_pretrained` to correctly
+    instantiate the custom model classes when `trust_remote_code=True`.
+
+    Args:
+        model (PreTrainedModel): The model instance to be saved.
+        auto_map (dict[str, object]): A dictionary mapping auto class names
+            (e.g., "AutoModelForCausalLM") to the corresponding custom class objects.
+        save_directory (str or os.PathLike): The directory where the model and
+            custom code files will be saved.
+        **kwargs: Additional keyword arguments to be passed to the
+            `model.save_pretrained` method.
+
+    Example:
+        ```python
+        # Assuming `model` is an instance of `SmileQwen2ForCausalLM`
+        # and `SmileQwen2Config`, `SmileQwen2Model`, `SmileQwen2ForCausalLM`
+        # are custom classes defined in your project.
+
+        save_pretrained_with_remote_code(
+            model,
+            auto_map={
+                "AutoConfig": SmileQwen2Config,
+                "AutoModel": SmileQwen2Model,
+                "AutoModelForCausalLM": SmileQwen2ForCausalLM,
+            },
+            save_directory="./my-custom-model",
+        )
+
+        # The model can then be loaded with `trust_remote_code=True`:
+        # from transformers import AutoModelForCausalLM
+        # loaded_model = AutoModelForCausalLM.from_pretrained(
+        #     "./my-custom-model", trust_remote_code=True
+        # )
+        ```
+    """
+    auto_map_files = {}
+    auto_map_strs = {}
+    for key, obj in auto_map.items():
+        auto_map_files[key] = inspect.getfile(obj)
+
+    for key, obj in auto_map.items():
+        auto_map_strs[key] = (
+            f"{(inspect.getmodule(obj).__name__).split('.')[-1]}.{obj.__name__}"
+        )
+
+    model.config.auto_map = auto_map_strs
+
+    # save model to `save_directory`
+    model.save_pretrained(save_directory=save_directory, **kwargs)
+
+    # copy source files to `save_directory`
+    for key, file_path in auto_map_files.items():
+        shutil.copy(
+            src=file_path, dst=os.path.join(save_directory, os.path.basename(file_path))
+        )
+    # construct `__init__.py`
+    init_file = os.path.join(save_directory, "__init__.py")
+    with open(init_file, "w") as f:
+        for key, file_name in auto_map_files.items():
+            base_name = os.path.basename(file_name).split(".")[0]
+            f.write(f"from .{base_name} import {auto_map[key].__name__}\n")
+
+
+def generate_readme_head(
+    models: list[str] | BaseModelPool,
+    library_name: str = "transformers",
+    tags: list[str] = ["fusion-bench", "merge"],
+):
+    text = "---\nbase_model:\n"
+    for model_name in models:
+        text += f"- {model_name}\n"
+    if library_name:
+        text += f"library_name: {library_name}\n"
+    text += "tags:\n"
+    for tag in tags:
+        text += f"- {tag}\n"
+    text += "---\n"
+    return text
+
+
+def generate_readme_body(
+    algorithm: BaseAlgorithm,
+    models_or_modelpool: Optional[list[str] | BaseModelPool] = None,
+    models: list[str] = None,
+):
+    text = """\
+# Merge
+
+This is a merge of pre-trained language models created using [fusion-bench](https://github.com/tanganke/fusion_bench).
+
+"""
+
+    if models is not None:
+        text += """
+## Models Merged
+
+The following models were included in the merge:
+
+"""
+        for model_name in models:
+            text += f"- {model_name}\n"
+        text += "\n"
+
+        try:
+            text += f"""\
+    ## Configuration
+
+    The following YAML configuration was used to produce this model:
+
+    ```yaml
+    {OmegaConf.to_yaml(algorithm.config, resolve=True, sort_keys=True)}
+    ```
+    """
+        except Exception as e:
+            return (
+                text  # If the algorithm config cannot be converted to YAML, we skip it.
+            )
+
+    if isinstance(models_or_modelpool, BaseModelPool):
+        try:
+            text += f"""
+```yaml
+{OmegaConf.to_yaml(models_or_modelpool.config, resolve=True, sort_keys=True)}
+```
+"""
+        except Exception as e:
+            pass  # If the model pool config cannot be converted to YAML, we skip it.
+    return text
+
+
+def generate_complete_readme(
+    algorithm: BaseAlgorithm, modelpool: BaseModelPool, models: list[str]
+):
+    # Generate the complete README content
+    text = generate_readme_head(
+        [modelpool.get_model_path(m) for m in modelpool.model_names]
+    )
+    readme_body = generate_readme_body(
+        algorithm,
+        models_or_modelpool=modelpool,
+        models=[modelpool.get_model_path(m) for m in modelpool.model_names],
+    )
+    complete_readme = text + "\n" + readme_body
+    return complete_readme

fusion_bench/models/linearized/linearized_model_utils.py

@@ -1,7 +1,7 @@
 import logging
 from collections import OrderedDict
 from copy import deepcopy
-from typing import Optional
+from typing import Any, Dict, Optional, Tuple
 
 import torch.nn as nn
 from torch.func import functional_call, jvp
@@ -9,7 +9,7 @@ from torch.func import functional_call, jvp
 log = logging.getLogger(__name__)
 
 
-def dict_params_to_tuple(dict_params: dict):
+def dict_params_to_tuple(dict_params: dict) -> Tuple:
     return tuple(v for k, v in dict_params.items())
 
 
@@ -33,7 +33,7 @@ class LinearizedModelWraper(nn.Module):
         for p in self.params0_values:
             p.requires_grad_(False)
 
-    def tuple_params_to_dict(self, tuple_params):
+    def tuple_params_to_dict(self, tuple_params) -> Dict[str, Any]:
         """
         Converts a tuple of parameters to a dictionary with keys corresponding to the parameter names.
 
@@ -50,7 +50,7 @@ class LinearizedModelWraper(nn.Module):
             state_dict[k] = p
         return state_dict
 
-    def forward(self, *args, **kwargs):
+    def forward(self, *args: Any, **kwargs: Any) -> Any:
         """
         Computes the linearized model output using a first-order Taylor decomposition.
 

fusion_bench/models/linearized/vision_model.py

@@ -70,7 +70,7 @@ def load_lora_vision_model_hf(
     peft_name: str,
     merge_and_unload: bool = False,
     return_vison_model=True,
-):
+) -> PeftModel:
     """
     Load a LoRA (Low-Rank Adaptation) vision model from Hugging Face.
 

fusion_bench/models/modeling_deepseek_v2/__init__.py

@@ -4,12 +4,12 @@ This is a direct copy of the DeepSeek-V2-Lite model from HuggingFace https://hug
 
 from .configuration_deepseek import DeepseekV2Config
 from .modeling_deepseek import (
+    DeepseekV2DecoderLayer,
     DeepseekV2ForCausalLM,
     DeepseekV2ForSequenceClassification,
     DeepseekV2MLP,
     DeepseekV2Model,
     DeepseekV2MoE,
-    DeepseekV2DecoderLayer,
 )
 from .modeling_deepseek import MoEGate as DeepseekV2MoEGate
 from .tokenization_deepseek_fast import DeepseekTokenizerFast

fusion_bench/models/modeling_deepseek_v2/modeling_deepseek.py

@@ -17,17 +17,18 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-""" PyTorch DeepSeek model."""
+"""PyTorch DeepSeek model."""
 import math
 import warnings
 from typing import List, Optional, Tuple, Union
 
+import numpy as np
 import torch
+import torch.distributed as dist
 import torch.nn.functional as F
 import torch.utils.checkpoint
 from torch import nn
 from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
-
 from transformers.activations import ACT2FN
 from transformers.cache_utils import Cache, DynamicCache
 from transformers.modeling_attn_mask_utils import (
@@ -54,9 +55,8 @@ from transformers.utils import (
     replace_return_docstrings,
 )
 from transformers.utils.import_utils import is_torch_fx_available
+
 from .configuration_deepseek import DeepseekV2Config
-import torch.distributed as dist
-import numpy as np
 
 if is_flash_attn_2_available():
     from flash_attn import flash_attn_func, flash_attn_varlen_func

fusion_bench/models/modeling_deepseek_v2/tokenization_deepseek_fast.py

@@ -1,6 +1,5 @@
 from typing import List, Optional, Union
 
-
 from transformers.models.llama import LlamaTokenizerFast
 
 

fusion_bench/models/modeling_smile_gemma2/__init__.py

@@ -0,0 +1,9 @@
+from . import register
+from .configuration_smile_gemma2 import SmileGemma2Config
+from .modeling_smile_gemma2 import (
+    SmileGemma2ForCausalLM,
+    SmileGemma2ForSequenceClassification,
+    SmileGemma2ForTokenClassification,
+    SmileGemma2Model,
+    SmileGemma2PreTrainedModel,
+)

fusion_bench/models/modeling_smile_gemma2/configuration_smile_gemma2.py

@@ -0,0 +1,20 @@
+from transformers.models.gemma2.configuration_gemma2 import Gemma2Config
+
+
+class SmileGemma2Config(Gemma2Config):
+    model_type = "smile_gemma2"
+
+    def __init__(
+        self,
+        num_experts_per_tok: int = 1,
+        rank_of_router: int = None,
+        rank_of_expert: int = None,
+        num_local_experts: int = None,
+        **kwargs,
+    ):
+        self.num_experts_per_tok = num_experts_per_tok
+        self.rank_of_router = rank_of_router
+        self.rank_of_expert = rank_of_expert
+        self.num_local_experts = num_local_experts
+
+        super().__init__(**kwargs)