From 6762842b0cbef9e55b73cc3fa2de3ccacff76769 Mon Sep 17 00:00:00 2001
From: ViperEkura <3081035982@qq.com>
Date: Sun, 20 Jul 2025 09:59:40 +0800
Subject: [PATCH] =?UTF-8?q?refactor(belle=5Fsft):=20=E9=87=8D=E6=9E=84?=
 =?UTF-8?q?=E5=A4=84=E7=90=86=E6=95=B0=E6=8D=AE=E9=9B=86=E7=9A=84=E9=80=BB?=
 =?UTF-8?q?=E8=BE=91?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 belle_sft.py | 39 ++++++++++++++++++++++++++++++++++-----
 utils.py     |  4 ++--
 2 files changed, 36 insertions(+), 7 deletions(-)

diff --git a/belle_sft.py b/belle_sft.py
index 4d6a521..8818f4a 100644
--- a/belle_sft.py
+++ b/belle_sft.py
@@ -1,10 +1,39 @@
 from datasets import load_dataset
 from utils import process_dataset
 
+
+def build_prompt(query, history) -> str:
+    ret_prompt = ""
+    if len(history) > 0:
+        for his_query, his_response in history:
+            ret_prompt += f"<|user|> {his_query} <|system|> <bos>{his_response}<eos>\n"
+    if query is not None:
+        ret_prompt += f"<|user|> {query} <|system|> <bos>"
+    return ret_prompt
+
+
+def process_func(input_dict: dict):
+    conversations = input_dict["conversations"]
+    n = len(conversations) // 2
+    examples = []
+    
+    for i in range(n):
+        user_msg = conversations[2*i]["value"]
+        assistant_msg = conversations[2*i+1]["value"]
+        examples.append((user_msg, assistant_msg))
+    
+    content = {
+        "text": build_prompt(None, examples)
+    }
+    print(content)
+    return content
+
+
 if __name__ == "__main__":
     dataset = load_dataset("BelleGroup/train_3.5M_CN")
-    # process_dataset(
-    #     dataset_dict=dataset,
-    #     output_subdir="belle_sft",
-    #     max_chunk_size=5,
-    # )
\ No newline at end of file
+    process_dataset(
+        dataset_dict=dataset,
+        output_subdir="belle_sft",
+        chunk_size=5,
+        process_func=process_func,
+    )
\ No newline at end of file
diff --git a/utils.py b/utils.py
index 8dc740d..7330a26 100644
--- a/utils.py
+++ b/utils.py
@@ -1,4 +1,4 @@
-from typing import List, Callable
+from typing import List, Callable, Union
 from datasets import DatasetDict
 from tokenizer import BpeTokenizer
 from tqdm import tqdm
@@ -64,7 +64,7 @@ def process_dataset(
     chunk_size: int = 1000000,
     split_name: str = "train",
     column_name: str = "text",
-    process_func: Callable[[dict], dict] = None,
+    process_func: Callable[[Union[dict, List[dict]]], dict] = None,
     normalization_func=comprehensive_normalization,
 ):
     train_dataset = dataset_dict[split_name]