bergson 2.2.4__tar.gz → 2.3.1__tar.gz

{bergson-2.2.4 → bergson-2.3.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bergson
-Version: 2.2.4
+Version: 2.3.1
 Summary: Tracing the memory of neural nets with data attribution
 License: MIT License
 Keywords: interpretability,explainable-ai

{bergson-2.2.4 → bergson-2.3.1}/bergson/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "2.2.4"
+__version__ = "2.3.1"
 
 import logging
 

{bergson-2.2.4 → bergson-2.3.1}/bergson/config/config.py

@@ -52,8 +52,13 @@ class DataConfig(Serializable):
     When specified, gradients are calculated using the policy
     gradient loss from Dr. GRPO. https://arxiv.org/abs/2503.20783"""
 
+    advantage_column: str = ""
+    """Optional column in the dataset that contains per-example advantages, e.g.
+    those logged by a GRPO trainer. When specified, gradients are calculated using
+    the Dr. GRPO policy gradient loss."""
+
     skip_nan_rewards: bool = False
-    """Whether to skip examples with NaN rewards."""
+    """Whether to skip examples with NaN rewards or advantages."""
 
     truncation: bool = False
     """Whether to truncate long documents to fit the token budget."""
@@ -73,6 +78,10 @@ class DataConfig(Serializable):
     sequences of this length. Incompatible with truncation and format_template."""
 
     def __post_init__(self):
+        if self.reward_column and self.advantage_column:
+            raise ValueError(
+                "reward_column and advantage_column cannot both be specified"
+            )
         if self.chunk_length > 0:
             if self.truncation:
                 raise ValueError("chunk_length and truncation cannot both be True")

{bergson-2.2.4 → bergson-2.3.1}/bergson/data.py

@@ -18,6 +18,7 @@ from datasets import (
     Dataset,
     DatasetDict,
     IterableDatasetDict,
+    ReadInstruction,
     concatenate_datasets,
     load_dataset,
     load_from_disk,
@@ -452,6 +453,7 @@ def load_data_string(
     data_kwargs: str = "",
 ) -> Dataset:
     """Load a dataset from a string identifier or path."""
+    from_disk = False
     if data_str.endswith(".csv"):
         ds = load_dataset("csv", data_files=data_str, split=split)
     elif data_str.endswith(".json") or data_str.endswith(".jsonl"):
@@ -461,6 +463,7 @@ def load_data_string(
         or (Path(data_str) / "dataset_dict.json").exists()
     ):
         ds = load_from_disk(data_str, keep_in_memory=False)
+        from_disk = True
     else:
         try:
             kwargs = simple_parse_kwargs_string(data_kwargs)
@@ -474,21 +477,40 @@ def load_data_string(
             # Automatically use load_from_disk if appropriate
             if "load_from_disk" in str(e):
                 ds = load_from_disk(data_str, keep_in_memory=False)
+                from_disk = True
             else:
                 raise e
 
-    if isinstance(ds, DatasetDict):
-        if "[" in split or "+" in split:
-            raise NotImplementedError(
-                f"Split slicing/concatenation expressions are not supported "
-                f"for load_from_disk paths: {split!r}"
-            )
-        ds = ds[split]
+    if from_disk:
+        ds = select_split(ds, split)
 
     ds = assert_type(Dataset, ds)
     return ds
 
 
+def select_split(ds: Dataset | DatasetDict, split: str) -> Dataset:
+    """The rows ``split`` names, in ``load_dataset``'s split syntax (such as
+    ``train[:100]`` or ``train[:10%]+test``), of a dataset from ``load_from_disk``.
+
+    A single saved dataset has no split names, so each part of ``split`` slices
+    the whole of it.
+    """
+    if "[" not in split and "+" not in split:
+        return ds[split] if isinstance(ds, DatasetDict) else ds
+
+    names = [part.split("[")[0].strip() for part in split.split("+")]
+    lengths = {
+        name: len(ds[name] if isinstance(ds, DatasetDict) else ds) for name in names
+    }
+    parts = [
+        (ds[i.splitname] if isinstance(ds, DatasetDict) else ds).select(
+            range(i.from_, i.to)
+        )
+        for i in ReadInstruction.from_spec(split).to_absolute(lengths)
+    ]
+    return parts[0] if len(parts) == 1 else concatenate_datasets(parts)
+
+
 def load_gradients(root_dir: Path | str) -> np.memmap:
     """Map the gradients stored in `root_dir` into memory as a flat
     ``(num_grads, total_grad_dim)`` array, with modules laid out in

{bergson-2.2.4 → bergson-2.3.1}/bergson/hessians/pipeline.py

@@ -17,7 +17,7 @@ from ..config.config_io import save_run_config
 from ..data import create_index, load_gradients
 from ..distributed import launch_distributed_run
 from ..score.score import score_dataset
-from ..utils.step_state import partial_path, prepare_step
+from ..utils.step_state import partial_path, prepare_step, promote_step
 from ..utils.worker_utils import validate_run_path
 from .apply_hessian import EkfacConfig, apply_worker
 from .astra import AstraPaths, astra_worker
@@ -139,7 +139,7 @@ def hessian_pipeline(
             index_cfg.distributed,
         )
         if index_cfg.distributed.rank == 0:
-            shutil.move(ekfac_cfg.run_path, transformed_query_path)
+            promote_step(transformed_query_path)
 
     astra_cfg = hessian_pipeline_cfg.astra
     if astra_cfg.num_steps > 0:
@@ -217,4 +217,4 @@ def _run_astra(
         ],
         index_cfg.distributed,
     )
-    shutil.move(part, out_path)
+    promote_step(out_path)

{bergson-2.2.4 → bergson-2.3.1}/bergson/utils/worker_utils.py

@@ -347,7 +347,9 @@ def filter_by_max_tokens(ds: Dataset, cfg: AttributionConfig) -> Dataset:
     return ds
 
 
-def max_tokens_for_model(tokenizer, model_str: str, revision: str | None) -> int:
+def max_tokens_for_model(
+    tokenizer, model_str: str, revision: str | None, trust_remote_code: bool = False
+) -> int:
     # You might think model_max_length should always be the same as
     # max_position_embeddings, but some models (e.g. Pythia!) have a smaller
     # max_position_embeddings than model_max_length, so we need to check both.
@@ -360,7 +362,9 @@ def max_tokens_for_model(tokenizer, model_str: str, revision: str | None) -> int
     except ValueError:
         pass
 
-    model_cfg = AutoConfig.from_pretrained(model_str, revision=revision)
+    model_cfg = AutoConfig.from_pretrained(
+        model_str, revision=revision, trust_remote_code=trust_remote_code
+    )
     model_max_length = getattr(tokenizer, "model_max_length", BIG_NUM)
     max_pos_emb = getattr(model_cfg, "max_position_embeddings", BIG_NUM)
     return min(model_max_length, max_pos_emb)
@@ -377,7 +381,12 @@ def setup_data_pipeline(
         data_cfg.dataset, data_cfg.split, data_cfg.subset, data_cfg.data_kwargs
     )
     tokenizer = AutoTokenizer.from_pretrained(cfg.tokenizer or cfg.model)
-    max_model_length = max_tokens_for_model(tokenizer, cfg.model, cfg.revision)
+    trust_remote_code = bool(
+        simple_parse_kwargs_string(cfg.model_kwargs).get("trust_remote_code", False)
+    )
+    max_model_length = max_tokens_for_model(
+        tokenizer, cfg.model, cfg.revision, trust_remote_code
+    )
 
     if data_cfg.chunk_length > 0:
         # Sanity check
@@ -449,23 +458,28 @@ def setup_data_pipeline(
                 f"Consider increasing --token_batch_size or using --truncation."
             )
 
-    if data_cfg.reward_column:
+    value_column = data_cfg.reward_column or data_cfg.advantage_column
+    if value_column:
         assert isinstance(ds, Dataset), "Dataset required for advantage estimation"
 
-        rewards = np.array(ds[data_cfg.reward_column], dtype=np.float64)
-        nan_mask = np.isnan(rewards)
+        values = np.array(ds[value_column], dtype=np.float64)
+        nan_mask = np.isnan(values)
         if nan_mask.any():
             if data_cfg.skip_nan_rewards:
-                print(f"Warning: Filtering out {nan_mask.sum()} rows with NaN rewards")
+                print(f"Warning: Filtering out {nan_mask.sum()} rows with NaN values")
                 ds = ds.filter(lambda _, idx: not nan_mask[idx], with_indices=True)
             else:
-                raise ValueError(
-                    f"Reward column '{data_cfg.reward_column}' contains NaN values"
-                )
+                raise ValueError(f"Column '{value_column}' contains NaN values")
 
+        if data_cfg.reward_column:
+            advantages = estimate_advantage(ds, data_cfg)
+        else:
+            advantages = ds[data_cfg.advantage_column][:]
+        if "advantage" in ds.column_names:
+            ds = ds.remove_columns("advantage")
         ds = ds.add_column(
             "advantage",
-            estimate_advantage(ds, data_cfg),
+            advantages,
             new_fingerprint="advantage",  # type: ignore
         )
 
@@ -475,6 +489,8 @@ def setup_data_pipeline(
 
     # Remove extraneous columns
     keep = {"length", "input_ids", "labels"}
+    if value_column:
+        keep.add("advantage")
     remove_columns -= keep
     remove_columns &= set(ds.column_names)
     if remove_columns:

{bergson-2.2.4 → bergson-2.3.1}/bergson.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bergson
-Version: 2.2.4
+Version: 2.3.1
 Summary: Tracing the memory of neural nets with data attribution
 License: MIT License
 Keywords: interpretability,explainable-ai

{bergson-2.2.4 → bergson-2.3.1}/bergson.egg-info/SOURCES.txt

@@ -159,6 +159,7 @@ tests/test_save_mode_final.py
 tests/test_score.py
 tests/test_score_plot.py
 tests/test_score_writer_dist.py
+tests/test_select_split.py
 tests/test_shard_load.py
 tests/test_source_fisher_normalization.py
 tests/test_source_optimizer_variants.py

{bergson-2.2.4 → bergson-2.3.1}/pyproject.toml

@@ -37,7 +37,7 @@ dependencies = [
     "bitsandbytes",
     "huggingface-hub>=1.13.0",
 ]
-version = "2.2.4"
+version = "2.3.1"
 [project.optional-dependencies]
 dev = [
     "pre-commit",

{bergson-2.2.4 → bergson-2.3.1}/tests/test_advantages.py

@@ -1,5 +1,6 @@
 import json
 
+import pytest
 from datasets import Dataset
 
 from bergson.config import DataConfig, IndexConfig
@@ -83,3 +84,57 @@ def test_advantages_computed_with_drop_columns(tmp_path):
     # "foo bar" group mean = 15.0: advantages = [-5.0, 5.0]
     assert abs(advantages[2] - (-5.0)) < 1e-9
     assert abs(advantages[3] - 5.0) < 1e-9
+
+
+def _pipeline_advantages(tmp_path, rows, **data_kwargs):
+    path = tmp_path / "rows.json"
+    path.write_text("\n".join(json.dumps(r) for r in rows))
+    cfg = IndexConfig(
+        run_path=str(tmp_path / "run"),
+        model="EleutherAI/pythia-14m",
+        token_batch_size=2048,
+        data=DataConfig(dataset=str(path), **data_kwargs),
+    )
+    ds, _ = setup_data_pipeline(cfg)
+    return ds["advantage"]
+
+
+def test_advantage_column_used_as_given(tmp_path):
+    rows = [
+        {"text": "hello world", "adv": 0.5},
+        {"text": "hello world", "adv": 2.0},
+        {"text": "foo bar", "adv": -1.5},
+    ]
+    assert _pipeline_advantages(tmp_path, rows, advantage_column="adv") == [
+        0.5,
+        2.0,
+        -1.5,
+    ]
+
+
+def test_existing_advantage_column(tmp_path):
+    """A dataset column named "advantage" can be the advantage column, and is
+    replaced by the estimate when rewards are given."""
+    rows = [
+        {"text": "hello world", "advantage": 0.5, "reward": 1.0},
+        {"text": "hello world", "advantage": 2.0, "reward": 3.0},
+    ]
+    assert _pipeline_advantages(tmp_path, rows, advantage_column="advantage") == [
+        0.5,
+        2.0,
+    ]
+    assert _pipeline_advantages(tmp_path, rows, reward_column="reward") == [-1.0, 1.0]
+
+
+def test_nan_advantages(tmp_path):
+    rows = [{"text": "hello world", "adv": 1.0}, {"text": "foo bar", "adv": None}]
+    with pytest.raises(ValueError, match="Column 'adv' contains NaN"):
+        _pipeline_advantages(tmp_path, rows, advantage_column="adv")
+    assert _pipeline_advantages(
+        tmp_path, rows, advantage_column="adv", skip_nan_rewards=True
+    ) == [1.0]
+
+
+def test_conflicting_advantage_config():
+    with pytest.raises(ValueError, match="cannot both be specified"):
+        DataConfig(reward_column="reward", advantage_column="adv")

bergson-2.3.1/tests/test_select_split.py

@@ -0,0 +1,39 @@
+from datasets import Dataset, DatasetDict
+
+from bergson.data import load_data_string
+
+
+def _saved(tmp_path, ds):
+    path = str(tmp_path / "ds")
+    ds.save_to_disk(path)
+    return path
+
+
+def test_saved_dataset_takes_split_slices(tmp_path):
+    path = _saved(tmp_path, Dataset.from_dict({"x": list(range(10))}))
+
+    assert load_data_string(path, "train")["x"] == list(range(10))
+    assert load_data_string(path, "train[:3]")["x"] == [0, 1, 2]
+    assert load_data_string(path, "train[50%:60%]")["x"] == [5]
+    assert load_data_string(path, "train[8:]+train[:1]")["x"] == [8, 9, 0]
+
+
+def test_saved_dataset_dict_takes_split_slices(tmp_path):
+    path = _saved(
+        tmp_path,
+        DatasetDict(
+            train=Dataset.from_dict({"x": list(range(10))}),
+            test=Dataset.from_dict({"x": list(range(100, 104))}),
+        ),
+    )
+
+    assert load_data_string(path, "test")["x"] == [100, 101, 102, 103]
+    assert load_data_string(path, "test[:2]")["x"] == [100, 101]
+    assert load_data_string(path, "train[:1]+test[-1:]")["x"] == [0, 103]
+
+
+def test_loaded_split_is_sliced_once(tmp_path):
+    path = tmp_path / "rows.csv"
+    path.write_text("x\n" + "\n".join(str(i) for i in range(8)) + "\n")
+
+    assert load_data_string(str(path), "train[:50%]")["x"] == [0, 1, 2, 3]