File size: 4,632 Bytes
0e83a2b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 | """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()
|