| """Explicit, resumable, revision-pinned upstream download with hash verification.""" |
| import argparse |
| import os |
| import hashlib |
| import json |
| from pathlib import Path |
| os.environ.setdefault("HF_HUB_DISABLE_XET", "1") |
| from huggingface_hub import HfApi, hf_hub_download |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| VARIANTS = ("5M", "28M", "123M", "209M") |
|
|
|
|
| def sha256(path): |
| h = hashlib.sha256() |
| with Path(path).open("rb") as f: |
| for chunk in iter(lambda: f.read(8 * 1024 * 1024), b""): |
| h.update(chunk) |
| return h.hexdigest() |
|
|
|
|
| def download(variant): |
| repo = f"mistralai/Timestral-{variant}" |
| api = HfApi() |
| pins_path = ROOT / "pins.json" |
| pins = json.loads(pins_path.read_text()) if pins_path.exists() else {} |
| revision = pins.get(variant, {}).get("revision") |
| info = api.model_info(repo, revision=revision, files_metadata=True) |
| path = ROOT / "models" / "source" / variant |
| path.mkdir(parents=True, exist_ok=True) |
| print(f"Downloading {repo}@{info.sha} to {path}", flush=True) |
| inventory = {} |
| for name in ("config.json", "model.safetensors", "LICENSE", "README.md"): |
| dest = Path(hf_hub_download(repo, name, revision=info.sha, local_dir=path)) |
| meta = next(s for s in info.siblings if s.rfilename == name) |
| if meta.lfs: |
| expected = meta.lfs.sha256 |
| actual = sha256(dest) |
| else: |
| expected = meta.blob_id |
| data = dest.read_bytes() |
| actual = hashlib.sha1(f"blob {len(data)}\0".encode() + data).hexdigest() |
| if actual != expected: |
| raise ValueError(f"Hub blob hash mismatch for {repo}/{name}") |
| inventory[name] = {"sha256": sha256(dest), "bytes": dest.stat().st_size, "hub_blob": expected} |
| pinned = pins.get(variant, {}).get("files", {}).get(name) |
| if pinned is not None and pinned != inventory[name]: |
| raise ValueError(f"pinned source hash mismatch: {repo}/{name}") |
| print(f" {name}: {dest.stat().st_size:,} bytes, verified", flush=True) |
| record = {"repo_id": repo, "revision": info.sha, "files": inventory} |
| (path / "source.json").write_text(json.dumps(record, indent=2) + "\n") |
| return path |
|
|
|
|
| if __name__ == "__main__": |
| p = argparse.ArgumentParser(description=__doc__) |
| p.add_argument("variants", nargs="*", choices=VARIANTS, default=list(VARIANTS)) |
| args = p.parse_args() |
| for size in args.variants: |
| download(size) |
|
|
