code/yunmo_next/protocol.py

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