ViperEkura/DataPipeline
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: 重构数据管道模块

Browse Source
This commit is contained in:
ViperEkura
2026-03-20 21:36:37 +08:00
parent 7045de9c32
commit 3fb2326115
20 changed files with 739 additions and 345 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+162 -47
View File
@@ -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")
```
modules/__init__.py
+17
View File
@@ -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
+13
View File
@@ -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
+77
View File
@@ -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
+41
View File
@@ -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
+252
View File
@@ -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
+110
View File
@@ -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
+27
View File
@@ -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
-234
View File
@@ -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
+4 -2
View File
@@ -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
+3 -2
View File
@@ -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
+4 -2
View File
@@ -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
+5 -3
View File
@@ -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
+4 -2
View File
@@ -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
-41
View File
@@ -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
+4 -2
View File
@@ -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
+6 -4
View File
@@ -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,
supervised_finetuning/sft_coder.py
+3 -2
View File
@@ -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
+3 -2
View File
@@ -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
+4 -2
View File
@@ -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,