File size: 2,248 Bytes
2bda233 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 | from __future__ import annotations
import unittest
from br_sovereign_llm.corpus import Split, SplitConfig, assign_split, split_bucket
class SplitAssignmentTests(unittest.TestCase):
def test_known_key_has_stable_bucket(self) -> None:
self.assertEqual(split_bucket("group:lei-9610"), 6756)
self.assertEqual(assign_split("group:lei-9610"), Split.TRAIN)
def test_same_leakage_group_always_has_same_split(self) -> None:
assignments = {assign_split("content-sha256:abc") for _ in range(20)}
self.assertEqual(assignments, {assign_split("content-sha256:abc")})
def test_all_partitions_are_reachable(self) -> None:
config = SplitConfig(
train_permyriad=8000,
validation_permyriad=1000,
test_permyriad=1000,
seed="coverage-test",
)
observed = {
assign_split(f"group-{index}", config=config) for index in range(500)
}
self.assertEqual(observed, {Split.TRAIN, Split.VALIDATION, Split.TEST})
def test_seed_is_part_of_assignment(self) -> None:
first = SplitConfig(
train_permyriad=5000,
validation_permyriad=2500,
test_permyriad=2500,
seed="first-seed",
)
second = SplitConfig(
train_permyriad=5000,
validation_permyriad=2500,
test_permyriad=2500,
seed="second-seed",
)
self.assertTrue(
any(
assign_split(f"group-{index}", config=first)
!= assign_split(f"group-{index}", config=second)
for index in range(100)
)
)
def test_split_proportions_must_sum_to_permyriad(self) -> None:
with self.assertRaises(ValueError):
SplitConfig(
train_permyriad=9800,
validation_permyriad=100,
test_permyriad=99,
)
def test_boolean_proportion_is_rejected(self) -> None:
with self.assertRaises(ValueError):
SplitConfig(
train_permyriad=True,
validation_permyriad=9999,
test_permyriad=0,
)
if __name__ == "__main__":
unittest.main()
|