#!/usr/bin/env python3
"""Convert native Gab weights (a headerless flat float32 buffer in MiniGPT.js
parameters() order) to model.safetensors with Transformers-style names.

MiniGPT.js stores matrices as [in, out] row-major; nn.Linear wants [out, in].
The per-head Wq/Wk/Wv/Wo are fused into one [768, 768] matrix each: Q/K/V by
stacking the heads' outputs, O by stacking the heads' inputs (summing per-head
output projections equals one projection of the concatenated heads).

Usage: convert.py gab1b.weights [config.json] [model.safetensors]
"""
import json
import sys

import numpy as np
from safetensors.numpy import save_file


def convert(weights_path, config_path="config.json", out_path="model.safetensors"):
    c = json.load(open(config_path))
    d, nh, hd = c["hidden_size"], c["num_attention_heads"], c["head_dim"]
    ne, h, v = c["num_experts"], c["intermediate_size"], c["vocab_size"]
    flat = np.fromfile(weights_path, dtype="<f4")
    pos = 0

    def take(rows, cols):
        nonlocal pos
        m = flat[pos:pos + rows * cols].reshape(rows, cols)
        pos += rows * cols
        return m

    t = {"model.embed_tokens.weight": take(v, d)}
    for b in range(c["num_hidden_layers"]):
        p = f"model.layers.{b}."
        t[p + "input_layernorm.weight"] = take(1, d)[0]
        q, k, val, o = [], [], [], []
        for _ in range(nh):
            q.append(take(d, hd)); k.append(take(d, hd)); val.append(take(d, hd)); o.append(take(hd, d))
        t[p + "self_attn.q_proj.weight"] = np.concatenate(q, axis=1).T
        t[p + "self_attn.k_proj.weight"] = np.concatenate(k, axis=1).T
        t[p + "self_attn.v_proj.weight"] = np.concatenate(val, axis=1).T
        t[p + "self_attn.o_proj.weight"] = np.concatenate(o, axis=0).T
        t[p + "post_attention_layernorm.weight"] = take(1, d)[0]
        t[p + "mlp.router.weight"] = take(d, ne).T
        for e in range(ne):
            q_ = f"{p}mlp.experts.{e}."
            t[q_ + "up_proj.weight"] = take(d, h).T
            t[q_ + "down_proj.weight"] = take(h, d).T
            t[q_ + "up_proj.bias"] = take(1, h)[0]
    t["model.norm.weight"] = take(1, d)[0]
    if pos != flat.size:
        sys.exit(f"Error: read {pos:,} floats but the file has {flat.size:,}; config doesn't match.")
    save_file({n: np.ascontiguousarray(a, dtype=np.float32) for n, a in t.items()}, out_path,
              metadata={"format": "pt"})
    print(f"{flat.size:,} parameters -> {out_path} ({len(t)} tensors)")


if __name__ == "__main__":
    convert(*sys.argv[1:])