refactor: 重构数据管道模块

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)