| """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." |
|
|
|
|
| |
| |
| |
| |
| _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) |
|
|
|
|
| |
| 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}" |
|
|
|
|
| |
| 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] |
| |
| 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 |
| 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 |
|
|
