#!/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

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

gmax = 0.0
gfile = ""
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"]
            assert meta["dtype"] == "BF16", f"{k}: dtype {meta['dtype']}"
            a = np.frombuffer(mm, dtype=np.uint16, count=(end - off) // 2,
                              offset=base + off)
            m = int((a & 0x7FFF).max())
            if m > gmax:
                gmax = m
                gfile = k
        mm.close()

# rebuild the BF16 value from the sign-cleared uint16 bit pattern
absmax = np.frombuffer(np.uint16(gmax).tobytes() + b"\x00\x00",
                       dtype=np.float32)[0]
cap = FP8_E4M3_MAX * SCALE
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")
sys.exit(0 if ok else 1)