[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Author: pre-commit-ci[bot]

README.md

@@ -5,7 +5,7 @@ Data attribution methods estimate the effect on a behavior of interest of removi
 
 ## Core features
 
-Per-token and per-sequence attribution is available everywhere. On-disk gradient stores and on-the-fly queries are supported. Almost every feature is available through both the CLI and a programmatic interface, which use a shared set of configuration dataclasses. Configuration dataclasses are always serialized to disk so commands can be reproduced in one line. To understand every available configuration option, [check out the documentation](https://bergson.readthedocs.io/en/latest/api.html#bergson.IndexConfig). 
+Per-token and per-sequence attribution is available everywhere. On-disk gradient stores and on-the-fly queries are supported. Almost every feature is available through both the CLI and a programmatic interface, which use a shared set of configuration dataclasses. Configuration dataclasses are always serialized to disk so commands can be reproduced in one line. To understand every available configuration option, [check out the documentation](https://bergson.readthedocs.io/en/latest/api.html#bergson.IndexConfig).
 
 Bergson uses FSDP2 or SimpleFSDP, BitsAndBytes, and low-level performance optimizations to support large models, datasets, and clusters. Bergson integrates with HuggingFace Transformers and Datasets, and also supports on-disk datasets in a variety of formats.
 

bergson/approx_unrolling/approx_unrolling_math.py

@@ -9,6 +9,8 @@
 import torch
 from torch import Tensor
 
+from bergson.cli.commands import Score
+from bergson.cli.config_io import save_run_config
 from bergson.config import (
     ApproxUnrollingConfig,
     DistributedConfig,
@@ -17,8 +19,6 @@
     ScoreConfig,
 )
 from bergson.data import load_scores
-from bergson.cli.commands import Score
-from bergson.cli.config_io import save_run_config
 from bergson.distributed import init_dist, launch_distributed_run
 from bergson.hessians.apply_hessian import EkfacApplicator, EkfacConfig
 from bergson.score.score import score_dataset

bergson/cli/commands.py

@@ -1,8 +1,8 @@
 """CLI command definitions.
 
 Each command is a thin dataclass that validates, persists its ``config.yaml``,
-and dispatches. They live here rather than in ``__main__`` so pipelines can 
-lazily import lower-level commands for step serialization without importing 
+and dispatches. They live here rather than in ``__main__`` so pipelines can
+lazily import lower-level commands for step serialization without importing
 the corresponding CLI entrypoint.
 """
 
@@ -82,7 +82,10 @@ def execute(self):
                 "if skip_index is True HessianConfig.method must be provided"
             )
 
-        if self.hessian_cfg is not None and self.hessian_cfg.method != "autocorrelation":
+        if (
+            self.hessian_cfg is not None
+            and self.hessian_cfg.method != "autocorrelation"
+        ):
             raise ValueError(
                 f"build only supports autocorrelation Hessians, got "
                 f"'{self.hessian_cfg.method}'. Use the `hessian` command for "
@@ -224,7 +227,6 @@ def execute(self):
         score_dataset(self.index_cfg, self.score_cfg, self.preprocess_cfg)
 
 
-
 @dataclass
 class Trackstar(Serializable):
     """Run hessians, build, and score as a single pipeline."""

bergson/cli/config_io.py

@@ -1,16 +1,14 @@
 """IO for Bergson config files (run with ``bergson path_to.yaml``.)"""
 
-from typing import TypeVar, Protocol, cast
 import subprocess
 from datetime import datetime, timezone
 from importlib.metadata import PackageNotFoundError
 from importlib.metadata import version as _pkg_version
 from pathlib import Path
-from typing import Any
+from typing import Any, Protocol, TypeVar, cast
 
 import yaml
 
-
 CONFIG_FILENAME = "config.yaml"
 
 
@@ -58,10 +56,7 @@ def _write(
     run_path: str | Path | None = None,
 ) -> Path:
     """Write a ``{[run_path], steps, metadata}`` document, metadata last."""
-    doc = {
-        "steps": steps,
-        "metadata": make_metadata()
-    }
+    doc = {"steps": steps, "metadata": make_metadata()}
     if run_path is not None:
         doc["run_path"] = str(run_path)
 
@@ -71,9 +66,7 @@ def _write(
     return path
 
 
-def save_run_config(
-    command: Any, run_dir: str | Path, *, name: str | None = None
-):
+def save_run_config(command: Any, run_dir: str | Path, *, name: str | None = None):
     """Write a one-step component ``config.yaml`` for ``command`` into ``run_dir``.
 
     It can be run using ``bergson <run_dir>/config.yaml``.
@@ -164,6 +157,5 @@ def load_subconfig(
     sub = cmd_dict.get(field)
     if sub is None:
         return default
-    
-    return cast(T, config_cls.from_dict(sub))
 
+    return cast(T, config_cls.from_dict(sub))

bergson/cli/config_runner.py

@@ -30,14 +30,11 @@ def _parse_steps(
                 f"Unknown command '{cmd_name}'. "
                 f"Valid commands: {sorted(command_registry)}."
             ) from None
-        
+
         # Hydrate config
         parsed_step = cmd_cls.from_dict(cmd_dict or {}, drop_extra_fields=False)
 
-        parsed.append((
-            cmd_name, 
-            parsed_step
-        ))
+        parsed.append((cmd_name, parsed_step))
     return parsed
 
 
@@ -49,11 +46,11 @@ def run_config(config_path: str, command_registry: dict[str, type]) -> None:
     Each step also writes its own component ``config.yaml`` into its run directory.
     """
     doc = read_config(config_path)
-    
+
     # Top-level run path for a multi-step pipeline
     run_path = doc.get("run_path")
     steps = _parse_steps(doc["steps"], command_registry)
-    
+
     multi = len(steps) > 1
 
     if multi:
@@ -65,7 +62,7 @@ def run_config(config_path: str, command_registry: dict[str, type]) -> None:
                 "logging pipeline config to %s",
                 run_path,
             )
-        
+
         resolved_steps = [{name: cmd.to_dict()} for name, cmd in steps]
         save_pipeline_config(resolved_steps, run_path)
 

bergson/process_grads.py

@@ -49,8 +49,8 @@ def normalize_flat_grad(
 
 def assert_autocorrelation_hessian(path: Path) -> None:
     """Verify that ``path`` contains an autocorrelation hessian."""
-    hessian_cfg = load_subconfig(path, 'hessian', HessianConfig)
-    
+    hessian_cfg = load_subconfig(path, "hessian", HessianConfig)
+
     assert hessian_cfg is not None
     assert hessian_cfg.method == "autocorrelation", (
         f"Hessian at '{path}' was computed with method "

bergson/trackstar.py

@@ -136,7 +136,9 @@ def _validate(cfg: IndexConfig):
         trackstar_cfg.score_cfg.higher_is_better = True
         _validate(score_index_cfg)
         save_run_config(
-            Score(trackstar_cfg.score_cfg, score_index_cfg, trackstar_cfg.preprocess_cfg),
+            Score(
+                trackstar_cfg.score_cfg, score_index_cfg, trackstar_cfg.preprocess_cfg
+            ),
             score_index_cfg.partial_run_path,
         )
         score_dataset(

tests/ekfac_tests/test_fim_accuracy.py

@@ -150,9 +150,7 @@ def test_kfac_fim_accuracy(seq_lengths, num_batches, max_rel_error, sample, tmp_
         path=str(index_cfg.partial_run_path),
     )
 
-    hessian_cfg = HessianConfig(
-        method="autocorrelation", use_dataset_labels=not sample
-    )
+    hessian_cfg = HessianConfig(method="autocorrelation", use_dataset_labels=not sample)
 
     computer = CollectorComputer(
         model=model,