File size: 1,794 Bytes
20c251e | 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 | from __future__ import annotations
from pathlib import Path
from dovla_cil.data.sharding import iter_cil_records
from dovla_cil.generation.pipeline import generate_builtin_toy_dataset, generate_cil_dataset
from dovla_cil.tasks.library import built_in_toy_tasks
from dovla_cil.utils.io import iter_jsonl, read_json
def test_toy_generation_pipeline_writes_shards_and_group_index(tmp_path: Path) -> None:
summary = generate_cil_dataset(
backend="toy",
tasks=built_in_toy_tasks()[:2],
out_dir=tmp_path,
num_states_per_task=1,
k=4,
seed=11,
shard_size=4,
inline_observations=True,
)
manifest = read_json(tmp_path / "manifest.json")
group_index = list(iter_jsonl(tmp_path / "group_index.jsonl"))
records = []
for shard in manifest["shards"]:
records.extend(iter_cil_records(tmp_path / str(shard["path"])))
assert summary.num_groups == 2
assert summary.num_records == 8
assert manifest["record_count"] == 8
assert manifest["group_count"] == 2
assert manifest["group_index_path"] == "group_index.jsonl"
assert len(group_index) == 2
assert all(entry["num_records"] == 4 for entry in group_index)
assert all(record.rank_within_group is not None for record in records)
assert all(record.regret is not None for record in records)
assert any(record.reward.terminal_success for record in records)
assert "expert" in summary.candidate_type_distribution
def test_builtin_group_shortcut_does_not_cap_at_library_size(tmp_path: Path) -> None:
summary = generate_builtin_toy_dataset(
out_dir=tmp_path,
groups=12,
k=2,
seed=17,
shard_size=64,
)
assert summary.num_groups == 12
assert summary.num_records == 24
|