#!/usr/bin/env python3
"""Streaming absmax audit of the Swift1.5 PLE n-gram table (BF16).

The PLE fp8 encoding uses ONE fixed scale (1.9931793212890625e-04, the
official checkpoint's plateau value). fp8_e4m3 saturates at 448, so the
fixed scale is legal iff absmax <= 448 * SCALE. This walks every
ngram_embedding.shard_*.weight tensor in the checkpoint (mmap, no copy),
computes the absmax from the BF16 bit pattern (clear the sign bit, max
the remaining uint16 — BF16 magnitudes are monotone in their bits), and
reports the headroom. Exit 0 only when headroom >= 1.5x.

Run from any cwd EXCEPT /tmp on dankserver (/tmp/inspect.py shadows the
stdlib). Usage: verify-swift15-ple.py <ckpt-dir>
"""
import json
import mmap
import os
import struct
import sys

SCALE = 1.9931793212890625e-04  # bf16 0x3951, the official plateau scale
FP8_E4M3_MAX = 448.0

import numpy as np
gmax = 0.0
gfile = ""
over_cap = 0
total = 0

ckpt = sys.argv[1]
index = json.load(open(os.path.join(ckpt, "model.safetensors.index.json")))
shards: dict[str, list[str]] = {}
for k, f in index["weight_map"].items():
    if "ngram_embedding.shard_" in k and k.endswith(".weight"):
        shards.setdefault(f, []).append(k)

files = len(shards)
tensors = sum(len(v) for v in shards.values())
print(f"n-gram shards: {files} files, {tensors} tensors")

import math

cap = FP8_E4M3_MAX * SCALE
# largest bf16 <= cap: truncate the float32 bits (cap > 0); bf16
# magnitudes are monotone in their bits, so |w| > cap <=> bits > cap_bits
cap_bits = int((np.float32(cap).view(np.uint32) >> 16) & 0xFFFF)
for fname, keys in sorted(shards.items()):
    path = os.path.join(ckpt, fname)
    with open(path, "rb") as fh:
        hlen = struct.unpack("<Q", fh.read(8))[0]
        header = json.loads(fh.read(hlen))
        base = 8 + hlen
        mm = mmap.mmap(fh.fileno(), 0, prot=mmap.PROT_READ)
        for k in keys:
            meta = header[k]
            off, end = meta["data_offsets"]
            a = np.frombuffer(mm, dtype=np.uint16, count=(end - off) // 2,
                              offset=base + off)
            v = a & 0x7FFF
            m = int(v.max())
            over_cap += int((v > np.uint16(cap_bits)).sum())
            total += v.size
            del a
            if m > gmax:
                gmax = m
                gfile = k
        mm.close()

# rebuild the BF16 value: the 16 bits are the HIGH half of the float32
absmax = np.frombuffer(b"\x00\x00" + np.uint16(gmax).tobytes(),
                        dtype=np.float32)[0]

print(f"absmax        = {absmax!r}  (tensor: {gfile})")
print(f"2^log2        = 2^{math.log2(absmax):.3f}")
print(f"fp8 cap       = 448 * {SCALE:.8g} = {cap:.8g}")
print(f"headroom      = {cap / absmax:.3f}x")
ok = cap / absmax >= 1.5
print("VERDICT:", "PASS (>= 1.5x headroom over the fixed scale)" if ok
      else "FAIL — fixed scale does not cover the Swift PLE")
print(f"values > cap  = {over_cap} of {total} ({over_cap / total:.3e})")
nxt_bits = int((np.float32(absmax / FP8_E4M3_MAX).view(np.uint32)
                >> 16) & 0xFFFF) + 1
nxt_scale = np.frombuffer(b"\x00\x00" + np.uint16(nxt_bits).tobytes(),
                          dtype=np.float32)[0]
print(f"next-bf16 scale covering absmax: {nxt_scale:.8g} "
      f"(0x{nxt_bits:04x}, headroom {FP8_E4M3_MAX * nxt_scale / absmax:.4f}x)")
sys.exit(0 if ok else 1)