"""Versioned tokenizer protocol ABI shared by data, post-training and serving. The tokenizer owns 256 atomic IDs at the top of its fixed 65,536-entry vocabulary. Token strings and IDs are immutable after tokenizer freeze; semantic use is versioned here so consumers cannot silently invent a second mapping. Reserved Omni aliases intentionally keep their on-disk ``<|reserved_NNN|>`` strings: a structured processor splices their IDs and projected features replace only the corresponding embeddings, without a tokenizer or embedding resize. """ from __future__ import annotations import hashlib import json from pathlib import Path from typing import Any from yunmo_next.protocol_tokens import ( ACTIVE_PROTOCOL_TOKENS, PROTOCOL_SLOT_COUNT, PROTOCOL_TOKENS, ) CONTROL_BASE = 65_280 VOCAB_SIZE = 65_536 FINAL_TOKENIZER_SHA256 = ( "0d1b129c0e757043821030f8662d248f5d96d86a8199c9ad64bc698f14e71324" ) TOKEN_TO_ID = { token: CONTROL_BASE + offset for offset, token in enumerate(PROTOCOL_TOKENS) } ID_TO_TOKEN = {identifier: token for token, identifier in TOKEN_TO_ID.items()} BOS = TOKEN_TO_ID["<|bos|>"] EOS = TOKEN_TO_ID["<|eos|>"] PAD = TOKEN_TO_ID["<|pad|>"] UNK = TOKEN_TO_ID["<|unk|>"] START = TOKEN_TO_ID["<|start|>"] CHANNEL = TOKEN_TO_ID["<|channel|>"] MESSAGE = TOKEN_TO_ID["<|message|>"] END = TOKEN_TO_ID["<|end|>"] RETURN = TOKEN_TO_ID["<|return|>"] RECIPIENT = TOKEN_TO_ID["<|recipient|>"] CALL = TOKEN_TO_ID["<|call|>"] DOC_START = TOKEN_TO_ID["<|doc_start|>"] DOC_END = TOKEN_TO_ID["<|doc_end|>"] TEXT_ROLES = ("system", "developer", "user", "assistant", "tool") TEXT_CHANNELS = ("analysis", "final") # These aliases reserve stable embedding positions for a future modality # adapter. Raw tokenizers can emit the literal reserved strings if they appear # in text; production text encoders must reject resulting control IDs. The # aliases are not audio-codebook IDs. A future Omni branch may replace the # embedding at frame/patch positions with projected encoder features while # keeping the text model vocabulary and LM head unchanged. OMNI_RESERVED_ALIASES = { "image_start": TOKEN_TO_ID["<|reserved_000|>"], "image_patch": TOKEN_TO_ID["<|reserved_001|>"], "image_end": TOKEN_TO_ID["<|reserved_002|>"], "audio_input_start": TOKEN_TO_ID["<|reserved_003|>"], "audio_input_frame": TOKEN_TO_ID["<|reserved_004|>"], "audio_input_end": TOKEN_TO_ID["<|reserved_005|>"], "audio_output_start": TOKEN_TO_ID["<|reserved_006|>"], "audio_output_frame": TOKEN_TO_ID["<|reserved_007|>"], "audio_output_end": TOKEN_TO_ID["<|reserved_008|>"], "voice_reference_start": TOKEN_TO_ID["<|reserved_009|>"], "voice_reference_frame": TOKEN_TO_ID["<|reserved_010|>"], "voice_reference_end": TOKEN_TO_ID["<|reserved_011|>"], } REPOSITORY_ROOT = Path(__file__).resolve().parents[1] PRODUCTION_TOKENIZER_DIRECTORY = ( REPOSITORY_ROOT / "data/tokenizer/Yunmo_Next_Tokenizer" ) DEFAULT_TOKENIZER = PRODUCTION_TOKENIZER_DIRECTORY / "tokenizer.json" DEFAULT_PROTOCOL_ABI = ( PRODUCTION_TOKENIZER_DIRECTORY / "protocol_abi_v2.json" ) PROTOCOL_ABI_TOP_LEVEL_KEYS = { "format", "tokenizer_sha256", "vocabulary_size", "lexical_vocabulary_size", "protocol_slot_count", "protocol_first_id", "protocol_last_id", "active_token_id_map", "reserved_slot_scheme", "text_message_protocol", "omni_reserved_aliases", "omni_embedding_protocol", } TEXT_ACTION_TOKENS = ( "<|bos|>", "<|eos|>", "<|pad|>", "<|start|>", "<|channel|>", "<|message|>", "<|end|>", "<|return|>", "<|recipient|>", "<|call|>", ) def sha256(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as handle: for chunk in iter(lambda: handle.read(8 * 2**20), b""): digest.update(chunk) return digest.hexdigest() def expected_token_map() -> dict[str, int]: return dict(TOKEN_TO_ID) def validate_tokenizer_protocol(tokenizer: Any) -> None: size = int(tokenizer.get_vocab_size(with_added_tokens=True)) if size != VOCAB_SIZE: raise RuntimeError( f"tokenizer vocabulary changed: expected={VOCAB_SIZE} actual={size}" ) mismatches = { token: (identifier, tokenizer.token_to_id(token)) for token, identifier in TOKEN_TO_ID.items() if tokenizer.token_to_id(token) != identifier } if mismatches: first = next(iter(mismatches.items())) raise RuntimeError(f"tokenizer protocol mapping changed: {first}") def validate_tokenizer_file(path: Path) -> None: resolved = path.expanduser().resolve(strict=True) observed = sha256(resolved) if observed != FINAL_TOKENIZER_SHA256: raise RuntimeError( "final tokenizer file identity changed: " f"expected={FINAL_TOKENIZER_SHA256} actual={observed}" ) def _strict_json_object(pairs: list[tuple[str, Any]]) -> dict[str, Any]: value: dict[str, Any] = {} for key, item in pairs: if key in value: raise ValueError(f"duplicate JSON key: {key}") value[key] = item return value def _reject_json_constant(value: str) -> None: raise ValueError(f"non-finite JSON number: {value}") def _canonical_json(value: Any) -> bytes: return json.dumps( value, ensure_ascii=False, sort_keys=True, separators=(",", ":"), allow_nan=False, ).encode("utf-8") def _normalize_string_set(value: Any, *, field: str) -> list[str]: if ( not isinstance(value, list) or any(type(item) is not str for item in value) or len(set(value)) != len(value) ): raise RuntimeError(f"tokenizer protocol ABI {field} is invalid") return sorted(value) def _normalize_active_token_map(value: Any) -> list[dict[str, Any]]: if not isinstance(value, list): raise RuntimeError("tokenizer protocol ABI active token map is invalid") normalized = [] for item in value: if ( not isinstance(item, dict) or set(item) != {"id", "token"} or type(item["id"]) is not int or type(item["token"]) is not str ): raise RuntimeError("tokenizer protocol ABI active token map is invalid") normalized.append({"id": item["id"], "token": item["token"]}) if ( len({item["id"] for item in normalized}) != len(normalized) or len({item["token"] for item in normalized}) != len(normalized) ): raise RuntimeError("tokenizer protocol ABI active token map has duplicates") return sorted(normalized, key=lambda item: (item["id"], item["token"])) def _expected_protocol_abi() -> dict[str, Any]: return { "format": "yunmo-next-tokenizer-protocol-abi-v2", "tokenizer_sha256": FINAL_TOKENIZER_SHA256, "vocabulary_size": VOCAB_SIZE, "lexical_vocabulary_size": CONTROL_BASE, "protocol_slot_count": PROTOCOL_SLOT_COUNT, "protocol_first_id": CONTROL_BASE, "protocol_last_id": VOCAB_SIZE - 1, "active_token_id_map": [ {"id": TOKEN_TO_ID[token], "token": token} for token in ACTIVE_PROTOCOL_TOKENS ], "reserved_slot_scheme": { "count": PROTOCOL_SLOT_COUNT - len(ACTIVE_PROTOCOL_TOKENS), "first_id": CONTROL_BASE + len(ACTIVE_PROTOCOL_TOKENS), "last_id": VOCAB_SIZE - 1, "token_format": "<|reserved_{index:03d}|>", }, "text_message_protocol": { "roles": sorted(TEXT_ROLES), "channels": sorted(TEXT_CHANNELS), "prompt_message": "<|start|>{role}<|channel|>analysis<|message|>{content}<|end|>", "assistant_reasoning": "<|start|>assistant<|channel|>analysis<|message|>{reasoning}<|end|>", "assistant_final": "<|start|>assistant<|channel|>final<|message|>{answer}<|return|>", "assistant_tool_call": "<|start|>assistant<|channel|>analysis<|recipient|>{tool_name}<|message|>{json}<|call|>", "tool_result": "<|start|>tool<|channel|>analysis<|message|>{result}<|end|>", "content_protocol_ids_allowed": False, "supervision": { "assistant_body_and_terminal": True, "prompt_and_structural_prefix": False, "label_offset": -1, }, }, "omni_reserved_aliases": OMNI_RESERVED_ALIASES, "omni_embedding_protocol": { "structured_content_part_types": sorted( [ "text", "image", "audio_input", "audio_output", "voice_reference", ] ), "raw_text_protocol_ids_allowed": False, "feature_injection": "replace_placeholder_embeddings", "span_layout": { "order": ["start", "placeholder", "end"], "placeholder_repeats": "feature_count", }, "feature_override": { "positions": { "dtype": "int64", "rank": 2, "columns": ["batch", "token_position"], }, "values": { "dtype": "model_embedding_dtype", "rank": 2, "last_dimension": "hidden_size", }, "placeholder_ids": {"dtype": "int64", "rank": 1}, "positions_unique": True, "positions_in_current_chunk": True, "input_ids_match_placeholder_ids": True, }, "supervision": { "control_and_placeholder_targets": False, "label_offset": -1, }, }, } def _normalize_protocol_abi(value: dict[str, Any]) -> dict[str, Any]: normalized = dict(value) normalized["active_token_id_map"] = _normalize_active_token_map( value["active_token_id_map"] ) text = value.get("text_message_protocol") if not isinstance(text, dict): raise RuntimeError("tokenizer protocol ABI text message contract is invalid") normalized_text = dict(text) normalized_text["roles"] = _normalize_string_set( text.get("roles"), field="text roles" ) normalized_text["channels"] = _normalize_string_set( text.get("channels"), field="text channels" ) normalized["text_message_protocol"] = normalized_text omni = value.get("omni_embedding_protocol") if not isinstance(omni, dict): raise RuntimeError("tokenizer protocol ABI Omni embedding contract is invalid") normalized_omni = dict(omni) normalized_omni["structured_content_part_types"] = _normalize_string_set( omni.get("structured_content_part_types"), field="Omni content part types" ) normalized["omni_embedding_protocol"] = normalized_omni return normalized def load_protocol_abi(path: Path = DEFAULT_PROTOCOL_ABI) -> dict[str, Any]: resolved = path.resolve(strict=True) try: value = json.loads( resolved.read_text(encoding="utf-8"), object_pairs_hook=_strict_json_object, parse_constant=_reject_json_constant, ) except (json.JSONDecodeError, ValueError) as error: raise RuntimeError("tokenizer protocol ABI is not strict JSON") from error if not isinstance(value, dict) or set(value) != PROTOCOL_ABI_TOP_LEVEL_KEYS: raise RuntimeError("tokenizer protocol ABI schema changed") normalized = _normalize_protocol_abi(value) if _canonical_json(normalized) != _canonical_json(_expected_protocol_abi()): raise RuntimeError("tokenizer protocol ABI executable semantics changed") return normalized def _canonical_sha256(value: Any) -> str: return hashlib.sha256(_canonical_json(value)).hexdigest() def text_protocol_semantics(path: Path = DEFAULT_PROTOCOL_ABI) -> dict[str, Any]: """Return only semantics consumed by text serialization and objectives.""" value = load_protocol_abi(path) token_ids = { item["token"]: item["id"] for item in value["active_token_id_map"] } return { "tokenizer_sha256": value["tokenizer_sha256"], "vocabulary_size": value["vocabulary_size"], "lexical_vocabulary_size": value["lexical_vocabulary_size"], "token_ids": {token: token_ids[token] for token in TEXT_ACTION_TOKENS}, "message_protocol": value["text_message_protocol"], } def omni_protocol_semantics(path: Path = DEFAULT_PROTOCOL_ABI) -> dict[str, Any]: """Return only semantics consumed by structured Omni input construction.""" value = load_protocol_abi(path) return { "tokenizer_sha256": value["tokenizer_sha256"], "vocabulary_size": value["vocabulary_size"], "lexical_vocabulary_size": value["lexical_vocabulary_size"], "reserved_aliases": value["omni_reserved_aliases"], "embedding_protocol": value["omni_embedding_protocol"], } def text_protocol_semantic_sha256(path: Path = DEFAULT_PROTOCOL_ABI) -> str: return _canonical_sha256(text_protocol_semantics(path)) def omni_protocol_semantic_sha256(path: Path = DEFAULT_PROTOCOL_ABI) -> str: return _canonical_sha256(omni_protocol_semantics(path)) def active_protocol_count() -> int: return len(ACTIVE_PROTOCOL_TOKENS)