diff --git a/examples/tau-bench/token_delta.py b/examples/tau-bench/token_delta.py new file mode 100644 index 0000000000..b78023d8ce --- /dev/null +++ b/examples/tau-bench/token_delta.py @@ -0,0 +1,58 @@ +from typing import Any + + +def get_token_delta( + tokenizer: Any, + messages: list[dict[str, Any]], + *, + include_generation_prompt: bool = False, +) -> tuple[list[int], list[int]]: + """Return the tokens and loss mask contributed by the last chat message.""" + if not messages: + raise ValueError("Cannot calculate a token delta for an empty conversation") + + is_assistant = messages[-1]["role"] == "assistant" + curr = tokenizer.apply_chat_template(messages, add_generation_prompt=False, tokenize=False) + + if is_assistant: + prev = tokenizer.apply_chat_template(messages[:-1], add_generation_prompt=False, tokenize=False) + generation_prompt = tokenizer.apply_chat_template(messages[:-1], add_generation_prompt=True, tokenize=False) + if not generation_prompt.startswith(prev): + raise ValueError("Adding the assistant generation prompt rewrote the rendered conversation") + if not curr.startswith(generation_prompt): + raise ValueError("The assistant response does not extend its generation prompt") + + generation_prompt_text = generation_prompt[len(prev) :] + if not include_generation_prompt: + new_text = curr[len(generation_prompt) :] + new_tokens = tokenizer.encode(new_text, add_special_tokens=False) + return new_tokens, [1] * len(new_tokens) + + new_text = curr[len(prev) :] + new_tokens = tokenizer.encode(new_text, add_special_tokens=False) + generation_prompt_length = len(tokenizer.encode(generation_prompt_text, add_special_tokens=False)) + masked_prefix_length = min(generation_prompt_length, len(new_tokens)) + loss_mask = [0] * masked_prefix_length + loss_mask.extend([1] * (len(new_tokens) - masked_prefix_length)) + return new_tokens, loss_mask + + prev = tokenizer.apply_chat_template(messages[:-1], add_generation_prompt=False, tokenize=False) + + if curr.startswith(prev): + new_text = curr[len(prev) :] + elif messages[-1]["role"] == "user": + # Reasoning templates such as Qwen3 can rewrite history when a new user + # message arrives. Render that message independently instead of slicing + # the rewritten conversation at the old conversation length. + new_text = tokenizer.apply_chat_template( + [messages[-1]], + add_generation_prompt=False, + tokenize=False, + ) + if not curr.endswith(new_text): + raise ValueError("The latest user message is not a standalone suffix of the rendered conversation") + else: + raise ValueError("The chat template rewrote history while calculating a non-user token delta") + + new_tokens = tokenizer.encode(new_text, add_special_tokens=False) + return new_tokens, [0] * len(new_tokens) diff --git a/examples/tau-bench/trainable_agents.py b/examples/tau-bench/trainable_agents.py index 5cfc3f43d0..effc5d7a66 100644 --- a/examples/tau-bench/trainable_agents.py +++ b/examples/tau-bench/trainable_agents.py @@ -8,6 +8,7 @@ from tau_bench.agents.base import Agent from tau_bench.agents.tool_calling_agent import RESPOND_ACTION_NAME, ToolCallingAgent from tau_bench.types import Action, RunConfig +from token_delta import get_token_delta from transformers import AutoTokenizer from slime.rollout.sglang_rollout import GenerateState @@ -274,7 +275,11 @@ async def asolve( # Add assistant response to conversation messages.append({"role": "assistant", "content": response}) - assistant_token_ids, assistant_loss_mask = self._get_token_delta(state.tokenizer, messages) + assistant_token_ids, assistant_loss_mask = self._get_token_delta( + state.tokenizer, + messages, + include_generation_prompt=bool(response_token_ids), + ) response_token_ids.extend(assistant_token_ids) loss_masks.extend(assistant_loss_mask) @@ -331,7 +336,13 @@ async def asolve( res, total_reward, info, messages, loss_masks, prompt_token_ids, response_token_ids ) - def _get_token_delta(self, tokenizer: AutoTokenizer, messages: list[dict]) -> tuple[list[int], list[int]]: + def _get_token_delta( + self, + tokenizer: AutoTokenizer, + messages: list[dict], + *, + include_generation_prompt: bool = False, + ) -> tuple[list[int], list[int]]: """ Calculate token delta for multi-turn conversations. @@ -347,24 +358,11 @@ def _get_token_delta(self, tokenizer: AutoTokenizer, messages: list[dict]) -> tu Returns: Tuple of (token_ids, loss_mask) """ - curr = tokenizer.apply_chat_template(messages, add_generation_prompt=False, tokenize=False) - token_ids = [] - loss_mask = [] - - # Case 1: last message is an assistant response - if messages[-1]["role"] == "assistant": - prev = tokenizer.apply_chat_template(messages[:-1], add_generation_prompt=True, tokenize=False) - new_tokens = tokenizer.encode(curr[len(prev) :], add_special_tokens=False) - token_ids += new_tokens - loss_mask += [1] * len(new_tokens) # Mask only the new assistant tokens - else: - # Case 2: last message is a tool response or environment observation - prev = tokenizer.apply_chat_template(messages[:-1], add_generation_prompt=False, tokenize=False) - new_tokens = tokenizer.encode(curr[len(prev) :], add_special_tokens=False) - token_ids += new_tokens - loss_mask += [0] * len(new_tokens) # Don't mask environment/tool tokens - - return token_ids, loss_mask + return get_token_delta( + tokenizer, + messages, + include_generation_prompt=include_generation_prompt, + ) def _build_final_result( self, diff --git a/tests/test_tau_bench_token_delta.py b/tests/test_tau_bench_token_delta.py new file mode 100644 index 0000000000..2b88c41120 --- /dev/null +++ b/tests/test_tau_bench_token_delta.py @@ -0,0 +1,207 @@ +import importlib.util +import re +from pathlib import Path + +import pytest + + +TOKEN_DELTA_PATH = Path(__file__).parents[1] / "examples" / "tau-bench" / "token_delta.py" + + +def _load_get_token_delta(): + if not TOKEN_DELTA_PATH.exists(): + pytest.fail("tau-bench token_delta helper does not exist") + + spec = importlib.util.spec_from_file_location("tau_bench_token_delta", TOKEN_DELTA_PATH) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module.get_token_delta + + +class HistoryRewritingTokenizer: + """Minimal chat template that hides old reasoning after a new user turn.""" + + @staticmethod + def _strip_reasoning(content: str) -> str: + return re.sub(r".*?", "", content, flags=re.DOTALL) + + def apply_chat_template(self, messages, *, add_generation_prompt, tokenize): + assert tokenize is False + last_user = max((i for i, message in enumerate(messages) if message["role"] == "user"), default=-1) + rendered = [] + for i, message in enumerate(messages): + content = message["content"] + if message["role"] == "assistant" and i < last_user: + content = self._strip_reasoning(content) + rendered.append(f'<{message["role"]}>{content}') + if add_generation_prompt: + rendered.append("") + return "".join(rendered) + + @staticmethod + def encode(text, *, add_special_tokens): + assert add_special_tokens is False + return list(text.encode()) + + @staticmethod + def decode(token_ids): + return bytes(token_ids).decode() + + +class BoundaryMergingTokenizer(HistoryRewritingTokenizer): + """Tokenizer where the generation-prefix tail merges with a leading newline.""" + + @staticmethod + def encode(text, *, add_special_tokens): + assert add_special_tokens is False + raw = text.encode() + token_ids = [] + index = 0 + while index < len(raw): + if raw[index : index + 2] == b">\n": + token_ids.append(1000) + index += 2 + else: + token_ids.append(raw[index]) + index += 1 + return token_ids + + +@pytest.mark.unit +def test_new_user_delta_survives_history_rewrite(): + get_token_delta = _load_get_token_delta() + tokenizer = HistoryRewritingTokenizer() + messages = [ + {"role": "user", "content": "first user"}, + {"role": "assistant", "content": "first reasoningfirst answer"}, + {"role": "user", "content": "second user must remain complete"}, + ] + + token_ids, loss_mask = get_token_delta(tokenizer, messages) + + assert tokenizer.decode(token_ids) == "second user must remain complete" + assert loss_mask == [0] * len(token_ids) + + +@pytest.mark.unit +def test_assistant_delta_keeps_existing_append_only_behavior(): + get_token_delta = _load_get_token_delta() + tokenizer = HistoryRewritingTokenizer() + messages = [ + {"role": "user", "content": "user question"}, + {"role": "assistant", "content": "reasoninganswer"}, + ] + + token_ids, loss_mask = get_token_delta(tokenizer, messages) + + assert tokenizer.decode(token_ids) == "reasoninganswer" + assert loss_mask == [1] * len(token_ids) + + +@pytest.mark.unit +def test_accumulated_multiturn_tokens_keep_every_assistant_generation_prefix(): + get_token_delta = _load_get_token_delta() + tokenizer = HistoryRewritingTokenizer() + messages = [ + {"role": "user", "content": "first user"}, + {"role": "assistant", "content": "first reasoningfirst answer"}, + {"role": "user", "content": "second user"}, + {"role": "assistant", "content": "second reasoningsecond answer"}, + ] + + initial_prompt = tokenizer.apply_chat_template(messages[:1], add_generation_prompt=True, tokenize=False) + token_ids = tokenizer.encode(initial_prompt, add_special_tokens=False) + loss_mask = [0] * len(token_ids) + + for end in range(2, len(messages) + 1): + include_generation_prompt = messages[end - 1]["role"] == "assistant" and end > 2 + delta_ids, delta_mask = get_token_delta( + tokenizer, + messages[:end], + include_generation_prompt=include_generation_prompt, + ) + token_ids.extend(delta_ids) + loss_mask.extend(delta_mask) + + decoded = tokenizer.decode(token_ids) + assert decoded.count("") == 2 + assert decoded.count("") == 2 + assert "first reasoning" in decoded + assert "second reasoning" in decoded + assert "second user" in decoded + + second_prefix = decoded.rindex("") + assert loss_mask[second_prefix : second_prefix + len("")] == [0] * len("") + + +@pytest.mark.unit +def test_accumulated_six_real_user_turns_keep_every_user_and_assistant_boundary(): + get_token_delta = _load_get_token_delta() + tokenizer = HistoryRewritingTokenizer() + messages = [{"role": "user", "content": "user 1"}] + + initial_prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False) + token_ids = tokenizer.encode(initial_prompt, add_special_tokens=False) + loss_mask = [0] * len(token_ids) + + for turn in range(1, 7): + messages.append( + { + "role": "assistant", + "content": f"reasoning {turn}answer {turn}", + } + ) + delta_ids, delta_mask = get_token_delta( + tokenizer, + messages, + include_generation_prompt=turn > 1, + ) + token_ids.extend(delta_ids) + loss_mask.extend(delta_mask) + + if turn < 6: + messages.append({"role": "user", "content": f"user {turn + 1}"}) + delta_ids, delta_mask = get_token_delta(tokenizer, messages) + token_ids.extend(delta_ids) + loss_mask.extend(delta_mask) + + decoded = tokenizer.decode(token_ids) + assert decoded.count("") == 6 + assert decoded.count("") == 6 + + for turn in range(1, 7): + user_span = f"user {turn}" + user_start = decoded.index(user_span) + assert loss_mask[user_start : user_start + len(user_span)] == [0] * len(user_span) + + assistant_span = f"reasoning {turn}answer {turn}" + assistant_start = decoded.index(assistant_span) + assert loss_mask[assistant_start : assistant_start + len(assistant_span)] == [1] * len(assistant_span) + + prefix_start = decoded.rfind("", 0, assistant_start) + assert loss_mask[prefix_start:assistant_start] == [0] * (assistant_start - prefix_start) + + +@pytest.mark.unit +def test_later_assistant_allows_bpe_merge_across_generation_prefix_boundary(): + get_token_delta = _load_get_token_delta() + tokenizer = BoundaryMergingTokenizer() + messages = [ + {"role": "user", "content": "first user"}, + {"role": "assistant", "content": "first reasoningfirst answer"}, + {"role": "user", "content": "second user"}, + {"role": "assistant", "content": "\nsecond answer"}, + ] + + token_ids, loss_mask = get_token_delta( + tokenizer, + messages, + include_generation_prompt=True, + ) + + expected_text = "\nsecond answer" + expected_ids = tokenizer.encode(expected_text, add_special_tokens=False) + generation_prefix_length = len(tokenizer.encode("", add_special_tokens=False)) + assert token_ids == expected_ids + assert loss_mask == [0] * generation_prefix_length + [1] * (len(expected_ids) - generation_prefix_length)