| """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() |
|
|
