wehub-resource-sync
79 lines
2.2 KiB
Python
79 lines
2.2 KiB
Python
#!/usr/bin/env python3
|
|
"""Upload a prepared local dataset folder to a Hugging Face dataset repo."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import time
|
|
from pathlib import Path
|
|
|
|
from huggingface_hub import HfApi
|
|
from huggingface_hub.errors import HfHubHTTPError
|
|
|
|
|
|
DEFAULT_LOCAL_DIR = Path(
|
|
"/home/user/wiki-screenshot-training/hf_dataset_export_sharded/screenshot-training"
|
|
)
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--repo-id", default="Chrisyichuan/screenshot-training")
|
|
parser.add_argument("--local-dir", type=Path, default=DEFAULT_LOCAL_DIR)
|
|
parser.add_argument("--repo-type", default="dataset")
|
|
parser.add_argument("--private", action="store_true")
|
|
parser.add_argument(
|
|
"--skip-create",
|
|
action="store_true",
|
|
help="Skip create_repo (use if repo already exists)",
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
def retry_on_429(fn, max_retries=5, initial_wait=60):
|
|
"""Retry a function on 429 rate limit errors with exponential backoff."""
|
|
for attempt in range(max_retries):
|
|
try:
|
|
return fn()
|
|
except HfHubHTTPError as e:
|
|
if "429" in str(e) and attempt < max_retries - 1:
|
|
wait = initial_wait * (2**attempt)
|
|
print(
|
|
f"Rate limited (429). Waiting {wait}s before retry {attempt + 2}/{max_retries}..."
|
|
)
|
|
time.sleep(wait)
|
|
else:
|
|
raise
|
|
|
|
|
|
def main() -> int:
|
|
args = parse_args()
|
|
api = HfApi()
|
|
|
|
if not args.skip_create:
|
|
print("Creating repo (with retry)...")
|
|
retry_on_429(
|
|
lambda: api.create_repo(
|
|
repo_id=args.repo_id,
|
|
repo_type=args.repo_type,
|
|
exist_ok=True,
|
|
private=args.private,
|
|
)
|
|
)
|
|
print("Repo ready.")
|
|
|
|
print(f"Uploading {args.local_dir} ...")
|
|
api.upload_large_folder(
|
|
repo_id=args.repo_id,
|
|
repo_type=args.repo_type,
|
|
folder_path=str(args.local_dir),
|
|
)
|
|
print(
|
|
f"Uploaded {args.local_dir} -> https://huggingface.co/datasets/{args.repo_id}"
|
|
)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|