diff --git a/src/teich/audit.py b/src/teich/audit.py index 6b8d3da..30c71ed 100644 --- a/src/teich/audit.py +++ b/src/teich/audit.py @@ -5,6 +5,8 @@ from datasets import Dataset +from .protocol import chatml_turn_role_at_start + @dataclass class SFTAuditReport: @@ -31,6 +33,20 @@ def _as_list(value: Any) -> list[Any]: return list(value) +def _contiguous_supervised_runs(labels: list[int]) -> list[tuple[int, int]]: + runs: list[tuple[int, int]] = [] + run_start: int | None = None + for index, label in enumerate(labels): + if label != -100 and run_start is None: + run_start = index + elif label == -100 and run_start is not None: + runs.append((run_start, index)) + run_start = None + if run_start is not None: + runs.append((run_start, len(labels))) + return runs + + def _audit_training_row(row: dict[str, Any], tokenizer: Any, row_index: int) -> tuple[list[str], list[str], dict[str, Any]]: errors: list[str] = [] warnings: list[str] = [] @@ -69,12 +85,24 @@ def _audit_training_row(row: dict[str, Any], tokenizer: Any, row_index: int) -> supervised_text = _decode(tokenizer, supervised_ids) sample["supervised_preview"] = supervised_text[:500] + # A non-assistant header at the beginning of a contiguous label run is + # unambiguously context. Protocol-looking strings later in assistant output + # may be quoted transcript examples and cannot be classified from labels + # alone; Teich's formatter validates those against marker-derived spans. + for run_start, run_end in _contiguous_supervised_runs(labels): + run_text = _decode(tokenizer, input_ids[run_start:run_end]) + role = chatml_turn_role_at_start(run_text) + if role is not None and role != "assistant": + errors.append( + f"row {row_index}: supervised run begins at masked-context ChatML {role!r} header" + ) + break + masked_ids = [token_id for token_id, label in zip(input_ids, labels) if label == -100] masked_text = _decode(tokenizer, masked_ids[-200:]) if masked_ids else "" sample["masked_suffix_preview"] = masked_text[-500:] suspicious_masked_markers = ( - "<|im_start|>user", "<|start_header_id|>user<|end_header_id|>", "user", "<|start_of_role|>user<|end_of_role|>", @@ -111,14 +139,17 @@ def audit_sft_dataset(dataset: Dataset, tokenizer: Any, *, sample_size: int | No if dataset.num_rows == 0: return SFTAuditReport(ok=False, errors=["dataset contains no rows"]) - limit = dataset.num_rows if sample_size is None else min(max(sample_size, 0), dataset.num_rows) - if limit == 0: - warnings.append("sample_size is 0; no rows audited") + preview_limit = dataset.num_rows if sample_size is None else min(max(sample_size, 0), dataset.num_rows) + if preview_limit == 0: + warnings.append("sample_size is 0; no row previews retained (all rows were still audited)") - for row_index in range(limit): + # Sampling controls report size only. Correctness checks are deliberately + # exhaustive so an audit that returns ok=True is a dataset-wide gate. + for row_index in range(dataset.num_rows): row_errors, row_warnings, sample = _audit_training_row(dataset[row_index], tokenizer, row_index) errors.extend(row_errors) warnings.extend(row_warnings) - samples.append(sample) + if row_index < preview_limit: + samples.append(sample) return SFTAuditReport(ok=not errors, errors=errors, warnings=warnings, samples=samples) diff --git a/src/teich/formatter.py b/src/teich/formatter.py index bed9a7c..06aed24 100644 --- a/src/teich/formatter.py +++ b/src/teich/formatter.py @@ -13,6 +13,7 @@ from rich.console import Console from .converter import normalize_training_messages +from .protocol import chatml_header_is_source_anchored, first_chatml_structural_turn_boundary_end _GEMMA_TURN_START_PATTERN = re.compile(r"<\|turn>(model|user|system)\n") @@ -39,6 +40,18 @@ "", "<|start_of_role|>assistant<|end_of_role|>", ) +_TURN_ROLES = ("system", "developer", "user", "assistant", "tool", "ipython") +_TURN_BLOCK_STARTS = ( + *((f"<|im_start|>{role}\n", role) for role in _TURN_ROLES), + *((f"<|start_header_id|>{role}<|end_header_id|>\n\n", role) for role in _TURN_ROLES), + *((f"<|start_header_id|>{role}<|end_header_id|>", role) for role in _TURN_ROLES), + ("user\n", "user"), + ("model\n", "assistant"), + *((f"<|{role}|>\n", role) for role in _TURN_ROLES), + *((f"<|{role}|>", role) for role in _TURN_ROLES), + *((f"<{role}>", role) for role in _TURN_ROLES), + *((f"<|start_of_role|>{role}<|end_of_role|>", role) for role in _TURN_ROLES), +) _ASSISTANT_BLOCK_END_TOKENS = ( "<|im_end|>", "<|eot_id|>", @@ -46,7 +59,17 @@ "", "", "<|end_of_text|>", + "<|end|>", +) +_TURN_BLOCK_END_TOKENS = ( + *_ASSISTANT_BLOCK_END_TOKENS, + _GEMMA_TURN_END, + "", + "", + "", + "", ) +_TURN_LEADING_TOKENS = ("", "", "<|begin_of_text|>") _REASONING_BLOCK_PATTERNS = ( re.compile(r"\n.*?\n\n?", re.DOTALL), re.compile(r".*?", re.DOTALL), @@ -1506,55 +1529,83 @@ def _resolve_assistant_prompt_prefixes( def _assistant_block_bounds(text: str, start: int, end: int) -> tuple[int, int] | None: - block_start = -1 - for token in _ASSISTANT_BLOCK_START_TOKENS: - token_start = text.rfind(token, 0, start) - if token_start > block_start: - block_start = token_start - if block_start < 0: - return None - block_end = -1 - for token in _ASSISTANT_BLOCK_END_TOKENS: - token_end_start = text.find(token, end) - if token_end_start >= 0 and (block_end < 0 or token_end_start < block_end): - block_end = token_end_start + len(token) - if block_end < 0: - return None - while block_end < len(text) and text[block_end] in "\r\n": - block_end += 1 - return block_start, block_end + return _AssistantBlockIndex.build(text).bounds(start, end) + + +def _is_valid_turn_header( + text: str, + start: int, + role: str, + source_spans: Sequence[Mapping[str, Any]] | None = None, +) -> bool: + prefix_end = start + while prefix_end > 0 and text[prefix_end - 1].isspace(): + prefix_end -= 1 + + leading_cursor = 0 + while leading_cursor < prefix_end: + while leading_cursor < prefix_end and text[leading_cursor].isspace(): + leading_cursor += 1 + leading_token = next( + (token for token in _TURN_LEADING_TOKENS if text.startswith(token, leading_cursor)), + None, + ) + if leading_token is None: + break + leading_cursor += len(leading_token) + if leading_cursor == prefix_end: + return True + if not any( + prefix_end >= len(token) and text.startswith(token, prefix_end - len(token), prefix_end) + for token in _TURN_BLOCK_END_TOKENS + ): + return False + if source_spans is not None and text.startswith("<|im_start|>", start): + return chatml_header_is_source_anchored(text, start, role, source_spans) + return True @dataclass(slots=True) class _AssistantBlockIndex: text: str - start_thresholds: list[int] - start_prefix_maxima: list[int] + turn_starts: list[int] + turn_content_starts: list[int] + turn_roles: list[str] end_starts: list[int] end_positions: list[int] @classmethod - def build(cls, text: str) -> _AssistantBlockIndex: - start_events: list[tuple[int, int]] = [] - for token in _ASSISTANT_BLOCK_START_TOKENS: + def build( + cls, + text: str, + source_spans: Sequence[Mapping[str, Any]] | None = None, + ) -> _AssistantBlockIndex: + turn_by_start: dict[int, tuple[int, str]] = {} + for token, role in _TURN_BLOCK_STARTS: cursor = 0 while True: token_start = text.find(token, cursor) if token_start < 0: break - start_events.append((token_start + len(token), token_start)) + if _is_valid_turn_header(text, token_start, role, source_spans): + content_start = token_start + len(token) + existing = turn_by_start.get(token_start) + if existing is None or content_start > existing[0]: + turn_by_start[token_start] = (content_start, role) cursor = token_start + len(token) - start_events.sort() - start_thresholds: list[int] = [] - start_prefix_maxima: list[int] = [] - max_start = -1 - for threshold, token_start in start_events: - max_start = max(max_start, token_start) - start_thresholds.append(threshold) - start_prefix_maxima.append(max_start) + + # Gemma's parser advances from each real terminator, so role-like + # strings inside message content never become structural turn starts. + for match in _gemma_turn_matches(text): + role = "assistant" if match.group(1) == "model" else match.group(1) + turn_by_start[match.start()] = (match.end(), role) + ordered_turns = sorted( + (turn_start, content_start, role) + for turn_start, (content_start, role) in turn_by_start.items() + ) end_by_start: dict[int, int] = {} - for token in _ASSISTANT_BLOCK_END_TOKENS: + for token in _TURN_BLOCK_END_TOKENS: cursor = 0 while True: token_start = text.find(token, cursor) @@ -1565,23 +1616,50 @@ def build(cls, text: str) -> _AssistantBlockIndex: ordered_ends = sorted(end_by_start.items()) return cls( text=text, - start_thresholds=start_thresholds, - start_prefix_maxima=start_prefix_maxima, + turn_starts=[start for start, _, _ in ordered_turns], + turn_content_starts=[content_start for _, content_start, _ in ordered_turns], + turn_roles=[role for _, _, role in ordered_turns], end_starts=[start for start, _ in ordered_ends], end_positions=[end for _, end in ordered_ends], ) - def bounds(self, start: int, end: int) -> tuple[int, int] | None: - start_index = bisect_right(self.start_thresholds, start) - 1 - if start_index < 0: + def originating_assistant_turn(self, position: int) -> tuple[int, int, int] | None: + turn_index = bisect_right(self.turn_starts, position) - 1 + if turn_index < 0: return None - block_start = self.start_prefix_maxima[start_index] - end_index = bisect_left(self.end_starts, end) - if end_index >= len(self.end_starts): + block_start = self.turn_starts[turn_index] + content_start = self.turn_content_starts[turn_index] + if content_start > position or self.turn_roles[turn_index] != "assistant": + return None + next_turn_index = turn_index + 1 + turn_boundary = ( + self.turn_starts[next_turn_index] + if next_turn_index < len(self.turn_starts) + else len(self.text) + ) + end_index = bisect_left(self.end_starts, turn_boundary) - 1 + if end_index >= 0 and self.end_starts[end_index] >= content_start: + block_end = self.end_positions[end_index] + else: + block_end = turn_boundary + if not self.text.startswith(_GEMMA_ASSISTANT_TURN_PREFIX, block_start): + while block_end < turn_boundary and self.text[block_end] in "\r\n": + block_end += 1 + return block_start, content_start, block_end + + def is_gemma_assistant_turn(self, position: int) -> bool: + assistant_turn = self.originating_assistant_turn(position) + return ( + assistant_turn is not None + and self.text.startswith(_GEMMA_ASSISTANT_TURN_PREFIX, assistant_turn[0]) + ) + + def bounds(self, start: int, end: int) -> tuple[int, int] | None: + del end + assistant_turn = self.originating_assistant_turn(start) + if assistant_turn is None: return None - block_end = self.end_positions[end_index] - while block_end < len(self.text) and self.text[block_end] in "\r\n": - block_end += 1 + block_start, _, block_end = assistant_turn return block_start, block_end def following_end(self, position: int) -> tuple[int, int] | None: @@ -1623,6 +1701,8 @@ def _expand_supervised_span( if assistant_block is None: return span block_start, block_end = assistant_block + if assistant_blocks.is_gemma_assistant_turn(start): + return span if not assistant_prompt_prefixes: return block_start, block_end block_text = text[block_start:block_end] @@ -1661,7 +1741,7 @@ def _expand_typed_spans( assistant_prompt_prefixes: tuple[str, ...], ) -> list[dict[str, Any]]: span_kinds = {span.get("kind") for span in spans} - assistant_blocks = _AssistantBlockIndex.build(text) + assistant_blocks = _AssistantBlockIndex.build(text, spans) reasoning_spans = _ContainingSpanIndex.build(_reasoning_spans(text)) if _SPAN_KIND_REASONING in span_kinds else None tool_call_spans = _ContainingSpanIndex.build(_tool_call_spans(text)) if _SPAN_KIND_TOOL_CALL in span_kinds else None tool_response_spans = ( @@ -1691,11 +1771,12 @@ def _expand_typed_spans( elif kind == _SPAN_KIND_TOOL_CALL: if tool_call_spans is not None: expanded_start, expanded_end = tool_call_spans.containing((start, end)) - expanded_start, expanded_end = _extend_span_to_following_assistant_end( - text, - (expanded_start, expanded_end), - assistant_blocks, - ) + if not assistant_blocks.is_gemma_assistant_turn(start): + expanded_start, expanded_end = _extend_span_to_following_assistant_end( + text, + (expanded_start, expanded_end), + assistant_blocks, + ) elif kind == _SPAN_KIND_TOOL_RESPONSE: if tool_response_spans is not None: expanded_start, expanded_end = tool_response_spans.containing((start, end)) @@ -1704,9 +1785,9 @@ def _expand_typed_spans( updated["end"] = expanded_end updated.setdefault("source_start", start) updated.setdefault("source_end", end) - if expanded_start < expanded_end: + if updated["start"] < updated["end"]: expanded_spans.append(updated) - return expanded_spans + return _clamp_model_spans_to_assistant_turns(text, expanded_spans) def _span_kind_enabled( @@ -1745,6 +1826,54 @@ def _source_spans_for_kind(spans: list[dict[str, Any]], kind: str) -> list[tuple return _merge_spans(source_spans) +def _clamp_model_spans_to_assistant_turns( + text: str, + spans: list[dict[str, Any]], +) -> list[dict[str, Any]]: + """Keep model-authored metadata inside the turn containing its source start.""" + assistant_blocks = _AssistantBlockIndex.build(text, spans) + clamped: list[dict[str, Any]] = [] + for span in spans: + updated = dict(span) + if updated.get("kind") in { + _SPAN_KIND_REASONING, + _SPAN_KIND_FINAL_ANSWER, + _SPAN_KIND_TOOL_CALL, + }: + source_start = updated.get("source_start", updated.get("start")) + if isinstance(source_start, int): + # Clamp the final span itself at a structural ChatML boundary. + # This is intentionally independent of assistant-block inference: + # marker extraction and kind-specific expansion take different + # paths, but none may supervise the following turn. A bare or + # quoted <|im_start|> is not a boundary unless it immediately + # follows <|im_end|> and carries a real role header. + structural_end = first_chatml_structural_turn_boundary_end( + text, + source_start, + spans, + ) + if structural_end is not None: + updated["end"] = min(updated["end"], structural_end) + source_end = updated.get("source_end") + if isinstance(source_end, int): + updated["source_end"] = min(source_end, structural_end) + assistant_turn = assistant_blocks.originating_assistant_turn(source_start) + if assistant_turn is None: + if assistant_blocks.turn_starts: + continue + else: + _, content_start, block_end = assistant_turn + updated["start"] = max(updated["start"], content_start) + updated["end"] = min(updated["end"], block_end) + source_end = updated.get("source_end") + if isinstance(source_end, int): + updated["source_end"] = min(source_end, block_end) + if updated["start"] < updated["end"]: + clamped.append(updated) + return clamped + + def _select_supervised_spans( text: str, spans: list[dict[str, Any]], @@ -1757,6 +1886,7 @@ def _select_supervised_spans( train_on_developer: bool, train_on_tool_responses: bool, ) -> list[tuple[int, int]]: + spans = _clamp_model_spans_to_assistant_turns(text, spans) selected = _merge_spans( [ (span["start"], span["end"]) diff --git a/src/teich/protocol.py b/src/teich/protocol.py new file mode 100644 index 0000000..0540278 --- /dev/null +++ b/src/teich/protocol.py @@ -0,0 +1,102 @@ +from __future__ import annotations + +import re +from collections.abc import Iterator, Mapping, Sequence +from typing import Any + + +_CHATML_ROLES = ("system", "developer", "user", "assistant", "tool", "ipython") +_CHATML_STRUCTURAL_TURN_BOUNDARY_PATTERN = re.compile( + r"<\|im_end\|>(?:\r?\n)?" + r"(?=<\|im_start\|>(?P" + "|".join(_CHATML_ROLES) + r")(?:\r?\n))" +) +_CHATML_TURN_HEADER_PATTERN = re.compile( + r"<\|im_start\|>(?P" + "|".join(_CHATML_ROLES) + r")(?:\r?\n)" +) + + +def _span_range(span: Mapping[str, Any]) -> tuple[int, int] | None: + start = span.get("source_start", span.get("start")) + end = span.get("source_end", span.get("end")) + if isinstance(start, int) and isinstance(end, int) and start < end: + return start, end + return None + + +def _roles_match(rendered_role: str, span_role: object) -> bool: + normalized = "assistant" if span_role == "model" else span_role + if rendered_role == "ipython": + return normalized in {"ipython", "tool"} + return normalized == rendered_role + + +def chatml_header_is_source_anchored( + text: str, + header_start: int, + role: str, + source_spans: Sequence[Mapping[str, Any]], +) -> bool: + """Reject protocol-looking transcript snippets inside known message content.""" + has_context_spans = any( + span.get("role") not in {"assistant", "model"} + for span in source_spans + ) + if not has_context_spans: + # External/legacy metadata may contain only the model span. Retain the + # fail-closed raw protocol boundary in that case because there is no + # marker-derived context with which to disambiguate transcript text. + return True + containing_ranges = [ + span_range + for span in source_spans + if (span_range := _span_range(span)) is not None + and span_range[0] < header_start < span_range[1] + ] + if not containing_ranges: + return True + + header = _CHATML_TURN_HEADER_PATTERN.match(text, header_start) + if header is None or header.group("role") != role: + return False + turn_end = text.find("<|im_end|>", header.end()) + if turn_end < 0: + return False + return any( + _roles_match(role, span.get("role")) + and (span_range := _span_range(span)) is not None + and header.end() <= span_range[0] < turn_end + for span in source_spans + ) + + +def iter_chatml_structural_turn_boundaries( + text: str, + start: int = 0, + source_spans: Sequence[Mapping[str, Any]] | None = None, +) -> Iterator[re.Match[str]]: + """Yield ChatML turn boundaries, excluding quoted or standalone token names.""" + for match in _CHATML_STRUCTURAL_TURN_BOUNDARY_PATTERN.finditer(text, max(start, 0)): + if source_spans is not None and not chatml_header_is_source_anchored( + text, + match.end(), + match.group("role"), + source_spans, + ): + continue + yield match + + +def first_chatml_structural_turn_boundary_end( + text: str, + start: int = 0, + source_spans: Sequence[Mapping[str, Any]] | None = None, +) -> int | None: + """Return the end of the first terminator/new-turn boundary at or after start.""" + match = next(iter_chatml_structural_turn_boundaries(text, start, source_spans), None) + return match.end() if match is not None else None + + +def chatml_turn_role_at_start(text: str) -> str | None: + """Return a ChatML role only when its header begins the supplied text.""" + match = _CHATML_TURN_HEADER_PATTERN.match(text) + return match.group("role") if match is not None else None diff --git a/tests/test_audit.py b/tests/test_audit.py index 05fa7b6..803b30c 100644 --- a/tests/test_audit.py +++ b/tests/test_audit.py @@ -21,6 +21,10 @@ def __init__(self): 5: "", 6: "<|turn>user", 7: "<|tool_response>", + 8: "<|im_end|>\n", + 9: "<|im_start|>user\n", + 10: '["<|im_start|>", "<|im_end|>"]', + 11: "EXAMPLE<|im_end|>\n<|im_start|>user\nLITERAL<|im_end|>", } def decode(self, token_ids, skip_special_tokens=False, clean_up_tokenization_spaces=False): @@ -62,13 +66,13 @@ def test_audit_sft_dataset_rejects_label_input_mismatch(): assert "labels differ from input_ids" in report.errors[0] -def test_audit_sft_dataset_rejects_supervised_user_marker(): +def test_audit_sft_dataset_rejects_run_starting_at_chatml_user_header(): dataset = Dataset.from_list( [ { - "input_ids": [3, 2], - "attention_mask": [1, 1], - "labels": [3, 2], + "input_ids": [1, 8, 9, 2], + "attention_mask": [1, 1, 1, 1], + "labels": [-100, -100, 9, 2], } ] ) @@ -76,7 +80,41 @@ def test_audit_sft_dataset_rejects_supervised_user_marker(): report = audit_sft_dataset(dataset, TinyTokenizer()) assert not report.ok - assert "<|im_start|>user" in report.errors[0] + assert "begins at masked-context ChatML 'user' header" in report.errors[0] + + +def test_audit_sft_dataset_allows_quoted_chatml_token_names(): + dataset = Dataset.from_list( + [ + { + "input_ids": [1, 10, 4], + "attention_mask": [1, 1, 1], + "labels": [1, 10, 4], + } + ] + ) + + report = audit_sft_dataset(dataset, TinyTokenizer()) + + assert report.ok + assert report.errors == [] + + +def test_audit_sft_dataset_allows_literal_full_chatml_transcript_in_assistant_output(): + dataset = Dataset.from_list( + [ + { + "input_ids": [1, 11, 4], + "attention_mask": [1, 1, 1], + "labels": [1, 11, 4], + } + ] + ) + + report = audit_sft_dataset(dataset, TinyTokenizer()) + + assert report.ok + assert report.errors == [] def test_audit_sft_dataset_rejects_gemma_context_markers(): @@ -108,6 +146,18 @@ def test_audit_sft_dataset_checks_all_rows_by_default(): assert any("row 8" in error and "<|turn>user" in error for error in report.errors) +def test_audit_sample_size_limits_previews_not_correctness_checks(): + safe = {"input_ids": [1, 2, 4], "attention_mask": [1, 1, 1], "labels": [-100, 2, 4]} + leaked = {"input_ids": [6, 2], "attention_mask": [1, 1], "labels": [6, 2]} + dataset = Dataset.from_list([safe] * 8 + [leaked]) + + report = audit_sft_dataset(dataset, TinyTokenizer(), sample_size=1) + + assert not report.ok + assert len(report.samples) == 1 + assert any("row 8" in error and "<|turn>user" in error for error in report.errors) + + def test_teich_example_has_single_safe_training_flow(): source = Path("teich_example.py").read_text(encoding="utf-8") tree = ast.parse(source) diff --git a/tests/test_formatter.py b/tests/test_formatter.py index b751ef6..9f71a76 100644 --- a/tests/test_formatter.py +++ b/tests/test_formatter.py @@ -4684,6 +4684,264 @@ def test_expand_typed_spans_indexes_large_conversations_once(): assert (final_span["start"], final_span["end"]) == expected +def test_expand_typed_spans_clamps_corrupt_granite_spans_to_originating_turn(): + assistant_prefix = "<|im_start|>assistant\n" + assistant_body = ( + "inspect\n" + "I will inspect it.\n" + "\n\n\n\n" + ) + first_turn = assistant_prefix + assistant_body + "<|im_end|>\n" + user_turn = "<|im_start|>user\nDo something else.<|im_end|>\n" + second_turn = assistant_prefix + "Second answer.<|im_end|>\n" + text = first_turn + user_turn + second_turn + corrupt_end = text.index("Do something else.") + len("Do something else.") + spans = [ + { + "start": text.index("inspect"), + "end": corrupt_end, + "kind": "reasoning", + "role": "assistant", + }, + { + "start": text.index("I will inspect it."), + "end": corrupt_end, + "kind": "final_answer", + "role": "assistant", + }, + { + # Mirrors the observed malformed Granite metadata: its tool-call + # marker begins in assistant prose and ends in the next user turn. + "start": text.index("I will inspect it."), + "end": corrupt_end, + "kind": "tool_call", + "role": "assistant", + }, + ] + + expanded = _expand_typed_spans(text, spans, (assistant_prefix,)) + + assert len(expanded) == 3 + next_turn_start = text.index("<|im_start|>user") + for span in expanded: + assert span["end"] <= next_turn_start + assert span["source_end"] <= next_turn_start + assert "<|im_start|>user" not in text[span["start"] : span["end"]] + assert "Do something else." not in text[span["start"] : span["end"]] + + +def test_expand_typed_spans_ignores_literal_role_headers_in_assistant_content(): + assistant_prefix = "<|im_start|>assistant\n" + literal_headers = ( + 'Explain and print ["<|im_start|>", "<|im_end|>"] plus ' + "<|im_start|>user\n as plain text." + ) + first_turn = assistant_prefix + literal_headers + "<|im_end|>\n" + next_turn = "<|im_start|>user\nActual next turn.<|im_end|>\n" + text = first_turn + next_turn + content_start = text.index("Explain") + spans = [ + { + "start": content_start, + "end": content_start + len(literal_headers), + "kind": "final_answer", + "role": "assistant", + } + ] + + expanded = _expand_typed_spans(text, spans, (assistant_prefix,)) + + assert len(expanded) == 1 + supervised = text[expanded[0]["start"] : expanded[0]["end"]] + assert literal_headers in supervised + assert supervised.endswith("<|im_end|>\n") + assert "Actual next turn." not in supervised + + +def test_expand_typed_spans_preserves_literal_full_chatml_transcript(): + assistant_prefix = "<|im_start|>assistant\n" + literal_transcript = ( + "Explain this literal transcript:\n" + "<|im_end|>\n<|im_start|>user\nEXAMPLE_USER<|im_end|>\n" + "and continue the explanation." + ) + first_turn = assistant_prefix + literal_transcript + "<|im_end|>\n" + next_turn = "<|im_start|>user\nACTUAL_USER<|im_end|>\n" + text = first_turn + next_turn + assistant_start = text.index("Explain this") + user_start = text.index("ACTUAL_USER") + spans = [ + { + "start": assistant_start, + "end": assistant_start + len(literal_transcript), + "kind": "final_answer", + "role": "assistant", + }, + { + "start": user_start, + "end": user_start + len("ACTUAL_USER"), + "kind": "user", + "role": "user", + }, + ] + + expanded = _expand_typed_spans(text, spans, (assistant_prefix,)) + + assistant_span = next(span for span in expanded if span.get("kind") == "final_answer") + supervised = text[assistant_span["start"] : assistant_span["end"]] + assert literal_transcript in supervised + assert supervised.endswith("<|im_end|>\n") + assert "ACTUAL_USER" not in supervised + + +def test_expand_typed_spans_stops_before_llama_ipython_tool_result(): + assistant_prefix = "<|start_header_id|>assistant<|end_header_id|>\n\n" + first_turn = assistant_prefix + "MODEL_ANSWER<|eot_id|>" + tool_prefix = "<|start_header_id|>ipython<|end_header_id|>\n\n" + tool_turn = tool_prefix + "SECRET_TOOL_RESULT<|eot_id|>" + next_turn = assistant_prefix + "NEXT_ANSWER<|eot_id|>" + text = first_turn + tool_turn + next_turn + answer_start = text.index("MODEL_ANSWER") + tool_start = text.index("SECRET_TOOL_RESULT") + spans = [ + { + "start": answer_start, + "end": tool_start + len("SECRET_TOOL_RESULT"), + "source_start": answer_start, + "source_end": answer_start + len("MODEL_ANSWER"), + "kind": "final_answer", + "role": "assistant", + }, + { + "start": tool_start, + "end": tool_start + len("SECRET_TOOL_RESULT"), + "kind": "tool_response", + "role": "tool", + }, + ] + + expanded = _expand_typed_spans(text, spans, (assistant_prefix,)) + + answer_span = next(span for span in expanded if span.get("kind") == "final_answer") + supervised = text[answer_span["start"] : answer_span["end"]] + assert supervised == "MODEL_ANSWER<|eot_id|>" + assert "ipython" not in supervised + assert "SECRET_TOOL_RESULT" not in supervised + + +def test_expand_typed_spans_recognizes_phi_end_terminator(): + user_prefix = "<|user|>\n" + assistant_prefix = "<|assistant|>\n" + text = user_prefix + "QUESTION<|end|>\n" + assistant_prefix + "ANSWER<|end|>\n" + user_start = text.index("QUESTION") + answer_start = text.index("ANSWER") + spans = [ + { + "start": user_start, + "end": user_start + len("QUESTION"), + "kind": "user", + "role": "user", + }, + { + "start": answer_start, + "end": answer_start + len("ANSWER"), + "kind": "final_answer", + "role": "assistant", + }, + ] + + expanded = _expand_typed_spans(text, spans, (assistant_prefix,)) + + answer_span = next(span for span in expanded if span.get("kind") == "final_answer") + assert text[answer_span["start"] : answer_span["end"]] == "ANSWER<|end|>\n" + + +def test_expand_typed_spans_rejects_model_kind_starting_in_user_turn(): + user_turn = "<|im_start|>user\nUSER_CONTENT<|im_end|>\n" + assistant_turn = "<|im_start|>assistant\nASSISTANT_CONTENT<|im_end|>\n" + text = user_turn + assistant_turn + user_start = text.index("USER_CONTENT") + spans = [ + { + "start": user_start, + "end": user_start + len("USER_CONTENT"), + "kind": "final_answer", + "role": "assistant", + } + ] + + assert _expand_typed_spans(text, spans, ("<|im_start|>assistant\n",)) == [] + + +def test_expand_typed_spans_clamps_corrupt_gemma_span_to_model_turn(): + model_prefix = "<|turn>model\n" + model_turn = model_prefix + "MODEL_OUTPUT\n" + user_turn = "<|turn>user\nSECRET_USER_CONTENT\n" + text = model_turn + user_turn + spans = [ + { + "start": text.index("MODEL_OUTPUT"), + "end": text.index("SECRET_USER_CONTENT") + len("SECRET_USER_CONTENT"), + "kind": "tool_call", + "role": "assistant", + } + ] + + expanded = _expand_typed_spans(text, spans, (model_prefix,)) + + assert len(expanded) == 1 + supervised = text[expanded[0]["start"] : expanded[0]["end"]] + assert supervised == "MODEL_OUTPUT" + assert "SECRET_USER_CONTENT" not in supervised + + +def test_mask_data_defensively_clamps_external_granite_span_metadata(): + tokenizer = TrainerStyleTokenizer() + prior_user_turn = "<|im_start|>user\nPRIOR_USER_MESSAGE<|im_end|>\n" + first_turn = ( + "<|im_start|>assistant\nI will inspect it.\n" + "bash\n<|im_end|>\n" + ) + user_turn = "<|im_start|>user\nSECRET_USER_MESSAGE<|im_end|>\n" + text = prior_user_turn + first_turn + user_turn + encoded = tokenizer(text, add_special_tokens=False) + prepared = Dataset.from_list( + [ + { + "text": text, + "input_ids": encoded["input_ids"], + "attention_mask": encoded["attention_mask"], + "teich_supervised_spans": [ + { + "start": text.index("PRIOR_USER_MESSAGE"), + "end": text.index("SECRET_USER_MESSAGE") + len("SECRET_USER_MESSAGE"), + "source_start": text.index("I will inspect it."), + "source_end": text.index("SECRET_USER_MESSAGE") + len("SECRET_USER_MESSAGE"), + "kind": "tool_call", + "role": "assistant", + } + ], + } + ] + ) + trainer = SimpleNamespace( + train_dataset=prepared, + eval_dataset=None, + processing_class=tokenizer, + args=SimpleNamespace(dataset_text_field="text", packing=False, max_length=4096), + ) + + trainer = mask_data(trainer, tokenizer=tokenizer, audit=True, verbose=False) + + row = trainer.train_dataset[0] + supervised_text = tokenizer.decode( + [token for token in row["labels"] if token != -100] + ) + assert supervised_text == "I will inspect it.\nbash\n<|im_end|>\n" + assert "SECRET_USER_MESSAGE" not in supervised_text + assert "<|im_start|>user" not in supervised_text + + def test_actual_qwen_template_receives_normalized_mapping_tool_arguments(): jinja2 = pytest.importorskip("jinja2") template_path = Path("qwen3.6_chat_template.jinja")