#!/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(" 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()