File size: 6,546 Bytes
aa7758f | 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 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 | #!/usr/bin/env python3
"""Verify a Q-Prefer adapter without loading the 4B base model."""
from __future__ import annotations
import argparse
import hashlib
import json
from pathlib import Path
from safetensors import safe_open
from qprefer_reward.constants import (
BASE_MODEL_ID,
BASE_MODEL_REVISION,
EXPECTED_SPECIAL_TOKEN_IDS,
PUBLISHED_ADAPTER_SHA256,
PUBLISHED_ADAPTER_TENSOR_COUNT,
PUBLISHED_RM_HEAD_SHAPE,
PUBLISHED_SPECIAL_EMBEDDINGS_SHA256,
SPECIAL_TOKENS,
)
def sha256(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as handle:
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
digest.update(chunk)
return digest.hexdigest()
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("adapter", type=Path)
parser.add_argument("--allow-unknown-checkpoint", action="store_true")
return parser.parse_args()
def verify_manifest(root: Path) -> None:
manifest_path = root / "artifact_manifest.json"
if not manifest_path.is_file():
raise FileNotFoundError(f"required artifact manifest is missing: {manifest_path}")
manifest = json.loads(manifest_path.read_text())
expected_metadata = {
"format_version": 1,
"base_model": BASE_MODEL_ID,
"base_revision": BASE_MODEL_REVISION,
"special_tokens": list(SPECIAL_TOKENS),
"special_token_ids": list(EXPECTED_SPECIAL_TOKEN_IDS),
"supported_dimensions": ["visual_quality", "text_alignment"],
"unsupported_dimensions": ["motion_quality"],
}
mismatches = {
key: (expected, manifest.get(key))
for key, expected in expected_metadata.items()
if manifest.get(key) != expected
}
if mismatches:
raise RuntimeError(f"unexpected artifact manifest metadata: {mismatches}")
files = manifest.get("files")
if not isinstance(files, dict) or not files:
raise RuntimeError("artifact manifest has no file checksums")
required = {
"adapter_config.json",
"adapter_model.safetensors",
"chat_template.jinja",
"README.md",
"special_token_embeddings.safetensors",
"tokenizer.json",
"tokenizer_config.json",
}
missing = required.difference(files)
if missing:
raise RuntimeError(f"artifact manifest is missing required files: {sorted(missing)}")
for filename, record in files.items():
if Path(filename).name != filename or filename == manifest_path.name:
raise RuntimeError(f"invalid artifact filename in manifest: {filename!r}")
if not isinstance(record, dict):
raise RuntimeError(f"invalid manifest record for {filename!r}")
path = root / filename
if not path.is_file():
raise FileNotFoundError(f"manifest file is missing: {path}")
observed_bytes = path.stat().st_size
observed_sha = sha256(path)
if record.get("bytes") != observed_bytes or record.get("sha256") != observed_sha:
raise RuntimeError(
f"artifact integrity check failed for {filename}: "
f"expected bytes/sha256={record.get('bytes')}/{record.get('sha256')}, "
f"observed={observed_bytes}/{observed_sha}"
)
def main() -> None:
args = parse_args()
root = args.adapter.expanduser().resolve()
weights = root / "adapter_model.safetensors"
config_path = root / "adapter_config.json"
if not weights.is_file() or not config_path.is_file():
raise FileNotFoundError("adapter_model.safetensors or adapter_config.json is missing")
verify_manifest(root)
observed_sha = sha256(weights)
if observed_sha != PUBLISHED_ADAPTER_SHA256 and not args.allow_unknown_checkpoint:
raise RuntimeError(
"adapter checksum mismatch: "
f"expected={PUBLISHED_ADAPTER_SHA256}, observed={observed_sha}"
)
config = json.loads(config_path.read_text())
expected_config = {
"peft_type": "LORA",
"r": 64,
"lora_alpha": 128,
"lora_dropout": 0.05,
}
mismatches = {
key: (expected, config.get(key))
for key, expected in expected_config.items()
if config.get(key) != expected
}
if mismatches:
raise RuntimeError(f"unexpected adapter configuration: {mismatches}")
if "rm_head" not in config.get("modules_to_save", []):
raise RuntimeError("adapter config does not preserve rm_head")
with safe_open(weights, framework="pt", device="cpu") as handle:
tensor_keys = list(handle.keys())
reward_head_keys = [key for key in tensor_keys if key.endswith("rm_head.weight")]
if len(tensor_keys) != PUBLISHED_ADAPTER_TENSOR_COUNT:
raise RuntimeError(
"unexpected adapter tensor count: "
f"expected={PUBLISHED_ADAPTER_TENSOR_COUNT}, observed={len(tensor_keys)}"
)
if len(reward_head_keys) != 1:
raise RuntimeError(f"expected one rm_head.weight, found {reward_head_keys}")
shape = tuple(handle.get_tensor(reward_head_keys[0]).shape)
if shape != PUBLISHED_RM_HEAD_SHAPE:
raise RuntimeError(
f"unexpected rm_head shape: expected={PUBLISHED_RM_HEAD_SHAPE}, observed={shape}"
)
embeddings = root / "special_token_embeddings.safetensors"
if not embeddings.is_file():
raise FileNotFoundError(f"required special-token embeddings are missing: {embeddings}")
embedding_sha = sha256(embeddings)
if embedding_sha != PUBLISHED_SPECIAL_EMBEDDINGS_SHA256:
raise RuntimeError(
"special-token embedding checksum mismatch: "
f"expected={PUBLISHED_SPECIAL_EMBEDDINGS_SHA256}, observed={embedding_sha}"
)
with safe_open(embeddings, framework="pt", device="cpu") as handle:
embedding_shape = tuple(handle.get_tensor("special_token_embeddings").shape)
if embedding_shape != PUBLISHED_RM_HEAD_SHAPE:
raise RuntimeError(
"special-token embedding shape must equal (3, hidden_size): "
f"expected={PUBLISHED_RM_HEAD_SHAPE}, observed={embedding_shape}"
)
print("Q-Prefer artifact verified")
print(f"adapter: {root}")
print(f"sha256: {observed_sha}")
print(f"tensors: {len(tensor_keys)}")
print(f"rm_head: {shape}")
print(f"special embeddings sha256: {embedding_sha}")
if __name__ == "__main__":
main()
|