#!/usr/bin/env python3
"""Dequantize gpt-oss-20b MXFP4 expert weights -> a plain BF16 safetensors dir.

Only the expert MLPs are MXFP4 (`*_blocks` uint8 packed 4-bit + `*_scales` uint8
E8M0, block=32); everything else is already BF16. We decode experts to BF16 and
re-emit a standard HF-format dir so xserv's normal safetensors loader reads it
(keeps the first MoE pass free of any MXFP4 code in Rust).

Fused expert outputs (per layer i), matching HF `GptOssExperts` param shapes:
  model.layers.{i}.mlp.experts.gate_up_proj   [E, hidden, 2*inter]  bf16
  model.layers.{i}.mlp.experts.down_proj      [E, inter, hidden]    bf16
  (*_bias tensors pass through unchanged)

NOTE on transpose: the MXFP4 `_blocks` decode to [E, OUT, IN] (out-major, the
contraction dim = nblk*32 last). HF's nn.Parameter for these is [E, IN, OUT]
(it does `x @ gate_up_proj`). We emit [E, IN, OUT] (transpose last two dims) so
the names/shapes match HF exactly and xserv can treat them uniformly.

Run on the GPU host (torch + the model + disk):
  python3 tools/gptoss_dequant.py /opt/wjh/models/gpt-oss-20b /opt/wjh/models/gpt-oss-20b-bf16
"""
import sys, os, json, glob
import torch
from safetensors import safe_open
from safetensors.torch import save_file

# FP4 E2M1 code -> value (OCP MX). 16 entries.
FP4 = torch.tensor(
    [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0,
     -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0], dtype=torch.float32)

def dequant(blocks: torch.Tensor, scales: torch.Tensor) -> torch.Tensor:
    """blocks uint8 [..., nblk, 16], scales uint8 [..., nblk] -> bf16 [..., nblk*32]."""
    blocks = blocks.to(torch.int64)
    lo = blocks & 0xF
    hi = (blocks >> 4) & 0xF
    codes = torch.stack([lo, hi], dim=-1).reshape(*blocks.shape[:-1], blocks.shape[-1] * 2)
    vals = FP4[codes]                                   # [..., nblk, 32] f32
    scale = torch.exp2(scales.to(torch.float32) - 127.0)  # [..., nblk]
    out = vals * scale[..., None]
    out = out.reshape(*out.shape[:-2], out.shape[-2] * out.shape[-1])
    return out.to(torch.bfloat16)

def main():
    src, dst = sys.argv[1], sys.argv[2]
    os.makedirs(dst, exist_ok=True)
    wm = json.load(open(os.path.join(src, "model.safetensors.index.json")))["weight_map"]

    shards = {}
    for name, shard in wm.items():
        shards.setdefault(shard, []).append(name)

    out = {}
    for shard in sorted(shards):
        h = safe_open(os.path.join(src, shard), framework="pt")
        keys = set(h.keys())
        for name in shards[shard]:
            if name.endswith("_scales"):
                continue
            if name.endswith("_blocks"):
                base = name[:-len("_blocks")]                 # ...gate_up_proj / ...down_proj
                sc = base + "_scales"
                sc_h = h if sc in keys else safe_open(os.path.join(src, wm[sc]), framework="pt")
                deq = dequant(h.get_tensor(name), sc_h.get_tensor(sc))  # [E, OUT, IN]
                out[base] = deq.transpose(1, 2).contiguous()           # [E, IN, OUT] (HF param layout)
                print("dequant", base, tuple(out[base].shape), flush=True)
            else:
                out[name] = h.get_tensor(name)                # already bf16/other; pass through

    save_file(out, os.path.join(dst, "model.safetensors"), metadata={"format": "pt"})
    for f in glob.glob(os.path.join(src, "*.json")) + glob.glob(os.path.join(src, "*.jinja")):
        b = os.path.basename(f)
        if b == "model.safetensors.index.json":
            continue
        open(os.path.join(dst, b), "wb").write(open(f, "rb").read())
    print("DEQUANT_DONE ->", dst, flush=True)

if __name__ == "__main__":
    main()