reafactor: 重构项目
This commit is contained in:
+143
-10
@@ -1,23 +1,36 @@
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional, Callable
|
||||
"""File, HDF5, JSONL I/O operations."""
|
||||
import json
|
||||
import os
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional, Callable, Union, Any
|
||||
|
||||
import h5py
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from tqdm import tqdm
|
||||
from datasets import Dataset
|
||||
|
||||
from .utils import error_handler
|
||||
from pipeline.utils import error_handler
|
||||
from pipeline.processors import BaseProcessor
|
||||
from pipeline.packing import SequencePacker
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class IOHandler:
|
||||
"""File and HDF5 read/write operations."""
|
||||
|
||||
@staticmethod
|
||||
def fetch_files(directory: str) -> List[str]:
|
||||
return [
|
||||
def fetch_files(directory: str, suffix: Optional[str] = None) -> List[str]:
|
||||
files = [
|
||||
os.path.join(root, f)
|
||||
for root, _, files in os.walk(directory)
|
||||
for f in files
|
||||
]
|
||||
if suffix:
|
||||
files = [f for f in files if f.endswith(suffix)]
|
||||
return sorted(files)
|
||||
|
||||
@staticmethod
|
||||
def fetch_folders(root_dir: str, filter_func: Optional[Callable[[str], bool]] = None) -> List[str]:
|
||||
@@ -34,7 +47,7 @@ class IOHandler:
|
||||
def save_h5(output_dir: str, file_name: str, tensor_group: Dict[str, List[Tensor]]) -> None:
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
full_path = os.path.join(output_dir, f"{file_name}.h5")
|
||||
|
||||
|
||||
with h5py.File(full_path, 'w') as f:
|
||||
for key, tensors in tensor_group.items():
|
||||
grp = f.create_group(key)
|
||||
@@ -43,12 +56,12 @@ class IOHandler:
|
||||
|
||||
@staticmethod
|
||||
@error_handler()
|
||||
def load_h5(file_path: str, share_memory=True) -> Dict[str, List[Tensor]]:
|
||||
def load_h5(file_path: str, share_memory: bool = True) -> Dict[str, List[Tensor]]:
|
||||
tensor_group: Dict[str, List[Tensor]] = {}
|
||||
|
||||
root_path = Path(file_path)
|
||||
h5_files = list(root_path.rglob("*.h5")) + list(root_path.rglob("*.hdf5"))
|
||||
|
||||
|
||||
for h5_file in h5_files:
|
||||
with h5py.File(h5_file, 'r') as f:
|
||||
for key in f.keys():
|
||||
@@ -60,9 +73,129 @@ class IOHandler:
|
||||
if share_memory:
|
||||
tensor = tensor.share_memory_()
|
||||
dsets.append(tensor)
|
||||
|
||||
|
||||
if tensor_group.get(key) is None:
|
||||
tensor_group[key] = []
|
||||
tensor_group[key].extend(dsets)
|
||||
|
||||
return tensor_group
|
||||
return tensor_group
|
||||
|
||||
|
||||
# ── Stage 1: Export HuggingFace Dataset to JSONL ──────────────────────────
|
||||
|
||||
|
||||
@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
|
||||
|
||||
|
||||
# ── Stage 2: Tokenize JSONL and cache to HDF5 ────────────────────────────
|
||||
|
||||
|
||||
@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
|
||||
|
||||
IOHandler.save_h5(output_dir, file_name, output)
|
||||
h5_path = os.path.join(output_dir, f"{file_name}.h5")
|
||||
output_files.append(h5_path)
|
||||
logger.info(f"Saved {h5_path}")
|
||||
|
||||
return output_files
|
||||
|
||||
Reference in New Issue
Block a user