File size: 1,914 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 | #!/usr/bin/env python
"""Push a packed wds dataset dir (shard-*.tar + manifest.json + tasks.json) to
the Hub as a private dataset repo, and verify the listing afterwards.
Usage:
python hf_push_wds.py --src /data/wds/tartandrive --repo AlexWortega/wds-tartandrive
"""
from __future__ import annotations
import argparse
import json
from pathlib import Path
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--src", type=Path, required=True)
ap.add_argument("--repo", required=True)
ap.add_argument("--public", action="store_true")
args = ap.parse_args()
from huggingface_hub import HfApi
# single packed dataset (manifest at top) or a bundle of per-task subdirs
if (args.src / "manifest.json").exists():
dirs = {"": json.loads((args.src / "manifest.json").read_text())}
else:
dirs = {p.name + "/": json.loads((p / "manifest.json").read_text())
for p in sorted(args.src.iterdir())
if p.is_dir() and (p / "manifest.json").exists()}
if not dirs:
raise SystemExit(f"{args.src}: no manifest.json anywhere")
api = HfApi()
api.create_repo(args.repo, repo_type="dataset", private=not args.public, exist_ok=True)
api.upload_large_folder(repo_id=args.repo, repo_type="dataset",
folder_path=str(args.src))
remote = set(api.list_repo_files(args.repo, repo_type="dataset"))
missing, total = [], 0
for pfx, mf in dirs.items():
missing += [pfx + s for s in mf["shards"] if pfx + s not in remote]
missing += [pfx + f for f in ("manifest.json", "tasks.json") if pfx + f not in remote]
total += mf["samples"]
if missing:
raise SystemExit(f"upload incomplete, missing: {missing[:10]}")
print(f"{args.repo}: {len(dirs)} dataset dir(s), {total:,} samples, all shards verified")
if __name__ == "__main__":
main()
|