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

fusion_bench/utils/parameters.py

@@ -129,7 +129,6 @@ def human_readable(num: int) -> str:
     Converts a number into a human-readable string with appropriate magnitude suffix.
 
     Examples:
-
         ```python
         print(human_readable(1500))
         # Output: '1.50K'
@@ -201,7 +200,6 @@ def count_parameters(module: nn.Module, non_zero_only: bool = False) -> tuple[in
         tuple: A tuple containing the number of trainable parameters and the total number of parameters.
 
     Examples:
-
         ```python
         # Count the parameters
         trainable_params, all_params = count_parameters(model)

fusion_bench/utils/rich_utils.py

@@ -188,17 +188,21 @@ if __name__ == "__main__":
     display_available_styles()
 
 
-def setup_colorlogging(force=False, **config_kwargs):
+def setup_colorlogging(
+    force=False,
+    level=logging.INFO,
+    **kwargs,
+):
     """
     Sets up color logging for the application.
     """
     FORMAT = "%(message)s"
 
     logging.basicConfig(
-        level=logging.INFO,
+        level=level,
         format=FORMAT,
         datefmt="[%X]",
         handlers=[RichHandler()],
         force=force,
-        **config_kwargs,
+        **kwargs,
     )

fusion_bench/utils/timer.py

@@ -6,38 +6,120 @@ log = logging.getLogger(__name__)
 
 class timeit_context:
     """
-    Usage:
+    A context manager for measuring and logging execution time of code blocks.
 
-    ```python
-    with timeit_context() as timer:
-        ... # code block to be measured
-    ```
+    This context manager provides precise timing measurements with automatic logging
+    of elapsed time. It supports nested timing contexts with proper indentation
+    for hierarchical timing analysis, making it ideal for profiling complex
+    operations with multiple sub-components.
+
+    Args:
+        msg (str, optional): Custom message to identify the timed code block.
+            If provided, logs "[BEGIN] {msg}" at start and includes context
+            in the final timing report. Defaults to None.
+        loglevel (int, optional): Python logging level for output messages.
+            Uses standard logging levels (DEBUG=10, INFO=20, WARNING=30, etc.).
+            Defaults to logging.INFO.
+
+    Example:
+        Basic usage:
+        ```python
+        with timeit_context("data loading"):
+            data = load_large_dataset()
+        # Logs: [BEGIN] data loading
+        # Logs: [END]   Elapsed time: 2.34s
+        ```
+
+        Nested timing:
+        ```python
+        with timeit_context("model training"):
+            with timeit_context("data preprocessing"):
+                preprocess_data()
+            with timeit_context("forward pass"):
+                model(data)
+        # Output shows nested structure:
+        # [BEGIN] model training
+        #   [BEGIN] data preprocessing
+        #   [END]   Elapsed time: 0.15s
+        #   [BEGIN] forward pass
+        #   [END]   Elapsed time: 0.89s
+        # [END]   Elapsed time: 1.04s
+        ```
+
+        Custom log level:
+        ```python
+        with timeit_context("debug operation", loglevel=logging.DEBUG):
+            debug_function()
+        ```
     """
 
     nest_level = -1
 
     def _log(self, msg):
+        """
+        Internal method for logging messages with appropriate stack level.
+
+        This helper method ensures that log messages appear to originate from
+        the caller's code rather than from internal timer methods, providing
+        more useful debugging information.
+
+        Args:
+            msg (str): The message to log at the configured log level.
+        """
         log.log(self.loglevel, msg, stacklevel=3)
 
     def __init__(self, msg: str = None, loglevel=logging.INFO) -> None:
+        """
+        Initialize a new timing context with optional message and log level.
+
+        Args:
+            msg (str, optional): Descriptive message for the timed operation.
+                If provided, will be included in the begin/end log messages
+                to help identify what is being timed. Defaults to None.
+            loglevel (int, optional): Python logging level for timer output.
+                Common values include:
+                - logging.DEBUG (10): Detailed debugging information
+                - logging.INFO (20): General information (default)
+                - logging.WARNING (30): Warning messages
+                - logging.ERROR (40): Error messages
+                Defaults to logging.INFO.
+        """
         self.loglevel = loglevel
         self.msg = msg
 
     def __enter__(self) -> None:
         """
-        Sets the start time and logs an optional message indicating the start of the code block execution.
+        Enter the timing context and start the timer.
 
-        Args:
-            msg: str, optional message to log
+        This method is automatically called when entering the 'with' statement.
+        It records the current timestamp, increments the nesting level for
+        proper log indentation, and optionally logs a begin message.
+
+        Returns:
+            None: This context manager doesn't return a value to the 'as' clause.
+                  All timing information is handled internally and logged automatically.
         """
         self.start_time = time.time()
         timeit_context.nest_level += 1
         if self.msg is not None:
             self._log("  " * timeit_context.nest_level + "[BEGIN] " + str(self.msg))
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
         """
-        Calculates the elapsed time and logs it, along with an optional message indicating the end of the code block execution.
+        Exit the timing context and log the elapsed time.
+
+        This method is automatically called when exiting the 'with' statement,
+        whether through normal completion or exception. It calculates the total
+        elapsed time and logs the results with proper nesting indentation.
+
+        Args:
+            exc_type (type): Exception type if an exception occurred, None otherwise.
+            exc_val (Exception): Exception instance if an exception occurred, None otherwise.
+            exc_tb (traceback): Exception traceback if an exception occurred, None otherwise.
+
+        Returns:
+            None: Does not suppress exceptions (returns None/False implicitly).
+                  Any exceptions that occurred in the timed block will propagate normally.
         """
         end_time = time.time()
         elapsed_time = end_time - self.start_time

{fusion_bench-0.2.22.dist-info → fusion_bench-0.2.24.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: fusion_bench
-Version: 0.2.22
+Version: 0.2.24
 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
@@ -23,12 +23,19 @@ Requires-Dist: rich
 Requires-Dist: scipy
 Requires-Dist: h5py
 Requires-Dist: pytest
+Requires-Dist: joblib
+Requires-Dist: bidict
 Requires-Dist: transformers!=4.49
 Requires-Dist: pillow!=11.2.1
 Provides-Extra: lm-eval-harness
 Requires-Dist: lm-eval; extra == "lm-eval-harness"
 Requires-Dist: immutabledict; extra == "lm-eval-harness"
 Requires-Dist: langdetect; extra == "lm-eval-harness"
+Requires-Dist: rich-run; extra == "lm-eval-harness"
+Provides-Extra: docs
+Requires-Dist: mkdocs; extra == "docs"
+Requires-Dist: mkdocs-material; extra == "docs"
+Requires-Dist: mkdocstrings[python]; extra == "docs"
 Dynamic: license-file
 
 <div align='center'>
@@ -151,7 +158,7 @@ This will install the latest version of fusion-bench and the dependencies requir
 Documentation for using LM-Eval Harness within FusionBench framework can be found at [this online documentation](https://tanganke.github.io/fusion_bench/taskpool/lm_eval_harness) or in the [`docs/taskpool/lm_eval_harness.md`](docs/taskpool/lm_eval_harness.md) markdown file.
 
 > [!TIP]
-> Documentation for merging large language models using FusionBench can be found at [this online documentation](https://tanganke.github.io/fusion_bench/modelpool/causal_lm) or in the [`docs/modelpool/causal_lm.md`](docs/modelpool/causal_lm.md) markdown file.
+> Documentation for merging large language models using FusionBench can be found at [this online documentation](https://tanganke.github.io/fusion_bench/modelpool/llm) or in the [`docs/modelpool/llm/index.md`](docs/modelpool/llm/index.md) markdown file.
 
 ## Introduction to Deep Model Fusion
 
@@ -179,7 +186,7 @@ The project is structured as follows:
   - `taskpool`: configuration files for the task pool.
   - `model`: configuration files for the models.
   - `dataset`: configuration files for the datasets.
-- `docs/`: documentation for the benchmark. We use [mkdocs](https://www.mkdocs.org/) to generate the documentation. Start the documentation server locally with `mkdocs serve`. The required packages can be installed with `pip install -r mkdocs-requirements.txt`.
+- `docs/`: documentation for the benchmark. We use [mkdocs](https://www.mkdocs.org/) to generate the documentation. Start the documentation server locally with `mkdocs serve`. The required packages can be installed with `pip install -e ".[docs]"`.
 - `examples/`: example scripts for running some of the experiments.
   > **naming convention**: `examples/{method_name}/` contains the files such as bash scripts and jupyter notebooks for the specific method.
 - `tests/`: unit tests for the benchmark.