Buckets:
| import json | |
| import tempfile | |
| import unittest | |
| from pathlib import Path | |
| import torch | |
| from shared.data import ( | |
| answer_reference, | |
| PackedTrajectoryCache, | |
| query_indices_for_update, | |
| rollout_indices_for_update, | |
| ) | |
| class SharedDataTest(unittest.TestCase): | |
| def test_cache_uses_manifest_index_as_stable_dataset_index(self): | |
| with tempfile.TemporaryDirectory() as directory: | |
| root = Path(directory) | |
| (root / "READY").write_text("complete\n") | |
| rows = [ | |
| { | |
| "sample_id": "a", | |
| "split": "train", | |
| "cache_shard": "one.pt", | |
| "cache_index": 0, | |
| }, | |
| { | |
| "sample_id": "b", | |
| "split": "validation", | |
| "cache_shard": "one.pt", | |
| "cache_index": 1, | |
| }, | |
| ] | |
| (root / "manifest.jsonl").write_text( | |
| "".join(json.dumps(row) + "\n" for row in rows) | |
| ) | |
| torch.save( | |
| { | |
| "schema": "zebra_uniar_cache_v1", | |
| "items": [{"sample_id": "a"}, {"sample_id": "b"}], | |
| }, | |
| root / "one.pt", | |
| ) | |
| cache = PackedTrajectoryCache(root) | |
| self.assertEqual(cache.indices("train"), (0,)) | |
| self.assertEqual(cache.sample_id(1), "b") | |
| self.assertEqual(cache.load(0)["sample_id"], "a") | |
| def test_answer_reference_requires_one_nonempty_answer(self): | |
| item = { | |
| "trajectory": { | |
| "assistant": [ | |
| {"role": "thought", "text": "reason"}, | |
| {"role": "answer", "text": "B"}, | |
| ] | |
| } | |
| } | |
| self.assertEqual( | |
| answer_reference(item), | |
| {"answer": "B", "schema": "answer_v1"}, | |
| ) | |
| with self.assertRaisesRegex(ValueError, "exactly one"): | |
| answer_reference({"trajectory": {"assistant": []}}) | |
| def test_query_schedule_is_fixed_budget_deterministic_and_epoch_covering(self): | |
| universe = tuple(range(10)) | |
| first = query_indices_for_update( | |
| universe, batch_size=4, seed=17, update=0 | |
| ) | |
| again = query_indices_for_update( | |
| universe, batch_size=4, seed=17, update=0 | |
| ) | |
| batches = [ | |
| query_indices_for_update(universe, batch_size=4, seed=17, update=i) | |
| for i in range(3) | |
| ] | |
| self.assertEqual(first, again) | |
| self.assertTrue(all(len(batch) == 4 for batch in batches)) | |
| self.assertEqual(len(set(batches[0] + batches[1])), 8) | |
| self.assertEqual(len(set(index for batch in batches for index in batch)), 10) | |
| def test_rollout_schedule_name_is_a_direct_alias(self): | |
| self.assertIs(rollout_indices_for_update, query_indices_for_update) | |
| def test_unshuffled_schedule_preserves_manifest_order(self): | |
| self.assertEqual( | |
| query_indices_for_update( | |
| (4, 8, 9), | |
| batch_size=5, | |
| seed=17, | |
| update=0, | |
| shuffle=False, | |
| ), | |
| (4, 8, 9, 4, 8), | |
| ) | |
| if __name__ == "__main__": | |
| unittest.main() | |
Xet Storage Details
- Size:
- 3.29 kB
- Xet hash:
- 092050ad7b3b3a36701172b69d0a64c9603224561b25aba7b9df0ce056ff287a
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.