From fcb5c2bcc8a6cf30495d4c0e9796cf3097616834 Mon Sep 17 00:00:00 2001
From: ViperEkura <3081035982@qq.com>
Date: Sun, 20 Jul 2025 11:21:18 +0800
Subject: [PATCH] =?UTF-8?q?feat(dataset):=20=E6=9B=B4=E6=96=B0=E6=95=B0?=
 =?UTF-8?q?=E6=8D=AE=E9=9B=86=E5=A4=84=E7=90=86=E5=92=8C=E4=BF=9D=E5=AD=98?=
 =?UTF-8?q?=E9=80=BB=E8=BE=91?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 belle_sft.py |  2 +-
 sft_file.py  | 15 +++++++++++++++
 utils.py     |  2 +-
 3 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/belle_sft.py b/belle_sft.py
index c783072..50aab61 100644
--- a/belle_sft.py
+++ b/belle_sft.py
@@ -32,6 +32,6 @@ if __name__ == "__main__":
     dataset = load_dataset("BelleGroup/train_3.5M_CN")
     process_dataset(
         dataset_dict=dataset,
-        output_subdir="belle_sft",
+        output_subdir="belle-sft",
         process_func=process_func,
     )
\ No newline at end of file
diff --git a/sft_file.py b/sft_file.py
index e69de29..66579ef 100644
--- a/sft_file.py
+++ b/sft_file.py
@@ -0,0 +1,15 @@
+from utils import dump_pkl_files, fetch_files
+from tokenizer import BpeTokenizer
+import os
+
+if __name__ == "__main__":
+    tokenizer = BpeTokenizer("tokenizer.json")
+    base_dir = [
+        os.path.join("dataset", "belle-sft"),
+    ]
+    base_out_dir = "pkl_output"
+    files = []
+    for dir_path in base_dir:
+        files.extend(fetch_files(dir_path))
+    
+    dump_pkl_files(tokenizer, files, base_out_dir, packing_size=2048)
\ No newline at end of file
diff --git a/utils.py b/utils.py
index ebb1f6a..f4701ab 100644
--- a/utils.py
+++ b/utils.py
@@ -38,7 +38,7 @@ def dump_pkl_files(
 ):  
     def process_line(line: str) -> Tensor:
             line = json.loads(line)[key]
-            processed_line = encoder(line)
+            processed_line = encoder(line) if encoder else line
             ids = tokenizer.encode(processed_line)
             arrow = torch.tensor(ids, dtype=torch.int32)
             return arrow