Adding data_template_path, and using chat_template_path for tokenizer template

Author: dasoto

src/maxtext/configs/post_train/rl.yml

@@ -208,7 +208,7 @@ reasoning_start_token: '<reasoning>'
 reasoning_end_token: '</reasoning>'
 solution_start_token: '<answer>'
 solution_end_token: '</answer>'
-chat_template_path: 'maxtext/examples/chat_templates/gsm8k_rl.json'
+data_template_path: 'maxtext/examples/chat_templates/gsm8k_rl.json'
 skip_jax_distributed_system: true
 
 # ====== Dataset Configuration ======

src/maxtext/configs/types.py

@@ -1081,9 +1081,9 @@ class Tokenizer(BaseModel):
       "",
       description="Chat template to use with HF tokenizers. It should be a valid Jinja2-formatted template.",
   )
-  tokenizer_chat_template_path: str = Field(
+  data_template_path: str = Field(
       "",
-      description="Path to a chat template file to be loaded into the tokenizer if missing.",
+      description="Path to a chat template file to be used when tokenizing the dataset. Used in RL workloads to provide the conversation.",
   )
   tokenize_train_data: bool = Field(True, description="If False, assumes the training dataset is pre-tokenized.")
   tokenize_eval_data: bool = Field(True, description="If False, assumes the evaluation dataset is pre-tokenized.")

src/maxtext/trainers/post_train/rl/train_rl.py

@@ -248,10 +248,10 @@ def prepare_datasets(
     model_tokenizer: AutoTokenizer,
 ) -> tuple[grain.IterDataset, grain.IterDataset | None]:
   """Setup and return train and test datasets."""
-  template_config = load_data_template_from_file(trainer_config.chat_template_path)
+  template_config = load_data_template_from_file(trainer_config.data_template_path)
   if template_config is None:
     raise ValueError(
-        f"Chat template is required for processing dataset but failed to load from {trainer_config.chat_template_path}"
+        f"Chat template is required for processing dataset but failed to load from {trainer_config.data_template_path}"
     )
 
   # Prepare train and test data from training data for certain datasets
@@ -548,10 +548,10 @@ def _reward_fn(**kwargs):
         epsilon_high=trainer_config.rl.epsilon_high,
     )
     # Instantiate the custom MaxText chat parser
-    template_config = load_data_template_from_file(trainer_config.chat_template_path)
+    template_config = load_data_template_from_file(trainer_config.data_template_path)
     if template_config is None:
       raise ValueError(
-          f"Chat template is required for AgenticGRPOLearner but failed to load from {trainer_config.chat_template_path}"
+          f"Chat template is required for AgenticGRPOLearner but failed to load from {trainer_config.data_template_path}"
       )
     chat_parser = utils_rl.MaxTextChatParser(
         model_tokenizer=model_tokenizer, template_config=template_config, tmvp_config=trainer_config
@@ -588,20 +588,20 @@ def configure_tokenizer_chat_template(model_tokenizer: Any, trainer_config: Any)
   if getattr(model_tokenizer, "chat_template", None) is None:
     if getattr(trainer_config, "chat_template", None):
       model_tokenizer.chat_template = trainer_config.chat_template
-    elif getattr(trainer_config, "tokenizer_chat_template_path", None):
+    elif getattr(trainer_config, "chat_template_path", None):
       from maxtext.input_pipeline.instruction_data_processing import (
           load_chat_template_from_file,
       )
       model_tokenizer.chat_template = load_chat_template_from_file(
-          trainer_config.tokenizer_chat_template_path
+          trainer_config.chat_template_path
       )
     else:
       raise ValueError(
           f"Tokenizer {getattr(trainer_config, 'tokenizer_path', None)!r} has no chat_template "
-          "and config.chat_template / config.tokenizer_chat_template_path "
+          "and config.chat_template / config.chat_template_path "
           "are both empty. Either pick an instruction-tuned tokenizer that "
           "ships with a chat_template, set config.chat_template to a Jinja "
-          "string, or set config.tokenizer_chat_template_path to a JSON file "
+          "string, or set config.chat_template_path to a JSON file "
           "with a 'chat_template' key."
       )
 

tests/post_training/unit/train_rl_test.py

@@ -334,7 +334,7 @@ def get_filtered_data_side_effect(dataset_name, model_tokenizer, template_config
         eval_split="eval",
         hf_train_files=None,
         hf_eval_files=None,
-        chat_template_path="maxtext/examples/chat_templates/gsm8k_rl.json",
+        data_template_path="maxtext/examples/chat_templates/gsm8k_rl.json",
         data_shuffle_seed=42,
         max_prefill_predict_length=10,
         batch_size=2,
@@ -389,7 +389,7 @@ def test_prepare_datasets_with_split(self, mock_load):
         eval_dataset_name="open-r1/OpenR1-Math-220k",
         train_split="train",
         hf_train_files="hf://open-r1/OpenR1-Math-220k/data/dummy.parquet",
-        chat_template_path="maxtext/examples/chat_templates/gsm8k_rl.json",
+        data_template_path="maxtext/examples/chat_templates/gsm8k_rl.json",
         data_shuffle_seed=42,
         num_batches=1,
         batch_size=5,
@@ -435,7 +435,7 @@ def test_prepare_datasets_without_split(self, mock_load):
         eval_split="test",
         hf_train_files="hf://openai/gsm8k/data/dummy.parquet",
         hf_eval_files="hf://openai/gsm8k/data/dummy.parquet",
-        chat_template_path="maxtext/examples/chat_templates/gsm8k_rl.json",
+        data_template_path="maxtext/examples/chat_templates/gsm8k_rl.json",
         data_shuffle_seed=42,
         num_batches=1,
         batch_size=5,
@@ -482,7 +482,7 @@ def test_chat_template_populated_from_config_string(self):
     mock_tokenizer.chat_template = None
     trainer_config = SimpleNamespace(
         chat_template="{{ messages[0].content }}",
-        tokenizer_chat_template_path=None,
+        chat_template_path=None,
         tokenizer_path="dummy-base-model",
     )
     train_rl.configure_tokenizer_chat_template(mock_tokenizer, trainer_config)
@@ -491,13 +491,13 @@ def test_chat_template_populated_from_config_string(self):
   @pytest.mark.cpu_only
   @mock.patch("maxtext.input_pipeline.instruction_data_processing.load_chat_template_from_file")
   def test_chat_template_populated_from_config_file(self, mock_load):
-    """Test that chat_template is loaded from tokenizer_chat_template_path when tokenizer lacks one."""
+    """Test that chat_template is loaded from chat_template_path when tokenizer lacks one."""
     mock_tokenizer = mock.MagicMock()
     mock_tokenizer.chat_template = None
     mock_load.return_value = "{% for message in messages %}{{ message.content }}{% endfor %}"
     trainer_config = SimpleNamespace(
         chat_template=None,
-        tokenizer_chat_template_path="/path/to/jinja_template.json",
+        chat_template_path="/path/to/jinja_template.json",
         tokenizer_path="dummy-base-model",
     )
     train_rl.configure_tokenizer_chat_template(mock_tokenizer, trainer_config)
@@ -513,7 +513,7 @@ def test_chat_template_raises_value_error_when_empty(self):
     mock_tokenizer.chat_template = None
     trainer_config = SimpleNamespace(
         chat_template=None,
-        tokenizer_chat_template_path=None,
+        chat_template_path=None,
         tokenizer_path="dummy-base-model",
     )
     with self.assertRaisesRegex(ValueError, "Tokenizer 'dummy-base-model' has no chat_template"):
@@ -526,7 +526,7 @@ def test_chat_template_unchanged_when_already_exists(self):
     mock_tokenizer.chat_template = "{{ existing_template }}"
     trainer_config = SimpleNamespace(
         chat_template="{{ overridden_template }}",
-        tokenizer_chat_template_path=None,
+        chat_template_path=None,
         tokenizer_path="dummy-instruction-tuned-model",
     )
     train_rl.configure_tokenizer_chat_template(mock_tokenizer, trainer_config)
@@ -549,7 +549,7 @@ def apply_chat_template(self, conversation, tokenize=False):
     tokenizer = DummyTokenizer()
     trainer_config = SimpleNamespace(
         chat_template="{{ messages[0].content }}",
-        tokenizer_chat_template_path=None,
+        chat_template_path=None,
         tokenizer_path="dummy-base-model",
     )
     # Initially, apply_chat_template fails (simulating HF tokenizer crash when chat_template is None)