[fix][data] Apply chat-template kwargs during filtering

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",

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

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

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

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)

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")]

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

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."""