File size: 2,791 Bytes
ce235ae | 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 | """Fetch and load the exact NanoJev checkpoint selected for conversion."""
from __future__ import annotations
import hashlib
import importlib.util
import json
from pathlib import Path
import torch
from huggingface_hub import snapshot_download
from safetensors.torch import load_file
from transformers import AutoConfig, AutoModel, AutoTokenizer
ROOT = Path(__file__).resolve().parent
LOCK = json.loads((ROOT / "assets.lock.json").read_text())
def sha256(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as source:
for chunk in iter(lambda: source.read(8 * 1024 * 1024), b""):
digest.update(chunk)
return digest.hexdigest()
def snapshot(*, with_weights: bool) -> Path:
patterns = list(LOCK["source_files"])
if with_weights:
patterns.append(LOCK["checkpoint"]["selected_weight"])
path = snapshot_download(
LOCK["checkpoint"]["repo"],
revision=LOCK["checkpoint"]["revision"],
allow_patterns=patterns,
token=False,
)
return Path(path)
def upstream_module(root: Path, name: str):
path = root / "source/scripts" / f"{name}.py"
spec = importlib.util.spec_from_file_location(f"nanojev_pinned_{name}", path)
if spec is None or spec.loader is None:
raise RuntimeError(f"Cannot load audited NanoJev source: {path}")
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module
def load_model() -> tuple[Path, AutoTokenizer, torch.nn.Module]:
root = snapshot(with_weights=True)
weights = root / LOCK["checkpoint"]["selected_weight"]
if sha256(weights) != LOCK["checkpoint"]["sha256"]:
raise ValueError("NanoJev trained checkpoint SHA256 mismatch")
run_config = json.loads((root / "config.json").read_text())
if run_config["set_head"] != "attention":
raise ValueError("Expected trained attention set head")
if run_config["resolved_model_revision"] != LOCK["upstream_base"]["revision"]:
raise ValueError("NanoJev base revision mismatch")
tokenizer = AutoTokenizer.from_pretrained(root / "tokenizer", local_files_only=True, trust_remote_code=False)
if tokenizer.pad_token_id is None:
tokenizer.pad_token = tokenizer.eos_token
config = AutoConfig.from_pretrained(root / "backbone_config", local_files_only=True, trust_remote_code=False)
config.use_cache = False
backbone = AutoModel.from_config(config, attn_implementation="sdpa", trust_remote_code=False).float()
model = upstream_module(root, "train_toy_decisions").DecisionModel(backbone, run_config["set_head"])
parameters = load_file(str(weights), device="cpu")
model.load_state_dict(parameters, strict=True)
del parameters
model.eval()
return root, tokenizer, model
|