refactor(belle_sft): 重构处理数据集的逻辑

2025-07-20 09:59:40 +08:00
parent faf9f81e7c
commit 6762842b0c
2 changed files with 36 additions and 7 deletions
@@ -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]