155 lines
4.8 KiB
Python
155 lines
4.8 KiB
Python
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()
|