#!/usr/bin/env python3
"""hf_do_ckpt.py (Hart, 09.11) — wagi HF Slayer1B 1.07 -> ckpt trenera r24.2 dla --init-from (bieg Slayer1B 1.08, 8xB200).

Eksport (eksport_hf_1b.py): stan modelu fp32 z kluczami trenera, BEZ head.weight (wiązanie z tok.weight),
config.json = ckpt["config"] + metadane. Tu odwrotność: head.weight = tok.weight, zapis {"model", "config", "step"}.
Bez sieci, bez optymalizatora (--init-from i tak bierze tylko wagi: r24.2 :1128-1140).

Kontrole (STOP przy rozjeździe):
  1. sha256 model.safetensors / config.json / trenera = piny;
  2. wiązanie w config.json = {"tok.weight": "head.weight"} i brak head.weight w safetensors;
  3. strict load do klasy GPT trenera (jak nauczyciel KD, r24.2 :1160-1161) + te same klucze configu, które sprawdza --init-from;
  4. po zapisie: ponowny torch.load(weights_only=True) i każdy tensor == safetensors bit w bit.
Użycie (węzeł): python hf_do_ckpt.py --pkg <katalog HF 1.07> --trainer train_gpt_ref_r24_2.py --out init_107.pt
"""
import argparse
import hashlib
import importlib.util
import json
import os
import sys
import types

import torch
from safetensors.torch import load_file

SHA_MODEL = "d826b4a11bf4eaccee10fe43ec19a4f6d515dfc5f3e5c859adb631f1a0b76ad5"   # HF Slayer1B-1.07 model.safetensors
SHA_CONFIG = "9e2062501a87d0b87d7ef75ba8bc5484f3b1afdf42081e727a74dc02f5a4afce"  # jego config.json
SHA_TRENER = "46f83b51fcec18ffc3273d80fb9eab29e182c52d8667800dab27b80f7ff18afc"  # train_gpt_ref_r24_2.py
KLUCZE_INIT = ("vocab", "n_layer", "n_embd", "n_head")                            # r24.2 :1131-1133
KLUCZE_DFLT = (("loop", 1), ("xsa", 0), ("rope_dims", 0), ("ffn", "gelu"), ("imha_t", 0), ("kv_heads", 0))  # :1134-1138


def stop(msg):
    print(f"STOP: {msg}", file=sys.stderr)
    sys.exit(2)


def sha(path):
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 24), b""):
            h.update(chunk)
    return h.hexdigest()


def main():
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("--pkg", required=True, help="katalog z model.safetensors i config.json")
    ap.add_argument("--trainer", required=True, help="train_gpt_ref_r24_2.py")
    ap.add_argument("--out", required=True)
    ap.add_argument("--sha-model", default=SHA_MODEL)
    ap.add_argument("--sha-config", default=SHA_CONFIG)
    ap.add_argument("--sha-trener", default=SHA_TRENER)
    a = ap.parse_args()

    if os.path.exists(a.out):
        stop(f"{a.out} istnieje")
    model_p, config_p = os.path.join(a.pkg, "model.safetensors"), os.path.join(a.pkg, "config.json")
    for path, want in ((model_p, a.sha_model), (config_p, a.sha_config), (a.trainer, a.sha_trener)):
        got = sha(path)
        if got != want:
            stop(f"{path}: sha {got[:12]} != pin {want[:12]}")

    sd = load_file(model_p, device="cpu")
    with open(config_p, encoding="utf-8") as f:
        cfg = json.load(f)
    if cfg.get("tied_weights") != {"tok.weight": "head.weight"} or "head.weight" in sd or "tok.weight" not in sd:
        stop("nieoczekiwane wiązanie wag (oczekiwane tied_weights tok.weight -> head.weight, bez head.weight w pliku)")
    for key in KLUCZE_INIT:
        if cfg.get(key) is None:
            stop(f"config.json bez '{key}' (wymagane przez --init-from)")
    sd["head.weight"] = sd["tok.weight"]

    spec = importlib.util.spec_from_file_location("trener_r24_2", a.trainer)
    trener = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(trener)  # moduł ma straż __main__: import nie startuje treningu
    with torch.random.fork_rng(devices=[]):
        model = trener.GPT(cfg["vocab"], cfg["n_layer"], cfg["n_embd"], cfg["n_head"], cfg["block"],
                           types.SimpleNamespace(**cfg))
    model.load_state_dict(sd, strict=True)

    state = {"model": model.state_dict(), "config": cfg, "step": int(cfg.get("step", 0)),
             "hf_init": {"model_sha256": a.sha_model, "config_sha256": a.sha_config}}
    tmp = a.out + ".tmp"
    torch.save(state, tmp)
    try:
        ck = torch.load(tmp, map_location="cpu", weights_only=True)
        c0 = ck["config"]
        for key in KLUCZE_INIT:
            if c0[key] != cfg[key]:
                stop(f"config po zapisie: {key} {c0[key]} != {cfg[key]}")
        dflt = {key: c0.get(key, d) for key, d in KLUCZE_DFLT}
        if set(ck["model"]) != set(sd):
            stop(f"klucze po zapisie != safetensors+head ({len(ck['model'])} vs {len(sd)})")
        rozne = [k for k, v in sd.items() if not torch.equal(ck["model"][k], v)]
        if rozne:
            stop(f"tensory po zapisie różne od safetensors: {rozne[:5]}")
    except SystemExit:
        os.remove(tmp)
        raise
    os.replace(tmp, a.out)
    n = sum(v.numel() for k, v in sd.items() if k != "head.weight")
    print(f"HF->CKPT OK {a.out} sha {sha(a.out)} | krok {state['step']} | P {n:,} | klucze {len(sd)} | "
          f"config --init-from {dict((k, cfg[k]) for k in KLUCZE_INIT)} {dflt}", flush=True)


if __name__ == "__main__":
    main()