Adds ability to pass chat_template_path argument into MaxText SFT, loading a separate chat_template from the tokenizer that's provided.

PiperOrigin-RevId: 900245511

Author: Google-ML-Automation

src/maxtext/input_pipeline/hf_data_processing.py

@@ -249,6 +249,10 @@ def preprocessing_pipeline(
   dataset = dataset.select_columns(data_column_names)
 
   if use_sft:
+    if not chat_template and chat_template_path:
+      chat_template = instruction_data_processing.get_chat_template_from_path(
+          chat_template_path
+      )
     data_processing_utils.validate_and_configure_sft_columns(data_column_names, tokenizer, chat_template)
 
     # convert instruction dataset to conversational format

src/maxtext/input_pipeline/instruction_data_processing.py

@@ -33,6 +33,37 @@ def load_template_from_file(template_path):
   return template_config
 
 
+def get_chat_template_from_path(chat_template_path):
+  """Loads a chat template from a file."""
+  if not chat_template_path:
+    return None
+
+  current_dir = os.path.dirname(os.path.abspath(__file__))
+  repo_root = os.path.abspath(os.path.join(current_dir, "..", ".."))
+  full_path = os.path.join(repo_root, chat_template_path)
+
+  if not os.path.isfile(full_path):
+    return None
+
+  if full_path.endswith((".jinja", ".j2", ".txt")):
+    with open(full_path, "r", encoding="utf-8") as f:
+      return f.read()
+
+  if full_path.endswith(".json"):
+    with open(full_path, "r", encoding="utf-8") as f:
+      try:
+        template_config = json.load(f)
+        if (
+            isinstance(template_config, dict)
+            and "chat_template" in template_config
+        ):
+          return template_config["chat_template"]
+      except json.JSONDecodeError:
+        return None
+
+  return None
+
+
 def get_template_placeholders(template):
   """Dynamically extracts the format keys (placeholders) from a template string."""
   # Finds all names inside {...}

tests/unit/instruction_data_processing_test.py

@@ -14,6 +14,9 @@
 
 """Instruction data processing test."""
 
+import json
+import os
+import tempfile
 import unittest
 
 from maxtext.input_pipeline import instruction_data_processing
@@ -139,6 +142,43 @@ def test_map_qa_data_to_conversation_with_no_templates(self):
       if data["role"] == "assistant":
         self.assertEqual(data["content"], "The capital of Germany is Berlin.")
 
+  def test_get_chat_template_from_path(self):
+    with tempfile.TemporaryDirectory() as tmpdir:
+      # Test .jinja file
+      jinja_path = os.path.join(tmpdir, "test.jinja")
+      with open(jinja_path, "w", encoding="utf-8") as f:
+        f.write("test jinja template")
+      self.assertEqual(
+          instruction_data_processing.get_chat_template_from_path(jinja_path),
+          "test jinja template",
+      )
+
+      # Test .json file with chat_template
+      json_path = os.path.join(tmpdir, "test.json")
+      with open(json_path, "w", encoding="utf-8") as f:
+        json.dump({"chat_template": "test json template"}, f)
+      self.assertEqual(
+          instruction_data_processing.get_chat_template_from_path(json_path),
+          "test json template",
+      )
+
+      # Test .json file without chat_template
+      json_no_key_path = os.path.join(tmpdir, "no_key.json")
+      with open(json_no_key_path, "w", encoding="utf-8") as f:
+        json.dump({"other_key": "other_value"}, f)
+      self.assertIsNone(
+          instruction_data_processing.get_chat_template_from_path(
+              json_no_key_path
+          )
+      )
+
+      # Test non-existent file
+      self.assertIsNone(
+          instruction_data_processing.get_chat_template_from_path(
+              "non_existent.jinja"
+          )
+      )
+
 
 if __name__ == "__main__":
   unittest.main()