diff --git a/skyrl/train/dataset/dataset.py b/skyrl/train/dataset/dataset.py index 386b1eb5ca..20599391dc 100644 --- a/skyrl/train/dataset/dataset.py +++ b/skyrl/train/dataset/dataset.py @@ -1,14 +1,21 @@ import os -from typing import List +from typing import Any, List, Optional import datasets from loguru import logger from transformers import PreTrainedTokenizerBase +from skyrl.utils.chat_template import apply_chat_template -def _prompt_not_too_long(doc, tokenizer, prompt_key, max_length): - tokens = tokenizer.apply_chat_template( - doc[prompt_key], add_generation_prompt=True, return_dict=False, tokenize=True + +def _prompt_not_too_long(doc, tokenizer, prompt_key, max_length, chat_template_kwargs): + tokens = apply_chat_template( + tokenizer, + doc[prompt_key], + chat_template_kwargs=chat_template_kwargs, + add_generation_prompt=True, + return_dict=False, + tokenize=True, ) return len(tokens) <= max_length @@ -22,11 +29,13 @@ def __init__( num_workers: int = 8, prompt_key: str = "prompt", env_class_key: str = "env_class", + chat_template_kwargs: Optional[dict[str, Any]] = None, ): self.tokenizer = tokenizer self.max_prompt_length = max_prompt_length self.prompt_key = prompt_key self.env_class_key = env_class_key + self.chat_template_kwargs = chat_template_kwargs or {} self.num_workers = num_workers self.datasets = datasets @@ -76,6 +85,7 @@ def _read_files_and_tokenize(self): "tokenizer": self.tokenizer, "prompt_key": self.prompt_key, "max_length": self.max_prompt_length, + "chat_template_kwargs": self.chat_template_kwargs, }, num_proc=self.num_workers, desc=f"Filtering prompts longer than {self.max_prompt_length} tokens", diff --git a/skyrl/train/entrypoints/main_base.py b/skyrl/train/entrypoints/main_base.py index b50efdbe47..65895edbc0 100644 --- a/skyrl/train/entrypoints/main_base.py +++ b/skyrl/train/entrypoints/main_base.py @@ -80,6 +80,7 @@ def get_train_dataset(self) -> PromptDataset: tokenizer=self.tokenizer, max_prompt_length=self.cfg.trainer.max_prompt_length, num_workers=8, + chat_template_kwargs=self.cfg.generator.chat_template_kwargs, ) # make sure the dataset is large enough to train on assert ( @@ -99,6 +100,7 @@ def get_eval_dataset(self) -> Optional[PromptDataset]: tokenizer=self.tokenizer, max_prompt_length=self.cfg.trainer.max_prompt_length, num_workers=8, + chat_template_kwargs=self.cfg.generator.chat_template_kwargs, ) return prompts_dataset return None diff --git a/skyrl/train/generators/skyrl_gym_generator.py b/skyrl/train/generators/skyrl_gym_generator.py index 26b1102e6c..9272b74415 100644 --- a/skyrl/train/generators/skyrl_gym_generator.py +++ b/skyrl/train/generators/skyrl_gym_generator.py @@ -32,6 +32,7 @@ TrajectoryID, ) from skyrl.train.generators.utils import ( + apply_chat_template, apply_overlong_filtering, get_custom_chat_template, get_generation_prompt_ids, @@ -185,12 +186,13 @@ def __init__( {"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "I am a user."}, ] - self.base_conversation_token_ids = tokenizer.apply_chat_template( + self.base_conversation_token_ids = apply_chat_template( + tokenizer, self.base_conversation, + chat_template_kwargs=self.generator_cfg.chat_template_kwargs, add_generation_prompt=False, tokenize=True, return_dict=False, - **self.generator_cfg.chat_template_kwargs, ) # We remove tokens after the last EOS token so that it can be captured in `observation_ids`. # For details, see https://docs.skyrl.ai/docs/tutorials/skyrl_gym_generator#multi-turn-tokenization-and-ti-to @@ -343,15 +345,16 @@ async def agent_loop( # init() returns the first prompt to be given to the model, and optional metadata dict chat_history, _ = await self._run_in_executor_if_available(env.init, chat_history) initial_chat_history_length = len(chat_history) - initial_input_ids = self.tokenizer.apply_chat_template( + initial_input_ids = apply_chat_template( + self.tokenizer, chat_history, + chat_template_kwargs=self.generator_cfg.chat_template_kwargs, # If retokenize_chat_history==True, avoid including the generation prompt in both the # prompt_ids and response_ids due to how `response_encodings["input_ids"]` works. add_generation_prompt=not retokenize_chat_history, chat_template=self.custom_chat_template if retokenize_chat_history else None, tokenize=True, return_dict=False, - **self.generator_cfg.chat_template_kwargs, ) initial_prompt_length = len(initial_input_ids) @@ -390,13 +393,14 @@ async def agent_loop( # 1. Generate output if is_step_wise or retokenize_chat_history: # re-apply whole chat template so length check is correct - agent_loop_state.input_ids = self.tokenizer.apply_chat_template( + agent_loop_state.input_ids = apply_chat_template( + self.tokenizer, chat_history, + chat_template_kwargs=self.generator_cfg.chat_template_kwargs, chat_template=self.custom_chat_template if retokenize_chat_history else None, add_generation_prompt=True, tokenize=True, return_dict=False, - **self.generator_cfg.chat_template_kwargs, ) agent_loop_state.loss_mask = [] agent_loop_state.rollout_logprobs = None @@ -530,16 +534,17 @@ async def agent_loop( # Note that during the agent loop, we still add the final observation messages/ tokens because we terminate the agent loop if the input length # exceeds the maximum if retokenize_chat_history: - response_encodings = self.tokenizer.apply_chat_template( + response_encodings = apply_chat_template( + self.tokenizer, agent_loop_state.chat_history[ initial_chat_history_length : len(agent_loop_state.chat_history) - len(new_obs) ], + chat_template_kwargs=self.generator_cfg.chat_template_kwargs, chat_template=self.custom_chat_template, add_generation_prompt=False, return_dict=True, return_assistant_tokens_mask=True, tokenize=True, - **self.generator_cfg.chat_template_kwargs, ) loss_mask = response_encodings["assistant_masks"] response_ids = response_encodings["input_ids"] @@ -661,12 +666,13 @@ def get_obs_ids_from_obs(self, new_obs: ConversationType, is_done: bool) -> List if len(new_obs) > 0: # For Qwen, this will generate `\n<|user|>Some observation<|im_end|>\n`. Note that the # first `\n` is generated since we stripped it in ``base_conversation_token_ids``. - obs_ids_to_add = self.tokenizer.apply_chat_template( + obs_ids_to_add = apply_chat_template( + self.tokenizer, [*self.base_conversation, *new_obs], + chat_template_kwargs=self.generator_cfg.chat_template_kwargs, add_generation_prompt=not is_done, tokenize=True, return_dict=False, - **self.generator_cfg.chat_template_kwargs, )[len(self.base_conversation_token_ids) :] elif not is_done: obs_ids_to_add = self.generation_prompt_ids diff --git a/skyrl/train/generators/utils.py b/skyrl/train/generators/utils.py index c84ea8a5aa..d5a3eb4fd1 100644 --- a/skyrl/train/generators/utils.py +++ b/skyrl/train/generators/utils.py @@ -17,6 +17,7 @@ TrainingPhase, TrajectoryID, ) +from skyrl.utils.chat_template import apply_chat_template from skyrl_gym.metrics import aggregate_for_environment @@ -165,15 +166,21 @@ def get_generation_prompt_ids(tokenizer, tokenizer_kwargs: Optional[dict] = None tokenizer_kwargs or {} ), "enable_thinking is not supported in get_generation_prompt_ids; use encode_messages_subset instead" kwargs = tokenizer_kwargs.copy() if tokenizer_kwargs else {} - empty_user = tokenizer.apply_chat_template( - [{"role": "user", "content": ""}], tokenize=True, return_dict=False, **kwargs + empty_user = apply_chat_template( + tokenizer, + [{"role": "user", "content": ""}], + chat_template_kwargs=kwargs, + add_generation_prompt=False, + tokenize=True, + return_dict=False, ) - empty_user_with_generation_prompt = tokenizer.apply_chat_template( + empty_user_with_generation_prompt = apply_chat_template( + tokenizer, [{"role": "user", "content": ""}], + chat_template_kwargs=kwargs, add_generation_prompt=True, tokenize=True, return_dict=False, - **kwargs, ) generation_prompt_ids = empty_user_with_generation_prompt[len(empty_user) :] @@ -605,21 +612,23 @@ def encode_messages_subset(messages: ConversationType, tokenizer, tokenizer_kwar {"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "I am a user."}, ] - base_conversation_token_ids = tokenizer.apply_chat_template( + base_conversation_token_ids = apply_chat_template( + tokenizer, base_conversation, + chat_template_kwargs=kwargs, add_generation_prompt=False, tokenize=True, return_dict=False, - **kwargs, ) full_conversation = base_conversation + messages - full_conversation_token_ids = tokenizer.apply_chat_template( + full_conversation_token_ids = apply_chat_template( + tokenizer, full_conversation, + chat_template_kwargs=kwargs, add_generation_prompt=False, tokenize=True, return_dict=False, - **kwargs, ) conversation_token_ids = full_conversation_token_ids[len(base_conversation_token_ids) :] return conversation_token_ids diff --git a/skyrl/utils/chat_template.py b/skyrl/utils/chat_template.py new file mode 100644 index 0000000000..3c04d27bf5 --- /dev/null +++ b/skyrl/utils/chat_template.py @@ -0,0 +1,5 @@ +def apply_chat_template(tokenizer, messages, *, chat_template_kwargs=None, **kwargs): + """Apply a tokenizer chat template with per-call controls taking precedence.""" + apply_kwargs = dict(chat_template_kwargs or {}) + apply_kwargs.update(kwargs) + return tokenizer.apply_chat_template(messages, **apply_kwargs) diff --git a/tests/train/dataset/test_dataset.py b/tests/train/dataset/test_dataset.py index 55e3fa0b72..c9a5e6c1a4 100644 --- a/tests/train/dataset/test_dataset.py +++ b/tests/train/dataset/test_dataset.py @@ -53,6 +53,33 @@ def test_prompt_dataset_filtering(mock_load_dataset, mock_tokenizer, sample_data assert extra == {"answer": "a1"} +@patch("datasets.load_dataset") +def test_prompt_filter_uses_generation_chat_template_kwargs(mock_load_dataset): + class ThinkingTokenizer: + def apply_chat_template(self, _prompt, *, enable_thinking=False, **kwargs): + assert kwargs["add_generation_prompt"] is True + assert kwargs["return_dict"] is False + assert kwargs["tokenize"] is True + return [1, 2, 3, 4] if enable_thinking else [1] + + mock_load_dataset.return_value = {"train": Dataset.from_dict({"prompt": [[{"role": "user", "content": "hi"}]]})} + + dataset = PromptDataset( + datasets=["dummy.parquet"], + tokenizer=ThinkingTokenizer(), + max_prompt_length=2, + num_workers=1, + chat_template_kwargs={ + "enable_thinking": True, + "add_generation_prompt": False, + "return_dict": True, + "tokenize": False, + }, + ) + + assert len(dataset) == 0 + + def test_collate_fn(): dataset = PromptDataset.__new__(PromptDataset) # Bypass __init__ sample_data = [("prompt 1", "env", {"answer": "a1"}, "1"), ("prompt 2", "env", {"answer": "a2"}, "2")] diff --git a/tests/train/generators/test_skyrl_gym_generator.py b/tests/train/generators/test_skyrl_gym_generator.py index 6d03d937b6..91ea5af8da 100644 --- a/tests/train/generators/test_skyrl_gym_generator.py +++ b/tests/train/generators/test_skyrl_gym_generator.py @@ -142,6 +142,32 @@ def mock_env_cfg(): return cfg +def test_generator_chat_template_controls_override_kwargs(mock_tokenizer, mock_llm, mock_env_cfg, generator_cfg): + generator_cfg.batched = False + generator_cfg.use_conversation_multi_turn = False + generator_cfg.chat_template_kwargs = { + "add_generation_prompt": True, + "return_dict": True, + "tokenize": False, + "custom_option": "preserved", + } + + SkyRLGymGenerator( + generator_cfg=generator_cfg, + skyrl_gym_cfg=mock_env_cfg, + inference_engine_client=mock_llm, + tokenizer=mock_tokenizer, + ) + + kwargs = mock_tokenizer.apply_chat_template.call_args.kwargs + assert kwargs == { + "add_generation_prompt": False, + "return_dict": False, + "tokenize": True, + "custom_option": "preserved", + } + + def validate_generator_input(input_batch: GeneratorInput) -> bool: """Validate that input_batch conforms to GeneratorInput TypedDict interface.""" # Check that input_batch has the required keys diff --git a/tests/train/generators/test_utils.py b/tests/train/generators/test_utils.py index 56ccab6c99..bcef8b9aa2 100644 --- a/tests/train/generators/test_utils.py +++ b/tests/train/generators/test_utils.py @@ -32,6 +32,39 @@ ) +def test_get_generation_prompt_ids_explicit_controls_override_kwargs(): + calls = [] + + class Tokenizer: + def apply_chat_template(self, _messages, **kwargs): + calls.append(kwargs) + return [1, 2] if not kwargs["add_generation_prompt"] else [1, 2, 3] + + tokenizer_kwargs = { + "add_generation_prompt": False, + "return_dict": True, + "tokenize": False, + "custom_option": "preserved", + } + generation_prompt_ids = get_generation_prompt_ids(Tokenizer(), tokenizer_kwargs=tokenizer_kwargs) + + assert generation_prompt_ids == [3] + assert calls == [ + { + "add_generation_prompt": False, + "return_dict": False, + "tokenize": True, + "custom_option": "preserved", + }, + { + "add_generation_prompt": True, + "return_dict": False, + "tokenize": True, + "custom_option": "preserved", + }, + ] + + @pytest.fixture def qwen3_acc_thinking_template(): """Load the qwen3_acc_thinking.jinja2 template."""