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

fusion_bench/utils/json.py

@@ -1,31 +1,72 @@
 import json
 from pathlib import Path
-from typing import Any, Union
+from typing import TYPE_CHECKING, Any, Union
 
+if TYPE_CHECKING:
+    from pyarrow.fs import FileSystem
 
-def save_to_json(obj, path: Union[str, Path]):
+
+def save_to_json(obj, path: Union[str, Path], filesystem: "FileSystem" = None):
     """
     save an object to a json file
 
     Args:
         obj (Any): the object to save
         path (Union[str, Path]): the path to save the object
+        filesystem (FileSystem, optional): PyArrow FileSystem to use for writing.
+            If None, uses local filesystem via standard Python open().
+            Can also be an s3fs.S3FileSystem or fsspec filesystem.
     """
-    with open(path, "w") as f:
-        json.dump(obj, f)
+    if filesystem is not None:
+        json_str = json.dumps(obj)
+        # Check if it's an fsspec-based filesystem (like s3fs)
+        if hasattr(filesystem, "open"):
+            # Direct fsspec/s3fs usage - more reliable for some endpoints
+            path_str = str(path)
+            with filesystem.open(path_str, "w") as f:
+                f.write(json_str)
+        else:
+            # Use PyArrow filesystem
+            path_str = str(path)
+            with filesystem.open_output_stream(path_str) as f:
+                f.write(json_str.encode("utf-8"))
+    else:
+        # Use standard Python file operations
+        with open(path, "w") as f:
+            json.dump(obj, f)
 
 
-def load_from_json(path: Union[str, Path]) -> Union[dict, list]:
+def load_from_json(
+    path: Union[str, Path], filesystem: "FileSystem" = None
+) -> Union[dict, list]:
     """load an object from a json file
 
     Args:
         path (Union[str, Path]): the path to load the object
+        filesystem (FileSystem, optional): PyArrow FileSystem to use for reading.
+            If None, uses local filesystem via standard Python open().
+            Can also be an s3fs.S3FileSystem or fsspec filesystem.
 
     Returns:
-        dict: the loaded object
+        Union[dict, list]: the loaded object
     """
-    with open(path, "r") as f:
-        return json.load(f)
+    if filesystem is not None:
+        # Check if it's an fsspec-based filesystem (like s3fs)
+        if hasattr(filesystem, "open"):
+            # Direct fsspec/s3fs usage
+            path_str = str(path)
+            with filesystem.open(path_str, "r") as f:
+                return json.load(f)
+        else:
+            # Use PyArrow filesystem
+            path_str = str(path)
+            with filesystem.open_input_stream(path_str) as f:
+                json_data = f.read().decode("utf-8")
+                return json.loads(json_data)
+    else:
+        # Use standard Python file operations
+        with open(path, "r") as f:
+            return json.load(f)
 
 
 def _is_list_of_dict(obj) -> bool:

fusion_bench/utils/state_dict_arithmetic.py

@@ -6,10 +6,13 @@ import torch
 from torch import Tensor
 from tqdm.auto import tqdm
 
+from fusion_bench.utils.type import TorchModelType
+
 from .type import BoolStateDictType, StateDictType
 
 __all__ = [
     "ArithmeticStateDict",
+    "load_state_dict_with_prefix",
     "state_dicts_check_keys",
     "state_dict_to_device",
     "num_params_of_state_dict",
@@ -646,6 +649,48 @@ def _validate_list_lengths_equal(
             pass
 
 
+def load_state_dict_with_prefix(
+    model: TorchModelType,
+    state_dict: StateDictType,
+    strict: bool = True,
+    assign: bool = False,
+    key_prefix: str = "model.",
+    operation: Literal["add", "remove"] = "remove",
+) -> TorchModelType:
+    """
+    Load a state dict into a model, adding or removing a prefix from the keys.
+
+    This is useful when loading state dicts saved with DataParallel, pytorch lightning or similar wrappers.
+
+    Args:
+        model: The model to load the state dict into.
+        state_dict: The state dictionary to load.
+        key_prefix: The prefix to add or remove from the keys.
+        operation: 'add' to add the prefix, 'remove' to remove it.
+
+    Returns:
+        The model with the loaded state dict.
+    """
+    if operation not in ("add", "remove"):
+        raise ValueError("operation must be either 'add' or 'remove'")
+
+    modified_state_dict = OrderedDict()
+    for key, value in state_dict.items():
+        if operation == "add":
+            new_key = f"{key_prefix}{key}"
+        else:  # operation == "remove"
+            if key.startswith(key_prefix):
+                new_key = key[len(key_prefix) :]
+            else:
+                raise ValueError(
+                    f"Key '{key}' does not start with prefix '{key_prefix}'"
+                )
+        modified_state_dict[new_key] = value
+
+    model.load_state_dict(modified_state_dict, strict=strict, assign=assign)
+    return model
+
+
 def state_dict_to_device(
     state_dict: StateDictType,
     device: Union[torch.device, str],
@@ -851,22 +896,48 @@ def state_dict_add_scalar(state_dict: StateDictType, scalar: Number) -> StateDic
     return OrderedDict((key, tensor + scalar) for key, tensor in state_dict.items())
 
 
-def state_dict_mul(state_dict: StateDictType, scalar: float) -> StateDictType:
+def state_dict_mul(
+    state_dict: StateDictType,
+    scalar: float,
+    *,
+    keep_dtype_when_zero: bool = True,
+    show_pbar: bool = False,
+) -> StateDictType:
     """
     Multiply all parameters in a state dict by a scalar.
 
     Args:
         state_dict: The state dict to multiply.
-        scalar: The scalar value to multiply each parameter by.
+        scalar (float): The scalar value to multiply each parameter by.
+        keep_dtype_when_zero (bool): Whether to keep the original data type of the tensors if either the tensor is all zeros or the scalar is zero.
+        show_pbar (bool): Whether to show a progress bar during computation.
 
     Returns:
         A new state dict with each parameter multiplied by the scalar.
     """
-    return OrderedDict((key, scalar * tensor) for key, tensor in state_dict.items())
+    new_state_dict = OrderedDict()
+    for key, tensor in (
+        state_dict.items()
+        if not show_pbar
+        else tqdm(state_dict.items(), desc="Multiplying state dict")
+    ):
+        if (
+            keep_dtype_when_zero
+            and not tensor.is_floating_point()  # when tensor is not floating point, multiplication by 0 keeps dtype
+            and (scalar == 0 or torch.all(tensor == 0))
+        ):
+            new_state_dict[key] = tensor.clone()
+        else:
+            new_state_dict[key] = scalar * tensor
+    return new_state_dict
 
 
 def state_dict_div(
-    state_dict: StateDictType, scalar: float, show_pbar: bool = False
+    state_dict: StateDictType,
+    scalar: float,
+    *,
+    keep_dtype_when_zero: bool = True,
+    show_pbar: bool = False,
 ) -> StateDictType:
     """
     Divide all parameters in a state dict by a scalar.
@@ -874,6 +945,7 @@ def state_dict_div(
     Args:
         state_dict: The state dict to divide.
         scalar: The scalar value to divide each parameter by.
+        keep_dtype_when_zero: Whether to keep the original data type of the tensors if the tensor is all zeros.
         show_pbar: Whether to show a progress bar during computation.
 
     Returns:
@@ -885,12 +957,21 @@ def state_dict_div(
     if scalar == 0:
         raise ZeroDivisionError("Cannot divide state dict by zero")
 
-    keys_iter = (
-        tqdm(state_dict.keys(), desc="Dividing state dict")
-        if show_pbar
-        else state_dict.keys()
-    )
-    return OrderedDict((key, state_dict[key] / scalar) for key in keys_iter)
+    new_state_dict = OrderedDict()
+    for key, tensor in (
+        state_dict.items()
+        if not show_pbar
+        else tqdm(state_dict.items(), desc="Dividing state dict")
+    ):
+        if (
+            keep_dtype_when_zero
+            and not tensor.is_floating_point()  # when tensor is not floating point, division by any scalar keeps dtype
+            and torch.all(tensor == 0)  # only check tensor for zero
+        ):
+            new_state_dict[key] = tensor.clone()
+        else:
+            new_state_dict[key] = tensor / scalar
+    return new_state_dict
 
 
 def state_dict_power(state_dict: StateDictType, p: float) -> StateDictType:

{fusion_bench-0.2.26.dist-info → fusion_bench-0.2.28.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
-Name: fusion_bench
-Version: 0.2.26
+Name: fusion-bench
+Version: 0.2.28
 Summary: A Comprehensive Benchmark of Deep Model Fusion
 Author-email: Anke Tang <tang.anke@foxmail.com>
 Project-URL: Repository, https://github.com/tanganke/fusion_bench