refactor: 修改项目结构

pipeline/__init__.py

@@ -0,0 +1,18 @@
+from .tokenizer import BpeTokenizer
+from .text import TextNormalizer
+from .packing import SequencePacker
+from .io import IOHandler
+from .processors import ProcessorFactory, BaseProcessor
+from .export import export_dataset
+from .cache import cache_jsonl
+
+__all__ = [
+    'BpeTokenizer',
+    'TextNormalizer',
+    'SequencePacker',
+    'IOHandler',
+    'ProcessorFactory',
+    'BaseProcessor',
+    'export_dataset',
+    'cache_jsonl',
+]

pipeline/cache.py

@@ -0,0 +1,64 @@
+"""将 JSONL 文件 tokenize 后打包存储为 HDF5"""
+import json
+import os
+from typing import List
+from pathlib import Path
+
+from tqdm import tqdm
+
+from .processors import BaseProcessor
+from .packing import SequencePacker
+from .io import IOHandler
+
+
+def cache_jsonl(
+    files: List[str],
+    output_dir: str,
+    processor: BaseProcessor,
+    *,
+    pack_size: int = -1,
+    pad_value: int = 1,
+) -> List[str]:
+    """
+    将 JSONL 文件 tokenize 后打包存储为 HDF5。
+
+    Args:
+        files: JSONL 文件路径列表
+        output_dir: H5 输出目录
+        processor: 已初始化的 Processor 实例
+        pack_size: 打包长度，<=0 表示不打包
+        pad_value: 填充值
+
+    Returns:
+        生成的 H5 文件路径列表
+    """
+    os.makedirs(output_dir, exist_ok=True)
+    output_files: List[str] = []
+
+    for file_path in files:
+        file_name = Path(file_path).stem
+
+        with open(file_path, "r", encoding="utf-8") as f:
+            lines = f.readlines()
+
+        arrows = []
+        for line in tqdm(lines, desc=f"Processing {file_name}", leave=False):
+            arrow = processor.process(json.loads(line))
+            if arrow is not None:
+                arrows.append(arrow)
+
+        package = {key: [a[key] for a in arrows] for key in processor.output_keys}
+
+        output = {}
+        for key in processor.output_keys:
+            if pack_size > 0:
+                output[key] = SequencePacker(pack_size, pad_value).pack(package[key])
+            else:
+                output[key] = package[key]
+
+        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}")
+
+    return output_files

pipeline/export.py

@@ -0,0 +1,54 @@
+"""将 HuggingFace Dataset 分块导出为 JSONL 文件"""
+import json
+import os
+from typing import Callable, Optional, List, Union
+
+
+def export_dataset(
+    dataset,
+    output_dir: str,
+    output_prefix: str,
+    *,
+    chunk_size: int = 1_000_000,
+    max_chunks: Optional[int] = None,
+    process_func: Optional[Callable] = None,
+    column: str = "text",
+) -> List[str]:
+    """
+    将 HuggingFace Dataset 分块导出为 JSONL 文件。
+
+    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 时使用）
+
+    Returns:
+        生成的文件路径列表
+    """
+    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")
+        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}")
+
+    return output_files

pipeline/io.py

@@ -0,0 +1,56 @@
+from pathlib import Path
+from typing import Dict, List
+import os
+import h5py
+import torch
+from torch import Tensor
+
+
+class IOHandler:
+    """文件和 HDF5 读写"""
+
+    @staticmethod
+    def fetch_files(directory: str) -> List[str]:
+        return [
+            os.path.join(root, f)
+            for root, _, files in os.walk(directory)
+            for f in files
+        ]
+
+    @staticmethod
+    def fetch_folders(root_dir: str, filter_func=None) -> List[str]:
+        folders = []
+        for root, dirs, _ in os.walk(root_dir):
+            for dir_name in dirs:
+                folder_path = os.path.join(root, dir_name)
+                if filter_func is None or filter_func(folder_path):
+                    folders.append(folder_path)
+        return folders
+
+    @staticmethod
+    def save_h5(file_path: str, file_name: str, tensor_group: Dict[str, List[Tensor]]) -> None:
+        os.makedirs(file_path, exist_ok=True)
+        full_path = os.path.join(file_path, f"{file_name}.h5")
+        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())
+
+    @staticmethod
+    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():
+                    grp = f[key]
+                    tensors = [
+                        (torch.from_numpy(dset[:]).share_memory_() if share_memory
+                         else torch.from_numpy(dset[:]))
+                        for dset_name in grp.keys()
+                        for dset in [grp[dset_name]]
+                    ]
+                    tensor_group.setdefault(key, []).extend(tensors)
+        return tensor_group

pipeline/packing.py

@@ -0,0 +1,35 @@
+from typing import List
+import torch
+from torch import Tensor
+
+
+class SequencePacker:
+    """序列打包（bin-packing）"""
+
+    def __init__(self, pack_size: int, pad_value: int = 0):
+        self.pack_size = pack_size
+        self.pad_value = pad_value
+
+    def pack(self, sequences: List[Tensor]) -> List[Tensor]:
+        packages = []
+        sequences.sort(key=lambda x: x.numel(), reverse=True)
+
+        current_pack = torch.full((self.pack_size,), self.pad_value, dtype=torch.int32)
+        current_pos = 0
+
+        for tensor in sequences:
+            tensor = tensor[:self.pack_size] if tensor.numel() > self.pack_size else tensor
+            tensor_size = tensor.numel()
+
+            if current_pos + tensor_size > self.pack_size:
+                packages.append(current_pack)
+                current_pack = torch.full((self.pack_size,), self.pad_value, dtype=torch.int32)
+                current_pos = 0
+
+            current_pack[current_pos:current_pos + tensor_size] = tensor
+            current_pos += tensor_size
+
+        if current_pos > 0:
+            packages.append(current_pack)
+
+        return packages

pipeline/processors.py

@@ -0,0 +1,89 @@
+from abc import ABC, abstractmethod
+from typing import Dict, List
+import torch
+
+
+class BaseProcessor(ABC):
+    """处理器抽象基类"""
+
+    @abstractmethod
+    def process(self, input_dict: dict) -> dict:
+        pass
+
+    @property
+    @abstractmethod
+    def output_keys(self) -> List[str]:
+        pass
+
+
+class PreTrainProcessor(BaseProcessor):
+    """预训练数据处理器"""
+
+    def __init__(self, tokenizer):
+        self.tokenizer = tokenizer
+
+    def process(self, input_dict: dict) -> dict:
+        segment = input_dict["text"]
+        tokens = self.tokenizer.encode(f"{segment}<eos>")
+        return {'sequence': torch.tensor(tokens, dtype=torch.int32)}
+
+    @property
+    def output_keys(self) -> List[str]:
+        return ["sequence"]
+
+
+class SFTProcessor(BaseProcessor):
+    """监督微调数据处理器"""
+
+    def __init__(self, tokenizer):
+        self.tokenizer = tokenizer
+
+    def process(self, input_dict: dict) -> dict:
+        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"
+        )
+        a = self.tokenizer.encode(f"{response}<|im_end|>\n<eos>")
+        tokens = torch.tensor(q + a, dtype=torch.int32)
+        loss_mask = torch.zeros_like(tokens, dtype=torch.bool)
+        loss_mask[len(q):] = True
+        return {"sequence": tokens, "loss_mask": loss_mask}
+
+    @property
+    def output_keys(self) -> List[str]:
+        return ["sequence", "loss_mask"]
+
+
+class DPOProcessor(BaseProcessor):
+    """DPO 偏好学习数据处理器"""
+
+    def __init__(self, tokenizer):
+        self.tokenizer = tokenizer
+
+    def process(self, input_dict: dict) -> dict:
+        # TODO: 实现 DPO 处理逻辑
+        return None
+
+    @property
+    def output_keys(self) -> List[str]:
+        return ["chosen", "chosen_mask", "rejected", "rejected_mask"]
+
+
+class ProcessorFactory:
+    """处理器工厂"""
+
+    _processors = {
+        "pt": PreTrainProcessor,
+        "sft": SFTProcessor,
+        "dpo": DPOProcessor,
+    }
+
+    @classmethod
+    def create(cls, processor_type: str, tokenizer) -> BaseProcessor:
+        if processor_type not in cls._processors:
+            raise ValueError(f"Invalid processor type: {processor_type}")
+        return cls._processors[processor_type](tokenizer)
+
+    @classmethod
+    def register(cls, processor_type: str, processor_class: type):
+        cls._processors[processor_type] = processor_class

pipeline/text.py

@@ -0,0 +1,21 @@
+import re
+from typing import Dict
+
+
+class TextNormalizer:
+    """文本规范化"""
+
+    DEFAULT_REPLACEMENTS = {
+        "\\[": "$$", "\\]": "$$", "\\(": "$", "\\)": "$",
+        '\u2018': "'", '\u2019': "'", '\u0060': "'",
+        '\u201C': '"', '\u201D': '"',
+        '\u2013': '-', '\u2014': '--', '\u2212': '-',
+        '\u00A0': ' ', '\u2026': '...'
+    }
+
+    def __init__(self, custom_rules: 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))
+
+    def normalize(self, text: str) -> str:
+        return self._pattern.sub(lambda m: self.replacements[m.group()], text)

pipeline/tokenizer.py

@@ -0,0 +1,106 @@
+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
+
+
+class BpeTokenizer:
+    def __init__(self, path=None):
+        self._control_tokens = ["<bos>", "<eos>", "<pad>"]
+        self._special_tokens = ["<|im_start|>", "<|im_end|>"]
+        
+        model = BPE()
+        self._tokenizer = Tokenizer(model)
+        self._tokenizer.normalizer = normalizers.Sequence([
+            normalizers.NFC(),
+            normalizers.Strip() 
+        ])
+        
+        self._tokenizer.pre_tokenizer = pre_tokenizers.Sequence([
+            pre_tokenizers.UnicodeScripts(),
+            pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=True)
+        ])
+        
+        self._tokenizer.decoder = decoders.ByteLevel()
+        self._tokenizer.post_processor = processors.ByteLevel(trim_offsets=True)
+        
+        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:
+        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))
+        
+        alphabet = pre_tokenizers.ByteLevel.alphabet()
+        min_size = len(alphabet) + len(self._control_tokens)
+        assert detail_vocab_size > min_size
+        
+        trainer = BpeTrainer(
+            vocab_size=detail_vocab_size,
+            min_frequency=min_freq,
+            limit_alphabet=detail_vocab_size // 6,
+            max_token_length=max_token_length,
+            special_tokens=self._control_tokens,
+            initial_alphabet=alphabet,
+            show_progress=True,
+        )
+        
+        return trainer, detail_vocab_size, reserved_tokens
+
+    def train(self, files, vocab_size, min_freq, reserved_token_size=100):
+        trainer, _, reserved_tokens = self._prepare_trainer(
+            vocab_size=vocab_size,
+            min_freq=min_freq,
+            reserved_token_size=reserved_token_size
+        )
+        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):
+        trainer, _, reserved_tokens = self._prepare_trainer(
+            vocab_size=vocab_size,
+            min_freq=min_freq,
+            reserved_token_size=reserved_token_size
+        )
+        self._tokenizer.train_from_iterator(iterator=iterator, trainer=trainer)
+        self._tokenizer.add_special_tokens(self._special_tokens + reserved_tokens)
+            
+    def save(self, path):
+        self._tokenizer.save(path)
+        
+    def load(self, path):
+        self._tokenizer = Tokenizer.from_file(path)
+
+    def encode(self, tokens: Union[str, List[str]], out_ids: bool=True, add_special_tokens: bool=False) -> List:
+        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
+        elif isinstance(tokens, list):
+            encoded_list: List[Encoding] = self._tokenizer.encode_batch(tokens, add_special_tokens=add_special_tokens)
+            return [encoded.ids if out_ids else encoded.tokens for encoded in encoded_list]
+
+    def decode(self, tokens: List[int], skip_special_tokens: bool=True) -> str:
+        return self._tokenizer.decode(tokens, skip_special_tokens=skip_special_tokens)
+    
+    def __len__(self) -> int:
+        return self._tokenizer.get_vocab_size()
+    
+    @property
+    def stop_ids(self) -> List[int]:
+        stop_token = self._control_tokens + self._special_tokens
+        stop_ids = [self._tokenizer.token_to_id(token) for token in stop_token]
+        return stop_ids
+    
+    @property
+    def bos_id(self) -> int:
+        return self._tokenizer.token_to_id("<bos>")
+    
+    @property
+    def eos_id(self) -> int:
+        return self._tokenizer.token_to_id("<eos>")
+    
+    @property
+    def pad_id(self) -> int:
+        return self._tokenizer.token_to_id("<pad>")