Fix chosen-token regularization

utils/test_dataset_helpers.py

@@ -0,0 +1,154 @@
+import json
+import tempfile
+import unittest
+from collections import Counter
+from pathlib import Path
+
+import numpy as np
+
+from utils.dataset_helpers import (
+    _chosen_target_quotas,
+    _trim_chosen_to_quotas,
+    load_ftpo_multi_dataset,
+)
+
+
+def _counts(rows):
+    return Counter(
+        token
+        for row in rows
+        for token in row["multi_chosen_decoded"]
+    )
+
+
+class _Tokenized:
+    def __init__(self, input_ids):
+        self.input_ids = input_ids
+
+
+class _FakeTokenizer:
+    truncation_side = "right"
+
+    def __call__(self, text, **_kwargs):
+        def encode(value):
+            if value.startswith("context-"):
+                return [1, 2]
+            return [100 + sum(value.encode("utf-8"))]
+
+        if isinstance(text, list):
+            return _Tokenized([encode(value) for value in text])
+        return _Tokenized(encode(text))
+
+
+class ChosenRegularisationTests(unittest.TestCase):
+    def test_zero_strength_disables_all_trimming(self):
+        counts = Counter({f"token-{i}": 20 - i for i in range(12)})
+
+        self.assertEqual(_chosen_target_quotas(counts, 0), dict(counts))
+
+    def test_positive_strength_caps_and_regularises_outliers(self):
+        counts = Counter({f"token-{i}": 100 - 5 * i for i in range(12)})
+
+        quotas = _chosen_target_quotas(counts, 0.2)
+
+        tenth_highest = sorted(counts.values(), reverse=True)[9]
+        self.assertLess(quotas["token-0"], counts["token-0"])
+        self.assertLessEqual(quotas["token-0"], tenth_highest)
+        self.assertEqual(quotas["token-11"], counts["token-11"])
+
+    def test_trimming_enforces_quotas_and_keeps_raw_fields_aligned(self):
+        rows = [
+            {
+                "multi_chosen_decoded": [" common", " rare-a", " common"],
+                "multi_chosen_raw": ["raw-common-1", "raw-rare-a", "raw-common-2"],
+            },
+            {
+                "multi_chosen_decoded": [" common", " rare-b"],
+                "multi_chosen_raw": ["raw-common-3", "raw-rare-b"],
+            },
+        ]
+        quotas = {" common": 2, " rare-a": 1, " rare-b": 1}
+
+        trimmed = _trim_chosen_to_quotas(
+            rows, quotas, np.random.default_rng(3407)
+        )
+
+        self.assertEqual(_counts(trimmed), Counter(quotas))
+        for row in trimmed:
+            self.assertEqual(
+                len(row["multi_chosen_decoded"]),
+                len(row["multi_chosen_raw"]),
+            )
+        # The helper does not mutate the source rows.
+        self.assertEqual(rows[0]["multi_chosen_decoded"].count(" common"), 2)
+
+    def test_trimming_is_reproducible(self):
+        rows = [
+            {
+                "multi_chosen_decoded": [" common", f" unique-{i}"],
+                "multi_chosen_raw": [" common", f" unique-{i}"],
+            }
+            for i in range(20)
+        ]
+        quotas = {" common": 5, **{f" unique-{i}": 1 for i in range(20)}}
+
+        first = _trim_chosen_to_quotas(
+            rows, quotas, np.random.default_rng(123)
+        )
+        second = _trim_chosen_to_quotas(
+            rows, quotas, np.random.default_rng(123)
+        )
+
+        self.assertEqual(first, second)
+        self.assertEqual(_counts(first)[" common"], 5)
+
+    def test_minimum_filter_applies_after_trimming(self):
+        rows = [
+            {"multi_chosen_decoded": [" common", " keep"]},
+            {"multi_chosen_decoded": [" common", " other"]},
+        ]
+
+        trimmed = _trim_chosen_to_quotas(
+            rows,
+            {" common": 1, " keep": 1, " other": 1},
+            np.random.default_rng(7),
+        )
+        surviving = [
+            row for row in trimmed if len(row["multi_chosen_decoded"]) >= 2
+        ]
+
+        self.assertEqual(len(surviving), 1)
+
+    def test_loader_filters_rows_after_applying_chosen_quotas(self):
+        rows = [
+            {
+                "context_with_chat_template": f"context-{i}",
+                "rejected_decoded": " reject",
+                "multi_chosen_decoded": [" common", f" unique-{i}"],
+                "multi_chosen_raw": ["raw-common", f"raw-unique-{i}"],
+            }
+            for i in range(4)
+        ]
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            path = Path(tmp_dir) / "ftpo.jsonl"
+            with path.open("w", encoding="utf-8") as handle:
+                for row in rows:
+                    handle.write(json.dumps(row) + "\n")
+
+            dataset = load_ftpo_multi_dataset(
+                path,
+                _FakeTokenizer(),
+                chosen_reg_strength=1.0,
+                min_chosen_tokens=2,
+                num_proc=1,
+            )
+
+        # " common" is trimmed from four occurrences to one, so only its
+        # containing row still meets the two-chosen-token minimum.
+        self.assertEqual(len(dataset), 1)
+        self.assertEqual(len(dataset[0]["chosen_ids"]), 2)
+
+
+if __name__ == "__main__":
+    unittest.main()