63 lines
2.0 KiB
Python
63 lines
2.0 KiB
Python
"""Shard downloader — fetches or stubs a model shard from HuggingFace Hub.
|
|
|
|
Cache layout: ~/.cache/meshnet/shards/<model>/layers_<start>-<end>/
|
|
|
|
For "stub-model" (no HF repo), a placeholder JSON file is written so the
|
|
test suite never touches the network.
|
|
"""
|
|
|
|
import json
|
|
from pathlib import Path
|
|
|
|
_DEFAULT_CACHE = Path.home() / ".cache" / "meshnet" / "shards"
|
|
|
|
|
|
def download_shard(
|
|
model: str,
|
|
shard_start: int,
|
|
shard_end: int,
|
|
cache_dir: Path = _DEFAULT_CACHE,
|
|
hf_repo: str | None = None,
|
|
progress: bool = True,
|
|
) -> Path:
|
|
"""Ensure the shard is present in *cache_dir* and return its local path.
|
|
|
|
When *hf_repo* is None (or *model* is ``"stub-model"``), a placeholder
|
|
weights file is created locally — no network access required. This keeps
|
|
the test suite hermetic while the real download path is exercised by
|
|
passing a non-stub *hf_repo*.
|
|
"""
|
|
shard_dir = cache_dir / model / f"layers_{shard_start}-{shard_end}"
|
|
shard_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
if hf_repo is None or model == "stub-model":
|
|
stub_file = shard_dir / "weights.json"
|
|
if not stub_file.exists():
|
|
stub_file.write_text(json.dumps({
|
|
"model": model,
|
|
"shard_start": shard_start,
|
|
"shard_end": shard_end,
|
|
"stub": True,
|
|
}))
|
|
if progress:
|
|
print(f" [stub] shard placeholder written to {stub_file}", flush=True)
|
|
else:
|
|
if progress:
|
|
print(f" [stub] shard already cached at {shard_dir}", flush=True)
|
|
return shard_dir
|
|
|
|
from huggingface_hub import snapshot_download # type: ignore[import]
|
|
|
|
if progress:
|
|
print(
|
|
f" Downloading layers {shard_start}-{shard_end} from {hf_repo} ...",
|
|
flush=True,
|
|
)
|
|
|
|
local_dir = snapshot_download(
|
|
repo_id=hf_repo,
|
|
cache_dir=str(cache_dir),
|
|
local_dir=str(shard_dir),
|
|
)
|
|
return Path(local_dir)
|