File size: 3,940 Bytes
cd94e1f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | """Cache the architecture of licence-gated models the demo cannot read.
Some of the best-known LLMs -- Llama and Gemma -- are licence-gated on the
Hugging Face Hub, so an anonymous demo cannot read their `config.json` and has
to refuse them. Serving them with a maintainer's own token would work, but it
would mean using one person's licence acceptance on behalf of every visitor.
Instead this script reads the architecture from a public, unquantised mirror of
the same weights and caches the handful of integers WattGPU actually needs.
Those numbers -- layer count, hidden size, head counts, parameter count -- are
published in the model cards and papers, so caching them redistributes nothing.
Each entry records the mirror it came from, so any value can be traced and
re-checked. Run this again to refresh:
python scripts/cache_gated_models.py
"""
from __future__ import annotations
import argparse
import json
import os
import sys
REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, REPO_ROOT)
from wattgpu_demo.hf_models import ( # noqa: E402
HF_ENDPOINT,
_get_json,
detect_quantisation,
llm_info_from_config,
)
# Gated model -> a public repository holding the same architecture at full
# precision. Mirrors are checked for quantisation before anything is cached.
MIRRORS = {
"meta-llama/Llama-3.2-1B": "unsloth/Llama-3.2-1B",
"meta-llama/Llama-3.2-1B-Instruct": "unsloth/Llama-3.2-1B-Instruct",
"meta-llama/Meta-Llama-3-8B-Instruct": "NousResearch/Meta-Llama-3-8B-Instruct",
"google/gemma-3-1b-it": "unsloth/gemma-3-1b-it",
"google/gemma-3-270m": "unsloth/gemma-3-270m",
}
DEFAULT_OUT = os.path.join(REPO_ROOT, "data", "gated_llms.json")
def fetch_entry(canonical: str, mirror: str) -> dict:
config = _get_json(f"{HF_ENDPOINT}/{mirror}/raw/main/config.json", mirror)
quantisation = detect_quantisation(mirror, config)
if quantisation:
raise ValueError(
f"{mirror} is quantised ({quantisation}); it cannot stand in for {canonical}")
info = _get_json(f"{HF_ENDPOINT}/api/models/{mirror}", mirror)
total = ((info or {}).get("safetensors") or {}).get("total")
if not isinstance(total, (int, float)) or total <= 0:
raise ValueError(f"{mirror} publishes no safetensors index")
# Reuse the same adapter the live path uses, so a cached model and a
# fetched one go through identical validation.
llm = llm_info_from_config(canonical, config, float(total) / 1e9, "published config")
return {
"model_type": llm.model_type,
"num_layers": llm.num_layers,
"hidden_size": llm.hidden_size,
"num_attention_heads": llm.num_attention_heads,
"num_key_value_heads": llm.num_key_value_heads,
"total_b_params": round(llm.total_b_params, 6),
"architectures": llm.architectures,
"max_position_embeddings": llm.max_position_embeddings,
"torch_dtype": llm.torch_dtype,
"source_mirror": mirror,
}
def main() -> int:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--out", default=DEFAULT_OUT)
args = parser.parse_args()
entries: dict[str, dict] = {}
for canonical, mirror in MIRRORS.items():
try:
entries[canonical] = fetch_entry(canonical, mirror)
e = entries[canonical]
print(f" {canonical}: {e['total_b_params']:.3f} B, "
f"{e['num_layers']} layers, via {mirror}")
except Exception as exc: # noqa: BLE001 - report and keep going
print(f" SKIPPED {canonical}: {exc}")
os.makedirs(os.path.dirname(args.out), exist_ok=True)
with open(args.out, "w") as fh:
json.dump(entries, fh, indent=2, sort_keys=True)
print(f"\nwrote {len(entries)} gated-model architectures to {args.out}")
return 0 if entries else 1
if __name__ == "__main__":
sys.exit(main())
|