refactor: 重构流水线架构，添加Pipeline抽象并拆分IOHandler
This commit is contained in:
@@ -0,0 +1,19 @@
|
||||
"""I/O module for file operations, HDF5 storage, and dataset export.
|
||||
|
||||
This module provides:
|
||||
- FileScanner: File and directory scanning utilities
|
||||
- HDF5Handler: Tensor data persistence
|
||||
- export_dataset: HuggingFace Dataset to JSONL export
|
||||
- cache_jsonl: JSONL to HDF5 tokenization and caching
|
||||
"""
|
||||
|
||||
from pipeline.io.file_scanner import FileScanner
|
||||
from pipeline.io.hdf5_handler import HDF5Handler
|
||||
from pipeline.io.export import export_dataset, cache_jsonl
|
||||
|
||||
__all__ = [
|
||||
"FileScanner",
|
||||
"HDF5Handler",
|
||||
"export_dataset",
|
||||
"cache_jsonl",
|
||||
]
|
||||
@@ -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
|
||||
@@ -0,0 +1,98 @@
|
||||
"""File system scanning utilities."""
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Callable, List, Optional
|
||||
|
||||
|
||||
class FileScanner:
|
||||
"""Utility for scanning and filtering files in a directory tree.
|
||||
|
||||
Example::
|
||||
|
||||
scanner = FileScanner()
|
||||
files = scanner.scan("./data", suffix=".jsonl")
|
||||
folders = scanner.scan_folders("./data", filter_func=is_valid_dir)
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def scan(
|
||||
directory: str,
|
||||
suffix: Optional[str] = None,
|
||||
recursive: bool = True,
|
||||
relative: bool = False,
|
||||
) -> List[str]:
|
||||
"""Scan directory for files matching criteria.
|
||||
|
||||
Args:
|
||||
directory: Root directory to scan.
|
||||
suffix: Filter files by extension (e.g., ".jsonl").
|
||||
recursive: Whether to search subdirectories.
|
||||
relative: Return relative paths if True.
|
||||
|
||||
Returns:
|
||||
Sorted list of file paths.
|
||||
"""
|
||||
root_path = Path(directory)
|
||||
if not root_path.exists():
|
||||
return []
|
||||
|
||||
if recursive:
|
||||
pattern = "**/*" + suffix if suffix else "**/*"
|
||||
files = [str(p) for p in root_path.glob(pattern) if p.is_file()]
|
||||
else:
|
||||
pattern = "*" + suffix if suffix else "*"
|
||||
files = [str(p) for p in root_path.glob(pattern) if p.is_file()]
|
||||
|
||||
if relative:
|
||||
files = [str(Path(f).relative_to(root_path)) for f in files]
|
||||
|
||||
return sorted(files)
|
||||
|
||||
@staticmethod
|
||||
def scan_folders(
|
||||
directory: str,
|
||||
filter_func: Optional[Callable[[str], bool]] = None,
|
||||
recursive: bool = True,
|
||||
) -> List[str]:
|
||||
"""Scan directory for subdirectories.
|
||||
|
||||
Args:
|
||||
directory: Root directory to scan.
|
||||
filter_func: Optional predicate to filter directories.
|
||||
recursive: Whether to search subdirectories.
|
||||
|
||||
Returns:
|
||||
Sorted list of directory paths.
|
||||
"""
|
||||
root_path = Path(directory)
|
||||
if not root_path.exists():
|
||||
return []
|
||||
|
||||
if recursive:
|
||||
folders = [str(p) for p in root_path.rglob("*") if p.is_dir()]
|
||||
else:
|
||||
folders = [str(p) for p in root_path.glob("*") if p.is_dir()]
|
||||
|
||||
if filter_func:
|
||||
folders = [f for f in folders if filter_func(f)]
|
||||
|
||||
return sorted(folders)
|
||||
|
||||
@staticmethod
|
||||
def group_by_extension(files: List[str]) -> dict[str, List[str]]:
|
||||
"""Group files by their extension.
|
||||
|
||||
Args:
|
||||
files: List of file paths.
|
||||
|
||||
Returns:
|
||||
Dictionary mapping extension to list of files.
|
||||
"""
|
||||
groups: dict[str, List[str]] = {}
|
||||
for f in files:
|
||||
ext = Path(f).suffix
|
||||
if ext not in groups:
|
||||
groups[ext] = []
|
||||
groups[ext].append(f)
|
||||
return groups
|
||||
@@ -0,0 +1,133 @@
|
||||
"""HDF5 storage operations for tensor data."""
|
||||
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
import h5py
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
from pipeline.utils import error_handler
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class HDF5Handler:
|
||||
"""Handler for reading and writing tensor data to HDF5 files.
|
||||
|
||||
Example::
|
||||
|
||||
handler = HDF5Handler()
|
||||
handler.save(output_dir, "data", {"input_ids": [tensor1, tensor2]})
|
||||
|
||||
loaded = handler.load("./output/data.h5")
|
||||
for tensor in loaded["input_ids"]:
|
||||
print(tensor.shape)
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
@error_handler()
|
||||
def save(
|
||||
output_dir: str,
|
||||
file_name: str,
|
||||
tensor_group: Dict[str, List[Tensor]],
|
||||
extension: str = ".h5",
|
||||
) -> str:
|
||||
"""Save tensor groups to HDF5 file.
|
||||
|
||||
Args:
|
||||
output_dir: Output directory path.
|
||||
file_name: Base name for the output file (without extension).
|
||||
tensor_group: Dictionary mapping group names to tensor lists.
|
||||
extension: File extension (default: ".h5").
|
||||
|
||||
Returns:
|
||||
Path to the saved file.
|
||||
"""
|
||||
import os
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
full_path = os.path.join(output_dir, f"{file_name}{extension}")
|
||||
|
||||
with h5py.File(full_path, "w") as f:
|
||||
for key, tensors in tensor_group.items():
|
||||
grp = f.create_group(key)
|
||||
for idx, tensor in enumerate(tensors):
|
||||
grp.create_dataset(f"data_{idx}", data=tensor.cpu().numpy())
|
||||
|
||||
logger.info(f"Saved HDF5 file: {full_path}")
|
||||
return full_path
|
||||
|
||||
@staticmethod
|
||||
@error_handler()
|
||||
def load(
|
||||
file_path: str,
|
||||
share_memory: bool = True,
|
||||
device: Optional[torch.device] = None,
|
||||
) -> Dict[str, List[Tensor]]:
|
||||
"""Load tensor groups from HDF5 file.
|
||||
|
||||
Args:
|
||||
file_path: Path to HDF5 file or directory containing HDF5 files.
|
||||
share_memory: Whether to use shared memory for tensors.
|
||||
device: Target device for tensors (default: CPU).
|
||||
|
||||
Returns:
|
||||
Dictionary mapping group names to tensor lists.
|
||||
"""
|
||||
root_path = Path(file_path)
|
||||
h5_files = []
|
||||
|
||||
if root_path.is_file():
|
||||
h5_files = [root_path]
|
||||
else:
|
||||
h5_files = list(root_path.rglob("*.h5")) + list(root_path.rglob("*.hdf5"))
|
||||
|
||||
if not h5_files:
|
||||
logger.warning(f"No HDF5 files found at: {file_path}")
|
||||
return {}
|
||||
|
||||
tensor_group: Dict[str, List[Tensor]] = {}
|
||||
|
||||
for h5_file in h5_files:
|
||||
with h5py.File(h5_file, "r") as f:
|
||||
for key in f.keys():
|
||||
grp = f[key]
|
||||
dsets = []
|
||||
for dset_name in grp.keys():
|
||||
dset = grp[dset_name]
|
||||
tensor = torch.from_numpy(dset[:])
|
||||
|
||||
if device is not None:
|
||||
tensor = tensor.to(device)
|
||||
elif share_memory:
|
||||
tensor = tensor.share_memory_()
|
||||
|
||||
dsets.append(tensor)
|
||||
|
||||
if tensor_group.get(key) is None:
|
||||
tensor_group[key] = []
|
||||
tensor_group[key].extend(dsets)
|
||||
|
||||
logger.info(f"Loaded HDF5: {len(tensor_group)} groups, "
|
||||
f"{sum(len(v) for v in tensor_group.values())} total tensors")
|
||||
|
||||
return tensor_group
|
||||
|
||||
@staticmethod
|
||||
def get_metadata(file_path: str) -> Dict[str, int]:
|
||||
"""Get metadata about an HDF5 file without loading full data.
|
||||
|
||||
Args:
|
||||
file_path: Path to HDF5 file.
|
||||
|
||||
Returns:
|
||||
Dictionary with group names and tensor counts.
|
||||
"""
|
||||
metadata: Dict[str, int] = {}
|
||||
|
||||
with h5py.File(file_path, "r") as f:
|
||||
for key in f.keys():
|
||||
metadata[key] = len(f[key].keys())
|
||||
|
||||
return metadata
|
||||
Reference in New Issue
Block a user