refactor: 重构流水线架构，添加Pipeline抽象并拆分IOHandler

pipeline/io/export.py

@@ -0,0 +1,141 @@
+"""Dataset export and caching utilities."""
+
+import json
+import logging
+import os
+from pathlib import Path
+from typing import Any, Callable, Dict, List, Optional, Union
+
+from datasets import Dataset
+from tqdm import tqdm
+
+from pipeline.io.file_scanner import FileScanner
+from pipeline.io.hdf5_handler import HDF5Handler
+from pipeline.processors import BaseProcessor
+from pipeline.packing import SequencePacker
+from pipeline.utils import error_handler
+
+logger = logging.getLogger(__name__)
+
+
+@error_handler()
+def export_dataset(
+    dataset: Dataset,
+    output_dir: str,
+    output_prefix: str,
+    *,
+    chunk_size: int = 1_000_000,
+    max_chunks: Optional[int] = None,
+    process_func: Optional[
+        Callable[[Dict[str, Any]], Union[Dict[str, Any], List[Dict[str, Any]]]]
+    ] = None,
+    column: str = "text",
+) -> List[str]:
+    """Export HuggingFace Dataset to JSONL files in chunks.
+
+    Args:
+        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)
+    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")
+        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
+
+
+@error_handler()
+def cache_jsonl(
+    files: List[str],
+    output_dir: str,
+    processor: BaseProcessor,
+    *,
+    pack_size: int = -1,
+    pad_value: int = 1,
+) -> List[str]:
+    """Tokenize JSONL files and pack them into HDF5 storage.
+
+    Args:
+        files: List of JSONL file paths.
+        output_dir: H5 output directory.
+        processor: Initialized Processor instance.
+        pack_size: Packing length, <=0 means no packing.
+        pad_value: Padding value.
+
+    Returns:
+        List of generated H5 file paths.
+    """
+    os.makedirs(output_dir, exist_ok=True)
+    output_files: List[str] = []
+    output_keys = processor.output_keys
+
+    for file_path in files:
+        file_name = Path(file_path).stem
+
+        arrows: Dict[str, List] = {key: [] for key in output_keys}
+
+        with open(file_path, "r", encoding="utf-8") as f:
+            for line_num, line in enumerate(
+                tqdm(f, desc=f"Processing {file_name}", leave=False), start=1
+            ):
+                try:
+                    result = processor.process(json.loads(line))
+                    if result is not None:
+                        for key in output_keys:
+                            arrows[key].append(result[key])
+                except json.JSONDecodeError as e:
+                    logger.warning(
+                        f"JSON decode error in {file_path} line {line_num}: {e}. Skipping line."
+                    )
+                    continue
+                except Exception as e:
+                    logger.warning(
+                        f"Unexpected error processing line {line_num} in {file_path}: {e}. Skipping line."
+                    )
+                    continue
+
+        if pack_size > 0:
+            output = {}
+            for key in output_keys:
+                packer = SequencePacker(pack_size, pad_value)
+                output[key] = packer.pack(arrows[key])
+        else:
+            output = arrows
+
+        h5_path = HDF5Handler.save(output_dir, file_name, output)
+        output_files.append(h5_path)
+        logger.info(f"Saved {h5_path}")
+
+    return output_files