diff --git a/README.md b/README.md index 1a75ddb..091055f 100644 --- a/README.md +++ b/README.md @@ -1,59 +1,174 @@ -## KHAOSZ-dataset +# DataPipeline -用于训练的KHAOSZ数据集 -### 项目结构 +用于训练KHAOSZ模型的数据集处理工具 -``` bash -. -│ .gitignore -│ dump_pt_file.py -│ dump_sft_file.py -│ README.md -│ run.py -│ tokenizer.json -│ -├───modules -│ │ tokenizer.py -│ └───utils.py -│ -├───pre_train -│ 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 +## 项目结构 + +``` +khaosz_dataset/ +├── modules/ # 核心模块 +│ ├── tokenizer.py # BPE Tokenizer +│ └── datapipeline/ # 数据管道模块 +│ ├── pipeline.py # 主数据管道 +│ ├── processors.py # 数据处理器(策略模式) +│ ├── io.py # 文件IO操作 +│ ├── packing.py # 序列打包 +│ └── text.py # 文本规范化 +├── tokenizer.json # Tokenizer配置 +└── pyproject.toml # 项目依赖 ``` +## 架构设计 -### 数据集特性 -- 支持多语言混合训练(中/英文) -- 包含以下预训练数据源: - - Chinese-C4 - - Chinese-Cosmopedia - - English-Fineweb - - English-Wiki -- 支持监督微调数据集: - - Ling-Coder-SFT - - Chinese-Instruct - - BelleGroup - - Magpie-Pro-300K +本项目采用模块化设计,应用了多种设计模式: + +### 设计模式 +1. **策略模式** - 不同类型的数据处理器(PT/SFT/DPO) +2. **工厂模式** - 处理器工厂统一创建实例 +3. **模板方法模式** - 数据管道流程标准化 + +### 核心组件 + +#### DataPipeline(数据管道) +主数据管道,负责数据集的分块处理、格式转换和存储。 + +```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") ``` \ No newline at end of file diff --git a/modules/__init__.py b/modules/__init__.py index e69de29..2c9ab1e 100644 --- a/modules/__init__.py +++ b/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' +] diff --git a/modules/datapipeline/__init__.py b/modules/datapipeline/__init__.py new file mode 100644 index 0000000..f38d5a1 --- /dev/null +++ b/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' +] diff --git a/modules/datapipeline/io.py b/modules/datapipeline/io.py new file mode 100644 index 0000000..30a060f --- /dev/null +++ b/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) diff --git a/modules/datapipeline/packing.py b/modules/datapipeline/packing.py new file mode 100644 index 0000000..7736ce4 --- /dev/null +++ b/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) diff --git a/modules/datapipeline/pipeline.py b/modules/datapipeline/pipeline.py new file mode 100644 index 0000000..3a588b1 --- /dev/null +++ b/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 + ) diff --git a/modules/datapipeline/processors.py b/modules/datapipeline/processors.py new file mode 100644 index 0000000..ae6b433 --- /dev/null +++ b/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}") + 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") + + 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 diff --git a/modules/datapipeline/text.py b/modules/datapipeline/text.py new file mode 100644 index 0000000..9e2030d --- /dev/null +++ b/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) diff --git a/modules/utils.py b/modules/utils.py deleted file mode 100644 index bc428fd..0000000 --- a/modules/utils.py +++ /dev/null @@ -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}") - 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") - 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}") diff --git a/pre_train/chinese-c4.py b/pre_train/chinese-c4.py index 76449c7..c6099e2 100644 --- a/pre_train/chinese-c4.py +++ b/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" ) \ No newline at end of file diff --git a/pre_train/chinese-cosmopedia.py b/pre_train/chinese-cosmopedia.py index b0d1805..927ac63 100644 --- a/pre_train/chinese-cosmopedia.py +++ b/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, diff --git a/pre_train/english-fineweb.py b/pre_train/english-fineweb.py index 23e7055..4097496 100644 --- a/pre_train/english-fineweb.py +++ b/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", ) \ No newline at end of file diff --git a/pre_train/english-wiki.py b/pre_train/english-wiki.py index 2e7ca2b..96f04c9 100644 --- a/pre_train/english-wiki.py +++ b/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, ) \ No newline at end of file diff --git a/reforce_learning/dpp_chinese_dpo_pairs.py b/reforce_learning/dpp_chinese_dpo_pairs.py index 2421dfa..7a2ca31 100644 --- a/reforce_learning/dpp_chinese_dpo_pairs.py +++ b/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 diff --git a/run.py b/run.py deleted file mode 100644 index b3729d4..0000000 --- a/run.py +++ /dev/null @@ -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() \ No newline at end of file diff --git a/supervised_finetuning/sft_belle.py b/supervised_finetuning/sft_belle.py index b01ff1f..34919b2 100644 --- a/supervised_finetuning/sft_belle.py +++ b/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, diff --git a/supervised_finetuning/sft_chinese_instruct.py b/supervised_finetuning/sft_chinese_instruct.py index ccbef6c..751a15b 100644 --- a/supervised_finetuning/sft_chinese_instruct.py +++ b/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) - resp = comprehensive_normalization(resp) + normalizer = TextNormalizer() + 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, diff --git a/supervised_finetuning/sft_coder.py b/supervised_finetuning/sft_coder.py index 0c6d8dd..3ca13a6 100644 --- a/supervised_finetuning/sft_coder.py +++ b/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 diff --git a/supervised_finetuning/sft_firefly-1.1m-rephrased.py b/supervised_finetuning/sft_firefly-1.1m-rephrased.py index 673850e..b09eaad 100644 --- a/supervised_finetuning/sft_firefly-1.1m-rephrased.py +++ b/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 diff --git a/supervised_finetuning/sft_magpie-pro-300k.py b/supervised_finetuning/sft_magpie-pro-300k.py index 50d24f2..0f67ba4 100644 --- a/supervised_finetuning/sft_magpie-pro-300k.py +++ b/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,