| """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() |
|
|
