| |
| from __future__ import annotations |
|
|
| import argparse |
| import sys |
| from collections import Counter |
| from pathlib import Path |
|
|
| PROJECT_ROOT = Path(__file__).resolve().parents[1] |
| if str(PROJECT_ROOT) not in sys.path: |
| sys.path.insert(0, str(PROJECT_ROOT)) |
|
|
| from dovla_cil.data.index import ShardIndex |
| from dovla_cil.data.schema import CILRecord |
| from dovla_cil.data.sharding import ShardReader, group_records, iter_cil_records |
|
|
|
|
| def main(argv: list[str] | None = None) -> int: |
| parser = argparse.ArgumentParser(description="Inspect a CIL dataset, metadata file, or shard.") |
| parser.add_argument("path", type=Path) |
| parser.add_argument("--group-id", default=None, help="Group to display. Defaults to first group.") |
| parser.add_argument( |
| "--max-rows", |
| type=int, |
| default=32, |
| help="Maximum rows to print in the ranking table.", |
| ) |
| args = parser.parse_args(argv) |
|
|
| if _looks_like_dataset(args.path): |
| _inspect_dataset(args.path, group_id=args.group_id, max_rows=args.max_rows) |
| else: |
| _inspect_single_shard(args.path, group_id=args.group_id, max_rows=args.max_rows) |
| return 0 |
|
|
|
|
| def _inspect_dataset(path: Path, *, group_id: str | None, max_rows: int) -> None: |
| reader = ShardReader(path) |
| index = reader.index |
| _print_summary(index) |
| entries = index.load_group_index() |
| if not entries: |
| print("group_index: missing or empty") |
| return |
| selected_group_id = group_id or str(entries[0]["group_id"]) |
| records = reader.load_group(selected_group_id) |
| _print_group_index_entry(_entry_by_group(entries, selected_group_id)) |
| _print_ranking_table(records, max_rows=max_rows) |
|
|
|
|
| def _inspect_single_shard(path: Path, *, group_id: str | None, max_rows: int) -> None: |
| records = list(iter_cil_records(path)) |
| groups = group_records(records) |
| candidate_counts = Counter(record.candidate_type for record in records) |
| print(f"shard: {path}") |
| print(f"num_records: {len(records)}") |
| print(f"num_groups: {len(groups)}") |
| print(f"candidate_type_counts: {dict(sorted(candidate_counts.items()))}") |
| if not groups: |
| return |
| selected_group_id = group_id or next(iter(groups)) |
| if selected_group_id not in groups: |
| raise KeyError(f"Unknown group {selected_group_id!r} in shard {path}") |
| _print_ranking_table(groups[selected_group_id], max_rows=max_rows) |
|
|
|
|
| def _print_summary(index: ShardIndex) -> None: |
| metadata = index.metadata |
| print(f"dataset: {metadata.get('dataset_name', index.dataset_dir.name)}") |
| print(f"path: {index.dataset_dir}") |
| print(f"version: {metadata.get('version', 'unknown')}") |
| print(f"schema_version: {metadata.get('schema_version', 'unknown')}") |
| print(f"backend: {metadata.get('backend', 'unknown')}") |
| print(f"created_at: {metadata.get('created_at', 'unknown')}") |
| print(f"num_groups: {index.group_count}") |
| print(f"num_records: {index.record_count}") |
| print(f"k: {metadata.get('k', 'unknown')}") |
| print(f"task_count: {metadata.get('task_count', 'unknown')}") |
| print(f"seed: {metadata.get('seed', 'unknown')}") |
| print(f"shard_count: {len(index.shards)}") |
| print(f"group_index: {index.group_index_path}") |
| print(f"record_index: {index.record_index_path}") |
|
|
|
|
| def _print_group_index_entry(entry: dict[str, object]) -> None: |
| print(f"sample_group: {entry.get('group_id')}") |
| print(f"sample_group_task: {entry.get('task_id')}") |
| print(f"sample_group_shard: {entry.get('shard_path')}") |
| print(f"sample_group_records: {entry.get('num_records')}") |
| print(f"sample_group_max_reward: {entry.get('max_reward')}") |
| print(f"sample_group_success_count: {entry.get('success_count')}") |
| print(f"sample_group_candidate_type_counts: {entry.get('candidate_type_counts')}") |
|
|
|
|
| def _print_ranking_table(records: list[CILRecord], *, max_rows: int) -> None: |
| ordered = sorted( |
| records, |
| key=lambda record: ( |
| record.rank_within_group if record.rank_within_group is not None else 10**9, |
| record.record_id, |
| ), |
| ) |
| print("ranking:") |
| print("record_id\tcandidate_type\treward.progress\tsuccess\tregret\trank\tfailure.type") |
| for record in ordered[:max_rows]: |
| print( |
| "\t".join( |
| [ |
| record.record_id, |
| record.candidate_type, |
| f"{record.reward.progress:.6g}", |
| str(record.reward.terminal_success), |
| "" if record.regret is None else f"{record.regret:.6g}", |
| "" |
| if record.rank_within_group is None |
| else str(record.rank_within_group), |
| record.failure.type if record.failure else "none", |
| ] |
| ) |
| ) |
| if len(ordered) > max_rows: |
| print(f"... {len(ordered) - max_rows} more records") |
|
|
|
|
| def _entry_by_group(entries: list[dict[str, object]], group_id: str) -> dict[str, object]: |
| for entry in entries: |
| if entry.get("group_id") == group_id: |
| return entry |
| raise KeyError(f"Unknown group {group_id!r}") |
|
|
|
|
| def _looks_like_dataset(path: Path) -> bool: |
| if path.is_dir(): |
| return (path / "metadata.json").exists() or (path / "manifest.json").exists() |
| return path.name in {"metadata.json", "manifest.json"} |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
