83 lines
2.6 KiB
Python
83 lines
2.6 KiB
Python
"""Download fixed upstream revisions. No remote dataset or model code is executed."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import shutil
|
|
from datetime import UTC, datetime
|
|
from pathlib import Path
|
|
|
|
from micro_scout.data import DATASET, REVISION, SOURCE_SHA256
|
|
from micro_scout.encoder import BASE_MODEL, BASE_REVISION
|
|
from micro_scout.io import atomic_json
|
|
|
|
|
|
def download(output: Path) -> None:
|
|
from huggingface_hub import HfApi, hf_hub_download, snapshot_download
|
|
|
|
output.mkdir(parents=True, exist_ok=True)
|
|
api = HfApi()
|
|
for name, info in (
|
|
("dataset", api.dataset_info(DATASET)),
|
|
("model", api.model_info(BASE_MODEL)),
|
|
):
|
|
atomic_json(
|
|
output / f"{name}_metadata.json",
|
|
{
|
|
"id": info.id,
|
|
"current_head_sha": info.sha,
|
|
"downloads": info.downloads,
|
|
"likes": info.likes,
|
|
"author": info.author,
|
|
"retrieved_at": datetime.now(UTC).isoformat(),
|
|
"used_revision": REVISION if name == "dataset" else BASE_REVISION,
|
|
},
|
|
)
|
|
for split in ("train", "validation", "test"):
|
|
destination = output / f"{split}.parquet"
|
|
if destination.exists():
|
|
if hashlib.sha256(destination.read_bytes()).hexdigest() != SOURCE_SHA256[split]:
|
|
raise ValueError(f"Existing source checksum mismatch: {destination}")
|
|
print(json.dumps({"existing": str(destination)}), flush=True)
|
|
continue
|
|
source = hf_hub_download(
|
|
DATASET,
|
|
f"python/{split}-00000-of-00001.parquet",
|
|
repo_type="dataset",
|
|
revision=REVISION,
|
|
)
|
|
temporary = destination.with_suffix(".parquet.tmp")
|
|
shutil.copyfile(source, temporary)
|
|
if hashlib.sha256(temporary.read_bytes()).hexdigest() != SOURCE_SHA256[split]:
|
|
temporary.unlink()
|
|
raise ValueError(f"Downloaded source checksum mismatch: {split}")
|
|
temporary.replace(destination)
|
|
snapshot_download(
|
|
BASE_MODEL,
|
|
revision=BASE_REVISION,
|
|
allow_patterns=[
|
|
"config.json",
|
|
"model.safetensors",
|
|
"tokenizer.json",
|
|
"tokenizer_config.json",
|
|
"vocab.txt",
|
|
"special_tokens_map.json",
|
|
"README.md",
|
|
"LICENSE",
|
|
],
|
|
max_workers=2,
|
|
)
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--output", type=Path, default=Path("data/source"))
|
|
args = parser.parse_args()
|
|
download(args.output)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|