special_tokens.py

"""
special_tokens.py -- central, *append-only* registry of special tokens for the
SpikeWhale length-max tokenizer.

WHY THIS FILE EXISTS
--------------------
The base vocab (tokenizer.json) is 16384 contiguous ids:
    0..3   -> <pad> <unk> <bos> <eos>
    4..259 -> the 256 raw bytes
    260..  -> learned byte-merges
Adding tokens "without breaking the model" has exactly one rule:

    ***APPEND ONLY. NEVER REORDER OR REMOVE AN EXISTING ID.***

Every existing id keeps pointing at the same embedding row and the same logit
column, so the model's behaviour on already-seen tokens is bit-for-bit
unchanged. New tokens are appended at ids >= 16384 and their embedding / lm_head
/ mtp rows are freshly initialised (near-zero contribution) so they are
no-ops until you train them.

To stay tensor-core friendly the final vocab is padded up to a multiple of
`VOCAB_MULTIPLE` (128) with `<|reserved_N|>` slots. Those reserves let you name
*future* tokens later by editing the registry WITHOUT another model resize, as
long as the total stays <= the padded size.

HOW TO ADD MORE LATER
---------------------
Append new names to NAMED_SPECIAL_TOKENS (at the END), then either:
  * if you still have <|reserved_*|> slots free, just rename a reserved id in
    tokenizer.json (no model change needed), or
  * re-run add_special_tokens.py to grow + re-pad the vocab (model resized).
"""

# Tensor-core / matmul friendly vocab alignment. 16384 is already 128*128.
VOCAB_MULTIPLE = 128

# ---------------------------------------------------------------------------
# The universal named set. ORDER IS PERMANENT -- append only, never reorder.
# Mixing the common conventions so the same model can do chat, reasoning,
# agentic tool use, and code infilling.
# ---------------------------------------------------------------------------
NAMED_SPECIAL_TOKENS = [
    # ChatML turn framing
    "<|im_start|>",
    "<|im_end|>",
    # Reasoning / scratchpad
    "<think>",
    "</think>",
    # Explicit solution block
    "<begin_solution>",
    "<end_solution>",
    # Agentic tool calling
    "<tool_call>",
    "</tool_call>",
    "<tool_response>",
    "</tool_response>",
    # Role markers (usable standalone or inside an im_start header)
    "<|system|>",
    "<|user|>",
    "<|assistant|>",
    # Fill-in-the-middle (code)
    "<|fim_prefix|>",
    "<|fim_middle|>",
    "<|fim_suffix|>",
    # Generic document separator
    "<|endoftext|>",
]


def build_special_token_list(base_vocab_size: int,
                             multiple: int = VOCAB_MULTIPLE):
    """
    Return the ordered list of tokens to APPEND after `base_vocab_size`:
    the named set followed by enough <|reserved_N|> slots to pad the final
    vocab size up to the next multiple of `multiple`.

    The returned list's element i gets id (base_vocab_size + i).
    """
    tokens = list(NAMED_SPECIAL_TOKENS)
    target = base_vocab_size + len(tokens)
    # round up to the next multiple (or stay put if already aligned)
    padded = ((target + multiple - 1) // multiple) * multiple
    n_reserved = padded - target
    tokens += [f"<|reserved_{i}|>" for i in range(n_reserved)]
    return tokens