From 7baa3ea0c3d92291a476f3af5e274061d2b333da Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Mon, 30 Mar 2026 16:28:56 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E5=A2=9E=E5=8A=A0=E6=97=A5=E5=BF=97?= =?UTF-8?q?=E7=AE=A1=E7=90=86?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- pipeline/__init__.py | 5 ++++ pipeline/cache.py | 36 ++++++++++++++++--------- pipeline/export.py | 53 +++++++++++++++++++++--------------- pipeline/io.py | 10 ++++--- pipeline/packing.py | 5 +++- pipeline/processors.py | 30 +++++++++++---------- pipeline/text.py | 6 ++--- pipeline/tokenizer.py | 16 +++++------ pipeline/utils.py | 61 ++++++++++++++++++++++++++++++++++++++++++ 9 files changed, 159 insertions(+), 63 deletions(-) create mode 100644 pipeline/utils.py diff --git a/pipeline/__init__.py b/pipeline/__init__.py index c57f032..8ab4c2e 100644 --- a/pipeline/__init__.py +++ b/pipeline/__init__.py @@ -1,3 +1,4 @@ +import logging from .tokenizer import BpeTokenizer from .text import TextNormalizer from .packing import SequencePacker @@ -5,6 +6,10 @@ from .io import IOHandler from .processors import ProcessorFactory, BaseProcessor from .export import export_dataset from .cache import cache_jsonl +from .utils import setup_logging + +# 配置项目级日志记录 +setup_logging() __all__ = [ 'BpeTokenizer', diff --git a/pipeline/cache.py b/pipeline/cache.py index 78a9ce9..3cfb80c 100644 --- a/pipeline/cache.py +++ b/pipeline/cache.py @@ -1,6 +1,7 @@ -"""将 JSONL 文件 tokenize 后打包存储为 HDF5""" +"""Tokenize JSONL files and pack them into HDF5 storage.""" import json import os +import logging from typing import List from pathlib import Path @@ -9,8 +10,12 @@ from tqdm import tqdm from .processors import BaseProcessor from .packing import SequencePacker from .io import IOHandler +from .utils import error_handler + +logger = logging.getLogger(__name__) +@error_handler() def cache_jsonl( files: List[str], output_dir: str, @@ -20,17 +25,17 @@ def cache_jsonl( pad_value: int = 1, ) -> List[str]: """ - 将 JSONL 文件 tokenize 后打包存储为 HDF5。 + Tokenize JSONL files and pack them into HDF5 storage. Args: - files: JSONL 文件路径列表 - output_dir: H5 输出目录 - processor: 已初始化的 Processor 实例 - pack_size: 打包长度，<=0 表示不打包 - pad_value: 填充值 + 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: - 生成的 H5 文件路径列表 + List of generated H5 file paths """ os.makedirs(output_dir, exist_ok=True) output_files: List[str] = [] @@ -40,8 +45,15 @@ def cache_jsonl( arrows = [] with open(file_path, "r", encoding="utf-8") as f: - for line in tqdm(f, desc=f"Processing {file_name}", leave=False): - arrow = processor.process(json.loads(line)) + for line_num, line in enumerate(tqdm(f, desc=f"Processing {file_name}", leave=False), start=1): + try: + arrow = processor.process(json.loads(line)) + 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 arrow is not None: arrows.append(arrow) @@ -50,7 +62,7 @@ def cache_jsonl( output = {} for key in processor.output_keys: if pack_size > 0: - packer = SequencePacker(pack_size, pad_value) # 每个键独立实例 + packer = SequencePacker(pack_size, pad_value) # independent instance per key output[key] = packer.pack(package[key]) else: output[key] = package[key] @@ -58,6 +70,6 @@ def cache_jsonl( IOHandler.save_h5(output_dir, file_name, output) h5_path = os.path.join(output_dir, f"{file_name}.h5") output_files.append(h5_path) - print(f"Saved {h5_path}") + logger.info(f"Saved {h5_path}") return output_files diff --git a/pipeline/export.py b/pipeline/export.py index 948adc9..6150f67 100644 --- a/pipeline/export.py +++ b/pipeline/export.py @@ -1,33 +1,40 @@ -"""将 HuggingFace Dataset 分块导出为 JSONL 文件""" +"""Export HuggingFace Dataset to JSONL files in chunks.""" import json import os -from typing import Callable, Optional, List, Union +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: Dataset, output_dir: str, output_prefix: str, *, chunk_size: int = 1_000_000, max_chunks: Optional[int] = None, - process_func: Optional[Callable] = None, + process_func: Optional[Callable[[Dict[str, Any]], Union[Dict[str, Any], List[Dict[str, Any]]]]] = None, column: str = "text", ) -> List[str]: """ - 将 HuggingFace Dataset 分块导出为 JSONL 文件。 + Export HuggingFace Dataset to JSONL files in chunks. Args: - dataset: HuggingFace Dataset 对象 - output_dir: 输出目录 - output_prefix: 输出文件名前缀，如 "chinese-c4-pretrain" - chunk_size: 每个文件的最大样本数 - max_chunks: 最多处理几个 chunk（用于调试） - process_func: 单条样本的转换函数 (dict) -> dict | list[dict] - column: 默认提取的文本列名（仅在 process_func 为 None 时使用） + 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) @@ -41,14 +48,16 @@ def export_dataset( chunk = dataset.select(range(start, end)) path = os.path.join(output_dir, f"{output_prefix}_chunk_{i}.jsonl") - 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) - print(f"[{i + 1}/{lim}] Saved {path}") + 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 diff --git a/pipeline/io.py b/pipeline/io.py index de7d733..33aec2b 100644 --- a/pipeline/io.py +++ b/pipeline/io.py @@ -1,13 +1,15 @@ from pathlib import Path -from typing import Dict, List +from typing import Dict, List, Optional, Callable import os import h5py import torch from torch import Tensor +from .utils import error_handler + class IOHandler: - """文件和 HDF5 读写""" + """File and HDF5 read/write operations.""" @staticmethod def fetch_files(directory: str) -> List[str]: @@ -18,7 +20,7 @@ class IOHandler: ] @staticmethod - def fetch_folders(root_dir: str, filter_func=None) -> List[str]: + def fetch_folders(root_dir: str, filter_func: Optional[Callable[[str], bool]] = None) -> List[str]: folders = [] for root, dirs, _ in os.walk(root_dir): for dir_name in dirs: @@ -28,6 +30,7 @@ class IOHandler: return folders @staticmethod + @error_handler() 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") @@ -38,6 +41,7 @@ class IOHandler: grp.create_dataset(f'data_{idx}', data=tensor.cpu().numpy()) @staticmethod + @error_handler() def load_h5(file_path: str, share_memory=True) -> Dict[str, List[Tensor]]: tensor_group: Dict[str, List[Tensor]] = {} diff --git a/pipeline/packing.py b/pipeline/packing.py index fa2b9d2..1c46236 100644 --- a/pipeline/packing.py +++ b/pipeline/packing.py @@ -3,12 +3,14 @@ from typing import List import torch from torch import Tensor +from .utils import error_handler + logger = logging.getLogger(__name__) class SequencePacker: - def __init__(self, pack_size: int, pad_value: int = 0, dtype=torch.int32): + def __init__(self, pack_size: int, pad_value: int = 0, dtype: torch.dtype = torch.int32): self.pack_size = pack_size self.pad_value = pad_value self.dtype = dtype @@ -21,6 +23,7 @@ class SequencePacker: ) self._current_pos = 0 + @error_handler() def pack(self, sequences: List[Tensor]) -> List[Tensor]: """ Pack sequences into fixed-size packages. diff --git a/pipeline/processors.py b/pipeline/processors.py index 53d74b1..34cf892 100644 --- a/pipeline/processors.py +++ b/pipeline/processors.py @@ -1,13 +1,15 @@ from abc import ABC, abstractmethod -from typing import Dict, List +from typing import Dict, List, Any import torch +from torch import Tensor +from .tokenizer import BpeTokenizer class BaseProcessor(ABC): - """处理器抽象基类""" + """Abstract base class for processors.""" @abstractmethod - def process(self, input_dict: dict) -> dict: + def process(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]: pass @property @@ -17,12 +19,12 @@ class BaseProcessor(ABC): class PreTrainProcessor(BaseProcessor): - """预训练数据处理器""" + """Pre-training data processor.""" - def __init__(self, tokenizer): + def __init__(self, tokenizer: BpeTokenizer): self.tokenizer = tokenizer - def process(self, input_dict: dict) -> dict: + def process(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]: segment = input_dict["text"] tokens = self.tokenizer.encode(f"{segment}") return {'sequence': torch.tensor(tokens, dtype=torch.int32)} @@ -33,12 +35,12 @@ class PreTrainProcessor(BaseProcessor): class SFTProcessor(BaseProcessor): - """监督微调数据处理器""" + """Supervised fine-tuning data processor.""" - def __init__(self, tokenizer): + def __init__(self, tokenizer: BpeTokenizer): self.tokenizer = tokenizer - def process(self, input_dict: dict) -> dict: + def process(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]: query, response = input_dict["query"], input_dict["response"] q = self.tokenizer.encode( f"<|im_start|>user\n{query}<|im_end|>\n<|im_start|>assistant\n" @@ -55,12 +57,12 @@ class SFTProcessor(BaseProcessor): class DPOProcessor(BaseProcessor): - """DPO 偏好学习数据处理器""" + """DPO preference learning data processor.""" - def __init__(self, tokenizer): + def __init__(self, tokenizer: BpeTokenizer): self.tokenizer = tokenizer - def process(self, input_dict: dict) -> dict: + def process(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]: query = input_dict["query"] chosen_response = input_dict["chosen"] rejected_response = input_dict["rejected"] @@ -92,7 +94,7 @@ class DPOProcessor(BaseProcessor): class ProcessorFactory: - """处理器工厂""" + """Processor factory.""" _processors = { "pt": PreTrainProcessor, @@ -101,7 +103,7 @@ class ProcessorFactory: } @classmethod - def create(cls, processor_type: str, tokenizer) -> BaseProcessor: + def create(cls, processor_type: str, tokenizer: BpeTokenizer) -> BaseProcessor: if processor_type not in cls._processors: raise ValueError(f"Invalid processor type: {processor_type}") return cls._processors[processor_type](tokenizer) diff --git a/pipeline/text.py b/pipeline/text.py index 22daa12..ffdbdaf 100644 --- a/pipeline/text.py +++ b/pipeline/text.py @@ -1,9 +1,9 @@ import re -from typing import Dict +from typing import Dict, Optional class TextNormalizer: - """文本规范化""" + """Text normalization.""" DEFAULT_REPLACEMENTS = { "\\[": "$$", "\\]": "$$", "\\(": "$", "\\)": "$", @@ -13,7 +13,7 @@ class TextNormalizer: '\u00A0': ' ', '\u2026': '...' } - def __init__(self, custom_rules: Dict[str, str] = None): + def __init__(self, custom_rules: Optional[Dict[str, str]] = None): self.replacements = {**self.DEFAULT_REPLACEMENTS, **(custom_rules or {})} self._pattern = re.compile('|'.join(re.escape(k) for k in self.replacements)) diff --git a/pipeline/tokenizer.py b/pipeline/tokenizer.py index 689b5c2..e78f17a 100644 --- a/pipeline/tokenizer.py +++ b/pipeline/tokenizer.py @@ -2,11 +2,11 @@ from tokenizers import Tokenizer, Encoding from tokenizers import decoders, processors, normalizers, pre_tokenizers from tokenizers.models import BPE from tokenizers.trainers import BpeTrainer -from typing import List, Union +from typing import List, Union, Optional, Tuple, Iterator class BpeTokenizer: - def __init__(self, path=None): + def __init__(self, path: Optional[str] = None): self._control_tokens = ["", "", ""] self._special_tokens = ["<|im_start|>", "<|im_end|>"] @@ -28,7 +28,7 @@ class BpeTokenizer: if path is not None: self._tokenizer = Tokenizer.from_file(path) - def _prepare_trainer(self, vocab_size: int, min_freq: int, reserved_token_size: int, max_token_length=18) -> tuple: + def _prepare_trainer(self, vocab_size: int, min_freq: int, reserved_token_size: int, max_token_length: int = 18) -> Tuple[BpeTrainer, int, List[str]]: assert reserved_token_size > len(self._special_tokens) reserved_tokens = [f"<|reserve{i:02d}|>" for i in range(reserved_token_size - len(self._special_tokens))] detail_vocab_size = vocab_size - (len(reserved_tokens) + len(self._special_tokens)) @@ -49,7 +49,7 @@ class BpeTokenizer: return trainer, detail_vocab_size, reserved_tokens - def train(self, files, vocab_size, min_freq, reserved_token_size=100): + def train(self, files: List[str], vocab_size: int, min_freq: int, reserved_token_size: int = 100) -> None: trainer, _, reserved_tokens = self._prepare_trainer( vocab_size=vocab_size, min_freq=min_freq, @@ -58,7 +58,7 @@ class BpeTokenizer: self._tokenizer.train(files=files, trainer=trainer) self._tokenizer.add_special_tokens(self._special_tokens + reserved_tokens) - def train_from_iterator(self, iterator, vocab_size, min_freq, reserved_token_size=100): + def train_from_iterator(self, iterator: Iterator[str], vocab_size: int, min_freq: int, reserved_token_size: int = 100) -> None: trainer, _, reserved_tokens = self._prepare_trainer( vocab_size=vocab_size, min_freq=min_freq, @@ -67,13 +67,13 @@ class BpeTokenizer: self._tokenizer.train_from_iterator(iterator=iterator, trainer=trainer) self._tokenizer.add_special_tokens(self._special_tokens + reserved_tokens) - def save(self, path): + def save(self, path: str) -> None: self._tokenizer.save(path) - def load(self, path): + def load(self, path: str) -> None: self._tokenizer = Tokenizer.from_file(path) - def encode(self, tokens: Union[str, List[str]], out_ids: bool=True, add_special_tokens: bool=False) -> List: + def encode(self, tokens: Union[str, List[str]], out_ids: bool = True, add_special_tokens: bool = False) -> Union[List[int], List[str], List[List[int]], List[List[str]]]: if isinstance(tokens, str): encoded: Encoding = self._tokenizer.encode(tokens, add_special_tokens=add_special_tokens) return encoded.ids if out_ids else encoded.tokens diff --git a/pipeline/utils.py b/pipeline/utils.py new file mode 100644 index 0000000..6283d1d --- /dev/null +++ b/pipeline/utils.py @@ -0,0 +1,61 @@ +import functools +import logging +import sys +import os +from typing import Optional, Callable, Any + + +def error_handler( + logger: Optional[logging.Logger] = None, + reraise: bool = True, + log_level: int = logging.ERROR, + capture_keyboard_interrupt: bool = False, +): + def decorator(func: Callable) -> Callable: + @functools.wraps(func) + def wrapper(*args, **kwargs) -> Any: + try: + return func(*args, **kwargs) + except Exception as e: + if not capture_keyboard_interrupt and isinstance(e, KeyboardInterrupt): + raise + + nonlocal logger + log = logger or logging.getLogger(func.__module__) + log.log( + log_level, + f"Error in {func.__name__}: {e}", + exc_info=True, + ) + if reraise: + raise + return None + return wrapper + return decorator + + +def setup_logging(level: Optional[int] = None) -> None: + + if level is None: + level_name = os.getenv("LOG_LEVEL", "INFO").upper() + level = getattr(logging, level_name, logging.INFO) + + root_logger = logging.getLogger() + if root_logger.handlers: + root_logger.setLevel(level) + return + + console_handler = logging.StreamHandler(sys.stdout) + console_handler.setLevel(level) + + formatter = logging.Formatter( + fmt="%(asctime)s - %(name)s - %(levelname)s - %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", + ) + console_handler.setFormatter(formatter) + + root_logger.setLevel(level) + root_logger.addHandler(console_handler) + + logging.getLogger("h5py").setLevel(logging.WARNING) + logging.getLogger("torch").setLevel(logging.WARNING) \ No newline at end of file