refactor: 重构数据管道模块

README.md

@@ -1,59 +1,174 @@
-## KHAOSZ-dataset
+# DataPipeline
 
-用于训练的KHAOSZ数据集
+用于训练KHAOSZ模型的数据集处理工具
-### 项目结构
 
-``` bash
+## 项目结构
-.
+
-│   .gitignore
+```
-│   dump_pt_file.py
+khaosz_dataset/
-│   dump_sft_file.py
+├── modules/                    # 核心模块
-│   README.md
+│   ├── tokenizer.py           # BPE Tokenizer
-│   run.py
+│   └── datapipeline/          # 数据管道模块
-│   tokenizer.json
+│       ├── pipeline.py        # 主数据管道
-│
+│       ├── processors.py      # 数据处理器（策略模式）
-├───modules
+│       ├── io.py             # 文件IO操作
-│   │   tokenizer.py
+│       ├── packing.py        # 序列打包
-│   └───utils.py
+│       └── text.py           # 文本规范化
-│
+├── tokenizer.json             # Tokenizer配置
-├───pre_train
+└── pyproject.toml            # 项目依赖
-│       chinese-c4.py
-│       chinese-cosmopedia.py
-│       english-fineweb.py
-│       english-wiki.py
-│
-├───reforce_learning
-│       dpp_chinese_dpo_pairs.py
-│
-└───supervised_finetuning
-        sft_belle.py
-        sft_chinese_instruct.py
-        sft_coder.py
-        sft_magpie-pro-300k.py
-        sft_small_talk.py
 ```
 
+## 架构设计
 
-### 数据集特性
+本项目采用模块化设计，应用了多种设计模式：
-- 支持多语言混合训练（中/英文）
+
-- 包含以下预训练数据源：
+### 设计模式
-  - Chinese-C4
+1. **策略模式** - 不同类型的数据处理器（PT/SFT/DPO）
-  - Chinese-Cosmopedia
+2. **工厂模式** - 处理器工厂统一创建实例
-  - English-Fineweb
+3. **模板方法模式** - 数据管道流程标准化
-  - English-Wiki
+
-- 支持监督微调数据集：
+### 核心组件
-  - Ling-Coder-SFT
+
-  - Chinese-Instruct
+#### DataPipeline（数据管道）
-  - BelleGroup
+主数据管道，负责数据集的分块处理、格式转换和存储。
-  - Magpie-Pro-300K
+
+```python
+from modules.datapipeline import DataPipeline
+
+pipeline = DataPipeline(output_dir="./dataset")
+pipeline.process_dataset(
+    dataset_dict=dataset,
+    output_subdir="my-data"
+)
+```
+
+#### ProcessorFactory（处理器工厂）
+创建不同类型的数据处理器。
+
+```python
+from modules.datapipeline import ProcessorFactory
+from modules.tokenizer import BpeTokenizer
+
+tokenizer = BpeTokenizer("tokenizer.json")
+
+# 创建预训练处理器
+processor = ProcessorFactory.create("pt", tokenizer)
+
+# 创建SFT处理器
+processor = ProcessorFactory.create("sft", tokenizer)
+```
+
+#### TextNormalizer（文本规范化）
+文本预处理和规范化。
+
+```python
+from modules.datapipeline import TextNormalizer
+
+normalizer = TextNormalizer()
+normalized_text = normalizer.normalize(text)
+```
+
+## 使用说明
+
+### 安装依赖
 
-### 使用说明
-1. 安装依赖：
 ```bash
-pip install datasets tokenizers tqdm torch
+pip install datasets tokenizers tqdm torch h5py
 ```
-运行数据处理：
+
+### 运行数据处理
+
+运行所有数据处理脚本：
 
 ```bash
 python run.py
 ```
+
+运行特定脚本：
+
+```bash
+# 预训练数据处理
+python pre_train/english-wiki.py
+
+# SFT数据处理
+python supervised_finetuning/sft_belle.py
+
+# DPO数据处理
+python reforce_learning/dpp_chinese_dpo_pairs.py
+```
+
+### 自定义数据处理
+
+#### 1. 基础数据处理
+
+```python
+from datasets import load_dataset
+from modules.datapipeline import DataPipeline
+
+# 加载数据集
+dataset = load_dataset("your-dataset")
+
+# 创建管道
+pipeline = DataPipeline()
+
+# 处理数据
+pipeline.process_dataset(
+    dataset_dict=dataset,
+    output_subdir="output-dir",
+    process_func=lambda x: {"text": x["content"]}
+)
+```
+
+#### 2. 自定义处理器
+
+```python
+from modules.datapipeline.processors import BaseProcessor
+from modules.datapipeline import ProcessorFactory
+
+class MyProcessor(BaseProcessor):
+    def __init__(self, tokenizer):
+        self.tokenizer = tokenizer
+    
+    def process(self, input_dict: dict) -> dict:
+        # 自定义处理逻辑
+        return {"processed": data}
+    
+    @property
+    def output_keys(self) -> list:
+        return ["processed"]
+
+# 注册处理器
+ProcessorFactory.register("my_type", MyProcessor)
+```
+
+#### 3. 文本规范化
+
+```python
+from modules.datapipeline import TextNormalizer
+
+# 使用默认规则
+normalizer = TextNormalizer()
+text = normalizer.normalize(text)
+
+# 自定义规则
+custom_rules = {"旧词": "新词"}
+normalizer = TextNormalizer(custom_rules)
+```
+
+## 数据输出格式
+
+### JSONL格式
+每个数据块保存为JSONL文件：
+```
+{"text": "训练文本内容..."}
+{"query": "问题", "response": "答案"}
+```
+
+### H5格式
+打包后的张量数据保存为HDF5格式，支持高效加载：
+```python
+from modules.datapipeline import IOHandler
+
+# 加载H5数据
+data = IOHandler.load_h5("./cached_data")
+```

modules/__init__.py

@@ -0,0 +1,17 @@
+from .tokenizer import BpeTokenizer
+from .datapipeline import (
+    DataPipeline,
+    ProcessorFactory,
+    IOHandler,
+    TextNormalizer,
+    SequencePacker
+)
+
+__all__ = [
+    'BpeTokenizer',
+    'DataPipeline',
+    'ProcessorFactory',
+    'IOHandler',
+    'TextNormalizer',
+    'SequencePacker'
+]

modules/datapipeline/__init__.py

@@ -0,0 +1,13 @@
+from .pipeline import DataPipeline
+from .processors import ProcessorFactory
+from .io import IOHandler
+from .text import TextNormalizer
+from .packing import SequencePacker
+
+__all__ = [
+    'DataPipeline',
+    'ProcessorFactory', 
+    'IOHandler',
+    'TextNormalizer',
+    'SequencePacker'
+]

modules/datapipeline/io.py

@@ -0,0 +1,77 @@
+from pathlib import Path
+from typing import Dict, List, Union
+import os
+import h5py
+import torch
+from torch import Tensor
+
+
+class IOHandler:
+    """文件和H5数据存储处理器"""
+    
+    @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:
+        """保存张量组到H5文件"""
+        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]]:
+        """从H5文件加载张量组"""
+        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
+
+
+# 向后兼容的函数接口
+def fetch_files(directory: str) -> List[str]:
+    return IOHandler.fetch_files(directory)
+
+def fetch_folders(root_dir: str, filter_func=None) -> List[str]:
+    return IOHandler.fetch_folders(root_dir, filter_func)
+
+def save_h5(file_path: str, file_name: str, tensor_group: Dict[str, List[Tensor]]) -> None:
+    return IOHandler.save_h5(file_path, file_name, tensor_group)
+
+def load_h5(file_path: str, share_memory: bool = True) -> Dict[str, List[Tensor]]:
+    return IOHandler.load_h5(file_path, share_memory)

modules/datapipeline/packing.py

@@ -0,0 +1,41 @@
+from typing import List
+import torch
+from torch import Tensor
+
+
+class SequencePacker:
+    """序列打包策略"""
+    
+    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
+
+
+def pack_sequences(sequences: List[Tensor], pack_size: int, pad_value: int) -> List[Tensor]:
+    """向后兼容的函数接口"""
+    return SequencePacker(pack_size, pad_value).pack(sequences)

modules/datapipeline/pipeline.py

@@ -0,0 +1,252 @@
+from typing import Dict, List, Callable, Union
+from datasets import DatasetDict
+from tqdm import tqdm
+import torch
+import json
+import os
+
+from .processors import ProcessorFactory
+from .packing import SequencePacker
+from .io import IOHandler
+from .text import TextNormalizer
+
+
+class DataPipeline:
+    """数据处理管道 - 模板方法模式"""
+    
+    def __init__(self, output_dir: str = None):
+        self.output_dir = output_dir or os.path.join(os.getcwd(), "dataset")
+    
+    def process_dataset(
+        self,
+        dataset_dict: DatasetDict,
+        output_subdir: str,
+        max_chunk_num: int = None,
+        chunk_size: int = 1000000,
+        split_name: str = "train",
+        column_name: str = "text",
+        process_func: Union[Callable[[dict], dict], Callable[[List[dict]], List[dict]]] = None,
+        normalization_func: Callable[[str], str] = None,
+        output_dir: str = None,
+    ) -> None:
+        """处理数据集的主流程"""
+        dataset = dataset_dict[split_name]
+        total_samples = len(dataset)
+        num_chunks = (total_samples // chunk_size) + 1
+        lim_chunks = min(max_chunk_num, num_chunks) if max_chunk_num else num_chunks
+        
+        output_dir = output_dir or os.path.join(self.output_dir, output_subdir)
+        os.makedirs(output_dir, exist_ok=True)
+        
+        # 处理每个数据块
+        for i in range(lim_chunks):
+            self._process_chunk(
+                dataset=dataset,
+                chunk_idx=i,
+                chunk_size=chunk_size,
+                total_samples=total_samples,
+                output_dir=output_dir,
+                output_subdir=output_subdir,
+                column_name=column_name,
+                process_func=process_func,
+                normalization_func=normalization_func
+            )
+    
+    def _process_chunk(
+        self,
+        dataset,
+        chunk_idx: int,
+        chunk_size: int,
+        total_samples: int,
+        output_dir: str,
+        output_subdir: str,
+        column_name: str,
+        process_func,
+        normalization_func
+    ) -> None:
+        """处理单个数据块"""
+        start_idx = chunk_idx * chunk_size
+        end_idx = min((chunk_idx + 1) * chunk_size, total_samples)
+        chunk = dataset.select(range(start_idx, end_idx))
+        
+        output_path = os.path.join(output_dir, f"{output_subdir}_text_chunk_{chunk_idx}.jsonl")
+        
+        with open(output_path, "w", encoding="utf-8") as f:
+            for example in chunk:
+                processed = self._process_example(
+                    example, column_name, process_func, normalization_func
+                )
+                self._write_processed(f, processed)
+        
+        print(f"Saved text chunk {chunk_idx} to {output_path}")
+    
+    def _process_example(
+        self,
+        example: dict,
+        column_name: str,
+        process_func,
+        normalization_func
+    ) -> Union[dict, List[dict]]:
+        """处理单个样本"""
+        if process_func:
+            return process_func(example)
+        
+        text = example[column_name]
+        if normalization_func:
+            text = normalization_func(text)
+        return {column_name: text}
+    
+    def _write_processed(self, file, processed) -> None:
+        """写入处理后的数据"""
+        if isinstance(processed, dict):
+            file.write(json.dumps(processed, ensure_ascii=False) + "\n")
+        elif isinstance(processed, list):
+            for item in processed:
+                file.write(json.dumps(item, ensure_ascii=False) + "\n")
+    
+    def cache_files(
+        self,
+        tokenizer,
+        files: List[str],
+        base_out_dir: str,
+        cache_type: str,
+        packing_size: int = -1,
+        pad_value: int = 1
+    ) -> None:
+        """缓存文件到H5格式"""
+        processor = ProcessorFactory.create(cache_type, tokenizer)
+        self.dump_files(
+            files=files,
+            base_out_dir=base_out_dir,
+            process_func=processor.process,
+            output_keys=processor.output_keys,
+            packing_size=packing_size,
+            pad_value=pad_value
+        )
+    
+    def dump_files(
+        self,
+        files: List[str],
+        base_out_dir: str,
+        process_func: Callable[[dict], dict],
+        output_keys: List[str],
+        packing_size: int = -1,
+        pad_value: int = 0
+    ) -> None:
+        """转储文件到H5格式"""
+        for file_path in files:
+            self._dump_single_file(
+                file_path=file_path,
+                base_out_dir=base_out_dir,
+                process_func=process_func,
+                output_keys=output_keys,
+                packing_size=packing_size,
+                pad_value=pad_value
+            )
+    
+    def _dump_single_file(
+        self,
+        file_path: str,
+        base_out_dir: str,
+        process_func: Callable[[dict], dict],
+        output_keys: List[str],
+        packing_size: int,
+        pad_value: int
+    ) -> None:
+        """转储单个文件"""
+        os.makedirs(base_out_dir, exist_ok=True)
+        file_name = os.path.basename(file_path)
+        out_file_name = file_name.split(".")[0]
+        
+        # 读取和处理数据
+        with open(file_path, "r") as f:
+            lines = f.readlines()
+        
+        arrows: List[Dict[str, torch.Tensor]] = []
+        for line in tqdm(lines, desc=f"Processing {file_name}", leave=False):
+            line_dict = json.loads(line)
+            arrow = process_func(line_dict)
+            if arrow is not None:
+                arrows.append(arrow)
+        
+        # 组织输出数据
+        package: Dict[str, List[torch.Tensor]] = {
+            key: [arrow[key] for arrow in arrows]
+            for key in output_keys
+        }
+        
+        # 打包序列（如果需要）
+        output_package = {}
+        for key in output_keys:
+            if packing_size > 0:
+                print(f"Packaging key: '{key}'")
+                packer = SequencePacker(packing_size, pad_value)
+                output_package[key] = packer.pack(package[key])
+            else:
+                output_package[key] = package[key]
+        
+        # 保存到H5
+        IOHandler.save_h5(base_out_dir, out_file_name, output_package)
+
+
+# 向后兼容的函数接口
+def process_dataset(
+    dataset_dict: DatasetDict,
+    output_subdir: str,
+    max_chunk_num: int = None,
+    chunk_size: int = 1000000,
+    split_name: str = "train",
+    column_name: str = "text",
+    process_func: Union[Callable[[dict], dict], Callable[[List[dict]], List[dict]]] = None,
+    normalization_func: Callable[[str], str] = None,
+    output_dir: str = None,
+) -> None:
+    """向后兼容的函数接口"""
+    pipeline = DataPipeline(output_dir)
+    normalizer = TextNormalizer() if normalization_func is None else None
+    norm_func = normalization_func or (normalizer.normalize if normalizer else None)
+    
+    return pipeline.process_dataset(
+        dataset_dict=dataset_dict,
+        output_subdir=output_subdir,
+        max_chunk_num=max_chunk_num,
+        chunk_size=chunk_size,
+        split_name=split_name,
+        column_name=column_name,
+        process_func=process_func,
+        normalization_func=norm_func,
+        output_dir=output_dir
+    )
+
+
+def cache_files(tokenizer, files, base_out_dir, cache_type, packing_size: int = -1, pad_value: int = 1):
+    """向后兼容的函数接口"""
+    pipeline = DataPipeline()
+    return pipeline.cache_files(
+        tokenizer=tokenizer,
+        files=files,
+        base_out_dir=base_out_dir,
+        cache_type=cache_type,
+        packing_size=packing_size,
+        pad_value=pad_value
+    )
+
+
+def dump_files(
+    files: List[str],
+    base_out_dir: str,
+    process_func: Callable[[dict], dict],
+    output_keys: List[str],
+    packing_size: int = -1,
+    pad_value: int = 0
+):
+    """向后兼容的函数接口"""
+    pipeline = DataPipeline()
+    return pipeline.dump_files(
+        files=files,
+        base_out_dir=base_out_dir,
+        process_func=process_func,
+        output_keys=output_keys,
+        packing_size=packing_size,
+        pad_value=pad_value
+    )

modules/datapipeline/processors.py

@@ -0,0 +1,110 @@
+from abc import ABC, abstractmethod
+from typing import Dict, List, Callable
+import torch
+from torch import Tensor
+
+
+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
+
+
+# 向后兼容的函数接口
+def get_pt_processor(tokenizer) -> Callable:
+    processor = PreTrainProcessor(tokenizer)
+    return processor.process
+
+def get_sft_processor(tokenizer) -> Callable:
+    processor = SFTProcessor(tokenizer)
+    return processor.process
+
+def get_dpo_processor(tokenizer) -> Callable:
+    processor = DPOProcessor(tokenizer)
+    return processor.process

modules/datapipeline/text.py

@@ -0,0 +1,27 @@
+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)
+
+
+def comprehensive_normalization(text: str) -> str:
+    """向后兼容的函数接口"""
+    return TextNormalizer().normalize(text)

modules/utils.py

@@ -1,234 +0,0 @@
-from pathlib import Path
-from typing import Dict, List, Callable, Union
-from datasets import DatasetDict
-from modules.tokenizer import BpeTokenizer
-from tqdm import tqdm
-from torch import Tensor
-
-import torch
-import h5py
-import json
-import os
-import re
-
-
-def fetch_files(directory):
-    return [os.path.join(root, f) 
-            for root, _, files in os.walk(directory) for f in files]
-
-def fetch_folders(root_dir, filter_func=None):
-    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
-
-def save_h5(file_path: str, file_name: str, tensor_group: Dict[str, List[Tensor]]):
-    os.makedirs(file_path, exist_ok=True)
-    full_file_path = os.path.join(file_path, f"{file_name}.h5")
-    with h5py.File(full_file_path, 'w') as f:
-        for key, tensors in tensor_group.items():
-            grp = f.create_group(key)
-            for idx, tensor in enumerate(tensors):
-                arr = tensor.cpu().numpy()
-                grp.create_dataset(f'data_{idx}', data=arr)
-
-def load_h5(file_path: str, share_memory=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]
-                dsets = []
-                for dset_name in grp.keys():
-                    dset = grp[dset_name]
-                    tensor = torch.from_numpy(dset[:])
-                    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
-
-def comprehensive_normalization(text):
-    replacements = {
-        "\\[": "$$", "\\]": "$$", "\\(": "$", "\\)": "$",
-        '\u2018': "'", '\u2019': "'", '\u0060': "'", '\u201C': '"', '\u201D': '"', 
-        '\u2013': '-', '\u2014': '--', '\u2212': '-', '\u00A0': ' ', '\u2026': '...'
-    }
-    pattern = re.compile('|'.join(re.escape(k) for k in replacements))
-    return pattern.sub(lambda m: replacements[m.group()], text) 
-
-
-def pack_sequences(sequences: List[Tensor], pack_size: int, pad_value: int) -> List[Tensor]:
-    packages = []
-    sequences.sort(key=lambda x: x.numel(), reverse=True)
-    current_pack = torch.full((pack_size,), pad_value, dtype=torch.int32)
-    current_pos = 0
-    
-    for tensor in sequences:
-        if tensor.numel() > pack_size:
-            tensor = tensor[:pack_size]
-        
-        tensor_size = tensor.numel()
-        
-        if current_pos + tensor_size > pack_size:
-            packages.append(current_pack)
-            current_pack = torch.full((pack_size,), 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
-
-
-def dump_files(
-    files: List[str], 
-    base_out_dir: str,
-    process_func: Callable[[dict], dict],
-    output_keys: List[str],
-    packing_size: int = -1,
-    pad_value: int = 0
-):
-        
-    for file_path in files:
-        os.makedirs(base_out_dir, exist_ok=True)
-        file_name = os.path.basename(file_path)
-        out_file_name = file_name.split(".")[0]
-        
-        arrows: List[Dict[str, Tensor]] = []
-        
-        with open(file_path, "r") as f:    
-            lines = f.readlines()
-            
-        for line in tqdm(lines, desc=f"Processing {file_name}", leave=False):
-            line_dict = json.loads(line)
-            arrow = process_func(line_dict)
-            arrows.append(arrow)
-        
-        package: Dict[str, List[Tensor]] = {}
-        for key in output_keys:
-            list_tensor = [arrow[key] for arrow in arrows]
-            package[key] = list_tensor
-            
-        output_package: Dict[str, List[Tensor]] = {}
-        
-        for key in output_keys:
-            if packing_size > 0:
-                print(f"Packaging key: '{key}'")
-                sequence = pack_sequences(package[key], packing_size, pad_value)
-            else:
-                sequence = package[key]
-                
-            output_package[key] = sequence
-         
-        save_h5(base_out_dir, out_file_name, output_package)
-
-
-def get_pt_processor(tokenizer: BpeTokenizer):
-    def processor(intput_dict: dict) -> dict:
-        segment = intput_dict["text"]
-        tokens = tokenizer.encode(f"{segment}<eos>")
-        tokens = torch.tensor(tokens, dtype=torch.int32)
-        
-        return {'sequence': tokens}
-
-    return processor
-
-
-def get_sft_processor(tokenizer: BpeTokenizer):
-    def processor(input_dict: dict):
-        query, response = input_dict["query"], input_dict["response"]
-        q = tokenizer.encode(f"<|im_start|>user\n{query}<|im_end|>\n<|im_start|>assistant\n")
-        a = 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}
-    
-    return processor
-
-def get_dpo_processor(tokenizer: BpeTokenizer):
-    def processor(input_dict: dict):
-        return None
-    
-    return processor
-        
-
-def cache_files(tokenizer, files, base_out_dir, cache_type, packing_size: int = -1, pad_value: int = 1):
-    processor = None
-    keys = []
-    if cache_type == "pt":
-        processor = get_pt_processor(tokenizer)
-        keys = ["sequence"]
-    elif cache_type == "sft":
-        processor = get_sft_processor(tokenizer)
-        keys = ["sequence", "loss_mask"]
-    elif cache_type == "dpo":
-        processor = get_dpo_processor(tokenizer)
-        keys = ["chosen", "chosen_mask", "rejected", "rejected_mask"]
-    else:
-        raise ValueError("Invalid cache type")
-    
-    dump_files(files, base_out_dir, processor, keys, packing_size, pad_value)
-                          
-            
-def process_dataset(
-    dataset_dict: DatasetDict,
-    output_subdir: str,
-    max_chunk_num: int = None,
-    chunk_size: int = 1000000,
-    split_name: str = "train",
-    column_name: str = "text",
-    process_func: Union[Callable[[dict], dict], Callable[[List[dict]], List[dict]]] = None,
-    normalization_func=comprehensive_normalization,
-    output_dir: str = None,
-):
-    train_dataset = dataset_dict[split_name]
-    total_samples = len(train_dataset)
-    num_chunks = (total_samples // chunk_size) + 1
-    lim_chunks = min(max_chunk_num, num_chunks) if max_chunk_num else num_chunks
-
-    if output_dir is None:
-        output_dir = os.path.join(os.getcwd(), "dataset", output_subdir)
-        
-    os.makedirs(output_dir, exist_ok=True)
-
-    for i in range(lim_chunks):
-        start_idx = i * chunk_size
-        end_idx = min((i + 1) * chunk_size, total_samples)
-        chunk = train_dataset.select(range(start_idx, end_idx))
-
-        output_path = os.path.join(output_dir, f"{output_subdir}_text_chunk_{i}.jsonl")
-        with open(output_path, "w", encoding="utf-8") as f:
-            for example in chunk:
-                if process_func:
-                    processed_example = process_func(example)
-                else:
-                    text = example[column_name]
-                    if normalization_func:
-                        text = normalization_func(text)
-                    processed_example = {column_name: text}
-                
-                if isinstance(processed_example, dict):
-                    f.write(json.dumps(processed_example, ensure_ascii=False) + "\n")
-                elif isinstance(processed_example, list):
-                    for item in processed_example:
-                        f.write(json.dumps(item, ensure_ascii=False) + "\n")
-                        
-
-        print(f"Saved text chunk {i} to {output_path}")

pre_train/chinese-c4.py

@@ -1,9 +1,11 @@
 from datasets import load_dataset
-from modules.utils import process_dataset
+from modules.datapipeline import DataPipeline
 
 if __name__ == "__main__":
     dataset = load_dataset("shjwudp/chinese-c4")
-    process_dataset(
+    
+    pipeline = DataPipeline()
+    pipeline.process_dataset(
         dataset_dict=dataset,
         output_subdir="chinese-c4-pretrain"
     )

pre_train/chinese-cosmopedia.py

@@ -1,5 +1,5 @@
 from datasets import load_dataset
-from modules.utils import process_dataset
+from modules.datapipeline import DataPipeline
 
 if __name__ == "__main__":
     chunk_size = 1000000
@@ -9,7 +9,8 @@ if __name__ == "__main__":
         data_files={"train": [f"data/000{i:02d}.parquet" for i in range(25)]}
     )
     
-    process_dataset(
+    pipeline = DataPipeline()
+    pipeline.process_dataset(
         dataset_dict=dataset,
         output_subdir="chinese-wiki-pretrain",
         chunk_size=chunk_size,

pre_train/english-fineweb.py

@@ -1,9 +1,11 @@
 from datasets import load_dataset
-from modules.utils import process_dataset
+from modules.datapipeline import DataPipeline
 
 if __name__ == "__main__":
     dataset = load_dataset("HuggingFaceFW/fineweb", "sample-10BT")
-    process_dataset(
+    
+    pipeline = DataPipeline()
+    pipeline.process_dataset(
         dataset_dict=dataset,
         output_subdir="english-fineweb-pretrain",
     )

pre_train/english-wiki.py

@@ -1,10 +1,12 @@
 from datasets import load_dataset
-from modules.utils import process_dataset
+from modules.datapipeline import DataPipeline
 
 if __name__ == "__main__":
     dataset = load_dataset("Blaze7451/enwiki_structured_content")
-    process_dataset(
+    
+    pipeline = DataPipeline()
+    pipeline.process_dataset(
         dataset_dict=dataset,
         output_subdir="english-wiki-pretrain",
-        max_chunk_size=5,
+        max_chunk_num=5,
     )

reforce_learning/dpp_chinese_dpo_pairs.py

@@ -1,5 +1,5 @@
 from datasets import load_dataset
-from modules.utils import process_dataset
+from modules.datapipeline import DataPipeline
 
 def process_func(input_dict: dict):
     return {
@@ -11,7 +11,9 @@ def process_func(input_dict: dict):
 
 if __name__ == "__main__":
     dataset = load_dataset("wenbopan/Chinese-dpo-pairs")
-    process_dataset(
+    
+    pipeline = DataPipeline()
+    pipeline.process_dataset(
         dataset_dict=dataset,
         output_subdir="Chinese-dpo-pairs",
         process_func=process_func

run.py

@@ -1,41 +0,0 @@
-import os
-import sys
-import subprocess
-
-PROJECT_ROOT = os.path.dirname(os.path.abspath(__file__))
-sys.path.append(PROJECT_ROOT)
-
-def run_script(script_path):
-    if not os.path.exists(script_path):
-        print(f"[Warning] File does not exist: {script_path}")
-        return
-
-    print(f"\n{'='*50}")
-    print(f"Running: {script_path}")
-    print(f"{'='*50}")
-
-    try:
-        env = os.environ.copy()
-        env['PYTHONPATH'] = PROJECT_ROOT
-        subprocess.run(
-            [sys.executable, script_path], 
-            check=True, 
-            cwd=PROJECT_ROOT, 
-            env=env
-        )
-    except subprocess.CalledProcessError as e:
-        print(f"[Error] Script execution failed: {script_path}, Error code: {e.returncode}")
-        
-def run_scripts(project_root: str, directory: str):
-    pre_train_dir = os.path.join(project_root, directory)
-    for file in os.listdir(pre_train_dir):
-        if file.endswith('.py'):
-            script_path = os.path.join(pre_train_dir, file)
-            run_script(script_path)
-
-def main():
-    run_scripts(PROJECT_ROOT, 'pre_train')
-    run_scripts(PROJECT_ROOT, 'supervised_finetuning')
-
-if __name__ == "__main__":
-    main()

supervised_finetuning/sft_belle.py

@@ -1,5 +1,5 @@
 from datasets import load_dataset
-from modules.utils import process_dataset
+from modules.datapipeline import DataPipeline
 
 
 
@@ -21,7 +21,9 @@ def process_func(input_dict: dict):
 
 if __name__ == "__main__":
     dataset = load_dataset("BelleGroup/train_3.5M_CN")
-    process_dataset(
+    
+    pipeline = DataPipeline()
+    pipeline.process_dataset(
         dataset_dict=dataset,
         output_subdir="belle-sft",
         process_func=process_func,

supervised_finetuning/sft_chinese_instruct.py

@@ -1,14 +1,15 @@
 from datasets import DatasetDict
 from datasets import load_dataset, concatenate_datasets
-from modules.utils import process_dataset, comprehensive_normalization
+from modules.datapipeline import DataPipeline, TextNormalizer
 
 
 def process_func(input_dict: dict):
     query = input_dict["prompt"] if input_dict["prompt"] else ""
     resp = input_dict["response"] if input_dict["response"] else ""
     
-    query = comprehensive_normalization(query)
+    normalizer = TextNormalizer()
-    resp = comprehensive_normalization(resp)
+    query = normalizer.normalize(query)
+    resp = normalizer.normalize(resp)
     
     return {"query": query, "response": resp }
 
@@ -25,7 +26,8 @@ if __name__ == "__main__":
     
     combined_dataset = concatenate_datasets(datasets)
     
-    process_dataset(
+    pipeline = DataPipeline()
+    pipeline.process_dataset(
         dataset_dict=DatasetDict({"train": combined_dataset}),
         output_subdir="chinese-instruct-sft",
         process_func=process_func,

supervised_finetuning/sft_coder.py

@@ -1,6 +1,6 @@
 # inclusionAI/Ling-Coder-SFT
 from datasets import load_dataset
-from modules.utils import process_dataset
+from modules.datapipeline import DataPipeline
 
 
 def process_func(input_dict: dict) -> dict:
@@ -13,7 +13,8 @@ def process_func(input_dict: dict) -> dict:
 if __name__ == "__main__":
     dataset = load_dataset("inclusionAI/Ling-Coder-SFT")
     
-    process_dataset(
+    pipeline = DataPipeline()
+    pipeline.process_dataset(
         dataset_dict=dataset,
         output_subdir="Ling-Coder-sft",
         process_func=process_func

supervised_finetuning/sft_firefly-1.1m-rephrased.py

@@ -1,6 +1,6 @@
 # Mxode/Firefly-1.1M-Rephrased
 from datasets import load_dataset
-from modules.utils import process_dataset
+from modules.datapipeline import DataPipeline
 
 def process_func(input_dict: dict):
     instruction = input_dict["instruction"]
@@ -11,7 +11,8 @@ def process_func(input_dict: dict):
 if __name__ == "__main__":
     dataset = load_dataset("Mxode/Firefly-1.1M-Rephrased")
     
-    process_dataset(
+    pipeline = DataPipeline()
+    pipeline.process_dataset(
         dataset_dict=dataset,
         output_subdir="Firefly-1.1M-Rephrased",
         process_func=process_func

supervised_finetuning/sft_magpie-pro-300k.py

@@ -1,6 +1,6 @@
 # HuggingFaceTB/Magpie-Pro-300K-Filtered-H4
 from datasets import load_dataset
-from modules.utils import process_dataset
+from modules.datapipeline import DataPipeline
 
 
 def process_func(input_dict: dict):
@@ -22,7 +22,9 @@ def process_func(input_dict: dict):
 
 if __name__ == "__main__":
     dataset = load_dataset("HuggingFaceTB/Magpie-Pro-300K-Filtered-H4")
-    process_dataset(
+    
+    pipeline = DataPipeline()
+    pipeline.process_dataset(
         dataset_dict=dataset,
         output_subdir="Magpie-Pro-300K-sft",
         process_func=process_func,