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