| """Generate a tiny synthetic decision dataset for pipeline smoke tests. |
| |
| Deterministic toy task: the state describes a colored box; the question asks |
| which color / how heavy / whether fragile, and targets follow a rule the model |
| can actually learn (so loss drops in a few steps). |
| |
| Usage: |
| python scripts/make_tiny_data.py --out /tmp/groot_tiny --n 96 |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import random |
| import sys |
| from pathlib import Path |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) |
| from groot1.data import make_row, write_jsonl |
|
|
| COLORS = ["red", "blue", "green", "amber"] |
|
|
|
|
| def gen_row(rng: random.Random, i: int, split: str): |
| color = rng.choice(COLORS) |
| weight = rng.randint(1, 10) |
| fragile = rng.random() < 0.5 |
| state = { |
| "text": f"A {color} box weighing {weight}kg, marked {'fragile' if fragile else 'sturdy'}.", |
| "metadata": {"weight": weight, "fragile": fragile}, |
| } |
| kind = i % 3 |
| if kind == 0: |
| target = [1.0 if c == color else 0.01 for c in COLORS] |
| return make_row("tiny", i, state, "Which color is the box?", "choice", COLORS, target, split) |
| if kind == 1: |
| band = min(3, weight // 3) |
| target = [1.0 if j == band else 0.05 for j in range(4)] |
| return make_row("tiny", i, state, "Rate the weight", "score", |
| [f"level {j}" for j in range(4)], target, split) |
| target = [0.95, 0.05] if fragile else [0.05, 0.95] |
| return make_row("tiny", i, state, "The box is fragile", "noul", ["false", "true"], target, split) |
|
|
|
|
| def main() -> None: |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--out", required=True) |
| ap.add_argument("--n", type=int, default=96) |
| args = ap.parse_args() |
| rng = random.Random(7) |
| out = Path(args.out) |
| out.mkdir(parents=True, exist_ok=True) |
| for split, n in (("train", args.n), ("val", max(16, args.n // 4))): |
| rows = [gen_row(rng, i, "train") for i in range(n)] |
| path = out / f"{split}.jsonl" |
| write_jsonl(rows, path) |
| print(f"wrote {n} rows -> {path}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
