Spaces:
Running
Running
| """Rebuild the training episodes the executor is trained on (one folder per episode: rows.jsonl, the step screenshots, | |
| episode.json) from the Hugging Face release, in the layout `python -m miniwebagent.train` reads. | |
| python scripts/fetch_training_data.py # MiniWebGen's episodes -> data/webmix/replay_v4/ | |
| python scripts/fetch_training_data.py --table task_centric_steps --out data/webmix/taskgen/rows_27k | |
| python scripts/fetch_training_data.py --src data/release/miniweb-v1-hf # from a local copy of the release | |
| Rows that are not training targets (a recovery episode's injected mistake and its replayed prefix) are written with | |
| "inject": true, which the trainer skips, as in the original episodes. | |
| """ | |
| import argparse, glob, json, os | |
| import pyarrow.parquet as pq | |
| ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) | |
| REPO = "sparklabuofu/MiniWeb" | |
| OUT = {"training_steps": os.path.join("data", "webmix", "replay_v4"), | |
| "task_centric_steps": os.path.join("data", "webmix", "taskgen", "rows_27k")} | |
| COLS = ["episode", "variant", "parts", "task_id", "step", "segment", "primitive", "family", "op", "resolved", | |
| "memory_by", "messages", "target", "is_target", "image"] | |
| def main(): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--table", choices=list(OUT), default="training_steps") | |
| ap.add_argument("--src", help="a local copy of the release (default: download the table from the Hub)") | |
| ap.add_argument("--repo", default=REPO) | |
| ap.add_argument("--out") | |
| a = ap.parse_args() | |
| out = a.out or os.path.join(ROOT, OUT[a.table]) | |
| src = a.src | |
| if not src: | |
| from huggingface_hub import snapshot_download | |
| src = snapshot_download(a.repo, repo_type="dataset", allow_patterns=[f"data/{a.table}/*"]) | |
| files = sorted(glob.glob(os.path.join(src, "data", a.table, "*.parquet"))) | |
| episodes, rows = {}, 0 | |
| handles = {} | |
| try: | |
| for f in files: | |
| for batch in pq.ParquetFile(f).iter_batches(batch_size=200, columns=COLS): | |
| for r in batch.to_pylist(): | |
| d = os.path.join(out, r["episode"]) | |
| if r["episode"] not in episodes: | |
| os.makedirs(d, exist_ok=True) | |
| episodes[r["episode"]] = r["parts"] | |
| parts = r["parts"] if r["variant"] == "chain" else None # only chains list their parts, as the originals | |
| json.dump({"task_id": r["task_id"], "variant": r["variant"], "parts": parts}, | |
| open(os.path.join(d, "episode.json"), "w")) | |
| handles[r["episode"]] = open(os.path.join(d, "rows.jsonl"), "w") | |
| img = r["image"] | |
| if img and img.get("bytes"): | |
| open(os.path.join(d, img["path"]), "wb").write(img["bytes"]) | |
| row = {"task_id": r["task_id"], "step": r["step"], "segment": r["segment"], "macro": r["primitive"], | |
| "family": r["family"], "op": r["op"], "resolved": r["resolved"], "memory_by": r["memory_by"], | |
| "messages": json.loads(r["messages"]) if r["messages"] else None, "target": r["target"]} | |
| if not r["is_target"]: | |
| row["inject"] = True | |
| handles[r["episode"]].write(json.dumps(row, ensure_ascii=False) + "\n") | |
| rows += 1 | |
| finally: | |
| for h in handles.values(): | |
| h.close() | |
| print(f"{len(episodes)} episodes, {rows} rows -> {out}") | |
| if __name__ == "__main__": | |
| main() | |