DataPipeline/pipeline/io/export.py

"""Dataset export and caching utilities."""

import json
import logging
import os
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Union

import torch
from datasets import Dataset
from torch import Tensor
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 pack_tensors, BasePacker
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


def merge_tensors(
    tensors: List[Tensor],
    group_size: int,
) -> List[Tensor]:
    """Merge a list of tensors into fewer larger tensors.

    Concatenates every group_size consecutive tensors into one merged
    tensor. This reduces the number of shm blocks when loading.

    Args:
        tensors: List of 1D tensors.
        group_size: Number of tensors to merge into each group.

    Returns:
        List of merged tensors.
    """
    if not tensors:
        return []

    merged: List[Tensor] = []
    for i in range(0, len(tensors), group_size):
        merged.append(torch.cat(tensors[i : i + group_size]))

    return merged


@error_handler()
def cache_jsonl(
    files: List[str],
    output_dir: str,
    processor: BaseProcessor,
    *,
    pack_size: int = -1,
    pad_value: int = 0,
    group_size: int = 1_000,
    pack_algo: Optional[str] = None,
) -> 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.
        group_size: Merge every this many packed chunks into one tensor,
            <=0 means no merging.
        pack_algo: Packing algorithm: 'bfd' (default), 'ffd',
            'greedy'. Only used when pack_size > 0.

    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 not arrows[output_keys[0]]:
            logger.warning(f"No valid samples in {file_path}, skipping")
            continue

        if pack_size > 0:
            dtypes = (
                dict(processor.schema.output_fields)
                if processor.schema is not None
                else None
            )
            pad_values = {k: (0 if k == "position_ids" else (False if k.endswith("_mask") else pad_value)) for k in output_keys}
            output = pack_tensors(arrows, pack_size, pad_value, dtypes, pad_values=pad_values, algo=pack_algo)
        else:
            output = arrows

        if group_size > 0 and output[output_keys[0]]:
            output = {
                key: merge_tensors(tensors, group_size)
                for key, tensors in output.items()
            }

        h5_path = HDF5Handler.save(output_dir, file_name, output)
        output_files.append(h5_path)
        logger.info(f"Saved {h5_path}")

    return output_files