figrid-board 0.8.3

#!/usr/bin/env python3
"""Sanitize and verify the codebook shipped in public build artifacts.

The runtime loader only needs the model dimensions and numeric weights.  This
tool removes all training-only metadata while freezing both the loader-visible
f32 values and their deployed quantized representation.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import re
import struct
import sys
from pathlib import Path
from typing import Any


ASSET_RELATIVE = Path("models/gomoku_codebook_v1_swapclosed.json")
PUBLIC_FORMAT = "noru-pattern4-codebook-eval-v1"
PUBLIC_MODEL = "codebook-region-fm"
EXPECTED_DIM = 16
EXPECTED_FM_RANK = 8
EXPECTED_REGIONS = 9
EXPECTED_LENGTHS = {"embeddings": 68_256, "head": 144, "factors": 1_152}
EXPECTED_SEMANTIC_SHA256 = (
    "C88366A868FD7203944E64E6AE8056144E796ECA1EF1CDE913FAE08CD9EB96CA"
)
EXPECTED_QUANTIZED_SHA256 = (
    "BDEAF5A5156BA2B3AD2B29B7D33843E936763D844973F73A36E023CE83E5228A"
)

ROOT_KEYS = {
    "embedding_dim",
    "fm_rank",
    "format",
    "model",
    "regions",
    "weights",
}
WEIGHT_KEYS = {"bias", "embeddings", "factors", "head"}
FORBIDDEN_KEYS = {
    "metadata",
    "normalization",
    "objective",
    "target_mode",
    "pattern_id_migration",
    "pattern_id_migration_source",
}


class CheckError(RuntimeError):
    pass


def f32(value: Any) -> float:
    if isinstance(value, bool) or not isinstance(value, (int, float)):
        raise CheckError(f"weight is not a JSON number: {value!r}")
    converted = struct.unpack("<f", struct.pack("<f", float(value)))[0]
    if not math.isfinite(converted):
        raise CheckError(f"weight is not finite f32: {value!r}")
    return converted


def f32_bytes(value: Any) -> bytes:
    return struct.pack("<f", f32(value))


def dimension(root: dict[str, Any], key: str) -> int:
    value = root.get(key)
    if value is None and isinstance(root.get("metadata"), dict):
        value = root["metadata"].get(key)
    if isinstance(value, bool) or not isinstance(value, int):
        raise CheckError(f"missing integer dimension {key}")
    return value


def weights(root: dict[str, Any]) -> dict[str, Any]:
    value = root.get("weights")
    if not isinstance(value, dict):
        raise CheckError("missing weights object")
    for key, expected in EXPECTED_LENGTHS.items():
        array = value.get(key)
        if not isinstance(array, list) or len(array) != expected:
            got = len(array) if isinstance(array, list) else None
            raise CheckError(f"{key} length {got!r}, expected {expected}")
    f32(value.get("bias"))
    return value


def semantic_payload(root: dict[str, Any]) -> bytes:
    model_weights = weights(root)
    out = bytearray(b"figrid-codebook-semantic-v1\0")
    for key in ("embedding_dim", "fm_rank", "regions"):
        out.extend(struct.pack("<Q", dimension(root, key)))
    for key in ("embeddings", "head", "factors"):
        array = model_weights[key]
        out.extend(struct.pack("<Q", len(array)))
        for value in array:
            out.extend(f32_bytes(value))
    out.extend(f32_bytes(model_weights["bias"]))
    return bytes(out)


def rust_round(value: float) -> int:
    return math.floor(value + 0.5) if value >= 0.0 else math.ceil(value - 0.5)


def quantize(value: Any, scale: int) -> int:
    product = f32(f32(value) * f32(scale))
    return max(-32_768, min(32_767, rust_round(product)))


def quantized_payload(root: dict[str, Any]) -> bytes:
    model_weights = weights(root)
    out = bytearray(b"figrid-codebook-quantized-v1\0")
    for value in (
        dimension(root, "embedding_dim"),
        dimension(root, "fm_rank"),
        dimension(root, "regions"),
        32,
        64,
        64,
    ):
        out.extend(struct.pack("<q", value))
    for key, scale in (("embeddings", 32), ("head", 64), ("factors", 64)):
        array = model_weights[key]
        out.extend(struct.pack("<Q", len(array)))
        for value in array:
            out.extend(struct.pack("<h", quantize(value, scale)))
    out.extend(f32_bytes(model_weights["bias"]))
    return bytes(out)


def sha256(data: bytes) -> str:
    return hashlib.sha256(data).hexdigest().upper()


def sanitized_root(root: dict[str, Any]) -> dict[str, Any]:
    model_weights = weights(root)
    sanitized = {
        "embedding_dim": dimension(root, "embedding_dim"),
        "fm_rank": dimension(root, "fm_rank"),
        "format": PUBLIC_FORMAT,
        "model": PUBLIC_MODEL,
        "regions": dimension(root, "regions"),
        "weights": {
            "bias": model_weights["bias"],
            "embeddings": model_weights["embeddings"],
            "factors": model_weights["factors"],
            "head": model_weights["head"],
        },
    }
    return sanitized


def canonical_bytes(root: dict[str, Any]) -> bytes:
    return (
        json.dumps(
            sanitized_root(root),
            ensure_ascii=False,
            allow_nan=False,
            separators=(",", ":"),
        ).encode("utf-8")
        + b"\n"
    )


def reject_private_metadata(raw: bytes, root: dict[str, Any]) -> None:
    if set(root) != ROOT_KEYS:
        raise CheckError(f"root keys {sorted(root)}, expected {sorted(ROOT_KEYS)}")
    model_weights = root.get("weights")
    if not isinstance(model_weights, dict) or set(model_weights) != WEIGHT_KEYS:
        actual = sorted(model_weights) if isinstance(model_weights, dict) else None
        raise CheckError(f"weight keys {actual!r}, expected {sorted(WEIGHT_KEYS)}")
    forbidden_present = sorted(FORBIDDEN_KEYS.intersection(root))
    if forbidden_present:
        raise CheckError(f"forbidden root keys: {forbidden_present}")

    # Build sensitive names in pieces so this verifier does not itself trip
    # the release-wide literal scanner it is designed to complement.
    forbidden_literals = (
        "ra" + "pfi",
        "pe" + "la",
        "sea" + "son",
        "exper" + "iments",
        ".work" + "trees",
    )
    text = raw.decode("utf-8").lower()
    for literal in forbidden_literals:
        if literal in text:
            raise CheckError(f"forbidden literal present: {literal!r}")
    if re.search(r"\brq\d{3}[a-z0-9_-]*\b", text, flags=re.IGNORECASE):
        raise CheckError("internal numbered identifier present")
    if re.search(r"[a-z]:\\(?:users|documents)\\", text, flags=re.IGNORECASE):
        raise CheckError("absolute Windows path present")


def verify(asset: Path, raw: bytes, root: dict[str, Any]) -> None:
    if dimension(root, "embedding_dim") != EXPECTED_DIM:
        raise CheckError("embedding_dim changed")
    if dimension(root, "fm_rank") != EXPECTED_FM_RANK:
        raise CheckError("fm_rank changed")
    if dimension(root, "regions") != EXPECTED_REGIONS:
        raise CheckError("regions changed")
    reject_private_metadata(raw, root)
    if raw != canonical_bytes(root):
        raise CheckError("asset is not the deterministic public serialization")

    semantic = sha256(semantic_payload(root))
    quantized = sha256(quantized_payload(root))
    if semantic != EXPECTED_SEMANTIC_SHA256:
        raise CheckError(f"semantic f32 SHA-256 {semantic}, expected {EXPECTED_SEMANTIC_SHA256}")
    if quantized != EXPECTED_QUANTIZED_SHA256:
        raise CheckError(
            f"quantized SHA-256 {quantized}, expected {EXPECTED_QUANTIZED_SHA256}"
        )
    print(
        "PASS_PUBLIC_CODEBOOK_ASSET "
        f"path={asset} bytes={len(raw)} semantic_sha256={semantic} "
        f"quantized_sha256={quantized} asset_sha256={sha256(raw)}"
    )


def parse(raw: bytes) -> dict[str, Any]:
    def reject_duplicate_keys(pairs: list[tuple[str, Any]]) -> dict[str, Any]:
        result: dict[str, Any] = {}
        for key, value in pairs:
            if key in result:
                raise CheckError(f"duplicate JSON key: {key!r}")
            result[key] = value
        return result

    def reject_nonstandard_number(value: str) -> Any:
        raise CheckError(f"non-standard JSON number: {value}")

    try:
        root = json.loads(
            raw,
            object_pairs_hook=reject_duplicate_keys,
            parse_constant=reject_nonstandard_number,
        )
    except (UnicodeDecodeError, json.JSONDecodeError) as error:
        raise CheckError(f"invalid UTF-8 JSON: {error}") from error
    if not isinstance(root, dict):
        raise CheckError("asset root must be an object")
    return root


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--asset",
        type=Path,
        default=Path(__file__).resolve().parents[1] / ASSET_RELATIVE,
    )
    parser.add_argument(
        "--write",
        action="store_true",
        help="replace a full training artifact with the deterministic public form",
    )
    args = parser.parse_args()

    try:
        raw = args.asset.read_bytes()
        root = parse(raw)
        semantic_before = sha256(semantic_payload(root))
        quantized_before = sha256(quantized_payload(root))
        if semantic_before != EXPECTED_SEMANTIC_SHA256:
            raise CheckError("refusing to sanitize an asset with different f32 semantics")
        if quantized_before != EXPECTED_QUANTIZED_SHA256:
            raise CheckError("refusing to sanitize an asset with different quantized semantics")
        if args.write:
            args.asset.write_bytes(canonical_bytes(root))
            raw = args.asset.read_bytes()
            root = parse(raw)
        verify(args.asset, raw, root)
        return 0
    except (OSError, CheckError) as error:
        print(f"PUBLIC_CODEBOOK_ASSET_ERROR: {error}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())