| """Single source of truth for the chat format, shared by SFT / distill / DPO / |
| GRPO training AND their sampling code (previously seven duplicated copies). |
| |
| Token-level template — role boundaries are special tokens, not plain text: |
| |
| [BOS] <|system|> {system} <|user|> {q1} <|assistant|> {a1} [EOS] <|user|> {q2} ... |
| |
| Design rules: |
| - Role tokens delimit turns. BPE can never merge across a special token, so a |
| conversation encodes to the same ids whether it is built turn-by-turn during |
| training or as a generation prompt at inference (the old plain-text |
| "User: ...\\n\\n" template drifted at every segment boundary). |
| - [EOS] appears ONLY after assistant turns: it means "assistant finished, |
| stop generating" — the same stop semantics as document ends in pretraining. |
| System/user turns need no terminator; the next role token is the boundary. |
| - The supervised span is exactly each assistant turn's content tokens plus its |
| trailing [EOS] (so the model is explicitly taught to stop). |
| - Content is encoded with `tokenizer.encode_special_tokens = True` |
| (see `load_chat_tokenizer`), so literal "[EOS]" / "<|user|>" strings inside |
| user or corpus text are tokenized as plain text and can never inject real |
| control tokens. |
| """ |
|
|
| from __future__ import annotations |
|
|
| from typing import Any |
|
|
| from tokenizers import Tokenizer |
|
|
| from src.special_tokens import ( |
| ASSISTANT_ID, |
| BOS_ID, |
| EOS_ID, |
| PAD_ID, |
| SYSTEM_ID, |
| USER_ID, |
| assert_tokenizer_contract, |
| ) |
|
|
| DEFAULT_SYSTEM = "You are a helpful assistant." |
|
|
| IGNORE_INDEX = -100 |
|
|
|
|
| |
| |
| |
| def configure_chat_tokenizer(tok: Tokenizer) -> Tokenizer: |
| """Make `tok.encode` treat special-token strings in raw text as plain text. |
| |
| All special tokens in this pipeline are inserted by ID by the code below, |
| never parsed out of content — this closes the prompt-injection hole where a |
| document containing the literal string "[EOS]" would encode to the real |
| EOS id. |
| """ |
| tok.encode_special_tokens = True |
| return tok |
|
|
|
|
| def load_chat_tokenizer(tokenizer_path: str) -> Tokenizer: |
| """Load tokenizer.json, assert the hardcoded special-token IDs, and disable |
| special-token matching in raw text. Every chat-stage script should load its |
| tokenizer through this.""" |
| assert_tokenizer_contract(tokenizer_path) |
| return configure_chat_tokenizer(Tokenizer.from_file(tokenizer_path)) |
|
|
|
|
| |
| |
| |
| def norm_newlines(s: str) -> str: |
| return (s or "").replace("\r\n", "\n").replace("\r", "\n") |
|
|
|
|
| def clean_text(s: str) -> str: |
| |
| return norm_newlines(s).strip() |
|
|
|
|
| def clean_text_assistant(s: str) -> str: |
| |
| |
| return norm_newlines(s).rstrip() |
|
|
|
|
| def _normalized_messages( |
| messages: list[dict[str, str]], |
| default_system: str | None, |
| *, |
| expected_end: str | None, |
| ) -> list[dict[str, str]]: |
| """Clean and validate the canonical chat state machine. |
| |
| Normally a chat has an initial system turn, followed by alternating user |
| and assistant turns. Explicit default_system=None preserves source content |
| verbatim and permits a missing system turn without injecting one. ``expected_end`` is ``"user"`` for |
| a generation prompt and ``"assistant"`` for a complete training sample. |
| No malformed turn is silently skipped. |
| """ |
| out: list[dict[str, str]] = [] |
| for index, message in enumerate(messages or []): |
| if not isinstance(message, dict): |
| raise ValueError(f"message {index} must be an object") |
| role_value = message.get("role") |
| role = role_value.strip().lower() if isinstance(role_value, str) else "" |
| if role not in ("system", "user", "assistant"): |
| raise ValueError(f"message {index} has invalid role {role_value!r}") |
| raw = message.get("content") |
| if not isinstance(raw, str): |
| raise ValueError(f"message {index} content must be a string") |
| text = raw if default_system is None else ( |
| clean_text_assistant(raw) if role == "assistant" else clean_text(raw) |
| ) |
| if not text.strip(): |
| if index == 0 and role == "system" and default_system is not None: |
| fallback = clean_text(default_system) |
| if not fallback: |
| raise ValueError( |
| "chat requires a non-empty initial system turn or non-empty default_system" |
| ) |
| text = fallback |
| else: |
| raise ValueError(f"message {index} ({role}) content must be non-empty") |
| out.append({"role": role, "content": text}) |
|
|
| if (not out or out[0]["role"] != "system") and default_system is not None: |
| fallback = clean_text(default_system) |
| if not fallback: |
| raise ValueError( |
| "chat requires a non-empty initial system turn or non-empty default_system" |
| ) |
| out.insert(0, {"role": "system", "content": fallback}) |
|
|
| has_system = bool(out) and out[0]["role"] == "system" |
| if not out or (has_system and len(out) == 1): |
| raise ValueError("chat requires at least one non-empty user turn") |
| for index, message in enumerate(out): |
| turn_index = index - int(has_system) |
| expected = "system" if has_system and index == 0 else ( |
| "user" if turn_index % 2 == 0 else "assistant" |
| ) |
| if message["role"] != expected: |
| raise ValueError( |
| "invalid chat role order: initial system must be followed by " |
| f"alternating user/assistant turns (index {index}: expected " |
| f"{expected!r}, got {message['role']!r})" |
| ) |
| if expected_end is not None and out[-1]["role"] != expected_end: |
| raise ValueError( |
| f"chat must end with a non-empty {expected_end} turn; got {out[-1]['role']!r}" |
| ) |
| return out |
|
|
|
|
| def prepare_prompt_messages( |
| messages: list[dict[str, str]], |
| default_system: str | None = DEFAULT_SYSTEM, |
| ) -> list[dict[str, str]]: |
| """Explicitly turn a valid conversation/example into a USER-ending prompt. |
| |
| Callers sampling from a complete SFT example must opt in to removing its |
| final assistant answer. ``encode_prompt`` itself never drops turns. |
| """ |
| normalized = _normalized_messages(messages, default_system, expected_end=None) |
| if normalized[-1]["role"] == "assistant": |
| normalized = normalized[:-1] |
| if normalized[-1]["role"] != "user": |
| raise ValueError("generation prompt context must end with a user turn") |
| return normalized |
|
|
|
|
| |
| |
| |
| _ROLE_TOKEN_ID = {"system": SYSTEM_ID, "user": USER_ID, "assistant": ASSISTANT_ID} |
|
|
|
|
| def _encode_normalized_chat( |
| tok: Tokenizer, messages: list[dict[str, str]] |
| ) -> tuple[list[int], list[int]]: |
| ids: list[int] = [BOS_ID] |
| labels: list[int] = [IGNORE_INDEX] |
| for index, message in enumerate(messages): |
| role = message["role"] |
| content_ids = tok.encode(message["content"]).ids |
| if not content_ids: |
| raise ValueError(f"message {index} ({role}) must encode to at least one token") |
| ids.append(_ROLE_TOKEN_ID[role]) |
| labels.append(IGNORE_INDEX) |
| if role == "assistant": |
| ids.extend(content_ids) |
| labels.extend(content_ids) |
| ids.append(EOS_ID) |
| labels.append(EOS_ID) |
| else: |
| ids.extend(content_ids) |
| labels.extend([IGNORE_INDEX] * len(content_ids)) |
| return ids, labels |
|
|
|
|
| def encode_chat( |
| tok: Tokenizer, |
| messages: list[dict[str, str]], |
| default_system: str | None = DEFAULT_SYSTEM, |
| ) -> tuple[list[int], list[int]]: |
| """Encode a full conversation for training. |
| |
| Returns (ids, labels), same length. labels[i] == ids[i] on supervised |
| positions (every assistant turn's content + its trailing EOS) and |
| IGNORE_INDEX everywhere else (BOS, role tokens, system/user content). |
| """ |
| msgs = _normalized_messages(messages, default_system, expected_end="assistant") |
| return _encode_normalized_chat(tok, msgs) |
|
|
|
|
| def encode_prompt( |
| tok: Tokenizer, |
| messages: list[dict[str, str]], |
| default_system: str | None = DEFAULT_SYSTEM, |
| mode: str = "full_context", |
| ) -> list[int]: |
| """Encode a generation prompt: context ending in the ``<|assistant|>`` cue. |
| |
| mode: |
| - "full_context": the entire validated USER-ending context (earlier |
| assistant turns keep their [EOS]). |
| - "last_user": system turn + last user turn only. |
| |
| The returned ids start with BOS and end with ASSISTANT_ID, exactly matching |
| the training-time prefix for an assistant turn. |
| """ |
| msgs = _normalized_messages(messages, default_system, expected_end="user") |
|
|
| if mode == "last_user": |
| msgs = [msgs[0], msgs[-1]] if msgs[0]["role"] == "system" else [msgs[-1]] |
| elif mode != "full_context": |
| raise ValueError(f"unknown prompt mode: {mode}") |
|
|
| ids, _ = _encode_normalized_chat(tok, msgs) |
| ids.append(ASSISTANT_ID) |
| return ids |
|
|
|
|
| def encode_completion( |
| tok: Tokenizer, |
| messages: list[dict[str, str]], |
| completion: str, |
| default_system: str | None = DEFAULT_SYSTEM, |
| ) -> tuple[list[int], list[int]]: |
| """Encode (prompt context + one assistant completion) for DPO-style scoring. |
| |
| `messages` is the shared context (must end with a user turn); `completion` |
| is a plain assistant string. Returns (ids, labels) where the supervised |
| span is the completion tokens + trailing EOS — logps therefore include the |
| stop decision. |
| """ |
| ids = encode_prompt(tok, messages, default_system, mode="full_context") |
| labels = [IGNORE_INDEX] * len(ids) |
| if not isinstance(completion, str) or not completion.strip(): |
| raise ValueError("DPO completion must be a non-empty, non-whitespace string") |
| comp_ids = tok.encode(completion if default_system is None else clean_text_assistant(completion)).ids |
| if not comp_ids: |
| raise ValueError("DPO completion must encode to at least one token") |
| ids.extend(comp_ids) |
| labels.extend(comp_ids) |
| ids.append(EOS_ID) |
| labels.append(EOS_ID) |
| return ids, labels |
|
|
|
|
| def truncate_chat_sequence( |
| ids: list[int], |
| labels: list[int] | None, |
| max_len: int, |
| ) -> tuple[list[int], list[int] | None]: |
| """Validate and truncate only at complete user-turn boundaries. |
| |
| The prefix is BOS, plus the initial SYSTEM turn when present. It is always |
| retained; source-preserving conversations may begin BOS + USER. The remainder is the largest recent suffix that begins at a real |
| ``USER_ID`` marker and runs through the original sequence end. A training |
| sequence must end with an assistant EOS; a prompt must end with the final |
| assistant cue. Literal special-token text cannot create these boundaries |
| because chat tokenizers disable special-token matching for raw content. |
| |
| If the system prefix plus the latest complete user-led suffix cannot fit, |
| this function raises instead of cutting content, role markers, assistant |
| completions, or EOS targets. |
| """ |
| if max_len <= 0: |
| raise ValueError("max_len must be positive") |
| if labels is not None and len(labels) != len(ids): |
| raise ValueError("ids and labels must have identical lengths") |
| if len(ids) < 3 or ids[0] != BOS_ID or ids[1] not in {SYSTEM_ID, USER_ID}: |
| raise ValueError("chat sequence must start with BOS_ID and a SYSTEM or USER turn") |
|
|
| forbidden_content_ids = {PAD_ID, BOS_ID, EOS_ID, SYSTEM_ID, USER_ID, ASSISTANT_ID} |
| role_ids = {SYSTEM_ID, USER_ID, ASSISTANT_ID} |
| first_role = 1 |
| if ids[1] == SYSTEM_ID: |
| first_role = next((index for index in range(2, len(ids)) if ids[index] in role_ids), len(ids)) |
| if first_role == 2: |
| raise ValueError("initial system content must contain at least one token") |
| if any(token_id in forbidden_content_ids for token_id in ids[2:first_role]): |
| raise ValueError("system content contains a structural special-token ID") |
| if first_role == len(ids) or ids[first_role] != USER_ID: |
| raise ValueError("chat sequence must contain a user turn after the initial system turn") |
|
|
| user_starts: list[int] = [] |
| cursor = first_role |
| while cursor < len(ids): |
| if ids[cursor] != USER_ID: |
| raise ValueError("chat role sequence must alternate USER and ASSISTANT") |
| user_starts.append(cursor) |
| user_content_start = cursor + 1 |
| assistant_pos = next( |
| ( |
| index |
| for index in range(user_content_start, len(ids)) |
| if ids[index] in forbidden_content_ids |
| ), |
| len(ids), |
| ) |
| if assistant_pos == user_content_start: |
| raise ValueError("user content must contain at least one token") |
| if assistant_pos == len(ids) or ids[assistant_pos] != ASSISTANT_ID: |
| raise ValueError("each user turn must be followed by an assistant turn") |
|
|
| assistant_content_start = assistant_pos + 1 |
| if assistant_content_start == len(ids): |
| if labels is not None: |
| raise ValueError("training chat must end with assistant content and EOS") |
| cursor = len(ids) |
| break |
|
|
| eos_pos = next( |
| ( |
| index |
| for index in range(assistant_content_start, len(ids)) |
| if ids[index] in forbidden_content_ids |
| ), |
| len(ids), |
| ) |
| if eos_pos == assistant_content_start: |
| raise ValueError("assistant content must contain at least one token") |
| if eos_pos == len(ids) or ids[eos_pos] != EOS_ID: |
| raise ValueError("each assistant completion must end with EOS") |
| cursor = eos_pos + 1 |
| if cursor == len(ids): |
| if labels is None: |
| raise ValueError("generation prompt must end with an assistant cue") |
| break |
|
|
| if labels is not None: |
| if ids[-1] != EOS_ID: |
| raise ValueError("training chat must end with a supervised assistant EOS") |
| if labels[-1] != EOS_ID: |
| raise ValueError("final assistant EOS must be supervised") |
|
|
| prefix_end = first_role |
| chosen_start = next( |
| (start for start in user_starts if prefix_end + (len(ids) - start) <= max_len), |
| None, |
| ) |
| if chosen_start is None: |
| minimum = prefix_end + (len(ids) - user_starts[-1]) |
| raise ValueError( |
| "chat sequence does not fit without cutting the system prefix or latest " |
| f"user-led suffix (requires at least {minimum} tokens, max_len={max_len})" |
| ) |
|
|
| if chosen_start == prefix_end: |
| return list(ids), list(labels) if labels is not None else None |
|
|
| kept_ids = ids[:prefix_end] + ids[chosen_start:] |
| if labels is None: |
| return kept_ids, None |
| return kept_ids, labels[:prefix_end] + labels[chosen_start:] |
|
|
|
|
| def pad_or_truncate( |
| ids: list[int], |
| labels: list[int], |
| seq_len: int, |
| pad_id: int = PAD_ID, |
| ) -> tuple[list[int], list[int]]: |
| """Structure-aware fixed-length shaping plus right padding.""" |
| ids, maybe_labels = truncate_chat_sequence(ids, labels, seq_len) |
| assert maybe_labels is not None |
| labels = maybe_labels |
| pad_n = seq_len - len(ids) |
| return ids + [pad_id] * pad_n, labels + [IGNORE_INDEX] * pad_n |
|
|
|
|
| def count_chat_tokens( |
| tok: Tokenizer, |
| messages: list[dict[str, str]], |
| default_system: str | None = DEFAULT_SYSTEM, |
| ) -> int: |
| """Exact token count of the training encoding (for mix/token budgeting).""" |
| ids, _ = encode_chat(tok, messages, default_system) |
| return len(ids) |
|
|
|
|
| def extract_last_user_and_ref(messages: list[dict[str, str]]) -> tuple[str, str]: |
| """Return (last_user_text, last_assistant_text_if_any) — for sample logs.""" |
| last_user = "" |
| ref = "" |
| for m in reversed(messages or []): |
| if (m.get("role") or "").strip().lower() == "user": |
| last_user = clean_text(m.get("content", "")) |
| break |
| for m in reversed(messages or []): |
| if (m.get("role") or "").strip().lower() == "assistant": |
| ref = clean_text_assistant(m.get("content", "")) |
| break |
| return last_user, ref |
|
|
|
|
| def decode_completion(tok: Tokenizer, ids: list[int]) -> str: |
| """Decode generated completion ids (the tokens AFTER the prompt), dropping |
| a trailing EOS if present. Replaces the old fragile rfind('Assistant: ').""" |
| if ids and ids[-1] == EOS_ID: |
| ids = ids[:-1] |
| return tok.decode(ids).strip() if ids else "" |
|
|
|
|
| |
| |
| |
| def is_refusal_text(text: str, patterns: list[str]) -> bool: |
| """If assistant content contains any refusal-ish substring, treat as refusal.""" |
| t = (text or "").strip().lower() |
| if not t: |
| return False |
| for p in patterns: |
| p2 = p.strip().lower() |
| if p2 and p2 in t: |
| return True |
| return False |
|
|
|
|
| def compute_example_weight_from_messages( |
| messages: list[dict[str, str]], |
| refusal_downweight: float, |
| refusal_patterns: list[str], |
| refusal_mode: str, |
| ) -> float: |
| """Scalar loss weight for a training example; downweights refusal-looking |
| assistant turns (refusal_mode="contains_any").""" |
| if refusal_downweight >= 1.0: |
| return 1.0 |
| if refusal_downweight <= 0.0: |
| return 0.0 |
| if refusal_mode != "contains_any": |
| raise ValueError(f"unknown refusal_mode: {refusal_mode}") |
| for m in messages or []: |
| if (m.get("role") or "").strip().lower() == "assistant": |
| if is_refusal_text(m.get("content", ""), refusal_patterns): |
| return refusal_downweight |
| return 1.0 |
|
|
|
|
| def build_example( |
| ex: dict[str, Any], |
| tok: Tokenizer, |
| seq_len: int, |
| default_system: str | None, |
| refusal_downweight: float, |
| refusal_patterns: list[str], |
| refusal_mode: str, |
| pad_id: int = PAD_ID, |
| ): |
| """One SFT/distill training example -> (input_ids, labels, example_weight). |
| |
| Tensors are torch.long of length seq_len; labels use IGNORE_INDEX outside |
| the supervised assistant spans. Honors meta.bucket safety exemption and |
| meta.weight multipliers exactly as before. |
| """ |
| import torch |
|
|
| messages = ex.get("messages") or [] |
| if not messages: |
| raise ValueError("missing messages") |
|
|
| meta = ex.get("meta") or {} |
| bucket = str(meta.get("bucket", "")).strip() |
| |
| |
| refusal_dw_eff = 1.0 if bucket in ("D_safety", "D") else refusal_downweight |
| ex_weight = compute_example_weight_from_messages( |
| messages, refusal_dw_eff, refusal_patterns, refusal_mode |
| ) |
| w0 = meta.get("weight", None) if isinstance(meta, dict) else None |
| if isinstance(w0, (int, float)): |
| ex_weight *= float(w0) |
|
|
| ids, labels = encode_chat(tok, messages, default_system) |
| if default_system is None and len(ids) > seq_len: |
| raise ValueError("source-preserving SFT rejects overlength conversations") |
| ids, labels = pad_or_truncate(ids, labels, seq_len, pad_id) |
| if all(label == IGNORE_INDEX for label in labels): |
| raise ValueError("SFT example requires a retained non-empty assistant turn") |
| return ( |
| torch.tensor(ids, dtype=torch.long), |
| torch.tensor(labels, dtype=torch.long), |
| float(ex_weight), |
| ) |
|
|
