Upload special_tokens.py with huggingface_hub

special_tokens.py

@@ -0,0 +1,85 @@
+"""
+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