From 2f919e9243cb07d0449799bb0beea9e008d76928 Mon Sep 17 00:00:00 2001
From: ViperEkura <3081035982@qq.com>
Date: Thu, 2 Jul 2026 17:40:03 +0800
Subject: [PATCH] feat: add MetaMathQA SFT export script and fix ChatML newline
 encoding

---
 pipeline/strategies/chatml.py                   | 17 +++++++++--------
 scripts/supervised_finetuning/sft_metamathqa.py | 17 +++++++++++++++++
 2 files changed, 26 insertions(+), 8 deletions(-)
 create mode 100644 scripts/supervised_finetuning/sft_metamathqa.py

diff --git a/pipeline/strategies/chatml.py b/pipeline/strategies/chatml.py
index cd8411c..b4f1493 100644
--- a/pipeline/strategies/chatml.py
+++ b/pipeline/strategies/chatml.py
@@ -14,17 +14,18 @@ class ChatMLStrategy(PromptStrategy):
     def __init__(
         self,
         tokenizer: AutoTokenizer,
-        user_start: str = "<｜im▁start｜>user\n",
-        user_end: str = "<｜im▁end｜>\n",
-        assistant_start: str = "<｜im▁start｜>assistant\n",
-        assistant_end: str = "<｜im▁end｜>\n",
+        user_start: str = "<｜im▁start｜>user",
+        user_end: str = "<｜im▁end｜>",
+        assistant_start: str = "<｜im▁start｜>assistant",
+        assistant_end: str = "<｜im▁end｜>",
     ):
         super().__init__(tokenizer)
+        nl_id = tokenizer.encode("a\nb", add_special_tokens=False)[1]
 
-        self._user_start_ids = self._encode_format(user_start)
-        self._user_end_ids = self._encode_format(user_end)
-        self._assistant_start_ids = self._encode_format(assistant_start)
-        self._assistant_end_ids = self._encode_format(assistant_end)
+        self._user_start_ids = self._encode_format(user_start) + [nl_id]
+        self._user_end_ids = self._encode_format(user_end) + [nl_id]
+        self._assistant_start_ids = self._encode_format(assistant_start) + [nl_id]
+        self._assistant_end_ids = self._encode_format(assistant_end) + [nl_id]
 
     @property
     def name(self) -> str:
diff --git a/scripts/supervised_finetuning/sft_metamathqa.py b/scripts/supervised_finetuning/sft_metamathqa.py
new file mode 100644
index 0000000..a26ff59
--- /dev/null
+++ b/scripts/supervised_finetuning/sft_metamathqa.py
@@ -0,0 +1,17 @@
+from datasets import load_dataset
+from pipeline import export_dataset
+
+
+def process_func(sample: dict) -> dict:
+    return {"query": sample["query"], "response": sample["response"]}
+
+
+if __name__ == "__main__":
+    dataset = load_dataset("meta-math/MetaMathQA", split="train")
+    export_dataset(
+        dataset=dataset,
+        output_dir="./dataset",
+        output_prefix="MetaMathQA",
+        process_func=process_func,
+        chunk_size=1_000_000,
+    )