File size: 10,276 Bytes
9d614f1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 | """Tokenizer surgeon (PRIVATE to Training_Module).
Safely registers <think>, </think> and <|im_end|> (+ <|im_start|>) as single tokens, makes sure the
embedding matrix has rows for them, and initialises new rows from the mean of their sub-token rows.
Also owns the shared chat format used by both training stages.
"""
from __future__ import annotations
import re
from typing import Dict, Iterable, List, Optional, Sequence
THINK_OPEN = "<think>"
THINK_CLOSE = "</think>"
IM_START = "<|im_start|>"
IM_END = "<|im_end|>"
ENDOFTEXT = "<|endoftext|>"
MODES = ("thinking", "non_thinking", "max_thinking")
DEFAULT_SYSTEM = "You are RSS AI, a helpful assistant created by RSS. You understand Hindi, Hinglish and English."
# Active vocabulary policy: ASCII (English/Hinglish/code/math/punctuation), Devanagari
# (including Vedic/extended marks), tokenizer control tokens, byte-fallback tokens, and
# every ID that actually occurs in the SFT examples. The original embedding table is
# deliberately NOT resized or rewritten: Stage-1 checkpoint keys/shapes stay unchanged.
_BYTE_FALLBACK = re.compile(r"^<0x[0-9A-Fa-f]{2}>$")
_SPECIAL_SPELLING = re.compile(r"^(?:<\|.*?\|>|<[^<>]+>)$")
_INDIC_RANGES = ((0x0900, 0x097F), (0x1CD0, 0x1CFF), (0xA8E0, 0xA8FF))
_INDIC_JOINERS = {0x200C, 0x200D}
def _is_active_token_spelling(token: str) -> bool:
if not isinstance(token, str):
return False
if _BYTE_FALLBACK.match(token) or _SPECIAL_SPELLING.match(token):
return True
for char in token:
code = ord(char)
if code <= 0x7F or code in _INDIC_JOINERS:
return True
if any(lo <= code <= hi for lo, hi in _INDIC_RANGES):
return True
return False
def summarize_active_vocab(tokenizer, dataset_token_ids: Iterable[int], active_token_ids: Sequence[int],
vocab_rows: int) -> Dict[str, int]:
"""Count the retained token families and foreign-token pruning without guessing from ratios."""
vocab = tokenizer.get_vocab()
special = {int(x) for x in getattr(tokenizer, "all_special_ids", []) or [] if x is not None}
for attr in ("bos_token_id", "eos_token_id", "pad_token_id", "unk_token_id", "sep_token_id",
"cls_token_id", "mask_token_id"):
value = getattr(tokenizer, attr, None)
if value is not None:
special.add(int(value))
categories = {"special_control": set(), "byte_fallback": set(),
"english_ascii_code_math": set(), "devanagari": set(), "foreign": set()}
for spelling, raw_id in vocab.items():
token_id = int(raw_id)
if _BYTE_FALLBACK.match(spelling):
categories["byte_fallback"].add(token_id)
elif token_id in special or _SPECIAL_SPELLING.match(spelling):
categories["special_control"].add(token_id)
elif any(ord(ch) <= 0x7F for ch in spelling):
categories["english_ascii_code_math"].add(token_id)
elif any(ord(ch) in _INDIC_JOINERS or any(lo <= ord(ch) <= hi for lo, hi in _INDIC_RANGES)
for ch in spelling):
categories["devanagari"].add(token_id)
else:
categories["foreign"].add(token_id)
active = {int(x) for x in active_token_ids}
data_only_foreign = categories["foreign"].intersection(active)
return {"full_vocab_size": int(vocab_rows), "active_vocab_size": len(active),
"inactive_vocab_size": int(vocab_rows) - len(active),
"special_control_kept": len(categories["special_control"]),
"byte_fallback_kept": len(categories["byte_fallback"]),
"english_ascii_code_math_kept": len(categories["english_ascii_code_math"]),
"devanagari_kept": len(categories["devanagari"]),
"dataset_only_foreign_kept": len(data_only_foreign),
"foreign_vocab_total": len(categories["foreign"]),
"foreign_vocab_pruned": len(categories["foreign"] - data_only_foreign)}
def build_active_vocab_ids(tokenizer, dataset_token_ids: Iterable[int], vocab_rows: Optional[int] = None) -> List[int]:
"""Return sorted original token IDs permitted by the RSS Hindi/English/code policy.
This is an output-vocabulary slice only. Input embeddings keep their original row IDs,
preserving strict compatibility with an existing Stage-1 checkpoint. Dataset IDs are
always unioned in, so no token present in an encoded SFT example can be pruned.
"""
vocab = tokenizer.get_vocab()
active = {int(token_id) for token_id in getattr(tokenizer, "all_special_ids", []) or []
if token_id is not None and int(token_id) >= 0}
for attr in ("bos_token_id", "eos_token_id", "pad_token_id", "unk_token_id",
"sep_token_id", "cls_token_id", "mask_token_id"):
token_id = getattr(tokenizer, attr, None)
if token_id is not None and int(token_id) >= 0:
active.add(int(token_id))
for spelling, token_id in vocab.items():
if _is_active_token_spelling(spelling):
active.add(int(token_id))
active.update(int(token_id) for token_id in dataset_token_ids if int(token_id) >= 0)
if vocab_rows is None:
vocab_rows = max([int(v) for v in vocab.values()] + [-1]) + 1
invalid = sorted(token_id for token_id in active if token_id >= int(vocab_rows))
if invalid:
raise ValueError(f"active tokenizer IDs exceed model vocabulary rows ({vocab_rows}): {invalid[:8]}")
if not active:
raise RuntimeError("active vocabulary is empty; refusing to construct a sliced loss")
return sorted(active)
# ------------------------------------------------------------- chat format
def build_system(system: str, mode: str) -> str:
return f"{(system or DEFAULT_SYSTEM).strip()}\n[MODE: {mode}]"
def render_prompt(system: str, user: str, mode: str) -> str:
return (f"{IM_START}system\n{build_system(system, mode)}{IM_END}\n"
f"{IM_START}user\n{user}{IM_END}\n"
f"{IM_START}assistant\n")
def render_target(thinking: str, assistant: str, mode: str) -> str:
if mode == "non_thinking":
return f"{assistant}{IM_END}"
return f"{THINK_OPEN}\n{thinking}\n{THINK_CLOSE}\n{assistant}{IM_END}"
# ----------------------------------------------------------------- surgery
def prepare_tokenizer(tokenizer, recipe_override: Optional[Dict] = None) -> Dict:
"""Register special tokens on the tokenizer. Returns the surgery info/recipe.
recipe_override: {new_token_id: [sub-token ids]} saved by an earlier run. It is used when the tokens already
exist in the tokenizer (so there is nothing to compute) but the weights still need the row initialisation."""
vocab = tokenizer.get_vocab()
wanted = [THINK_OPEN, THINK_CLOSE, IM_START, IM_END]
missing = [t for t in wanted if t not in vocab]
# sub-token recipe MUST be computed before the tokens exist as single tokens
sources = {t: tokenizer.encode(t, add_special_tokens=False) for t in missing}
if missing:
existing: List[str] = list(getattr(tokenizer, "additional_special_tokens", []) or [])
tokenizer.add_special_tokens({"additional_special_tokens": existing + missing})
if tokenizer.pad_token is None:
tokenizer.pad_token = ENDOFTEXT if ENDOFTEXT in tokenizer.get_vocab() else tokenizer.eos_token
tokenizer.eos_token = IM_END
tokenizer.model_max_length = 10 ** 9 # we chunk ourselves; silence length warnings
tokenizer.padding_side = "right"
recipe = {tokenizer.convert_tokens_to_ids(t): ids for t, ids in sources.items()}
if not recipe and recipe_override:
recipe = {int(k): [int(x) for x in v] for k, v in recipe_override.items()}
eos_ids = [tokenizer.convert_tokens_to_ids(IM_END)]
if ENDOFTEXT in tokenizer.get_vocab():
eos_ids.append(tokenizer.convert_tokens_to_ids(ENDOFTEXT))
return {
"added": missing, "recipe": recipe, "resized": False,
"ids": {"think_open": tokenizer.convert_tokens_to_ids(THINK_OPEN),
"think_close": tokenizer.convert_tokens_to_ids(THINK_CLOSE),
"im_end": tokenizer.convert_tokens_to_ids(IM_END),
"im_start": tokenizer.convert_tokens_to_ids(IM_START),
"pad": tokenizer.pad_token_id, "eos_ids": eos_ids},
}
def finalize_model_embeddings(model, tokenizer, info: Dict) -> Dict:
"""Make sure the embedding matrix covers the tokenizer, init new rows, fix generation ids."""
import torch
needed = len(tokenizer)
rows = int(getattr(model.config, "vocab_size", 0) or model.get_input_embeddings().weight.shape[0])
ring = getattr(model, "_rss_ring", None)
if needed > rows:
if ring is not None:
raise RuntimeError(f"tokenizer needs {needed} rows but the checkpoint only has {rows}; "
"resizing is not possible in tier-3 ring mode (use tier 2 once to resize).")
model.resize_token_embeddings(needed, pad_to_multiple_of=8)
info["resized"] = True
print(f" [tokenizer] embedding matrix resized {rows} -> {model.get_input_embeddings().weight.shape[0]}")
else:
print(f" [tokenizer] {len(info['added'])} token(s) registered, embedding has {rows} rows -> no resize needed")
if ring is None and info["recipe"]:
with torch.no_grad():
embed = model.get_input_embeddings().weight
out = model.get_output_embeddings()
out_w = out.weight if out is not None and out.weight.data_ptr() != embed.data_ptr() else None
for new_id, src in info["recipe"].items():
if not src:
continue
idx = torch.tensor(src, dtype=torch.long, device=embed.device)
embed[new_id] = embed[idx].float().mean(0).to(embed.dtype)
if out_w is not None:
idx_o = torch.tensor(src, dtype=torch.long, device=out_w.device)
out_w[new_id] = out_w[idx_o].float().mean(0).to(out_w.dtype)
print(" [tokenizer] new token rows initialised from the mean of their sub-token embeddings")
gen = model.generation_config
gen.eos_token_id = info["ids"]["eos_ids"]
gen.pad_token_id = info["ids"]["pad"]
return info
|