feat: 增加日志管理

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',

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

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

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]] = {}
 

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.

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}<eos>")
         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)

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))
 

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 = ["<bos>", "<eos>", "<pad>"]
         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

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)