package_model.py

"""Convert a flat Mettle checkpoint into a Hugging Face safetensors artifact."""

import argparse
import hashlib
import json
import os
from collections import OrderedDict
from pathlib import Path

import torch
from safetensors.torch import save_file


EXPECTED_HEAD_KEYS = {
    "head.P",
    "head.Q",
    "head.mix.0.bias",
    "head.mix.0.weight",
    "head.mix.2.bias",
    "head.mix.2.weight",
    "head.shift.0.bias",
    "head.shift.0.weight",
    "head.shift.2.bias",
    "head.shift.2.weight",
}
RC2_SOURCE_SHA256 = (
    "ad79736b61638a58a72c5dedd7eb21582e94c6397564ae28c1c5da0af4d8a9be"
)


def sha256(path, chunk_size=16 * 1024 * 1024):
    digest = hashlib.sha256()
    with open(path, "rb") as stream:
        while chunk := stream.read(chunk_size):
            digest.update(chunk)
    return digest.hexdigest()


def convert_state_dict(source):
    """Prefix backbone tensors for :class:`MettleModel`."""
    head_keys = {key for key in source if key.startswith("head.")}
    if head_keys != EXPECTED_HEAD_KEYS:
        missing = sorted(EXPECTED_HEAD_KEYS - head_keys)
        unexpected = sorted(head_keys - EXPECTED_HEAD_KEYS)
        raise ValueError(
            "checkpoint refinement-head keys do not match the release contract; "
            f"missing={missing}, unexpected={unexpected}"
        )
    converted = OrderedDict()
    for key in sorted(source):
        target = key if key.startswith("head.") else f"backbone.{key}"
        converted[target] = source[key].contiguous()
    return converted


def validate_shapes(state):
    p_shape = tuple(state["head.P"].shape)
    q_shape = tuple(state["head.Q"].shape)
    if p_shape != (16, 1536, 8) or q_shape != p_shape:
        raise ValueError(
            "checkpoint refinement-head shape does not match RC2: "
            f"P={p_shape}, Q={q_shape}"
        )
    if tuple(state["head.mix.0.weight"].shape) != (256, 1536):
        raise ValueError("checkpoint refinement-head hidden size is not 256")
    if tuple(state["head.shift.2.weight"].shape) != (1536, 256):
        raise ValueError("checkpoint refinement-head output width is not 1536")


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("checkpoint", help="flat RC2 .pt checkpoint")
    parser.add_argument(
        "--output-dir",
        default=str(Path(__file__).resolve().parent),
        help="directory containing the Hugging Face repository files",
    )
    parser.add_argument(
        "--expected-sha256",
        default=RC2_SOURCE_SHA256,
        help="refuse a source checkpoint whose SHA-256 differs",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="validate and map the checkpoint without writing weights",
    )
    return parser.parse_args()


def main():
    args = parse_args()
    checkpoint = Path(args.checkpoint).resolve()
    output_dir = Path(args.output_dir).resolve()
    source_sha = sha256(checkpoint)
    if args.expected_sha256 and source_sha != args.expected_sha256.lower():
        raise SystemExit(
            f"source SHA-256 mismatch: expected {args.expected_sha256}, "
            f"found {source_sha}"
        )

    source = torch.load(checkpoint, map_location="cpu", weights_only=True)
    validate_shapes(source)
    converted = convert_state_dict(source)
    print(
        f"validated {len(source)} tensors; mapped {len(converted)} tensors; "
        f"source sha256={source_sha}"
    )
    if args.dry_run:
        return

    output_dir.mkdir(parents=True, exist_ok=True)
    output_path = output_dir / "model.safetensors"
    manifest_path = output_dir / "weights_manifest.json"
    if output_path.exists():
        raise SystemExit(f"refusing to overwrite {output_path}")
    if manifest_path.exists():
        raise SystemExit(f"refusing to overwrite {manifest_path}")
    temporary = output_dir / ".model.safetensors.incomplete"
    if temporary.exists():
        raise SystemExit(f"refusing to overwrite incomplete file {temporary}")
    save_file(converted, str(temporary), metadata={"format": "pt"})
    os.replace(temporary, output_path)
    packaged_sha = sha256(output_path)
    manifest = {
        "format": "safetensors",
        "model_file": output_path.name,
        "model_sha256": packaged_sha,
        "source_checkpoint_sha256": source_sha,
        "tensor_count": len(converted),
    }
    manifest_path.write_text(
        json.dumps(manifest, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    print(f"wrote {output_path} sha256={packaged_sha}")
    print(f"wrote {manifest_path}")


if __name__ == "__main__":
    main()