seeker
49 lines
1.4 KiB
Python
49 lines
1.4 KiB
Python
#!/usr/bin/env python3
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
|
|
from huggingface_hub import model_info, snapshot_download
|
|
|
|
from alpamayo_lite.config import load_config
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--config", default="configs/local_overfit.yaml")
|
|
args = parser.parse_args()
|
|
cfg = load_config(args.config)
|
|
output_dir = cfg.resolve(cfg.paths.model_dir)
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
info = model_info(cfg.vlm.repo_id, files_metadata=True)
|
|
total = info.safetensors.total if info.safetensors is not None else None
|
|
if total is not None and not (700_000_000 <= total <= 900_000_000):
|
|
raise RuntimeError(f"Unexpected model size for {cfg.vlm.repo_id}: {total:,} parameters")
|
|
|
|
snapshot_download(
|
|
repo_id=cfg.vlm.repo_id,
|
|
local_dir=output_dir,
|
|
allow_patterns=[
|
|
"*.json",
|
|
"*.txt",
|
|
"*.safetensors",
|
|
],
|
|
)
|
|
provenance = {
|
|
"repo_id": cfg.vlm.repo_id,
|
|
"revision": info.sha,
|
|
"safetensors_parameters": total,
|
|
"pipeline_tag": info.pipeline_tag,
|
|
"library_name": info.library_name,
|
|
}
|
|
(output_dir / "download_provenance.json").write_text(
|
|
json.dumps(provenance, indent=2, ensure_ascii=False) + "\n", encoding="utf-8"
|
|
)
|
|
print(json.dumps(provenance, indent=2, ensure_ascii=False))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|