64 lines
2.2 KiB
Python
64 lines
2.2 KiB
Python
"""Export HuggingFace Dataset to JSONL files in chunks."""
|
|
import json
|
|
import os
|
|
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,
|
|
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
|