| 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() | |