validate_repo.py

#!/usr/bin/env python3
"""Validate repository structure, config, tokenizer IDs, and Safetensors index."""

from __future__ import annotations

import json
import math
import re
import struct
from pathlib import Path


ROOT = Path(__file__).resolve().parent
INTERNAL_PATH_RE = re.compile(r"/(?:wbl" r"-fast|home)/|usr" r"s/|skt-" r"shared")
EXPECTED_SPECIAL_TOKENS = {
    163692: "<|pad|>",
    163702: "<|im_start|>",
    163704: "<|im_end|>",
    163840: "<|audio_start|>",
    163841: "<|audio_end|>",
    163842: "<|speaker_embedding_placeholder|>",
    163843: "<|audio_output_placeholder|>",
    163844: "<|audio_input_placeholder|>",
    163845: "<|audio_output_pad|>",
    163846: "<|audio_output_end_pad|>",
}
REQUIRED_FILES = {
    ".gitattributes",
    "README.md",
    "SHA256SUMS",
    "chat_template.jinja",
    "config.json",
    "configuration_axk1.py",
    "configuration_raon.py",
    "model.safetensors.index.json",
    "modeling_axk1.py",
    "modeling_raon.py",
    "provenance.json",
    "requirements.txt",
    "special_tokens_map.json",
    "tokenizer.json",
    "tokenizer_config.json",
    "validate_model.py",
}


def load_json(name: str) -> dict:
    with (ROOT / name).open(encoding="utf-8") as handle:
        return json.load(handle)


def iter_named_values(value: object, name: str = ""):
    if isinstance(value, dict):
        for key, child in value.items():
            child_name = f"{name}.{key}" if name else key
            yield from iter_named_values(child, child_name)
    elif isinstance(value, list):
        for index, child in enumerate(value):
            yield from iter_named_values(child, f"{name}[{index}]")
    else:
        yield name, value


def read_safetensors_header(path: Path) -> dict:
    with path.open("rb") as handle:
        raw_length = handle.read(8)
        if len(raw_length) != 8:
            raise AssertionError(f"{path.name}: missing Safetensors header length")
        (header_length,) = struct.unpack("<Q", raw_length)
        if header_length <= 0 or header_length > 1_000_000_000:
            raise AssertionError(f"{path.name}: invalid header length {header_length}")
        header = json.loads(handle.read(header_length))
    return header


def main() -> None:
    missing = sorted(name for name in REQUIRED_FILES if not (ROOT / name).is_file())
    assert not missing, f"Missing required files: {missing}"

    config = load_json("config.json")
    assert config["model_type"] == "raon"
    assert config["architectures"] == ["RaonModel"]
    assert config["auto_map"] == {
        "AutoConfig": "configuration_raon.RaonConfig",
        "AutoModel": "modeling_raon.RaonModel",
        "AutoModelForCausalLM": "modeling_raon.RaonModel",
    }
    assert config["text_model_config"]["model_type"] == "AXK1"

    for field_name, field_value in iter_named_values(config):
        if field_name.endswith("_name_or_path") or field_name == "text_model_path":
            assert field_value in ("", None), f"Internal model path remains at {field_name}: {field_value!r}"

    for path in ROOT.iterdir():
        if path.is_file() and path.suffix in {".json", ".md", ".py", ".txt", ".jinja"}:
            text = path.read_text(encoding="utf-8")
            assert not INTERNAL_PATH_RE.search(text), f"Internal path found in {path.name}"

    tokenizer = load_json("tokenizer.json")
    added_tokens = {int(item["id"]): item["content"] for item in tokenizer["added_tokens"]}
    for token_id, token_text in EXPECTED_SPECIAL_TOKENS.items():
        assert added_tokens.get(token_id) == token_text, (
            f"Tokenizer mismatch for ID {token_id}: "
            f"expected {token_text!r}, found {added_tokens.get(token_id)!r}"
        )

    index = load_json("model.safetensors.index.json")
    weight_map = index["weight_map"]
    expected_keys = set(weight_map)
    shard_names = sorted(set(weight_map.values()))
    assert len(shard_names) == 9, f"Expected 9 shards, found {len(shard_names)}"

    header_keys: set[str] = set()
    total_tensor_elements = 0
    total_tensor_bytes = 0
    routed_experts: dict[tuple[int, int], int] = {}
    routed_pattern = re.compile(r"^text_model\.layers\.(\d+)\.mlp\.experts\.(\d+)\.")
    for shard_name in shard_names:
        shard_path = ROOT / shard_name
        assert shard_path.is_file(), f"Missing shard: {shard_name}"
        header = read_safetensors_header(shard_path)
        for tensor_name, tensor_info in header.items():
            if tensor_name == "__metadata__":
                continue
            assert tensor_name not in header_keys, f"Duplicate tensor: {tensor_name}"
            assert weight_map.get(tensor_name) == shard_name, (
                f"Index maps {tensor_name} to {weight_map.get(tensor_name)!r}, "
                f"but it is stored in {shard_name}"
            )
            header_keys.add(tensor_name)
            tensor_elements = math.prod(tensor_info["shape"])
            total_tensor_elements += tensor_elements
            if match := routed_pattern.match(tensor_name):
                expert_key = (int(match.group(1)), int(match.group(2)))
                routed_experts[expert_key] = routed_experts.get(expert_key, 0) + tensor_elements
            start, end = tensor_info["data_offsets"]
            total_tensor_bytes += end - start

    assert header_keys == expected_keys, (
        f"Index/header key mismatch: "
        f"{len(expected_keys - header_keys)} missing, "
        f"{len(header_keys - expected_keys)} unexpected"
    )
    metadata = index["metadata"]
    assert metadata["total_parameters"] == 21_197_920_833
    # Safetensors may include persistent buffers, so total tensor elements can
    # be larger than the model parameter count recorded by save_pretrained.
    assert total_tensor_elements >= metadata["total_parameters"]
    assert total_tensor_bytes == metadata["total_size"], (
        f"Tensor byte mismatch: {total_tensor_bytes} != {metadata['total_size']}"
    )

    top_k = int(config["text_model_config"]["num_experts_per_tok"])
    routed_total = sum(routed_experts.values())
    active_routed = 0
    for layer in sorted({layer for layer, _ in routed_experts}):
        layer_experts = sorted(
            (value for (expert_layer, _), value in routed_experts.items() if expert_layer == layer),
            reverse=True,
        )
        active_routed += sum(layer_experts[:top_k])
    active_parameters = metadata["total_parameters"] - routed_total + active_routed
    assert active_parameters == 3_456_014_913

    print(
        "Repository validation passed: "
        f"{len(header_keys):,} tensors, "
        f"{metadata['total_parameters']:,} total parameters, "
        f"{active_parameters:,} active parameters, "
        f"{total_tensor_bytes:,} tensor bytes, "
        f"{len(shard_names)} shards."
    )


if __name__ == "__main__":
    main()