File size: 2,368 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 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | #!/usr/bin/env python
"""Fetch wds dataset repos (hf_push_wds.py output) to a local root for training.
Reads the same specs yaml train_fast.py uses (source.type: wds): every entry
with an hf_repo is downloaded to <root>/<repo basename> unless already complete
(all shards from manifest.json present).
Usage:
python fetch_wds.py --specs configs/wds_mix.yaml --root /data/wds
"""
from __future__ import annotations
import argparse
import json
from pathlib import Path
import yaml
def _complete(d: Path) -> bool:
mf_local = d / "manifest.json"
if not mf_local.exists():
return False
mf = json.loads(mf_local.read_text())
return all((d / s).exists() for s in mf["shards"])
def _dataset_dirs(root: Path):
"""A repo is either one packed dataset (manifest at top) or a bundle of
per-task subdirs each with its own manifest (wds-unitree)."""
if (root / "manifest.json").exists():
return [root]
return [p for p in sorted(root.iterdir()) if p.is_dir() and (p / "manifest.json").exists()]
def fetch(repo: str, dest: Path):
from huggingface_hub import snapshot_download
dest.mkdir(parents=True, exist_ok=True)
dirs = _dataset_dirs(dest)
if dirs and all(_complete(d) for d in dirs):
print(f"{repo}: already complete ({len(dirs)} dataset dir(s))")
return
snapshot_download(repo, repo_type="dataset", local_dir=str(dest))
dirs = _dataset_dirs(dest)
if not dirs:
raise SystemExit(f"{repo}: no manifest.json found after download")
for d in dirs:
if not _complete(d):
raise SystemExit(f"{repo}: incomplete after download: {d}")
mf = json.loads((d / "manifest.json").read_text())
print(f"{repo}::{d.name}: {len(mf['shards'])} shards, {mf['samples']:,} samples")
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--specs", type=Path, required=True)
ap.add_argument("--root", type=Path, required=True)
args = ap.parse_args()
specs = yaml.safe_load(args.specs.read_text())["datasets"]
done = set()
for s in specs:
repo = s.get("hf_repo")
if repo and "dir" not in s and repo not in done:
fetch(repo, args.root / repo.split("/")[-1])
done.add(repo)
elif not repo:
print(f"local: {s.get('dir')}")
if __name__ == "__main__":
main()
|