Spaces:
Sleeping
Sleeping
| #!/usr/bin/env python3 | |
| """Sync the private Hugging Face dataset into the Space local disk.""" | |
| from __future__ import annotations | |
| import argparse | |
| import os | |
| import time | |
| from pathlib import Path | |
| from huggingface_hub import snapshot_download | |
| def main() -> int: | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument('--repo-id', required=True) | |
| parser.add_argument('--local-dir', type=Path, required=True) | |
| parser.add_argument('--include', action='append', default=[]) | |
| parser.add_argument('--clean', action='store_true') | |
| args = parser.parse_args() | |
| args.local_dir.mkdir(parents=True, exist_ok=True) | |
| token = os.environ.get('HF_TOKEN') | |
| if not token: | |
| raise SystemExit('HF_TOKEN secret is required to sync the private dataset') | |
| if args.clean: | |
| for child in args.local_dir.iterdir(): | |
| if child.is_dir(): | |
| import shutil | |
| shutil.rmtree(child) | |
| else: | |
| child.unlink() | |
| for attempt in range(1, 5): | |
| try: | |
| kwargs = { | |
| 'repo_id': args.repo_id, | |
| 'repo_type': 'dataset', | |
| 'local_dir': str(args.local_dir), | |
| 'token': token, | |
| 'max_workers': 2, | |
| } | |
| if args.include: | |
| kwargs['allow_patterns'] = args.include | |
| snapshot_download(**kwargs) | |
| return 0 | |
| except Exception as exc: | |
| if attempt == 4: | |
| raise | |
| wait_seconds = 20 * attempt | |
| print(f'dataset sync failed on attempt {attempt}: {exc}', flush=True) | |
| print(f'retrying in {wait_seconds}s', flush=True) | |
| time.sleep(wait_seconds) | |
| return 1 | |
| if __name__ == '__main__': | |
| raise SystemExit(main()) | |