scripts/convert_evaluation_checkpoint.py

#!/usr/bin/env python3
"""Convert the frozen FP32 evaluation lens to deterministic Safetensors."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
from collections.abc import Mapping
from pathlib import Path
from typing import Any

import safetensors
import torch
from safetensors import safe_open

SOURCE_SHA256 = "8f752032a26a5196c1cb447ef63f01e8a29820ff0c57178dd80c6e26d32b12a9"
SOURCE_PROVENANCE_SHA256 = (
    "b6e5764fa1580a142403a425ecd03cafe55d4cc36ca1e6b90da7fc19a35aad36"
)
SOURCE_VALIDATION_SHA256 = (
    "4aea71008a10ef2d043129f7767e5d387372fe1fd5022550b59017a44e0b965d"
)
SOURCE_MATRIX_STATS_SHA256 = (
    "a21fe7c1f661fa9b5455d89c0a415beae794012eab38f782b667b27fac942401"
)
SOURCE_FIT_CHECKPOINT_SHA256 = (
    "a1236cfe5d04601575b3de150ffe50e3a67e755ede1197ecf74c205a31bdc258"
)
SOURCE_EXPORT_SCRIPT_SHA256 = (
    "d1d9e0b7afd1d69771907a839f62b30936995fdd63b17ea2ac80f724b3cdce17"
)
FP16_ARTIFACT_SHA256 = (
    "089d776979408f23e5377539c15aa8025d171633718ccdce709bcd3372e7942c"
)
MODEL_ID = "WeiboAI/VibeThinker-3B"
MODEL_REVISION = "77bd2cced09193c8b9a59a32bd8577bbd1f3e01c"
SOURCE_LAYERS = tuple(range(0, 36, 2))
TARGET_LAYER = 35
D_MODEL = 2048
N_PROMPTS = 1000
EXPECTED_TOP_LEVEL_KEYS = {"J", "n_prompts", "source_layers", "d_model"}
FORBIDDEN_PUBLIC_FRAGMENTS = (
    os.sep.join(("", "Users", "")),
    os.sep.join(("", "Volumes", "")),
    os.sep.join(("", "workspace")),
)


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(8 * 1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def tensor_storage_bytes(tensor: torch.Tensor) -> bytes:
    """Return C-contiguous little-endian FP32 bytes without value conversion."""

    array = tensor.detach().cpu().contiguous().view(torch.int32).numpy()
    return array.astype("<i4", copy=False).tobytes(order="C")


def tensor_sha256(tensor: torch.Tensor) -> str:
    return hashlib.sha256(tensor_storage_bytes(tensor)).hexdigest()


def write_json(path: Path, value: Any) -> None:
    path.write_text(
        json.dumps(value, indent=2, sort_keys=True, ensure_ascii=True) + "\n",
        encoding="utf-8",
    )


def require(condition: bool, message: str) -> None:
    if not condition:
        raise ValueError(message)


def validate_source(checkpoint: Any) -> dict[int, torch.Tensor]:
    require(isinstance(checkpoint, Mapping), "checkpoint must be a mapping")
    require(
        set(checkpoint) == EXPECTED_TOP_LEVEL_KEYS,
        f"unexpected checkpoint keys: {sorted(checkpoint)}",
    )
    require(checkpoint["n_prompts"] == N_PROMPTS, "unexpected n_prompts")
    require(checkpoint["d_model"] == D_MODEL, "unexpected d_model")
    require(
        tuple(checkpoint["source_layers"]) == SOURCE_LAYERS,
        "unexpected source_layers",
    )
    matrices = checkpoint["J"]
    require(isinstance(matrices, Mapping), "J must be a layer-to-tensor mapping")
    require(set(matrices) == set(SOURCE_LAYERS), "unexpected J layer keys")

    validated: dict[int, torch.Tensor] = {}
    for layer in SOURCE_LAYERS:
        tensor = matrices[layer]
        require(isinstance(tensor, torch.Tensor), f"J[{layer}] is not a tensor")
        require(tensor.device.type == "cpu", f"J[{layer}] is not on CPU")
        require(tensor.dtype == torch.float32, f"J[{layer}] is not FP32")
        require(
            tuple(tensor.shape) == (D_MODEL, D_MODEL),
            f"J[{layer}] shape mismatch",
        )
        require(tensor.is_contiguous(), f"J[{layer}] is not contiguous")
        require(
            bool(torch.isfinite(tensor).all()),
            f"J[{layer}] contains non-finite values",
        )
        validated[layer] = tensor
    return validated


def public_header() -> dict[str, str]:
    return {
        "artifact_kind": "jacobian_lens_evaluation_fp32",
        "d_model": str(D_MODEL),
        "format": "pt",
        "model_id": MODEL_ID,
        "model_revision": MODEL_REVISION,
        "n_prompts": str(N_PROMPTS),
        "schema_version": "1",
        "source_fit_checkpoint_sha256": SOURCE_FIT_CHECKPOINT_SHA256,
        "source_fp32_checkpoint_sha256": SOURCE_SHA256,
        "source_layers": json.dumps(SOURCE_LAYERS, separators=(",", ":")),
        "target_layer": str(TARGET_LAYER),
        "tensor_dtype": "float32",
        "tensor_key_pattern": "J.{source_layer}",
    }


def assert_public_header(metadata: Mapping[str, str]) -> None:
    encoded = json.dumps(dict(metadata), sort_keys=True)
    for fragment in FORBIDDEN_PUBLIC_FRAGMENTS:
        require(
            fragment not in encoded,
            f"private fragment found in Safetensors header: {fragment}",
        )


def save_deterministic_safetensors(
    tensors: Mapping[str, torch.Tensor],
    path: Path,
    metadata: Mapping[str, str],
) -> None:
    offset = 0
    header: dict[str, Any] = {
        "__metadata__": {key: metadata[key] for key in sorted(metadata)}
    }
    for key in sorted(tensors):
        tensor = tensors[key]
        require(tensor.dtype == torch.float32, f"{key} is not FP32")
        nbytes = tensor.numel() * tensor.element_size()
        header[key] = {
            "dtype": "F32",
            "shape": list(tensor.shape),
            "data_offsets": [offset, offset + nbytes],
        }
        offset += nbytes

    encoded_header = json.dumps(
        header,
        ensure_ascii=False,
        separators=(",", ":"),
    ).encode("utf-8")
    encoded_header += b" " * ((-len(encoded_header)) % 8)
    with path.open("wb") as handle:
        handle.write(len(encoded_header).to_bytes(8, "little", signed=False))
        handle.write(encoded_header)
        for key in sorted(tensors):
            handle.write(tensor_storage_bytes(tensors[key]))


def compare_fp16(
    matrices: Mapping[int, torch.Tensor],
    fp16_path: Path,
) -> dict[str, Any]:
    require(
        sha256_file(fp16_path) == FP16_ARTIFACT_SHA256,
        "FP16 companion artifact SHA-256 mismatch",
    )
    per_layer: dict[str, Any] = {}
    error_squared = 0.0
    reference_squared = 0.0
    with safe_open(fp16_path, framework="pt", device="cpu") as fp16_artifact:
        require(
            set(fp16_artifact.keys()) == {f"J.{layer}" for layer in SOURCE_LAYERS},
            "FP16 companion key set mismatch",
        )
        for layer in SOURCE_LAYERS:
            fp32 = matrices[layer]
            stored_fp16 = fp16_artifact.get_tensor(f"J.{layer}")
            cast_fp16 = fp32.to(torch.float16)
            require(
                torch.equal(cast_fp16.view(torch.int16), stored_fp16.view(torch.int16)),
                f"FP32-to-FP16 cast differs at J.{layer}",
            )
            error = cast_fp16.float() - fp32
            error_norm = float(torch.linalg.vector_norm(error))
            reference_norm = float(torch.linalg.vector_norm(fp32))
            layer_error_squared = error_norm**2
            layer_reference_squared = reference_norm**2
            error_squared += layer_error_squared
            reference_squared += layer_reference_squared
            per_layer[str(layer)] = {
                "cast_matches_model_safetensors_exactly": True,
                "max_absolute_error": float(error.abs().max().item()),
                "relative_frobenius_error": error_norm / reference_norm,
            }
    return {
        "schema_version": 1,
        "artifact_kind": "fp32_to_fp16_lens_compatibility",
        "fp32_source_checkpoint_sha256": SOURCE_SHA256,
        "fp16_artifact": "model.safetensors",
        "fp16_artifact_sha256": FP16_ARTIFACT_SHA256,
        "conversion": "IEEE_FP32_to_FP16_round_to_nearest_even",
        "all_layer_casts_match_exactly": True,
        "relative_frobenius_error": (error_squared / reference_squared) ** 0.5,
        "max_absolute_error": max(
            record["max_absolute_error"] for record in per_layer.values()
        ),
        "per_layer": per_layer,
    }


def convert(source: Path, output_dir: Path, overwrite: bool) -> dict[str, Any]:
    source_digest = sha256_file(source)
    require(source_digest == SOURCE_SHA256, "source checkpoint SHA-256 mismatch")
    checkpoint = torch.load(source, map_location="cpu", weights_only=True)
    matrices = validate_source(checkpoint)
    tensors = {f"J.{layer}": matrices[layer] for layer in SOURCE_LAYERS}

    output_dir.mkdir(parents=True, exist_ok=True)
    output_path = output_dir / "evaluation.safetensors"
    if output_path.exists() and not overwrite:
        raise FileExistsError(
            f"refusing to overwrite {output_path.name}; pass --overwrite"
        )
    metadata = public_header()
    assert_public_header(metadata)
    temporary_path = output_dir / ".evaluation.safetensors.tmp"
    save_deterministic_safetensors(tensors, temporary_path, metadata)
    os.replace(temporary_path, output_path)

    manifest_tensors: dict[str, Any] = {}
    exact_matches = 0
    with safe_open(output_path, framework="pt", device="cpu") as artifact:
        stored_metadata = artifact.metadata() or {}
        require(stored_metadata == metadata, "Safetensors metadata changed")
        assert_public_header(stored_metadata)
        require(set(artifact.keys()) == set(tensors), "Safetensors key set mismatch")
        for layer in SOURCE_LAYERS:
            key = f"J.{layer}"
            source_tensor = matrices[layer]
            output_tensor = artifact.get_tensor(key)
            require(output_tensor.dtype == torch.float32, f"{key} dtype mismatch")
            require(
                tuple(output_tensor.shape) == (D_MODEL, D_MODEL),
                f"{key} shape mismatch",
            )
            require(output_tensor.is_contiguous(), f"{key} is not contiguous")
            require(torch.equal(output_tensor, source_tensor), f"{key} value mismatch")
            require(
                torch.equal(
                    output_tensor.view(torch.int32),
                    source_tensor.view(torch.int32),
                ),
                f"{key} FP32 bit-pattern mismatch",
            )
            source_tensor_digest = tensor_sha256(source_tensor)
            require(
                tensor_sha256(output_tensor) == source_tensor_digest,
                f"{key} raw byte hash mismatch",
            )
            exact_matches += 1
            manifest_tensors[key] = {
                "dtype": "float32",
                "nbytes": source_tensor.numel() * source_tensor.element_size(),
                "numel": source_tensor.numel(),
                "sha256_c_contiguous_little_endian_bytes": source_tensor_digest,
                "shape": list(source_tensor.shape),
                "source_layer": layer,
            }

    output_digest = sha256_file(output_path)
    output_size = output_path.stat().st_size
    manifest = {
        "schema_version": 1,
        "artifact": "evaluation.safetensors",
        "artifact_sha256": output_digest,
        "artifact_size_bytes": output_size,
        "source_checkpoint_sha256": source_digest,
        "tensor_count": len(manifest_tensors),
        "tensor_storage_bytes": sum(
            record["nbytes"] for record in manifest_tensors.values()
        ),
        "tensors": manifest_tensors,
    }
    compatibility = compare_fp16(matrices, output_dir / "model.safetensors")
    compatibility["fp32_artifact"] = "evaluation.safetensors"
    compatibility["fp32_artifact_sha256"] = output_digest
    validation = {
        "schema_version": 1,
        "artifact": "evaluation.safetensors",
        "artifact_sha256": output_digest,
        "artifact_size_bytes": output_size,
        "source_checkpoint_sha256": source_digest,
        "exact_tensor_matches": exact_matches,
        "expected_tensor_matches": len(SOURCE_LAYERS),
        "tensor_values_changed": 0,
        "checks": {
            "all_source_tensors_contiguous": True,
            "all_source_tensors_finite": True,
            "all_source_tensors_fp32": True,
            "all_source_tensors_shape_2048x2048": True,
            "fp16_cast_matches_companion_artifact": True,
            "roundtrip_all_tensor_byte_hashes_equal": True,
            "roundtrip_all_tensor_dtypes_equal": True,
            "roundtrip_all_tensor_shapes_equal": True,
            "roundtrip_all_tensor_values_equal": True,
            "roundtrip_all_tensor_bit_patterns_equal": True,
            "roundtrip_key_set_exact": True,
            "safetensors_header_public_safe": True,
            "safetensors_header_roundtrip_exact": True,
            "source_checkpoint_sha256_exact": True,
            "source_metadata_exact": True,
            "source_top_level_key_set_exact": True,
        },
        "ok": exact_matches == len(SOURCE_LAYERS),
    }
    provenance = {
        "schema_version": 1,
        "artifact_kind": "jacobian_lens_evaluation_fp32_provenance",
        "artifact": {
            "filename": "evaluation.safetensors",
            "format": "safetensors",
            "sha256": output_digest,
            "size_bytes": output_size,
            "tensor_conversion": "lossless_fp32_reserialization",
        },
        "source_checkpoint": {
            "format": "pytorch",
            "sha256": source_digest,
        },
        "derivation": {
            "formula": "jacobian_sum[layer] / n_done",
            "n_done": N_PROMPTS,
            "source_fit_checkpoint_sha256": SOURCE_FIT_CHECKPOINT_SHA256,
            "source_export_script_sha256": SOURCE_EXPORT_SCRIPT_SHA256,
            "source_matrix_stats_sha256": SOURCE_MATRIX_STATS_SHA256,
            "source_provenance_sha256": SOURCE_PROVENANCE_SHA256,
            "source_validation_sha256": SOURCE_VALIDATION_SHA256,
        },
        "model": {
            "architecture": "Qwen2ForCausalLM",
            "d_model": D_MODEL,
            "id": MODEL_ID,
            "n_layers": 36,
            "revision": MODEL_REVISION,
            "revision_binding": "inferred_hub_head_unchanged_since_before_fit",
            "revision_last_modified": "2026-06-30T11:35:41+00:00",
            "tied_embeddings": True,
        },
        "lens": {
            "d_model": D_MODEL,
            "dtype": "float32",
            "hook_convention": "post_transformer_block_output_residual",
            "n_prompts": N_PROMPTS,
            "source_layers": list(SOURCE_LAYERS),
            "target_layer": TARGET_LAYER,
            "estimator": {
                "dim_batch": 8,
                "exclude_final_position": True,
                "fit_dtype": "bfloat16",
                "max_seq_len": 128,
                "name": "causal_all_current_and_future_targets_mean_jacobian",
                "prompt_aggregation": "equal_weight_mean_over_prompts",
                "skip_first": 16,
                "source_position_aggregation": "mean_over_valid_source_positions",
                "target_position_aggregation": (
                    "sum_over_valid_targets_at_or_after_source"
                ),
            },
        },
        "compatibility": {
            "file": "evaluation_compatibility.json",
            "fp16_artifact": "model.safetensors",
            "fp16_artifact_sha256": FP16_ARTIFACT_SHA256,
            "all_layer_casts_match_exactly": True,
        },
        "software": {
            "conversion_runtime": {
                "python_implementation": "CPython",
                "safetensors": safetensors.__version__,
                "torch": torch.__version__,
            },
            "anthropic_jacobian_lens_commit": (
                "581d398613e5602a5af361e1c34d3a92ea82ba8e"
            ),
        },
        "limitations": [
            "The original fit did not store the resolved Hugging Face commit; the revision binding was reconstructed from the Hub head and its last-modified timestamp.",
            "The recorded evaluation validates token readout and does not establish causal steering or a global workspace.",
        ],
    }

    write_json(output_dir / "evaluation_tensor_manifest.json", manifest)
    write_json(output_dir / "evaluation_compatibility.json", compatibility)
    write_json(output_dir / "evaluation_validation.json", validation)
    write_json(output_dir / "evaluation_provenance.json", provenance)
    return {
        "artifact": output_path.name,
        "artifact_sha256": output_digest,
        "artifact_size_bytes": output_size,
        "exact_tensor_matches": exact_matches,
        "fp16_cast_matches": compatibility["all_layer_casts_match_exactly"],
        "source_checkpoint_sha256": source_digest,
    }


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--source", required=True, type=Path)
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=Path(__file__).resolve().parents[1],
    )
    parser.add_argument("--overwrite", action="store_true")
    args = parser.parse_args()
    result = convert(args.source, args.output_dir.resolve(), args.overwrite)
    print(json.dumps(result, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()