File size: 3,276 Bytes
adc02fa | 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 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 | #!/usr/bin/env python
from __future__ import annotations
import argparse
import sys
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.generation.pipeline import ( # noqa: E402
generate_builtin_toy_dataset,
generate_cil_dataset,
load_task_specs,
print_generation_summary,
)
from dovla_cil.tasks.library import ToyTaskLibrary
def main(argv: list[str] | None = None) -> int:
parser = argparse.ArgumentParser(description="Generate DoVLA-CIL shards.")
parser.add_argument("--backend", choices=["toy"], default="toy")
parser.add_argument("--tasks", type=Path, default=None, help="TaskSpec JSONL or JSON path.")
parser.add_argument("--out", "--output", dest="out", type=Path, default=Path("outputs/cil_toy"))
parser.add_argument("--num-states-per-task", type=int, default=10)
parser.add_argument(
"--num-tasks",
type=int,
default=None,
help="Optional limit on loaded/generated tasks before CIL generation.",
)
parser.add_argument("--k", type=int, default=16)
parser.add_argument("--seed", type=int, default=0)
parser.add_argument("--shard-size", "--max-records-per-shard", dest="shard_size", type=int, default=1000)
parser.add_argument("--inline-observations", action="store_true")
parser.add_argument(
"--use-vlm-annotations",
action="store_true",
help="Ask the configured VLM for concise semantic failure explanations.",
)
parser.add_argument(
"--vlm-cache",
type=Path,
default=None,
help="JSON cache path for VLM failure annotations.",
)
parser.add_argument(
"--groups",
type=int,
default=None,
help="Legacy toy shortcut: generate this many built-in task groups.",
)
args = parser.parse_args(argv)
if args.tasks is None and args.groups is not None:
summary = generate_builtin_toy_dataset(
out_dir=args.out,
groups=args.groups,
k=args.k,
seed=args.seed,
shard_size=args.shard_size,
inline_observations=True,
use_vlm_annotations=args.use_vlm_annotations,
vlm_cache=args.vlm_cache,
)
else:
tasks = (
load_task_specs(args.tasks)
if args.tasks is not None
else ToyTaskLibrary().list(args.num_tasks)
)
if args.tasks is not None and args.num_tasks is not None:
if args.num_tasks <= 0:
raise ValueError("--num-tasks must be positive when provided")
tasks = tasks[: args.num_tasks]
summary = generate_cil_dataset(
backend=args.backend,
tasks=tasks,
out_dir=args.out,
num_states_per_task=args.num_states_per_task,
k=args.k,
seed=args.seed,
shard_size=args.shard_size,
inline_observations=args.inline_observations,
use_vlm_annotations=args.use_vlm_annotations,
vlm_cache=args.vlm_cache,
)
print_generation_summary(summary)
return 0
if __name__ == "__main__":
raise SystemExit(main())
|