| |
| """Download pinned datasets and model/tokenizer snapshots with cache resume.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| from datetime import datetime, timezone |
| import os |
| from pathlib import Path |
| import shutil |
| import sys |
|
|
| import yaml |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--root", type=Path, default=Path(__file__).resolve().parents[1]) |
| parser.add_argument("--config", type=Path) |
| parser.add_argument("--dry-run", action="store_true") |
| parser.add_argument("--allow-low-disk", action="store_true") |
| return parser.parse_args() |
|
|
|
|
| def load_configuration(root: Path, config_path: Path | None) -> dict[str, object]: |
| path = config_path or root / "artifacts" / "data.yaml" |
| return yaml.safe_load(path.read_text(encoding="utf-8")) |
|
|
|
|
| def target_lines(configuration: dict[str, object]) -> list[str]: |
| datasets = configuration["datasets"] |
| models = configuration["models"] |
| return [ |
| *[ |
| f"DATASET {name} {datasets[name]['repo_id']}@{datasets[name]['revision']}" |
| for name in ("lm1b", "owt") |
| ], |
| *[ |
| f"MODEL {name}@{models[name]}" |
| for name in ("bert-base-uncased", "gpt2", "gpt2-large") |
| ], |
| ] |
|
|
|
|
| def directory_size(path: Path) -> int: |
| return sum(candidate.stat().st_size for candidate in path.rglob("*") if candidate.is_file()) |
|
|
|
|
| def _relative_cache_path(path: Path, root: Path) -> str: |
| try: |
| return path.resolve().relative_to(root).as_posix() |
| except ValueError: |
| return path.resolve().as_posix() |
|
|
|
|
| def main() -> int: |
| args = parse_args() |
| root = args.root.resolve() |
| configuration = load_configuration(root, args.config) |
| for line in target_lines(configuration): |
| print(line) |
| if args.dry_run: |
| return 0 |
|
|
| hf_home = root / "data" / "raw" / "huggingface" |
| datasets_cache = hf_home / "datasets" |
| hub_cache = hf_home / "hub" |
| hf_home.mkdir(parents=True, exist_ok=True) |
| os.environ["HF_HOME"] = str(hf_home) |
| os.environ["HF_DATASETS_CACHE"] = str(datasets_cache) |
|
|
| from datasets import load_dataset |
| from huggingface_hub import snapshot_download |
|
|
| sys.path.insert(0, str(root / "src")) |
| from dlb.data import disk_preflight_evidence, validate_download_manifest |
| from dlb.io import atomic_json_write |
|
|
| dataset_records: dict[str, object] = {} |
| owt_disk_preflight = None |
| for name in ("lm1b", "owt"): |
| specification = configuration["datasets"][name] |
| if name == "owt": |
| free_bytes = shutil.disk_usage(root).free |
| owt_disk_preflight = disk_preflight_evidence( |
| free_bytes, allow_low_disk=args.allow_low_disk |
| ) |
| if owt_disk_preflight["override_used"]: |
| print( |
| f"WARNING low disk override: {free_bytes} bytes free, required 55 GiB", |
| file=sys.stderr, |
| ) |
| if name == "lm1b": |
| source_metadata_snapshot = snapshot_download( |
| repo_id=specification["repo_id"], |
| repo_type="dataset", |
| revision=specification["revision"], |
| cache_dir=hub_cache, |
| ) |
| dataset_snapshot = snapshot_download( |
| repo_id=specification["repo_id"], |
| repo_type="dataset", |
| revision=specification["parquet_revision"], |
| cache_dir=hub_cache, |
| allow_patterns="plain_text/**/*.parquet", |
| ) |
| snapshot_path = Path(dataset_snapshot) |
| loaded = load_dataset( |
| "parquet", |
| data_files={ |
| "train": [ |
| str(path) |
| for path in sorted( |
| (snapshot_path / "plain_text" / "train").glob("*.parquet") |
| ) |
| ], |
| "test": [ |
| str(path) |
| for path in sorted( |
| (snapshot_path / "plain_text" / "test").glob("*.parquet") |
| ) |
| ], |
| }, |
| cache_dir=str(datasets_cache), |
| download_mode="reuse_dataset_if_exists", |
| ) |
| archive_specification = configuration["lm1b_archive"] |
| if len(loaded["train"]) != int(archive_specification["train_documents"]): |
| raise RuntimeError("LM1B training document count differs from pinned metadata") |
| if len(loaded["test"]) != int(archive_specification["test_documents"]): |
| raise RuntimeError("LM1B test document count differs from pinned metadata") |
| else: |
| source_metadata_snapshot = None |
| dataset_snapshot = None |
| loaded = load_dataset( |
| specification["repo_id"], |
| specification["config"], |
| revision=specification["revision"], |
| cache_dir=str(datasets_cache), |
| download_mode="reuse_dataset_if_exists", |
| ) |
| if len(loaded["train"]) != int(specification["documents"]): |
| raise RuntimeError("OpenWebText document count differs from pinned metadata") |
| dataset_records[name] = { |
| "repo_id": specification["repo_id"], |
| "source_revision": specification["revision"], |
| "cache_files": { |
| split: [ |
| _relative_cache_path(Path(entry["filename"]), root) |
| for entry in value.cache_files |
| ] |
| for split, value in loaded.items() |
| }, |
| "split_rows": {split: len(value) for split, value in loaded.items()}, |
| "dataset_snapshot": ( |
| _relative_cache_path(Path(dataset_snapshot), root) if dataset_snapshot else None |
| ), |
| "source_metadata_snapshot": ( |
| _relative_cache_path(Path(source_metadata_snapshot), root) |
| if source_metadata_snapshot |
| else None |
| ), |
| "materialization_revision": specification.get("parquet_revision"), |
| } |
| print(f"DOWNLOADED DATASET {name} {dataset_records[name]['split_rows']}") |
|
|
| model_records: dict[str, object] = {} |
| for name in ("bert-base-uncased", "gpt2", "gpt2-large"): |
| revision = configuration["models"][name] |
| snapshot_path = Path( |
| snapshot_download( |
| repo_id=name, |
| revision=revision, |
| cache_dir=hub_cache, |
| ) |
| ) |
| model_records[name] = { |
| "repo_id": name, |
| "revision": revision, |
| "snapshot_path": snapshot_path.relative_to(root).as_posix(), |
| "size_bytes": directory_size(snapshot_path), |
| } |
| print(f"DOWNLOADED MODEL {name} {snapshot_path}") |
|
|
| manifest = { |
| "schema_version": 1, |
| "created_at": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"), |
| "hf_home": hf_home.relative_to(root).as_posix(), |
| "owt_disk_preflight": owt_disk_preflight, |
| "datasets": dataset_records, |
| "models": model_records, |
| } |
| for dataset_name in ("lm1b", "owt"): |
| validate_download_manifest(configuration, manifest, dataset_name) |
| atomic_json_write(root / "data" / "manifests" / "downloads.json", manifest) |
| print(f"RAW CACHE BYTES {directory_size(hf_home)}") |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
