diff --git a/docs/training.md b/docs/training.md index 482c66c..e41804f 100644 --- a/docs/training.md +++ b/docs/training.md @@ -226,7 +226,10 @@ For Gemma 4, Teich supervises exactly one closing `` for a completed model turn whenever reasoning, final-answer, or tool-call supervision is enabled for that turn. It does not add a second terminator inside a continuing tool-call chain. The terminator remains a target even when the final answer is -masked, so reasoning-only and tool-only fine-tunes still learn to stop. +masked, so reasoning-only and tool-only fine-tunes still learn to stop. Empty +final assistant messages and unresolved tool-result prefixes do not create +synthetic `final_answer` spans; the stopping token is attached only after an +actual enabled target is selected. ## Masking Policy diff --git a/src/teich/formatter.py b/src/teich/formatter.py index 77eeec8..b1aacfb 100644 --- a/src/teich/formatter.py +++ b/src/teich/formatter.py @@ -1695,13 +1695,14 @@ def _select_supervised_spans( selected = _merge_spans(selected) # A Gemma turn terminator is useful exactly when the same model turn still - # contains an enabled training target. It is initially part of the inferred - # final-answer span, so restore it when final answers are disabled but an - # enabled reasoning or tool-call target remains. Conversely, do not keep a - # reasoning-only row alive solely by labeling after reasoning was - # excluded. + # contains an enabled model-authored training target. Terminator-only spans + # are not final answers, so add the terminator from the rendered protocol + # instead of relying on degenerate final-answer metadata. Conversely, do + # not keep a row alive solely by labeling after every assistant + # target was excluded. orphaned_turn_ends: list[tuple[int, int]] = [] retained_turn_ends: list[tuple[int, int]] = [] + model_selected = _subtract_spans(selected, _tool_response_spans(text)) turn_matches = _gemma_turn_matches(text) for index, match in enumerate(turn_matches): if match.group(1) != "model": @@ -1711,16 +1712,12 @@ def _select_supervised_spans( if turn_end < 0: continue turn_end_span = (turn_end, turn_end + len(_GEMMA_TURN_END)) - turn_end_was_supervised = any( - span["start"] <= turn_end and span["end"] >= turn_end_span[1] - for span in spans - ) has_selected_content = any( text[max(start, match.end()) : min(end, turn_end)].strip() - for start, end in selected + for start, end in model_selected if start < turn_end and end > match.end() ) - if has_selected_content and turn_end_was_supervised: + if has_selected_content: retained_turn_ends.append(turn_end_span) elif not has_selected_content: orphaned_turn_ends.append(turn_end_span) @@ -1906,22 +1903,6 @@ def _supervised_text_and_spans( ) if inferred_spans: return formatted_text, _span_dicts(inferred_spans) - gemma_spans = _gemma_like_supervised_spans(formatted_text) - if gemma_spans: - gemma_spans = _subtract_spans(gemma_spans, _tool_call_spans(formatted_text)) - gemma_spans = _subtract_spans(gemma_spans, _tool_response_spans(formatted_text)) - gemma_spans = _subtract_spans(gemma_spans, _reasoning_spans(formatted_text)) - supervised_spans.extend( - { - "start": start, - "end": end, - "source_start": start, - "source_end": end, - "kind": _SPAN_KIND_FINAL_ANSWER, - "role": "assistant", - } - for start, end in gemma_spans - ) assistant_prompt_prefixes = _resolve_assistant_prompt_prefixes( renderer, messages, diff --git a/tests/test_formatter.py b/tests/test_formatter.py index 94287b4..2cbdfc4 100644 --- a/tests/test_formatter.py +++ b/tests/test_formatter.py @@ -1533,6 +1533,128 @@ def apply_chat_template(self, *args, **kwargs): assert supervised_text.endswith("") +def test_gemma_empty_final_message_does_not_create_a_degenerate_final_answer_span(): + tokenizer = GemmaLikeOffsetTokenizer() + tokenizer.name_or_path = "google/gemma-4-26B-A4B-it" + dataset = Dataset.from_list( + [ + { + "messages": [ + {"role": "user", "content": "list files"}, + { + "role": "assistant", + "content": "", + "reasoning_content": "inspect first", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "bash", "arguments": {"command": "ls"}}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "name": "bash", "content": "SECRET"}, + {"role": "assistant", "content": ""}, + ], + "tools": [ + { + "type": "function", + "function": { + "name": "bash", + "parameters": {"type": "object", "properties": {"command": {"type": "string"}}}, + }, + } + ], + } + ] + ) + + prepared = prepare_data(dataset, tokenizer, tokenize=True, strict=True, verbose=False) + spans = prepared[0]["teich_supervised_spans"] + assert not any(span.get("kind") == "final_answer" for span in spans) + + trainer = SimpleNamespace( + train_dataset=prepared, + eval_dataset=None, + args=SimpleNamespace(dataset_text_field="text", max_length=None, packing=False), + tokenizer=tokenizer, + ) + trainer = mask_data( + trainer, + tokenizer=tokenizer, + train_on_reasoning=True, + train_on_final_answers=True, + train_on_tools=True, + audit=True, + verbose=False, + ) + row = trainer.train_dataset[0] + supervised_text = tokenizer.decode([token for token in row["labels"] if token != -100]) + assert "inspect first" in supervised_text + assert '<|tool_call>call:bash{command:"ls"}' in supervised_text + assert "SECRET" not in supervised_text + assert supervised_text.endswith("") + assert supervised_text.count("") == 1 + + +def test_gemma_tool_only_mask_does_not_cross_two_tool_response_boundaries(): + tokenizer = GemmaLikeOffsetTokenizer() + tokenizer.name_or_path = "google/gemma-4-26B-A4B-it" + messages = [{"role": "user", "content": "inspect twice"}] + tools = [ + { + "type": "function", + "function": { + "name": "bash", + "parameters": {"type": "object", "properties": {"command": {"type": "string"}}}, + }, + } + ] + for index, command in enumerate(("ls", "pwd"), start=1): + messages.extend( + [ + { + "role": "assistant", + "content": f"narration {index}", + "tool_calls": [ + { + "id": f"call_{index}", + "type": "function", + "function": {"name": "bash", "arguments": {"command": command}}, + } + ], + }, + { + "role": "tool", + "tool_call_id": f"call_{index}", + "name": "bash", + "content": f"SECRET_RESPONSE_{index}", + }, + ] + ) + messages.append({"role": "assistant", "content": "final text"}) + + training_data = prepare_and_mask_for_test( + Dataset.from_list([{"messages": messages, "tools": tools}]), + tokenizer, + train_on_reasoning=False, + train_on_final_answers=False, + train_on_tools=True, + strict=True, + ) + row = training_data[0] + supervised_text = tokenizer.decode([token for token in row["labels"] if token != -100]) + + assert 'call:bash{command:"ls"}' in supervised_text + assert 'call:bash{command:"pwd"}' in supervised_text + assert "narration 1" not in supervised_text + assert "narration 2" not in supervised_text + assert "final text" not in supervised_text + assert "SECRET_RESPONSE_1" not in supervised_text + assert "SECRET_RESPONSE_2" not in supervised_text + assert supervised_text.count("") == 1 + + def test_prepare_and_mask_falls_back_when_gemma_drops_marker_boundaries_with_thinking_disabled(): class MarkerDroppingGemmaTokenizer(GemmaLikeOffsetTokenizer): def apply_chat_template(self, *args, **kwargs): @@ -4117,6 +4239,60 @@ def test_new_gemma4_template_strict_markers_tolerate_trimmed_embedded_thinking() assert "gh issue view 1123 --json title,body" in supervised_text +def test_new_gemma4_unresolved_tool_call_does_not_label_tool_response_prefix_as_answer(): + jinja2 = pytest.importorskip("jinja2") + template_path = Path("new_gemma_4_template.jinja") + if not template_path.exists(): + pytest.skip(f"{template_path} is not available") + + tokenizer = RealJinjaChatTemplateTokenizer(template_path, jinja2) + answer = "A" * 92 + tools = _real_template_tool_call_dataset()[0]["tools"] + + def row(content: str) -> dict[str, object]: + return { + "messages": [ + {"role": "user", "content": "Inspect the repository."}, + { + "role": "assistant", + "content": content, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "bash", "arguments": {"command": "ls"}}, + } + ], + }, + ], + "tools": tools, + } + + prepared = prepare_data( + Dataset.from_list([row(""), row(answer)]), + tokenizer, + tokenize=True, + strict=True, + verbose=False, + ) + + assert prepared[0]["text"].endswith("<|tool_response>") + assert not any( + span.get("kind") == "final_answer" + for span in prepared[0]["teich_supervised_spans"] + ) + + final_spans = [ + span + for span in prepared[1]["teich_supervised_spans"] + if span.get("kind") == "final_answer" + ] + assert len(final_spans) == 1 + assert prepared[1]["text"][final_spans[0]["start"] : final_spans[0]["end"]] == answer + assert prepared[1]["text"][final_spans[0]["end"] :].startswith("<|tool_call>") + assert prepared[1]["text"].endswith("<|tool_response>") + + def test_marker_boundary_whitespace_reconciliation_handles_insertions_and_deletions(): spans = [{"start": 5, "end": 12, "kind": "final_answer", "role": "assistant"}] deleted = _reconcile_marker_boundary_whitespace("start\n\nanswer", spans, "startanswer") diff --git a/tests/test_runner.py b/tests/test_runner.py index 53e0a38..d3993d6 100644 --- a/tests/test_runner.py +++ b/tests/test_runner.py @@ -43,9 +43,20 @@ def _free_tcp_port() -> int: return int(sock.getsockname()[1]) -def _wait_for_tcp_port(port: int, timeout: float = 5.0) -> None: +def _wait_for_tcp_port( + port: int, + timeout: float = 5.0, + *, + process: subprocess.Popen[str] | None = None, +) -> None: deadline = time.time() + timeout while time.time() < deadline: + if process is not None and process.poll() is not None: + stdout, stderr = process.communicate() + raise RuntimeError( + f"process exited with code {process.returncode} before port {port} opened" + f"\nstdout:\n{stdout}\nstderr:\n{stderr}" + ) try: with socket.create_connection(("127.0.0.1", port), timeout=0.1): return @@ -994,8 +1005,6 @@ def test_claude_openrouter_proxy_strips_decoded_compression_headers(tmp_path: Pa if shutil.which("node") is None: pytest.skip("node is required for the proxy smoke test") - upstream_port = _free_tcp_port() - proxy_port = _free_tcp_port() response_payload = b'{"ok":true}\n' seen: dict[str, object] = {} @@ -1018,7 +1027,11 @@ def do_POST(self) -> None: def log_message(self, format: str, *args: object) -> None: return - server = ThreadingHTTPServer(("127.0.0.1", upstream_port), Handler) + # Bind the upstream before selecting the proxy port so the two ephemeral + # allocations cannot resolve to the same released port under CI load. + server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + upstream_port = int(server.server_address[1]) + proxy_port = _free_tcp_port() server_thread = threading.Thread(target=server.serve_forever, daemon=True) server_thread.start() proxy_script = tmp_path / "claude_openrouter_proxy.js" @@ -1036,7 +1049,7 @@ def log_message(self, format: str, *args: object) -> None: text=True, ) try: - _wait_for_tcp_port(proxy_port) + _wait_for_tcp_port(proxy_port, process=process) request = Request( f"http://127.0.0.1:{proxy_port}/v1/messages?beta=true", data=json.dumps( @@ -1106,8 +1119,6 @@ def test_claude_openrouter_proxy_handles_requests_concurrently(tmp_path: Path): if shutil.which("node") is None: pytest.skip("node is required for the proxy smoke test") - upstream_port = _free_tcp_port() - proxy_port = _free_tcp_port() request_count = 5 response_delay = 0.35 stats = {"inflight": 0, "max_inflight": 0, "requests": 0} @@ -1137,7 +1148,12 @@ def do_POST(self) -> None: def log_message(self, format: str, *args: object) -> None: return - server = ThreadingHTTPServer(("127.0.0.1", upstream_port), Handler) + # Keep the upstream port reserved while selecting the proxy port. Calling + # _free_tcp_port twice can return the same released port and make Node exit + # with EADDRINUSE before the readiness probe observes it. + server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + upstream_port = int(server.server_address[1]) + proxy_port = _free_tcp_port() server_thread = threading.Thread(target=server.serve_forever, daemon=True) server_thread.start() proxy_script = tmp_path / "claude_openrouter_proxy.js" @@ -1166,7 +1182,7 @@ def send_request(index: int) -> bytes: return response.read() try: - _wait_for_tcp_port(proxy_port) + _wait_for_tcp_port(proxy_port, process=process) started = time.perf_counter() with concurrent.futures.ThreadPoolExecutor(max_workers=request_count) as executor: bodies = list(executor.map(send_request, range(request_count))) diff --git a/tests/test_tokenizer_smoke.py b/tests/test_tokenizer_smoke.py index 25c3118..b6f7361 100644 --- a/tests/test_tokenizer_smoke.py +++ b/tests/test_tokenizer_smoke.py @@ -268,6 +268,7 @@ def test_real_gemma4_supervises_exactly_one_turn_end_for_enabled_targets(model_i assert ("Careful reasoning." in supervised_text) is train_on_reasoning assert ("Final answer." in supervised_text) is train_on_final_answers + tool_prepared = prepare_data( _tool_call_dataset(), tokenizer, @@ -304,6 +305,69 @@ def test_real_gemma4_supervises_exactly_one_turn_end_for_enabled_targets(model_i assert tool_supervised_text.count("") == 1 assert tool_supervised_text.endswith("") + empty_final_messages = list(_tool_call_dataset()[0]["messages"]) + empty_final_messages[-1] = {"role": "assistant", "content": ""} + empty_final_prepared = prepare_data( + Dataset.from_list( + [ + { + "messages": empty_final_messages, + "tools": _tool_call_dataset()[0]["tools"], + } + ] + ), + tokenizer, + tokenize=True, + strict=True, + max_length=4096, + verbose=False, + ) + assert not any( + span.get("kind") == "final_answer" + for span in empty_final_prepared[0]["teich_supervised_spans"] + ) + + +@pytest.mark.integration +@pytest.mark.tokenizer_smoke +@pytest.mark.parametrize( + "model_id", + [ + pytest.param("google/gemma-4-E4B-it", id="gemma-4-e4b-unresolved-tool"), + pytest.param("google/gemma-4-26B-A4B-it", id="gemma-4-26b-a4b-unresolved-tool"), + pytest.param("google/gemma-4-31B-it", id="gemma-4-31b-unresolved-tool"), + ], +) +def test_real_gemma4_unresolved_tool_call_has_no_synthetic_final_answer(model_id: str): + if not _tokenizer_smokes_enabled(): + pytest.skip("Set TEICH_RUN_TOKENIZER_SMOKES=1 to run real Hugging Face tokenizer smokes.") + transformers = pytest.importorskip("transformers") + tokenizer = transformers.AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) + answer = "A" * 92 + source = _tool_call_dataset()[0] + messages = [ + source["messages"][1], + {**source["messages"][2], "content": answer, "reasoning_content": ""}, + ] + + prepared = prepare_data( + Dataset.from_list([{"messages": messages, "tools": source["tools"]}]), + tokenizer, + tokenize=True, + strict=True, + verbose=False, + ) + text = prepared[0]["text"] + final_spans = [ + span + for span in prepared[0]["teich_supervised_spans"] + if span.get("kind") == "final_answer" + ] + + assert text.endswith("<|tool_response>") + assert len(final_spans) == 1 + assert text[final_spans[0]["start"] : final_spans[0]["end"]] == answer + @pytest.mark.integration @pytest.mark.tokenizer_smoke