bergson 2.2.4__tar.gz → 2.3.0__tar.gz

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

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

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

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

{bergson-2.2.4 → bergson-2.3.0}/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.0}/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.0}/bergson.egg-info/PKG-INFO

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

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

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

{bergson-2.2.4 → bergson-2.3.0}/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")