Fix chat template pathing in RL config and template loading

PiperOrigin-RevId: 886940444

Author: SurbhiJainUSC

src/maxtext/configs/post_train/rl.yml

@@ -177,7 +177,7 @@ reasoning_start_token: '<reasoning>'
 reasoning_end_token: '</reasoning>'
 solution_start_token: '<answer>'
 solution_end_token: '</answer>'
-chat_template_path: 'src/maxtext/examples/chat_templates/gsm8k_rl.json'
+chat_template_path: 'maxtext/examples/chat_templates/gsm8k_rl.json'
 skip_jax_distributed_system: True
 
 # # TODO(@mazumdera): fix this

src/maxtext/input_pipeline/instruction_data_processing.py

@@ -24,6 +24,9 @@
 def load_template_from_file(template_path):
   """Loads a template from a file."""
   template_config = None
+  current_dir = os.path.dirname(os.path.abspath(__file__))
+  repo_root = os.path.abspath(os.path.join(current_dir, "..", ".."))
+  template_path = os.path.join(repo_root, template_path)
   if os.path.isfile(template_path) and template_path.endswith(".json"):
     with open(template_path, encoding="utf-8") as f:
       template_config = json.load(f)

tests/unit/instruction_data_processing_test.py

@@ -21,6 +21,22 @@
 
 class InstructionDataProcessingTest(unittest.TestCase):
 
+  def test_load_template_from_file(self):
+    template_config = instruction_data_processing.load_template_from_file("maxtext/examples/chat_templates/gsm8k_rl.json")
+    self.assertEqual(
+        template_config,
+        {
+            "SYSTEM_PROMPT": (
+                "You are given a problem. Think about the problem and provide"
+                " your reasoning. Place it between {reasoning_start_token} and"
+                " {reasoning_end_token}. Then, provide the final answer (i.e.,"
+                " just one numerical value) between {solution_start_token} and"
+                " {solution_end_token}."
+            ),
+            "TEMPLATE": ("<start_of_turn>user\n{system_prompt}\n\n{question}<end_of_turn>\n<start_of_turn>model"),
+        },
+    )
+
   def test_map_qa_data_to_conversation_with_prompt_completion_template(self):
     template_config = {
         "PROMPT_TEMPLATE": "This is a question: {question}",