feat: 增加日志管理

pipeline/export.py

@@ -1,33 +1,40 @@
-"""将 HuggingFace Dataset 分块导出为 JSONL 文件"""
+"""Export HuggingFace Dataset to JSONL files in chunks."""
 import json
 import os
-from typing import Callable, Optional, List, Union
+import logging
+from typing import Callable, Optional, List, Union, Dict, Any
+
+from datasets import Dataset
+from .utils import error_handler
+
+logger = logging.getLogger(__name__)
 
 
+@error_handler()
 def export_dataset(
-    dataset,
+    dataset: Dataset,
     output_dir: str,
     output_prefix: str,
     *,
     chunk_size: int = 1_000_000,
     max_chunks: Optional[int] = None,
-    process_func: Optional[Callable] = None,
+    process_func: Optional[Callable[[Dict[str, Any]], Union[Dict[str, Any], List[Dict[str, Any]]]]] = None,
     column: str = "text",
 ) -> List[str]:
     """
-    将 HuggingFace Dataset 分块导出为 JSONL 文件。
+    Export HuggingFace Dataset to JSONL files in chunks.
 
     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 时使用）
+        dataset: HuggingFace Dataset object
+        output_dir: Output directory
+        output_prefix: Output file name prefix, e.g., "chinese-c4-pretrain"
+        chunk_size: Maximum number of samples per file
+        max_chunks: Maximum number of chunks to process (for debugging)
+        process_func: Single sample transformation function (dict) -> dict | list[dict]
+        column: Default text column name (only used when process_func is None)
 
     Returns:
-        生成的文件路径列表
+        List of generated file paths
     """
     os.makedirs(output_dir, exist_ok=True)
     total = len(dataset)
@@ -41,14 +48,16 @@ def export_dataset(
         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}")
+        try:
+            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)
+            logger.info(f"[{i + 1}/{lim}] Saved {path}")
+        except (OSError, IOError) as e:
+            logger.error(f"Failed to write chunk {i} to {path}: {e}")
 
     return output_files