refactor: 修改项目结构

pipeline/export.py

@@ -0,0 +1,54 @@
+"""将 HuggingFace Dataset 分块导出为 JSONL 文件"""
+import json
+import os
+from typing import Callable, Optional, List, Union
+
+
+def export_dataset(
+    dataset,
+    output_dir: str,
+    output_prefix: str,
+    *,
+    chunk_size: int = 1_000_000,
+    max_chunks: Optional[int] = None,
+    process_func: Optional[Callable] = None,
+    column: str = "text",
+) -> List[str]:
+    """
+    将 HuggingFace Dataset 分块导出为 JSONL 文件。
+
+    Args:
+        dataset: HuggingFace Dataset 对象
+        output_dir: 输出目录
+        output_prefix: 输出文件名前缀，如 "chinese-c4-pretrain"
+        chunk_size: 每个文件的最大样本数
+        max_chunks: 最多处理几个 chunk（用于调试）
+        process_func: 单条样本的转换函数 (dict) -> dict | list[dict]
+        column: 默认提取的文本列名（仅在 process_func 为 None 时使用）
+
+    Returns:
+        生成的文件路径列表
+    """
+    os.makedirs(output_dir, exist_ok=True)
+    total = len(dataset)
+    num_chunks = (total + chunk_size - 1) // chunk_size
+    lim = min(max_chunks, num_chunks) if max_chunks else num_chunks
+
+    output_files: List[str] = []
+    for i in range(lim):
+        start = i * chunk_size
+        end = min(start + chunk_size, total)
+        chunk = dataset.select(range(start, end))
+
+        path = os.path.join(output_dir, f"{output_prefix}_chunk_{i}.jsonl")
+        with open(path, "w", encoding="utf-8") as f:
+            for example in chunk:
+                processed = process_func(example) if process_func else {column: example[column]}
+                items = processed if isinstance(processed, list) else [processed]
+                for item in items:
+                    f.write(json.dumps(item, ensure_ascii=False) + "\n")
+
+        output_files.append(path)
+        print(f"[{i + 1}/{lim}] Saved {path}")
+
+    return output_files