File size: 2,066 Bytes
b152c62 | 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 | #!/usr/bin/env python
"""Generate configs/wds_mix.yaml from converted wds dirs.
Scans <root> for dataset dirs (manifest.json at top level, or one level down
for bundles like unitree/<task>), weights sqrt(samples) — same rationale as
make_stream_specs.py: plain-proportional lets the biggest sets dominate.
Usage:
python make_wds_mix.py --root data/wds --hf-prefix AlexWortega \
> configs/wds_mix.yaml
"""
from __future__ import annotations
import argparse
import json
import math
from pathlib import Path
# local dir name -> hf repo basename (bundles keep subdir)
REPO_FOR = {"unitree": "wds-unitree", "go_stanford": "wds-go-stanford",
"recon": "wds-recon", "sacson": "wds-sacson",
"scand": "wds-scand", "tartandrive": "wds-tartandrive"}
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--root", type=Path, required=True)
ap.add_argument("--hf-prefix", default="AlexWortega")
args = ap.parse_args()
rows = []
for top in sorted(args.root.iterdir()):
if not top.is_dir():
continue
repo = f"{args.hf_prefix}/{REPO_FOR.get(top.name, 'wds-' + top.name.replace('_', '-'))}"
if (top / "manifest.json").exists():
rows.append((repo, None, json.loads((top / "manifest.json").read_text())))
else:
for sub in sorted(top.iterdir()):
if sub.is_dir() and (sub / "manifest.json").exists():
rows.append((repo, sub.name, json.loads((sub / "manifest.json").read_text())))
total = sum(m["samples"] for _, _, m in rows)
print(f"# autogenerated by make_wds_mix.py: {len(rows)} datasets, "
f"{total:,} samples. weight = sqrt(samples).")
print("datasets:")
for repo, sub, m in rows:
w = math.sqrt(m["samples"])
print(f" - hf_repo: {repo}")
if sub:
print(f" subdir: {sub}")
print(f" weight: {w:.1f} # {m['samples']:,} samples")
print(f" embodiment_id: {m['embodiment_id']}")
if __name__ == "__main__":
main()
|