modify process_data to generate separate user/system parts in prompts

Author: andytwigg

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

@@ -527,27 +527,28 @@ def _to_str(val):
     answer = extract_hash_answer(answer)
 
   return {
-      # passed to model forward pass
-      "prompts": model_tokenizer.apply_chat_template(
-          [
-              {
-                  "role": "user",
-                  "content": template_config["TEMPLATE"].format(
-                      system_prompt=template_config["SYSTEM_PROMPT"].format(
-                          reasoning_start_token=tmvp_config.reasoning_start_token,
-                          reasoning_end_token=tmvp_config.reasoning_end_token,
-                          solution_start_token=tmvp_config.solution_start_token,
-                          solution_end_token=tmvp_config.solution_end_token,
-                      ),
-                      question=question,
-                  ),
-              },
-          ],
-          tokenize=False,
-          add_generation_prompt=True,
-      ),
-      # passed to reward functions
-      "question": question,
-      # passed to reward functions
-      "answer": answer,
+    # passed to model forward pass
+    "prompts": model_tokenizer.apply_chat_template(
+      [
+        {
+          "role": "system",
+          "content": template_config["SYSTEM_PROMPT"].format(
+            reasoning_start_token=tmvp_config.reasoning_start_token,
+            reasoning_end_token=tmvp_config.reasoning_end_token,
+            solution_start_token=tmvp_config.solution_start_token,
+            solution_end_token=tmvp_config.solution_end_token,
+          ),
+        },
+        {
+          "role": "user",
+          "content": question,
+        }
+      ],
+      tokenize=False,
+      add_generation_prompt=True,
+    ),
+    # passed to reward functions
+    "question": question,
+    # passed to reward functions
+    "answer": answer,
   }