File size: 1,388 Bytes
d17e08e | 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 | """Load Slayer weights (model.safetensors + config.json) into the GPT class of train_gpt_ref_r23.py.
Usage:
import load_slayer
model = load_slayer.load("path/to/Slayer1B-1.00") # float32, eval mode, on CPU
logits, _ = model(token_ids) # token_ids: LongTensor [batch, seq], seq <= config["block"]
Tokenizer: tokenizers.Tokenizer.from_file("tokenizer.json"); id 0 is end-of-text.
"""
import importlib.util
import json
import os
from types import SimpleNamespace
from safetensors.torch import load_file
def _trainer(path):
spec = importlib.util.spec_from_file_location("train_gpt_ref_r23", os.path.join(path, "train_gpt_ref_r23.py"))
mod = importlib.util.module_from_spec(spec)
spec.loader.exec_module(mod)
return mod
def build(cfg, path):
"""Empty model with the architecture from config (no weights loaded)."""
c = SimpleNamespace(**cfg)
return _trainer(path).GPT(c.vocab, c.n_layer, c.n_embd, c.n_head, c.block, c)
def load(path, device="cpu"):
with open(os.path.join(path, "config.json"), encoding="utf-8") as f:
cfg = json.load(f)
model = build(cfg, path)
sd = load_file(os.path.join(path, "model.safetensors"))
for dst, src in cfg.get("tied_weights", {}).items():
sd.setdefault(src, sd[dst])
model.load_state_dict(sd, strict=True)
return model.to(device).float().eval()
|