File size: 2,068 Bytes
9204f7b | 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 | """Offline integrity and model-structure checks for the release directory."""
from __future__ import annotations
import hashlib
import json
from pathlib import Path
import torch
from safetensors import safe_open
from safetensors.torch import load_file
from modeling_anka import RMW3, RMW3Config, count_parameters
ROOT = Path(__file__).resolve().parent
EXPECTED_PARAMETERS = 48_944_657
def sha256(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as stream:
for chunk in iter(lambda: stream.read(8 * 1024 * 1024), b""):
digest.update(chunk)
return digest.hexdigest()
def main() -> None:
required = (
"model.safetensors",
"tokenizer.json",
"config.json",
"modeling_anka.py",
"inference.py",
)
for name in required:
if not (ROOT / name).is_file():
raise FileNotFoundError(name)
config_data = json.loads((ROOT / "config.json").read_text())
config = RMW3Config(
**{field: config_data[field] for field in RMW3Config.__dataclass_fields__}
)
model = RMW3(config).to(dtype=torch.bfloat16)
state = load_file(str(ROOT / "model.safetensors"), device="cpu")
model.load_state_dict(state, strict=True)
if count_parameters(model) != EXPECTED_PARAMETERS:
raise RuntimeError("Parameter-count check failed")
if model.lm_head.weight.data_ptr() != model.token_embeddings.weight.data_ptr():
raise RuntimeError("Weight-tying check failed")
if not torch.equal(state["lm_head.weight"], state["token_embeddings.weight"]):
raise RuntimeError("Serialized tied weights differ")
with safe_open(ROOT / "model.safetensors", framework="pt", device="cpu") as file:
if len(file.keys()) != 199:
raise RuntimeError("Unexpected tensor count")
print("ANKA-50M RMW3 release: PASS")
print(f"Parameters : {EXPECTED_PARAMETERS:,}")
print(f"Tensors : {len(state)}")
print(f"Weights SHA: {sha256(ROOT / 'model.safetensors')}")
if __name__ == "__main__":
main()
|