"""Nexum release-folder validation.""" from __future__ import annotations import hashlib import json import re from dataclasses import asdict, dataclass from pathlib import Path from typing import Any REQUIRED_FILES: tuple[str, ...] = ( "config.json", "tokenizer.json", "tokenizer_config.json", "chat_template.jinja", "generation_config.json", "state_config.json", "tensor_map.json", ) NUMBERED_TENSOR_RE = re.compile(r"^\d{6}\.safetensors$") SHA256_RE = re.compile(r"^[0-9a-f]{64}$") @dataclass(frozen=True) class BundleFile: name: str exists: bool bytes: int = 0 sha256: str = "" keys: int | None = None @dataclass(frozen=True) class BundleReport: ok: bool model_dir: str files: tuple[BundleFile, ...] missing: tuple[str, ...] tensor_map_ok: bool load_map_ok: bool tensor_count: int tensor_bytes: int tensor_hashes_verified: bool tokenizer_ok: bool errors: tuple[str, ...] def to_dict(self) -> dict[str, Any]: return { "ok": self.ok, "model_dir": Path(self.model_dir).name, "files": [asdict(item) for item in self.files], "missing": list(self.missing), "tensor_map_ok": self.tensor_map_ok, "load_map_ok": self.load_map_ok, "tensor_count": self.tensor_count, "tensor_bytes": self.tensor_bytes, "tensor_hashes_verified": self.tensor_hashes_verified, "tokenizer_ok": self.tokenizer_ok, "errors": list(self.errors), } 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 _count_tensor_keys(path: Path) -> int | None: if path.suffix != ".safetensors": return None from safetensors import safe_open with safe_open( # type: ignore[no-untyped-call] str(path), framework="pt", device="cpu" ) as handle: return len(handle.keys()) def _read_json(path: Path) -> dict[str, Any]: value = json.loads(path.read_text(encoding="utf-8")) if not isinstance(value, dict): raise ValueError(f"{path.name} must contain a JSON object") return value def _tensor_entries(map_path: Path) -> list[dict[str, Any]]: data = _read_json(map_path) tensors = data.get("tensors") if not isinstance(tensors, list): raise ValueError("tensor_map.json must contain a tensors list") return [item for item in tensors if isinstance(item, dict)] def _product_contract(release_cfg: dict[str, Any]) -> tuple[str, int, bool]: product = release_cfg.get("product") if not isinstance(product, dict): return "Nexum-Universal", 113, True variant = product.get("variant") expected = { "Nexum-Lite": (113, True), "Nexum-Universal": (113, True), "Nexum-Expanded": (113, True), } if variant not in expected: raise ValueError("config.json product variant is unsupported") package_count, complete_authority = expected[str(variant)] if ( product.get("tensor_package_count") != package_count or product.get("complete_none_authority") is not complete_authority ): raise ValueError("config.json product profile does not match its variant") return str(variant), package_count, complete_authority def release_artifact_sha256( model_dir: str | Path, *, validated_report: BundleReport ) -> str: """Return one identity for the exact validated numbered model bundle.""" root = Path(model_dir).expanduser().resolve() report_root = Path(validated_report.model_dir).expanduser().resolve() if report_root != root: raise ValueError("validated bundle report belongs to another model folder") if not validated_report.ok or not validated_report.tensor_hashes_verified: raise ValueError("release identity requires a deep validated model bundle") files = [ { "path": name, "bytes": (root / name).stat().st_size, "sha256": _sha256(root / name), } for name in REQUIRED_FILES ] tensors = [] for entry in _tensor_entries(root / "tensor_map.json"): tensors.append( { "id": str(entry.get("id") or ""), "file": str(entry.get("file") or ""), "bytes": int(entry.get("bytes") or 0), "sha256": str(entry.get("sha256") or ""), "binding_sha256": str(entry.get("binding_sha256") or ""), "descriptor_sha256": str(entry.get("descriptor_sha256") or ""), "key_count": int(entry.get("key_count") or 0), } ) canonical = json.dumps( { "schema": "nexum.release-artifact.v1", "files": files, "tensors": tensors, }, ensure_ascii=True, separators=(",", ":"), sort_keys=True, ).encode("utf-8") return hashlib.sha256(canonical).hexdigest() def validate_bundle(model_dir: str | Path, *, deep: bool = False) -> BundleReport: root = Path(model_dir).expanduser().resolve() errors: list[str] = [] missing: list[str] = [] files: list[BundleFile] = [] for name in REQUIRED_FILES: path = root / name if not path.exists(): missing.append(name) files.append(BundleFile(name=name, exists=False)) continue try: files.append( BundleFile( name=name, exists=True, bytes=path.stat().st_size, sha256=_sha256(path), keys=_count_tensor_keys(path) if deep else None, ) ) except Exception as exc: errors.append(f"{name}: {type(exc).__name__}: {exc}") files.append(BundleFile(name=name, exists=True, bytes=path.stat().st_size)) tensor_map_ok = False load_map_ok = False tensor_count = 0 tensor_bytes = 0 tensor_hashes_verified = False tensor_dir = root / "safetensors" if not tensor_dir.is_dir(): errors.append("missing safetensors directory") elif (root / "tensor_map.json").is_file(): try: entries = _tensor_entries(root / "tensor_map.json") tensor_count = len(entries) tensor_bytes = sum(int(item.get("bytes") or 0) for item in entries) indexed = { str(item.get("file")) for item in entries if isinstance(item.get("file"), str) } actual = { f"safetensors/{path.name}" for path in tensor_dir.glob("*.safetensors") } bad_names = sorted( path.name for path in tensor_dir.glob("*.safetensors") if not NUMBERED_TENSOR_RE.fullmatch(path.name) ) unknown = sorted(actual - indexed) absent = sorted(indexed - actual) if bad_names: errors.append(f"non-numbered safetensors files: {bad_names[:8]}") if unknown: errors.append(f"unmapped safetensors files: {unknown[:8]}") if absent: errors.append(f"missing mapped safetensors files: {absent[:8]}") ids = [str(item.get("id") or "") for item in entries] expected_ids = [f"{index:06d}" for index in range(1, len(entries) + 1)] if ids != expected_ids: errors.append( "tensor ids must be contiguous, ordered, and start at 000001" ) for item in entries: tensor_id = str(item.get("id") or "") expected_file = f"safetensors/{tensor_id}.safetensors" if item.get("file") != expected_file: errors.append(f"tensor {tensor_id} must map to {expected_file}") break path = root / expected_file if ( path.is_file() and int(item.get("bytes") or -1) != path.stat().st_size ): errors.append( f"tensor {tensor_id} byte size does not match tensor_map.json" ) break if not SHA256_RE.fullmatch(str(item.get("sha256") or "")): errors.append(f"tensor {tensor_id} has no valid sha256") break if not SHA256_RE.fullmatch( str(item.get("binding_sha256") or "") ) or not SHA256_RE.fullmatch(str(item.get("descriptor_sha256") or "")): errors.append(f"tensor {tensor_id} has no binding fingerprints") break if path.is_file(): from safetensors import safe_open with safe_open( # type: ignore[no-untyped-call] str(path), framework="pt", device="cpu" ) as handle: keys = list(handle.keys()) metadata = handle.metadata() expected_keys = [ f"{index:06d}" for index in range(1, len(keys) + 1) ] if keys != expected_keys: errors.append( f"tensor {tensor_id} keys must be contiguous and numeric" ) break if metadata: errors.append(f"tensor {tensor_id} metadata must be empty") break if int(item.get("key_count") or -1) != len(keys): errors.append( f"tensor {tensor_id} key count does not match tensor_map.json" ) break map_data = _read_json(root / "tensor_map.json") if map_data.get("schema") != "nexum.tensor_map.v4": errors.append("tensor_map.json schema is incompatible") release_cfg = _read_json(root / "config.json") _variant, expected_package_count, _complete_authority = ( _product_contract(release_cfg) ) if len(entries) != expected_package_count: errors.append( "tensor map package count does not match the product profile" ) load_map_ok = ( bool(entries) and ids == expected_ids and len(indexed) == len(entries) and indexed == actual and not bad_names and all( item.get("file") == f"safetensors/{str(item.get('id') or '')}.safetensors" for item in entries ) ) if not load_map_ok: errors.append( "ordered tensor entries must map every numbered file exactly once" ) if deep and not errors: for item in entries: tensor_id = str(item["id"]) path = root / str(item["file"]) expected_hash = str(item.get("sha256") or "") actual_hash = _sha256(path) if actual_hash != expected_hash: errors.append( f"tensor {tensor_id} sha256 does not match tensor_map.json" ) break try: key_count = _count_tensor_keys(path) except Exception as exc: errors.append( f"tensor {tensor_id}: {type(exc).__name__}: {exc}" ) break if not key_count: errors.append(f"tensor {tensor_id} contains no tensors") break tensor_hashes_verified = not errors if ( not bad_names and not unknown and not absent and entries and ids == expected_ids and load_map_ok ): tensor_map_ok = True except Exception as exc: errors.append(f"tensor_map: {type(exc).__name__}: {exc}") tokenizer_ok = False try: tokenizer = _read_json(root / "tokenizer.json") tokenizer_cfg = _read_json(root / "tokenizer_config.json") generation_cfg = _read_json(root / "generation_config.json") release_cfg = _read_json(root / "config.json") state_cfg = _read_json(root / "state_config.json") variant, expected_package_count, complete_authority = _product_contract( release_cfg ) architecture = release_cfg.get("runtime_config") if not isinstance(architecture, dict): raise ValueError("config.json has no runtime configuration object") if release_cfg.get("schema") != "nexum.release.v4": raise ValueError("config.json schema is incompatible") if release_cfg.get("model_type") != "nexum_native": raise ValueError("config.json model type is incompatible") if release_cfg.get("architectures") != ["NexumForCausalLM"]: raise ValueError("config.json architecture entry point is incompatible") required_architecture = ( "hidden_size", "num_hidden_layers", "vocab_size", "max_position_embeddings", "bos_token_id", "eos_token_id", "pad_token_id", ) if not all( isinstance(architecture.get(name), int) and int(architecture[name]) > 0 for name in required_architecture ): raise ValueError("config.json architecture values are incomplete") token_ids = { "bos_token_id": int(architecture["bos_token_id"]), "eos_token_id": int(architecture["eos_token_id"]), "pad_token_id": int(architecture["pad_token_id"]), } if any(generation_cfg.get(name) != value for name, value in token_ids.items()): raise ValueError("generation token ids do not match config.json") if ( tokenizer_cfg.get("model_max_length") != architecture["max_position_embeddings"] ): raise ValueError("tokenizer context length does not match config.json") state_bindings = state_cfg.get("bindings") if state_cfg.get("schema") != "nexum.state.v2" or not isinstance( state_bindings, list ): raise ValueError("state_config.json is incompatible") if complete_authority and not state_bindings: raise ValueError("complete product profile has no numbered state bindings") if not complete_authority and state_bindings: raise ValueError("compact product profile has unexpected state bindings") if int(expected_package_count) != tensor_count: raise ValueError("product package count does not match tensor map") identity = release_cfg.get("identity") if ( not isinstance(identity, dict) or identity.get("variant") not in (None, variant) ): raise ValueError("config.json identity does not match product variant") model_payload = tokenizer.get("model") base_vocab = ( model_payload.get("vocab") if isinstance(model_payload, dict) else None ) added_tokens = tokenizer.get("added_tokens") if not isinstance(base_vocab, dict) or not isinstance(added_tokens, list): raise ValueError("tokenizer vocabulary is incomplete") vocabulary_ids = { int(value) for value in base_vocab.values() if isinstance(value, int) } vocabulary_ids.update( int(item["id"]) for item in added_tokens if isinstance(item, dict) and isinstance(item.get("id"), int) ) expected_vocab_size = int(architecture["vocab_size"]) if vocabulary_ids != set(range(expected_vocab_size)): raise ValueError("tokenizer does not cover the complete model vocabulary") if deep: from tokenizers import Tokenizer # type: ignore[import-untyped] native_tokenizer = Tokenizer.from_file(str(root / "tokenizer.json")) for token_id in range(expected_vocab_size): if not native_tokenizer.id_to_token( token_id ) or not native_tokenizer.decode([token_id], skip_special_tokens=False): raise ValueError(f"tokenizer id {token_id} is not reversible") tokenizer_ok = "eos_token" in tokenizer_cfg except Exception as exc: errors.append(f"tokenizer: {type(exc).__name__}: {exc}") ok = not missing and not errors and tensor_map_ok and load_map_ok and tokenizer_ok return BundleReport( ok=ok, model_dir=str(root), files=tuple(files), missing=tuple(missing), tensor_map_ok=tensor_map_ok, load_map_ok=load_map_ok, tensor_count=tensor_count, tensor_bytes=tensor_bytes, tensor_hashes_verified=tensor_hashes_verified, tokenizer_ok=tokenizer_ok, errors=tuple(errors), )