reafactor: 重构项目
This commit is contained in:
@@ -1,111 +1,188 @@
|
||||
# DataPipeline
|
||||
|
||||
数据集处理工具，支持预训练 / SFT / DPO 三种训练范式。
|
||||
语言模型训练数据集处理工具包。支持 PT / SFT / DPO 范式。
|
||||
|
||||
## 项目结构
|
||||
|
||||
```
|
||||
pipeline/
|
||||
├── tokenizer.py # BPE 分词器
|
||||
├── text.py # 文本规范化
|
||||
├── packing.py # 序列打包
|
||||
├── io.py # 文件/HDF5 读写
|
||||
├── processors.py # PT / SFT / DPO 处理器
|
||||
├── export.py # Dataset → JSONL
|
||||
└── cache.py # JSONL → Tokenize → H5
|
||||
├── tokenizer.py # BPE tokenizer
|
||||
├── text.py # Text normalization
|
||||
├── packing.py # Sequence packing
|
||||
├── io.py # File / HDF5 I/O
|
||||
├── processors.py # PT / SFT / DPO processors
|
||||
├── export.py # Dataset -> JSONL
|
||||
├── cache.py # JSONL -> Tokenize -> H5
|
||||
├── utils.py # Logging, error handling
|
||||
└── strategies/ # Prompt strategy (strategy pattern)
|
||||
├── base.py # PromptStrategy ABC
|
||||
├── chatml.py # ChatML format (configurable tokens)
|
||||
├── alpaca.py # Alpaca format (configurable tokens)
|
||||
└── factory.py # StrategyFactory
|
||||
```
|
||||
|
||||
## 设计理念
|
||||
## 数据流
|
||||
|
||||
模块**独立可用**，通过磁盘文件解耦，按需组合：
|
||||
整体分为两个阶段，通过磁盘 JSONL 文件解耦：
|
||||
|
||||
```
|
||||
Dataset → export_dataset() → JSONL → cache_jsonl() → HDF5
|
||||
↑
|
||||
processors.py
|
||||
packing.py
|
||||
io.py
|
||||
Stage 1: Export Dataset Stage 2: Tokenize & Cache
|
||||
┌──────────────────────────────────────┐ ┌──────────────────────────────────────┐
|
||||
│ │ │ │
|
||||
│ HuggingFace Dataset │ │ JSONL files (one dict per line) │
|
||||
│ │ │ │ │ │
|
||||
│ ▼ process_func (optional) │ │ ▼ json.loads() │
|
||||
│ export_dataset() │ │ Processor.process(input_dict) │
|
||||
│ │ chunk by chunk_size │ │ │ │
|
||||
│ ▼ │ │ │ ┌───────────────────────┐ │
|
||||
│ ./dataset/prefix_chunk_0.jsonl ────────────────|> │ PT / SFT / DPO │ │
|
||||
│ ./dataset/prefix_chunk_1.jsonl │ │ │ │ processor details │ │
|
||||
│ ... │ │ │ └───────────────────────┘ │
|
||||
│ │ │ ▼ │
|
||||
│ modules: export.py │ │ List[Tensor] │
|
||||
│ │ │ │ │
|
||||
│ │ │ ▼ SequencePacker (optional) │
|
||||
│ │ │ Fixed-length packed tensors │
|
||||
│ │ │ │ │
|
||||
│ │ │ ▼ IOHandler.save_h5() │
|
||||
│ │ │ ./cached/chunk_0.h5 │
|
||||
│ │ │ ./cached/chunk_1.h5 │
|
||||
│ │ │ ... │
|
||||
│ │ │ │
|
||||
│ │ │ modules: io.py, processors.py, │
|
||||
│ │ │ packing.py, tokenizer.py │
|
||||
└──────────────────────────────────────┘ └──────────────────────────────────────┘
|
||||
```
|
||||
|
||||
## 使用方法
|
||||
### Processor 转换规则
|
||||
|
||||
### 1. 导出数据集
|
||||
**PT (Pre-training)**
|
||||
```
|
||||
Input: {"text": "Hello world"}
|
||||
Action: tokenizer.encode(text + "<eos>")
|
||||
Output: {"sequence": Tensor[int32]}
|
||||
```
|
||||
|
||||
**SFT (Supervised Fine-tuning)** — uses PromptStrategy
|
||||
```
|
||||
Input: {"query": "...", "response": "..."}
|
||||
Action:
|
||||
1. strategy.build_prompt(input_dict) -> "<|im_start|>user\n...\n<|im_start|>assistant\n"
|
||||
2. concat response + response_suffix
|
||||
3. tokenizer.encode full string
|
||||
4. build loss_mask: query part=False, response part=True
|
||||
Output: {"sequence": Tensor[int32], "loss_mask": Tensor[bool]}
|
||||
```
|
||||
|
||||
**DPO (Direct Preference Optimization)** — uses PromptStrategy
|
||||
```
|
||||
Input: {"query": "...", "chosen": "...", "rejected": "..."}
|
||||
Action:
|
||||
1. strategy.build_prompt(input_dict) -> shared query prompt
|
||||
2. encode chosen = query + response_start + chosen + response_suffix
|
||||
3. encode rejected = query + response_start + rejected + response_suffix
|
||||
4. build masks: query part=False, response part=True
|
||||
Output: {"chosen": Tensor, "chosen_mask": Tensor[bool],
|
||||
"rejected": Tensor, "rejected_mask": Tensor[bool]}
|
||||
```
|
||||
|
||||
### 序列打包
|
||||
|
||||
当 `pack_size > 0` 时启用，使用 First-Fit Decreasing 算法将变长序列填充到固定长度：
|
||||
|
||||
- 按序列长度降序排列
|
||||
- 逐个放入当前包，超长则截断
|
||||
- 当前包放不下时开新包
|
||||
- 不足部分用 `pad_value` 填充
|
||||
|
||||
## 快速开始
|
||||
|
||||
### 1. 导出数据集为 JSONL
|
||||
|
||||
```python
|
||||
from datasets import load_dataset
|
||||
from pipeline.export import export_dataset
|
||||
from pipeline import export_dataset
|
||||
|
||||
dataset = load_dataset("your-dataset")
|
||||
export_dataset(
|
||||
dataset=dataset["train"],
|
||||
output_dir="./data",
|
||||
output_prefix="train",
|
||||
process_func=lambda x: {"text": x["content"]}, # 可选
|
||||
process_func=lambda x: {"text": x["content"]},
|
||||
)
|
||||
```
|
||||
|
||||
### 2. Tokenize 并缓存
|
||||
### 2. 分词并缓存为 HDF5
|
||||
|
||||
```python
|
||||
from pipeline import BpeTokenizer, ProcessorFactory, cache_jsonl
|
||||
|
||||
tokenizer = BpeTokenizer("tokenizer.json")
|
||||
processor = ProcessorFactory.create("pt", tokenizer) # "pt" | "sft" | "dpo"
|
||||
processor = ProcessorFactory.create("pt", tokenizer)
|
||||
|
||||
cache_jsonl(
|
||||
files=["./data/train.jsonl"],
|
||||
output_dir="./cached",
|
||||
processor=processor,
|
||||
pack_size=4096, # <=0 不打包
|
||||
pad_value=1,
|
||||
pack_size=4096,
|
||||
)
|
||||
```
|
||||
|
||||
### 3. 处理器类型
|
||||
|
||||
| 类型 | key | 输入 | 输出 |
|
||||
|------|-----|------|------|
|
||||
| 预训练 | `"pt"` | `{"text": "..."}` | `["sequence"]` |
|
||||
| SFT | `"sft"` | `{"query": "...", "response": "..."}` | `["sequence", "loss_mask"]` |
|
||||
| DPO | `"dpo"` | `{"query": "...", "chosen": "...", "rejected": "..."}` | `["chosen", "chosen_mask", "rejected", "rejected_mask"]` |
|
||||
|
||||
## 参数参考
|
||||
|
||||
### export_dataset()
|
||||
|
||||
| 参数 | 类型 | 默认 | 说明 |
|
||||
|------|------|------|------|
|
||||
| `dataset` | Dataset | 必填 | HuggingFace Dataset |
|
||||
| `output_dir` | str | 必填 | 输出目录 |
|
||||
| `output_prefix` | str | 必填 | 文件名前缀 |
|
||||
| `chunk_size` | int | 1_000_000 | 每个文件的最大样本数 |
|
||||
| `max_chunks` | int | None | 最大 chunk 数量 |
|
||||
| `process_func` | callable | None | 样本转换函数 |
|
||||
| `column` | str | "text" | 默认文本列名 |
|
||||
|
||||
### cache_jsonl()
|
||||
|
||||
| 参数 | 类型 | 默认 | 说明 |
|
||||
|------|------|------|------|
|
||||
| `files` | List[str] | 必填 | JSONL 文件列表 |
|
||||
| `output_dir` | str | 必填 | 输出目录 |
|
||||
| `processor` | BaseProcessor | 必填 | 处理器实例 |
|
||||
| `pack_size` | int | -1 | 打包长度，<=0 不打包 |
|
||||
| `pad_value` | int | 1 | 填充值 |
|
||||
|
||||
### SequencePacker
|
||||
### 3. 使用策略模式（默认 token）
|
||||
|
||||
```python
|
||||
packer = SequencePacker(pack_size=4096, pad_value=0)
|
||||
packed = packer.pack([tensor1, tensor2, ...]) # → List[Tensor]
|
||||
from pipeline import StrategyFactory, ProcessorFactory
|
||||
|
||||
strategy = StrategyFactory.create("alpaca")
|
||||
processor = ProcessorFactory.create_with_strategy("sft", tokenizer, strategy)
|
||||
```
|
||||
|
||||
### IOHandler
|
||||
### 4. 使用策略模式（自定义 token）
|
||||
|
||||
```python
|
||||
# 保存
|
||||
IOHandler.save_h5("./out", "name", {"key": [tensor1, tensor2]})
|
||||
# 自定义 ChatML 的特殊 token
|
||||
strategy = StrategyFactory.create("chatml",
|
||||
user_start="<s>user\n",
|
||||
user_end="</s>\n",
|
||||
assistant_start="<s>assistant\n",
|
||||
assistant_end="</s>\n<eos>",
|
||||
)
|
||||
processor = ProcessorFactory.create_with_strategy("sft", tokenizer, strategy)
|
||||
```
|
||||
|
||||
# 加载
|
||||
data = IOHandler.load_h5("./out") # → {"key": [tensor1, ...]}
|
||||
## 策略格式
|
||||
|
||||
| Strategy | Key | Default Tokens |
|
||||
|-----------|------------|--------------------------------------------------------------------------------------------------|
|
||||
| ChatML | `"chatml"` | `<\|im_start\|>user`, `<\|im_end\|>`, `<\|im_start\|>assistant`, `<eos>` |
|
||||
| Alpaca | `"alpaca"` | `### Instruction:`, `### Response:`, `<eos>` |
|
||||
|
||||
所有策略的 token 均可通过构造函数参数自定义，同时支持通过 `StrategyFactory.register()` 注册新格式。
|
||||
|
||||
## 命令行工具
|
||||
|
||||
```bash
|
||||
# 缓存 JSONL 到 H5
|
||||
python scripts/cache_h5.py pt ./dataset/chinese-c4-pretrain
|
||||
python scripts/cache_h5.py sft ./dataset/belle-sft --pack-size 4096 --strategy alpaca
|
||||
```
|
||||
|
||||
## 脚本示例
|
||||
|
||||
```
|
||||
scripts/
|
||||
├── cache_h5.py # Stage 2: JSONL -> H5 (CLI tool)
|
||||
├── pre_train/
|
||||
│ ├── chinese-c4.py # Chinese pretrain data export
|
||||
│ ├── chinese-cosmopedia.py # Chinese pretrain data export
|
||||
│ ├── english-fineweb.py # English pretrain data export
|
||||
│ └── english-wiki.py # English pretrain data export
|
||||
├── supervised_finetuning/
|
||||
│ ├── sft_belle.py # Belle Chinese SFT data export
|
||||
│ ├── sft_chinese_instruct.py # Chinese instruct SFT (with TextNormalizer)
|
||||
│ ├── sft_coder.py # Code SFT data export
|
||||
│ ├── sft_firefly-1.1m-rephrased.py # Firefly SFT data export
|
||||
│ └── sft_magpie-pro-300k.py # Magpie Pro SFT data export
|
||||
└── reforce_learning/
|
||||
└── dpp_chinese_dpo_pairs.py # Chinese DPO preference pairs export
|
||||
```
|
||||
+185
@@ -0,0 +1,185 @@
|
||||
# Pipeline 设计文档
|
||||
|
||||
## 概述
|
||||
|
||||
本项目是一个语言模型训练数据集处理工具包，将 HuggingFace 数据集转换为模型可直接消费的 HDF5 格式。整体分为两个阶段，通过磁盘 JSONL 文件解耦。
|
||||
|
||||
## 数据流架构
|
||||
|
||||
```
|
||||
Stage 1: Export Stage 2: Tokenize & Cache
|
||||
┌──────────────────────────┐ ┌──────────────────────────────────┐
|
||||
│ │ │ │
|
||||
│ HuggingFace Dataset │ │ JSONL files (one dict/line) │
|
||||
│ │ │ │ │ │
|
||||
│ ▼ process_func │ │ ▼ json.loads() │
|
||||
│ export_dataset() │ │ Processor.process(input_dict) │
|
||||
│ │ chunk_size │ │ │ │
|
||||
│ ▼ │ │ ▼ SequencePacker (opt.) │
|
||||
│ prefix_chunk_0.jsonl ────────▶ IOHandler.save_h5() │
|
||||
│ prefix_chunk_1.jsonl │ │ │ │
|
||||
│ ... │ │ ▼ │
|
||||
│ │ │ chunk_0.h5 │
|
||||
│ module: export.py │ │ chunk_1.h5 │
|
||||
│ │ │ │
|
||||
│ │ │ modules: cache.py, │
|
||||
│ │ │ processors.py, packing.py, │
|
||||
│ │ │ io.py, tokenizer.py │
|
||||
└──────────────────────────┘ └──────────────────────────────────┘
|
||||
```
|
||||
|
||||
### 阶段 1: export_dataset()
|
||||
|
||||
**职责**: 将 HuggingFace Dataset 导出为 JSONL 文件。
|
||||
|
||||
**流程**:
|
||||
1. 按 `chunk_size`（默认 100 万条）将数据集分片
|
||||
2. 每条样本通过可选的 `process_func` 转换（字段映射、文本清洗等）
|
||||
3. `process_func` 返回单个 dict 或 list[dict]（支持一对多展开）
|
||||
4. 每片写为一个 JSONL 文件
|
||||
|
||||
**模块**: `pipeline/export.py`
|
||||
|
||||
### 阶段 2: cache_jsonl()
|
||||
|
||||
**职责**: 将 JSONL 文件分词、打包、保存为 HDF5 格式。
|
||||
|
||||
**流程**:
|
||||
1. 逐行读取 JSONL，解析为 dict
|
||||
2. 通过 `Processor.process(input_dict)` 转换为 tensor dict
|
||||
3. 可选通过 `SequencePacker` 打包为固定长度
|
||||
4. 通过 `IOHandler.save_h5()` 写入 HDF5
|
||||
|
||||
**模块**: `pipeline/io.py`（编排）+ `pipeline/processors.py`（转换）+ `pipeline/packing.py`（打包）+ `pipeline/tokenizer.py`（分词）
|
||||
|
||||
## 处理器设计
|
||||
|
||||
### ProcessorFactory
|
||||
|
||||
通过工厂模式创建不同类型的处理器：
|
||||
|
||||
| 方法 | 说明 |
|
||||
|------|------|
|
||||
| `create(type, tokenizer)` | 创建处理器，SFT/DPO 使用默认 ChatML 策略 |
|
||||
| `create_with_strategy(type, tokenizer, strategy)` | 创建带自定义策略的处理器 |
|
||||
| `create_with_strategy_name(type, tokenizer, name, **kwargs)` | 通过名称创建策略，`**kwargs` 透传用于自定义 token |
|
||||
| `register(type, processor_class)` | 注册自定义处理器 |
|
||||
|
||||
### 处理器类型
|
||||
|
||||
**PreTrainProcessor** (`"pt"`)
|
||||
```
|
||||
Input: {"text": "Hello world"}
|
||||
Action: tokenizer.encode(text + "<eos>")
|
||||
Output: {"sequence": Tensor[int32]}
|
||||
```
|
||||
|
||||
**SFTProcessor** (`"sft"`) — 注入 PromptStrategy
|
||||
```
|
||||
Input: {"query": "...", "response": "..."}
|
||||
Action:
|
||||
1. strategy.build_prompt(input_dict) -> query prompt with special tokens
|
||||
2. concat: prompt + response + response_suffix
|
||||
3. tokenizer.encode -> token ids
|
||||
4. build loss_mask: query=False, response=True
|
||||
Output: {"sequence": Tensor[int32], "loss_mask": Tensor[bool]}
|
||||
```
|
||||
|
||||
**DPOProcessor** (`"dpo"`) — 注入 PromptStrategy
|
||||
```
|
||||
Input: {"query": "...", "chosen": "...", "rejected": "..."}
|
||||
Action:
|
||||
1. strategy.build_prompt(input_dict) -> shared query prompt
|
||||
2. encode chosen = query + response_start + chosen + suffix
|
||||
3. encode rejected = query + response_start + rejected + suffix
|
||||
4. build masks: query=False, response=True
|
||||
Output: {"chosen": Tensor, "chosen_mask": Tensor[bool],
|
||||
"rejected": Tensor, "rejected_mask": Tensor[bool]}
|
||||
```
|
||||
|
||||
## 策略模式
|
||||
|
||||
### PromptStrategy
|
||||
|
||||
抽象 prompt/response 格式，所有特殊 token 可通过构造函数自定义配置。
|
||||
|
||||
**接口**:
|
||||
- `build_prompt(input_dict)` — 构建包含 query 的完整 prompt
|
||||
- `build_response_prefix()` — response 前缀（当前均返回空串）
|
||||
- `build_response_suffix()` — response 后缀（含 `<eos>`）
|
||||
- `response_start_token` — response 起始 token（用于 DPO 的 loss mask 定位）
|
||||
- `eos_tokens` — 结束 token
|
||||
|
||||
**内置实现**:
|
||||
|
||||
ChatML:
|
||||
|
||||
| Parameter | Default Token |
|
||||
|-----------------|--------------------------------------------|
|
||||
| user_start | `<\|im_start\|>user\n` |
|
||||
| user_end | `<\|im_end\|>\n` |
|
||||
| assistant_start | `<\|im_start\|>assistant\n` |
|
||||
| assistant_end | `<\|im_end\|>\n<eos>` |
|
||||
|
||||
Alpaca:
|
||||
|
||||
| Parameter | Default Token |
|
||||
|------------------|------------------------|
|
||||
| instruction_start | `### Instruction:\n` |
|
||||
| response_start | `### Response:\n` |
|
||||
| response_suffix | `\n<eos>` |
|
||||
|
||||
**自定义示例**:
|
||||
```python
|
||||
strategy = StrategyFactory.create("chatml",
|
||||
user_start="<s>user\n",
|
||||
user_end="</s>\n",
|
||||
assistant_start="<s>assistant\n",
|
||||
assistant_end="</s>\n<eos>",
|
||||
)
|
||||
```
|
||||
|
||||
**注册新格式**:
|
||||
```python
|
||||
from pipeline import StrategyFactory
|
||||
from pipeline.strategies import PromptStrategy
|
||||
|
||||
class MyStrategy(PromptStrategy):
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "my_format"
|
||||
# ... 实现抽象方法
|
||||
|
||||
StrategyFactory.register("my_format", MyStrategy)
|
||||
```
|
||||
|
||||
## 序列打包
|
||||
|
||||
`SequencePacker` 将变长序列打包为固定长度的 tensor，使用 First-Fit Decreasing 算法：
|
||||
|
||||
1. 按序列长度降序排列
|
||||
2. 逐个放入当前包，超长则截断并警告
|
||||
3. 当前包放不下时保存并开新包
|
||||
4. 不足部分用 `pad_value` 填充
|
||||
|
||||
打包在 `cache_jsonl()` 中按 `pack_size` 启用（`> 0` 时生效，`<= 0` 跳过打包）。
|
||||
|
||||
## 辅助模块
|
||||
|
||||
### TextNormalizer (`pipeline/text.py`)
|
||||
|
||||
文本标准化：统一标点符号（弯引号→直引号、破折号、省略号等），通过 `DEFAULT_REPLACEMENTS` 配置，支持 `custom_rules` 扩展。
|
||||
|
||||
### IOHandler (`pipeline/io.py`)
|
||||
|
||||
- `save_h5(output_dir, file_name, tensor_group)` — 按 key 分组存储 tensor 到 HDF5
|
||||
- `load_h5(file_path, share_memory)` — 加载 HDF5，支持共享内存（用于 DataLoader 多进程）
|
||||
- `fetch_files(directory)` / `fetch_folders(root_dir)` — 文件/目录遍历
|
||||
|
||||
### BpeTokenizer (`pipeline/tokenizer.py`)
|
||||
|
||||
基于 HuggingFace `tokenizers` 库的 BPE 分词器，支持从文件加载、训练、保存。内置 `<bos>`/`<eos>`/`<pad>` 控制符和 `<|im_start|>`/`<|im_end|>` 特殊 token。
|
||||
|
||||
## API 参考
|
||||
|
||||
各模块的详细 API 文档请参阅对应的源文件。
|
||||
+14
-9
@@ -1,17 +1,17 @@
|
||||
import logging
|
||||
from .tokenizer import BpeTokenizer
|
||||
from .text import TextNormalizer
|
||||
from .packing import SequencePacker
|
||||
from .io import IOHandler
|
||||
from .processors import ProcessorFactory, BaseProcessor
|
||||
from .export import export_dataset
|
||||
from .cache import cache_jsonl
|
||||
from .utils import setup_logging
|
||||
from pipeline.tokenizer import BpeTokenizer
|
||||
from pipeline.text import TextNormalizer
|
||||
from pipeline.packing import SequencePacker
|
||||
from pipeline.io import IOHandler, export_dataset, cache_jsonl
|
||||
from pipeline.processors import ProcessorFactory, BaseProcessor
|
||||
from pipeline.utils import setup_logging
|
||||
from pipeline.strategies import PromptStrategy, ChatMLStrategy, AlpacaStrategy, StrategyFactory
|
||||
|
||||
# 配置项目级日志记录
|
||||
# Configure project-level logging
|
||||
setup_logging()
|
||||
|
||||
__all__ = [
|
||||
# Core modules
|
||||
'BpeTokenizer',
|
||||
'TextNormalizer',
|
||||
'SequencePacker',
|
||||
@@ -20,4 +20,9 @@ __all__ = [
|
||||
'BaseProcessor',
|
||||
'export_dataset',
|
||||
'cache_jsonl',
|
||||
# Strategy pattern
|
||||
'PromptStrategy',
|
||||
'ChatMLStrategy',
|
||||
'AlpacaStrategy',
|
||||
'StrategyFactory',
|
||||
]
|
||||
|
||||
@@ -1,82 +0,0 @@
|
||||
"""Tokenize JSONL files and pack them into HDF5 storage."""
|
||||
import json
|
||||
import os
|
||||
import logging
|
||||
from typing import List, Dict
|
||||
from pathlib import Path
|
||||
|
||||
from tqdm import tqdm
|
||||
|
||||
from .processors import BaseProcessor
|
||||
from .packing import SequencePacker
|
||||
from .io import IOHandler
|
||||
from .utils import error_handler
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@error_handler()
|
||||
def cache_jsonl(
|
||||
files: List[str],
|
||||
output_dir: str,
|
||||
processor: BaseProcessor,
|
||||
*,
|
||||
pack_size: int = -1,
|
||||
pad_value: int = 1,
|
||||
) -> List[str]:
|
||||
"""
|
||||
Tokenize JSONL files and pack them into HDF5 storage.
|
||||
|
||||
Args:
|
||||
files: List of JSONL file paths
|
||||
output_dir: H5 output directory
|
||||
processor: Initialized Processor instance
|
||||
pack_size: Packing length, <=0 means no packing
|
||||
pad_value: Padding value
|
||||
|
||||
Returns:
|
||||
List of generated H5 file paths
|
||||
"""
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
output_files: List[str] = []
|
||||
# Cache output_keys to avoid repeated attribute access
|
||||
output_keys = processor.output_keys
|
||||
|
||||
for file_path in files:
|
||||
file_name = Path(file_path).stem
|
||||
|
||||
# Pre-allocate lists for each output key
|
||||
arrows: Dict[str, List] = {key: [] for key in output_keys}
|
||||
|
||||
# Read and process all lines
|
||||
with open(file_path, "r", encoding="utf-8") as f:
|
||||
for line_num, line in enumerate(tqdm(f, desc=f"Processing {file_name}", leave=False), start=1):
|
||||
try:
|
||||
result = processor.process(json.loads(line))
|
||||
if result is not None:
|
||||
# Batch append: add each key's tensor to corresponding list
|
||||
for key in output_keys:
|
||||
arrows[key].append(result[key])
|
||||
except json.JSONDecodeError as e:
|
||||
logger.warning(f"JSON decode error in {file_path} line {line_num}: {e}. Skipping line.")
|
||||
continue
|
||||
except Exception as e:
|
||||
logger.warning(f"Unexpected error processing line {line_num} in {file_path}: {e}. Skipping line.")
|
||||
continue
|
||||
|
||||
# Convert lists to tensors once per key
|
||||
if pack_size > 0:
|
||||
output = {}
|
||||
for key in output_keys:
|
||||
packer = SequencePacker(pack_size, pad_value)
|
||||
output[key] = packer.pack(arrows[key])
|
||||
else:
|
||||
# No packing: directly use the arrow tensors
|
||||
output = arrows
|
||||
|
||||
IOHandler.save_h5(output_dir, file_name, output)
|
||||
h5_path = os.path.join(output_dir, f"{file_name}.h5")
|
||||
output_files.append(h5_path)
|
||||
logger.info(f"Saved {h5_path}")
|
||||
|
||||
return output_files
|
||||
@@ -1,63 +0,0 @@
|
||||
"""Export HuggingFace Dataset to JSONL files in chunks."""
|
||||
import json
|
||||
import os
|
||||
import logging
|
||||
from typing import Callable, Optional, List, Union, Dict, Any
|
||||
|
||||
from datasets import Dataset
|
||||
from .utils import error_handler
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@error_handler()
|
||||
def export_dataset(
|
||||
dataset: Dataset,
|
||||
output_dir: str,
|
||||
output_prefix: str,
|
||||
*,
|
||||
chunk_size: int = 1_000_000,
|
||||
max_chunks: Optional[int] = None,
|
||||
process_func: Optional[Callable[[Dict[str, Any]], Union[Dict[str, Any], List[Dict[str, Any]]]]] = None,
|
||||
column: str = "text",
|
||||
) -> List[str]:
|
||||
"""
|
||||
Export HuggingFace Dataset to JSONL files in chunks.
|
||||
|
||||
Args:
|
||||
dataset: HuggingFace Dataset object
|
||||
output_dir: Output directory
|
||||
output_prefix: Output file name prefix, e.g., "chinese-c4-pretrain"
|
||||
chunk_size: Maximum number of samples per file
|
||||
max_chunks: Maximum number of chunks to process (for debugging)
|
||||
process_func: Single sample transformation function (dict) -> dict | list[dict]
|
||||
column: Default text column name (only used when process_func is None)
|
||||
|
||||
Returns:
|
||||
List of generated file paths
|
||||
"""
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
total = len(dataset)
|
||||
num_chunks = (total + chunk_size - 1) // chunk_size
|
||||
lim = min(max_chunks, num_chunks) if max_chunks else num_chunks
|
||||
|
||||
output_files: List[str] = []
|
||||
for i in range(lim):
|
||||
start = i * chunk_size
|
||||
end = min(start + chunk_size, total)
|
||||
chunk = dataset.select(range(start, end))
|
||||
|
||||
path = os.path.join(output_dir, f"{output_prefix}_chunk_{i}.jsonl")
|
||||
try:
|
||||
with open(path, "w", encoding="utf-8") as f:
|
||||
for example in chunk:
|
||||
processed = process_func(example) if process_func else {column: example[column]}
|
||||
items = processed if isinstance(processed, list) else [processed]
|
||||
for item in items:
|
||||
f.write(json.dumps(item, ensure_ascii=False) + "\n")
|
||||
output_files.append(path)
|
||||
logger.info(f"[{i + 1}/{lim}] Saved {path}")
|
||||
except (OSError, IOError) as e:
|
||||
logger.error(f"Failed to write chunk {i} to {path}: {e}")
|
||||
|
||||
return output_files
|
||||
+139
-6
@@ -1,23 +1,36 @@
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional, Callable
|
||||
"""File, HDF5, JSONL I/O operations."""
|
||||
import json
|
||||
import os
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional, Callable, Union, Any
|
||||
|
||||
import h5py
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from tqdm import tqdm
|
||||
from datasets import Dataset
|
||||
|
||||
from .utils import error_handler
|
||||
from pipeline.utils import error_handler
|
||||
from pipeline.processors import BaseProcessor
|
||||
from pipeline.packing import SequencePacker
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class IOHandler:
|
||||
"""File and HDF5 read/write operations."""
|
||||
|
||||
@staticmethod
|
||||
def fetch_files(directory: str) -> List[str]:
|
||||
return [
|
||||
def fetch_files(directory: str, suffix: Optional[str] = None) -> List[str]:
|
||||
files = [
|
||||
os.path.join(root, f)
|
||||
for root, _, files in os.walk(directory)
|
||||
for f in files
|
||||
]
|
||||
if suffix:
|
||||
files = [f for f in files if f.endswith(suffix)]
|
||||
return sorted(files)
|
||||
|
||||
@staticmethod
|
||||
def fetch_folders(root_dir: str, filter_func: Optional[Callable[[str], bool]] = None) -> List[str]:
|
||||
@@ -43,7 +56,7 @@ class IOHandler:
|
||||
|
||||
@staticmethod
|
||||
@error_handler()
|
||||
def load_h5(file_path: str, share_memory=True) -> Dict[str, List[Tensor]]:
|
||||
def load_h5(file_path: str, share_memory: bool = True) -> Dict[str, List[Tensor]]:
|
||||
tensor_group: Dict[str, List[Tensor]] = {}
|
||||
|
||||
root_path = Path(file_path)
|
||||
@@ -66,3 +79,123 @@ class IOHandler:
|
||||
tensor_group[key].extend(dsets)
|
||||
|
||||
return tensor_group
|
||||
|
||||
|
||||
# ── Stage 1: Export HuggingFace Dataset to JSONL ──────────────────────────
|
||||
|
||||
|
||||
@error_handler()
|
||||
def export_dataset(
|
||||
dataset: Dataset,
|
||||
output_dir: str,
|
||||
output_prefix: str,
|
||||
*,
|
||||
chunk_size: int = 1_000_000,
|
||||
max_chunks: Optional[int] = None,
|
||||
process_func: Optional[Callable[[Dict[str, Any]], Union[Dict[str, Any], List[Dict[str, Any]]]]] = None,
|
||||
column: str = "text",
|
||||
) -> List[str]:
|
||||
"""
|
||||
Export HuggingFace Dataset to JSONL files in chunks.
|
||||
|
||||
Args:
|
||||
dataset: HuggingFace Dataset object
|
||||
output_dir: Output directory
|
||||
output_prefix: Output file name prefix, e.g., "chinese-c4-pretrain"
|
||||
chunk_size: Maximum number of samples per file
|
||||
max_chunks: Maximum number of chunks to process (for debugging)
|
||||
process_func: Single sample transformation function (dict) -> dict | list[dict]
|
||||
column: Default text column name (only used when process_func is None)
|
||||
|
||||
Returns:
|
||||
List of generated file paths
|
||||
"""
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
total = len(dataset)
|
||||
num_chunks = (total + chunk_size - 1) // chunk_size
|
||||
lim = min(max_chunks, num_chunks) if max_chunks else num_chunks
|
||||
|
||||
output_files: List[str] = []
|
||||
for i in range(lim):
|
||||
start = i * chunk_size
|
||||
end = min(start + chunk_size, total)
|
||||
chunk = dataset.select(range(start, end))
|
||||
|
||||
path = os.path.join(output_dir, f"{output_prefix}_chunk_{i}.jsonl")
|
||||
try:
|
||||
with open(path, "w", encoding="utf-8") as f:
|
||||
for example in chunk:
|
||||
processed = process_func(example) if process_func else {column: example[column]}
|
||||
items = processed if isinstance(processed, list) else [processed]
|
||||
for item in items:
|
||||
f.write(json.dumps(item, ensure_ascii=False) + "\n")
|
||||
output_files.append(path)
|
||||
logger.info(f"[{i + 1}/{lim}] Saved {path}")
|
||||
except (OSError, IOError) as e:
|
||||
logger.error(f"Failed to write chunk {i} to {path}: {e}")
|
||||
|
||||
return output_files
|
||||
|
||||
|
||||
# ── Stage 2: Tokenize JSONL and cache to HDF5 ────────────────────────────
|
||||
|
||||
|
||||
@error_handler()
|
||||
def cache_jsonl(
|
||||
files: List[str],
|
||||
output_dir: str,
|
||||
processor: BaseProcessor,
|
||||
*,
|
||||
pack_size: int = -1,
|
||||
pad_value: int = 1,
|
||||
) -> List[str]:
|
||||
"""
|
||||
Tokenize JSONL files and pack them into HDF5 storage.
|
||||
|
||||
Args:
|
||||
files: List of JSONL file paths
|
||||
output_dir: H5 output directory
|
||||
processor: Initialized Processor instance
|
||||
pack_size: Packing length, <=0 means no packing
|
||||
pad_value: Padding value
|
||||
|
||||
Returns:
|
||||
List of generated H5 file paths
|
||||
"""
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
output_files: List[str] = []
|
||||
output_keys = processor.output_keys
|
||||
|
||||
for file_path in files:
|
||||
file_name = Path(file_path).stem
|
||||
|
||||
arrows: Dict[str, List] = {key: [] for key in output_keys}
|
||||
|
||||
with open(file_path, "r", encoding="utf-8") as f:
|
||||
for line_num, line in enumerate(tqdm(f, desc=f"Processing {file_name}", leave=False), start=1):
|
||||
try:
|
||||
result = processor.process(json.loads(line))
|
||||
if result is not None:
|
||||
for key in output_keys:
|
||||
arrows[key].append(result[key])
|
||||
except json.JSONDecodeError as e:
|
||||
logger.warning(f"JSON decode error in {file_path} line {line_num}: {e}. Skipping line.")
|
||||
continue
|
||||
except Exception as e:
|
||||
logger.warning(f"Unexpected error processing line {line_num} in {file_path}: {e}. Skipping line.")
|
||||
continue
|
||||
|
||||
if pack_size > 0:
|
||||
output = {}
|
||||
for key in output_keys:
|
||||
packer = SequencePacker(pack_size, pad_value)
|
||||
output[key] = packer.pack(arrows[key])
|
||||
else:
|
||||
output = arrows
|
||||
|
||||
IOHandler.save_h5(output_dir, file_name, output)
|
||||
h5_path = os.path.join(output_dir, f"{file_name}.h5")
|
||||
output_files.append(h5_path)
|
||||
logger.info(f"Saved {h5_path}")
|
||||
|
||||
return output_files
|
||||
|
||||
+93
-72
@@ -1,109 +1,130 @@
|
||||
import logging
|
||||
from typing import List, Optional
|
||||
from typing import List
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
from .utils import error_handler
|
||||
from pipeline.utils import error_handler
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class SequencePacker:
|
||||
"""
|
||||
Packs variable-length sequences into fixed-size tensors, suitable for
|
||||
concatenating unequal-length training samples into uniform shapes
|
||||
for DataLoader / model training.
|
||||
|
||||
def __init__(self, pack_size: int, pad_value: int = 0, dtype: torch.dtype = torch.int32):
|
||||
Algorithm (Sorted Greedy Fill, based on First-Fit Decreasing heuristic):
|
||||
|
||||
Input: sequences = [A(len=5), B(len=2), C(len=3)], pack_size = 8
|
||||
|
||||
1. Validate & Normalize
|
||||
- Check 1D dimension, unify dtype, truncate overlong sequences with warning
|
||||
- Result: [(A,5), (B,2), (C,3)]
|
||||
|
||||
2. Sort by length descending (FFD)
|
||||
- Result: [(A,5), (C,3), (B,2)]
|
||||
|
||||
3. Greedy fill: write into a pre-allocated buffer sequentially, flush when full
|
||||
- Write A(5) -> buffer = [A A A A A _ _ _], pos=5
|
||||
- Write C(3) -> pos+3=8 <= 8 -> buffer = [A A A A A C C C], pos=8
|
||||
- Buffer full -> flush as package[0], reset buffer & pos=0
|
||||
- Write B(2) -> buffer = [B B _ _ _ _ _ _], pos=2
|
||||
- Loop ends -> flush tail -> package[1] = [B B 0 0 0 0 0 0]
|
||||
|
||||
Output: [package[0], package[1]]
|
||||
|
||||
Cross-group consistency:
|
||||
When packing different key groups (e.g. sequences and loss_masks)
|
||||
with separate pack() calls, tensors at the same index always have
|
||||
identical lengths, so the descending sort produces the exact same
|
||||
ordering. Element-level correspondence across groups is preserved.
|
||||
|
||||
Performance:
|
||||
- Pre-allocated buffer reused via fill_() to avoid repeated tensor creation
|
||||
- Attributes cached as local variables inside the loop to reduce lookup overhead
|
||||
"""
|
||||
|
||||
def __init__(self, pack_size: int, pad_value: int = 0, dtype: torch.dtype = None):
|
||||
self.pack_size = pack_size
|
||||
self.pad_value = pad_value
|
||||
self.dtype = dtype
|
||||
# Pre-allocate buffer for better performance
|
||||
self._buffer: Optional[Tensor] = None
|
||||
self._reset()
|
||||
|
||||
def _reset(self) -> None:
|
||||
"""Reset internal state for instance reuse."""
|
||||
# Reuse buffer instead of creating new tensors
|
||||
if self._buffer is None or self._buffer.shape[0] != self.pack_size:
|
||||
self._buffer = torch.full(
|
||||
(self.pack_size,), self.pad_value, dtype=self.dtype
|
||||
)
|
||||
else:
|
||||
self._buffer.fill_(self.pad_value)
|
||||
self._current_pos = 0
|
||||
self.dtype = dtype # None = follow input dtype
|
||||
self._buffer: Tensor | None = None
|
||||
self._pos = 0
|
||||
self._packages: List[Tensor] = []
|
||||
# Backward compatibility: maintain _current_pack reference
|
||||
self._current_pack = self._buffer
|
||||
|
||||
def reset(self) -> None:
|
||||
"""Reset packer state for instance reuse, unlocking dtype."""
|
||||
self.dtype = None
|
||||
self._buffer = None
|
||||
self._pos = 0
|
||||
self._packages = []
|
||||
|
||||
@error_handler()
|
||||
def pack(self, sequences: List[Tensor]) -> List[Tensor]:
|
||||
"""
|
||||
Pack sequences into fixed-size packages.
|
||||
Pack sequences into fixed-size packages using First-Fit Decreasing.
|
||||
|
||||
Sequences are sorted by length descending to minimize wasted padding.
|
||||
All tensor groups (e.g. sequences, loss_masks) with matching per-item
|
||||
lengths produce identical ordering, so cross-group correspondence is preserved.
|
||||
|
||||
Args:
|
||||
sequences: List of input tensors
|
||||
sequences: List of 1D input tensors.
|
||||
|
||||
Returns:
|
||||
List of packed tensors, each with length equal to pack_size
|
||||
List of packed tensors, each with length equal to pack_size.
|
||||
"""
|
||||
# Input validation
|
||||
if not sequences:
|
||||
return []
|
||||
|
||||
# Validate and cache tensor sizes in one pass
|
||||
tensor_sizes = []
|
||||
# --- validate & normalize in a single pass ---
|
||||
normalized: list[tuple[Tensor, int]] = []
|
||||
target_dtype = self.dtype if self.dtype is not None else sequences[0].dtype
|
||||
for i, seq in enumerate(sequences):
|
||||
if seq.dim() != 1:
|
||||
raise ValueError(
|
||||
f"Expected 1D tensor at index {i}, got {seq.dim()}D tensor with shape {seq.shape}"
|
||||
)
|
||||
tensor_sizes.append(seq.numel())
|
||||
if seq.dtype != self.dtype:
|
||||
logger.warning(
|
||||
f"Input tensor dtype {seq.dtype} does not match packer dtype {self.dtype}, "
|
||||
f"will be converted. This may affect packing efficiency."
|
||||
)
|
||||
if seq.dtype != target_dtype:
|
||||
seq = seq.to(target_dtype)
|
||||
length = seq.numel()
|
||||
if length > self.pack_size:
|
||||
seq = seq[: self.pack_size]
|
||||
length = self.pack_size
|
||||
normalized.append((seq, length))
|
||||
|
||||
# Reset state for new packing
|
||||
# --- reset internal state ---
|
||||
buf = self._buffer
|
||||
if buf is None or buf.dtype != target_dtype:
|
||||
buf = torch.full((self.pack_size,), self.pad_value, dtype=target_dtype)
|
||||
self._buffer = buf
|
||||
buf.fill_(self.pad_value)
|
||||
self._pos = 0
|
||||
self._packages = []
|
||||
self._reset()
|
||||
|
||||
# Combine sequences with their sizes for sorting
|
||||
indexed_seqs = list(zip(sequences, tensor_sizes))
|
||||
# Sort by size descending (First-Fit Decreasing algorithm)
|
||||
indexed_seqs.sort(key=lambda x: x[1], reverse=True)
|
||||
# --- sort by length descending (FFD heuristic) ---
|
||||
normalized.sort(key=lambda x: x[1], reverse=True)
|
||||
|
||||
for tensor, tensor_size in indexed_seqs:
|
||||
# Truncate sequences that exceed pack_size
|
||||
if tensor_size > self.pack_size:
|
||||
logger.warning(
|
||||
f"Sequence length {tensor_size} exceeds pack_size {self.pack_size}, truncating"
|
||||
)
|
||||
tensor_size = self.pack_size
|
||||
tensor = tensor[: self.pack_size]
|
||||
# --- greedy fill ---
|
||||
buf = self._buffer
|
||||
pos = self._pos
|
||||
packages = self._packages
|
||||
pack_size = self.pack_size
|
||||
pad_value = self.pad_value
|
||||
|
||||
# Current package is full, create a new one
|
||||
if self._current_pos + tensor_size > self.pack_size:
|
||||
# Finish current package (pad to pack_size)
|
||||
package = self._buffer.clone()
|
||||
self._packages.append(package)
|
||||
# Reset buffer for reuse
|
||||
self._buffer.fill_(self.pad_value)
|
||||
self._current_pos = 0
|
||||
for tensor, length in normalized:
|
||||
if pos + length > pack_size:
|
||||
# flush current package
|
||||
packages.append(buf.clone())
|
||||
buf.fill_(pad_value)
|
||||
pos = 0
|
||||
buf[pos : pos + length] = tensor
|
||||
pos += length
|
||||
|
||||
# Place tensor in current package
|
||||
self._buffer[self._current_pos : self._current_pos + tensor_size] = tensor
|
||||
self._current_pos += tensor_size
|
||||
|
||||
# Handle the last package (pad to pack_size)
|
||||
if self._current_pos > 0:
|
||||
package = self._buffer.clone()
|
||||
self._packages.append(package)
|
||||
|
||||
# Clear buffer and reset state for backward compatibility
|
||||
self._buffer = None
|
||||
self._current_pack = None
|
||||
self._current_pos = 0
|
||||
# flush the last (possibly partial) package
|
||||
if pos > 0:
|
||||
packages.append(buf.clone())
|
||||
|
||||
# write back state
|
||||
self._pos = pos
|
||||
return self._packages
|
||||
|
||||
def reset(self) -> None:
|
||||
"""Reset packer state for reuse. More efficient than creating a new instance."""
|
||||
self._reset()
|
||||
@@ -1,122 +0,0 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Dict, List, Any
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from .tokenizer import BpeTokenizer
|
||||
|
||||
|
||||
class BaseProcessor(ABC):
|
||||
"""Abstract base class for processors."""
|
||||
|
||||
@abstractmethod
|
||||
def process(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]:
|
||||
pass
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def output_keys(self) -> List[str]:
|
||||
pass
|
||||
|
||||
|
||||
class PreTrainProcessor(BaseProcessor):
|
||||
"""Pre-training data processor."""
|
||||
|
||||
def __init__(self, tokenizer: BpeTokenizer):
|
||||
self.tokenizer = tokenizer
|
||||
|
||||
def process(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]:
|
||||
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):
|
||||
"""Supervised fine-tuning data processor."""
|
||||
|
||||
def __init__(self, tokenizer: BpeTokenizer):
|
||||
self.tokenizer = tokenizer
|
||||
|
||||
def process(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]:
|
||||
query = input_dict["query"]
|
||||
response = 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>")
|
||||
|
||||
q_len = len(q)
|
||||
tokens = torch.tensor(q + a, dtype=torch.int32)
|
||||
loss_mask = torch.zeros(q_len + len(a), dtype=torch.bool)
|
||||
loss_mask[q_len:] = True
|
||||
return {"sequence": tokens, "loss_mask": loss_mask}
|
||||
|
||||
@property
|
||||
def output_keys(self) -> List[str]:
|
||||
return ["sequence", "loss_mask"]
|
||||
|
||||
|
||||
class DPOProcessor(BaseProcessor):
|
||||
"""DPO preference learning data processor."""
|
||||
|
||||
def __init__(self, tokenizer: BpeTokenizer):
|
||||
self.tokenizer = tokenizer
|
||||
|
||||
def process(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]:
|
||||
query = input_dict["query"]
|
||||
chosen_response = input_dict["chosen"]
|
||||
rejected_response = input_dict["rejected"]
|
||||
|
||||
q = self.tokenizer.encode(
|
||||
f"<|im_start|>user\n{query}<|im_end|>\n<|im_start|>assistant\n"
|
||||
)
|
||||
|
||||
chosen = self.tokenizer.encode(f"{chosen_response}<|im_end|>\n<eos>")
|
||||
q_len = len(q)
|
||||
chosen_len = len(chosen)
|
||||
|
||||
chosen_tokens = torch.tensor(q + chosen, dtype=torch.int32)
|
||||
chosen_mask = torch.zeros(q_len + chosen_len, dtype=torch.bool)
|
||||
chosen_mask[q_len:] = True
|
||||
|
||||
rejected = self.tokenizer.encode(f"{rejected_response}<|im_end|>\n<eos>")
|
||||
rejected_len = len(rejected)
|
||||
|
||||
rejected_tokens = torch.tensor(q + rejected, dtype=torch.int32)
|
||||
rejected_mask = torch.zeros(q_len + rejected_len, dtype=torch.bool)
|
||||
rejected_mask[q_len:] = True
|
||||
|
||||
return {
|
||||
"chosen": chosen_tokens,
|
||||
"chosen_mask": chosen_mask,
|
||||
"rejected": rejected_tokens,
|
||||
"rejected_mask": rejected_mask,
|
||||
}
|
||||
|
||||
@property
|
||||
def output_keys(self) -> List[str]:
|
||||
return ["chosen", "chosen_mask", "rejected", "rejected_mask"]
|
||||
|
||||
|
||||
class ProcessorFactory:
|
||||
"""Processor factory."""
|
||||
|
||||
_processors = {
|
||||
"pt": PreTrainProcessor,
|
||||
"sft": SFTProcessor,
|
||||
"dpo": DPOProcessor,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def create(cls, processor_type: str, tokenizer: BpeTokenizer) -> 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
|
||||
@@ -0,0 +1,18 @@
|
||||
"""Data processors with factory pattern.
|
||||
|
||||
Processor classes are registered at definition time via decorators and
|
||||
can be created through :class:`ProcessorFactory`.
|
||||
"""
|
||||
from pipeline.processors.base import BaseProcessor
|
||||
from pipeline.processors.factory import ProcessorFactory
|
||||
from pipeline.processors.pretrain import PreTrainProcessor
|
||||
from pipeline.processors.sft import SFTProcessor
|
||||
from pipeline.processors.dpo import DPOProcessor
|
||||
|
||||
__all__ = [
|
||||
"BaseProcessor",
|
||||
"ProcessorFactory",
|
||||
"PreTrainProcessor",
|
||||
"SFTProcessor",
|
||||
"DPOProcessor",
|
||||
]
|
||||
@@ -0,0 +1,31 @@
|
||||
"""Processor base class and shared utilities."""
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Dict, List, Any, Tuple
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
|
||||
def _encode_with_mask(
|
||||
prompt_tokens: List[int],
|
||||
response_tokens: List[int],
|
||||
) -> Tuple[Tensor, Tensor]:
|
||||
"""Concatenate token lists and build loss mask (prompt=False, response=True)."""
|
||||
q_len = len(prompt_tokens)
|
||||
combined = torch.tensor(prompt_tokens + response_tokens, dtype=torch.int32)
|
||||
mask = torch.zeros(q_len + len(response_tokens), dtype=torch.bool)
|
||||
mask[q_len:] = True
|
||||
return combined, mask
|
||||
|
||||
|
||||
class BaseProcessor(ABC):
|
||||
"""Abstract base class for processors."""
|
||||
|
||||
@abstractmethod
|
||||
def process(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]:
|
||||
pass
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def output_keys(self) -> List[str]:
|
||||
pass
|
||||
@@ -0,0 +1,51 @@
|
||||
"""DPO preference learning data processor."""
|
||||
from typing import Dict, List, Any, Optional
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
from pipeline.tokenizer import BpeTokenizer
|
||||
from pipeline.strategies import PromptStrategy, ChatMLStrategy
|
||||
from pipeline.processors.base import BaseProcessor, _encode_with_mask
|
||||
from pipeline.processors.factory import ProcessorFactory
|
||||
|
||||
|
||||
@ProcessorFactory.register("dpo")
|
||||
class DPOProcessor(BaseProcessor):
|
||||
"""DPO preference learning data processor.
|
||||
|
||||
Supports custom prompt strategy via constructor parameter.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
tokenizer: BpeTokenizer,
|
||||
strategy: Optional[PromptStrategy] = None,
|
||||
):
|
||||
self.tokenizer = tokenizer
|
||||
self.strategy = strategy or ChatMLStrategy(tokenizer)
|
||||
|
||||
def process(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]:
|
||||
query_tokens = self.tokenizer.encode(input_dict["query"])
|
||||
chosen_tokens = self.tokenizer.encode(input_dict["chosen"])
|
||||
rejected_tokens = self.tokenizer.encode(input_dict["rejected"])
|
||||
|
||||
prompt = self.strategy.assemble_prompt(query_tokens)
|
||||
|
||||
chosen_t, chosen_m = _encode_with_mask(
|
||||
prompt, self.strategy.assemble_response(chosen_tokens)
|
||||
)
|
||||
rejected_t, rejected_m = _encode_with_mask(
|
||||
prompt, self.strategy.assemble_response(rejected_tokens)
|
||||
)
|
||||
|
||||
return {
|
||||
"chosen": chosen_t,
|
||||
"chosen_mask": chosen_m,
|
||||
"rejected": rejected_t,
|
||||
"rejected_mask": rejected_m,
|
||||
}
|
||||
|
||||
@property
|
||||
def output_keys(self) -> List[str]:
|
||||
return ["chosen", "chosen_mask", "rejected", "rejected_mask"]
|
||||
@@ -0,0 +1,122 @@
|
||||
"""Factory for creating and registering processors."""
|
||||
from typing import Dict, List, Any, Optional, Type
|
||||
|
||||
from pipeline.processors.base import BaseProcessor
|
||||
from pipeline.tokenizer import BpeTokenizer
|
||||
from pipeline.strategies import PromptStrategy, StrategyFactory
|
||||
|
||||
|
||||
class ProcessorFactory:
|
||||
"""Registry and factory for BaseProcessor implementations.
|
||||
|
||||
Supports decorator-based registration for extensible processor types.
|
||||
|
||||
Example usage::
|
||||
|
||||
@ProcessorFactory.register("custom")
|
||||
class CustomProcessor(BaseProcessor):
|
||||
...
|
||||
|
||||
processor = ProcessorFactory.create(optimizer, "custom", **kwargs)
|
||||
"""
|
||||
|
||||
PROCESSOR_MAP: Dict[str, Type[BaseProcessor]] = {}
|
||||
|
||||
@classmethod
|
||||
def register(cls, name: str):
|
||||
"""Decorator to register a new processor class.
|
||||
|
||||
Args:
|
||||
name: Registration name for the processor.
|
||||
|
||||
Returns:
|
||||
Decorator function that registers the processor class.
|
||||
"""
|
||||
|
||||
def decorator(processor_cls: Type[BaseProcessor]) -> Type[BaseProcessor]:
|
||||
if not issubclass(processor_cls, BaseProcessor):
|
||||
raise TypeError(
|
||||
f"{processor_cls.__name__} must inherit from BaseProcessor"
|
||||
)
|
||||
cls.PROCESSOR_MAP[name] = processor_cls
|
||||
return processor_cls
|
||||
|
||||
return decorator
|
||||
|
||||
@classmethod
|
||||
def create(cls, processor_type: str, tokenizer: BpeTokenizer) -> BaseProcessor:
|
||||
"""Create a processor by type name (uses default ChatMLStrategy for SFT/DPO).
|
||||
|
||||
Args:
|
||||
processor_type: Registered processor name (e.g. ``"pt"``, ``"sft"``, ``"dpo"``).
|
||||
tokenizer: Tokenizer instance.
|
||||
|
||||
Returns:
|
||||
Processor instance.
|
||||
|
||||
Raises:
|
||||
ValueError: If processor_type is not registered.
|
||||
"""
|
||||
if processor_type not in cls.PROCESSOR_MAP:
|
||||
raise ValueError(
|
||||
f"Unknown processor type: '{processor_type}'. "
|
||||
f"Supported types: {sorted(cls.PROCESSOR_MAP.keys())}"
|
||||
)
|
||||
return cls.PROCESSOR_MAP[processor_type](tokenizer)
|
||||
|
||||
@classmethod
|
||||
def create_with_strategy(
|
||||
cls,
|
||||
processor_type: str,
|
||||
tokenizer: BpeTokenizer,
|
||||
strategy: PromptStrategy,
|
||||
) -> BaseProcessor:
|
||||
"""Create a processor with a custom strategy.
|
||||
|
||||
Only SFT and DPO processors accept a strategy; PreTrain ignores it.
|
||||
|
||||
Args:
|
||||
processor_type: Registered processor name.
|
||||
tokenizer: Tokenizer instance.
|
||||
strategy: Prompt strategy instance.
|
||||
|
||||
Returns:
|
||||
Processor instance configured with strategy.
|
||||
"""
|
||||
if processor_type not in cls.PROCESSOR_MAP:
|
||||
raise ValueError(
|
||||
f"Unknown processor type: '{processor_type}'. "
|
||||
f"Supported types: {sorted(cls.PROCESSOR_MAP.keys())}"
|
||||
)
|
||||
|
||||
processor_cls = cls.PROCESSOR_MAP[processor_type]
|
||||
if processor_type == "pt":
|
||||
return processor_cls(tokenizer)
|
||||
return processor_cls(tokenizer, strategy=strategy)
|
||||
|
||||
@classmethod
|
||||
def create_with_strategy_name(
|
||||
cls,
|
||||
processor_type: str,
|
||||
tokenizer: BpeTokenizer,
|
||||
strategy_name: str,
|
||||
**strategy_kwargs,
|
||||
) -> BaseProcessor:
|
||||
"""Create a processor with a strategy selected by name.
|
||||
|
||||
Args:
|
||||
processor_type: Registered processor name.
|
||||
tokenizer: Tokenizer instance.
|
||||
strategy_name: Registered strategy name (``"chatml"``, ``"alpaca"``, etc.).
|
||||
**strategy_kwargs: Forwarded to the strategy constructor.
|
||||
|
||||
Returns:
|
||||
Processor instance.
|
||||
"""
|
||||
strategy = StrategyFactory.create(strategy_name, tokenizer, **strategy_kwargs)
|
||||
return cls.create_with_strategy(processor_type, tokenizer, strategy)
|
||||
|
||||
@classmethod
|
||||
def available_types(cls) -> List[str]:
|
||||
"""Return list of registered processor type names."""
|
||||
return list(cls.PROCESSOR_MAP.keys())
|
||||
@@ -0,0 +1,26 @@
|
||||
"""Pre-training data processor."""
|
||||
from typing import Dict, List, Any
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
from pipeline.tokenizer import BpeTokenizer
|
||||
from pipeline.processors.base import BaseProcessor
|
||||
from pipeline.processors.factory import ProcessorFactory
|
||||
|
||||
|
||||
@ProcessorFactory.register("pt")
|
||||
class PreTrainProcessor(BaseProcessor):
|
||||
"""Pre-training data processor."""
|
||||
|
||||
def __init__(self, tokenizer: BpeTokenizer):
|
||||
self.tokenizer = tokenizer
|
||||
|
||||
def process(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]:
|
||||
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"]
|
||||
@@ -0,0 +1,40 @@
|
||||
"""Supervised fine-tuning data processor."""
|
||||
from typing import Dict, List, Any, Optional
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
from pipeline.tokenizer import BpeTokenizer
|
||||
from pipeline.strategies import PromptStrategy, ChatMLStrategy
|
||||
from pipeline.processors.base import BaseProcessor, _encode_with_mask
|
||||
from pipeline.processors.factory import ProcessorFactory
|
||||
|
||||
|
||||
@ProcessorFactory.register("sft")
|
||||
class SFTProcessor(BaseProcessor):
|
||||
"""Supervised fine-tuning data processor.
|
||||
|
||||
Supports custom prompt strategy via constructor parameter.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
tokenizer: BpeTokenizer,
|
||||
strategy: Optional[PromptStrategy] = None,
|
||||
):
|
||||
self.tokenizer = tokenizer
|
||||
self.strategy = strategy or ChatMLStrategy(tokenizer)
|
||||
|
||||
def process(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]:
|
||||
query_tokens = self.tokenizer.encode(input_dict["query"])
|
||||
response_tokens = self.tokenizer.encode(input_dict["response"])
|
||||
|
||||
prompt = self.strategy.assemble_prompt(query_tokens)
|
||||
response = self.strategy.assemble_response(response_tokens)
|
||||
|
||||
tokens, loss_mask = _encode_with_mask(prompt, response)
|
||||
return {"sequence": tokens, "loss_mask": loss_mask}
|
||||
|
||||
@property
|
||||
def output_keys(self) -> List[str]:
|
||||
return ["sequence", "loss_mask"]
|
||||
@@ -0,0 +1,14 @@
|
||||
"""Strategy pattern for prompt/response format abstraction."""
|
||||
from pipeline.strategies.base import PromptStrategy
|
||||
from pipeline.strategies.factory import StrategyFactory
|
||||
|
||||
# Import strategy implementations to trigger decorator registration
|
||||
from pipeline.strategies.chatml import ChatMLStrategy # noqa: F401
|
||||
from pipeline.strategies.alpaca import AlpacaStrategy # noqa: F401
|
||||
|
||||
__all__ = [
|
||||
"PromptStrategy",
|
||||
"StrategyFactory",
|
||||
"ChatMLStrategy",
|
||||
"AlpacaStrategy",
|
||||
]
|
||||
@@ -0,0 +1,43 @@
|
||||
"""Alpaca format strategy."""
|
||||
from typing import List
|
||||
|
||||
from pipeline.tokenizer import BpeTokenizer
|
||||
from pipeline.strategies.base import PromptStrategy
|
||||
from pipeline.strategies.factory import StrategyFactory
|
||||
|
||||
|
||||
@StrategyFactory.register("alpaca")
|
||||
class AlpacaStrategy(PromptStrategy):
|
||||
"""Alpaca format: ``### Instruction: ... \\n\\n### Response: ... <eos>``"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
tokenizer: BpeTokenizer,
|
||||
instruction_start: str = "### Instruction:\n",
|
||||
response_start: str = "### Response:\n",
|
||||
response_suffix: str = "\n<eos>",
|
||||
):
|
||||
super().__init__(tokenizer)
|
||||
self.instruction_start = instruction_start
|
||||
self.response_start = response_start
|
||||
self.response_suffix = response_suffix
|
||||
|
||||
self._instruction_start_ids = self._encode_format(instruction_start)
|
||||
self._separator_ids = self._encode_format("\n\n")
|
||||
self._response_start_ids = self._encode_format(response_start)
|
||||
self._response_suffix_ids = self._encode_format(response_suffix)
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "alpaca"
|
||||
|
||||
def assemble_prompt(self, query_tokens: List[int]) -> List[int]:
|
||||
return (
|
||||
self._instruction_start_ids
|
||||
+ query_tokens
|
||||
+ self._separator_ids
|
||||
+ self._response_start_ids
|
||||
)
|
||||
|
||||
def assemble_response(self, response_tokens: List[int]) -> List[int]:
|
||||
return response_tokens + self._response_suffix_ids
|
||||
@@ -0,0 +1,41 @@
|
||||
"""Abstract base class for prompt construction strategies."""
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import List
|
||||
|
||||
from pipeline.tokenizer import BpeTokenizer
|
||||
|
||||
|
||||
class PromptStrategy(ABC):
|
||||
"""Abstract base for prompt/response format strategies.
|
||||
|
||||
Strategies operate at the token level: the Processor tokenizes raw
|
||||
text (query, response, …) and passes token lists to the Strategy,
|
||||
which assembles them with pre-encoded format tokens.
|
||||
"""
|
||||
|
||||
def __init__(self, tokenizer: BpeTokenizer):
|
||||
self.tokenizer = tokenizer
|
||||
|
||||
def _encode_format(self, text: str) -> List[int]:
|
||||
"""Encode a format string that may contain special tokens."""
|
||||
return self.tokenizer.encode(text, add_special_tokens=False)
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def name(self) -> str:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def assemble_prompt(self, query_tokens: List[int]) -> List[int]:
|
||||
"""Assemble query tokens into a complete prompt with format tokens.
|
||||
|
||||
The prompt includes all tokens up to (and including) the response
|
||||
start marker, e.g. ``<|im_start|>assistant\n``.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def assemble_response(self, response_tokens: List[int]) -> List[int]:
|
||||
"""Wrap response tokens with format tokens (suffix, eos, etc)."""
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"{self.__class__.__name__}(name='{self.name}')"
|
||||
@@ -0,0 +1,41 @@
|
||||
"""ChatML format strategy."""
|
||||
from typing import List
|
||||
|
||||
from pipeline.tokenizer import BpeTokenizer
|
||||
from pipeline.strategies.base import PromptStrategy
|
||||
from pipeline.strategies.factory import StrategyFactory
|
||||
|
||||
|
||||
@StrategyFactory.register("chatml")
|
||||
class ChatMLStrategy(PromptStrategy):
|
||||
"""ChatML format: ``<|im_start|>user ... <|im_end|> <|im_start|>assistant ... <|im_end|> <eos>``"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
tokenizer: BpeTokenizer,
|
||||
user_start: str = "<|im_start|>user\n",
|
||||
user_end: str = "<|im_end|>\n",
|
||||
assistant_start: str = "<|im_start|>assistant\n",
|
||||
assistant_end: str = "<|im_end|>\n<eos>",
|
||||
):
|
||||
super().__init__(tokenizer)
|
||||
|
||||
self._user_start_ids = self._encode_format(user_start)
|
||||
self._user_end_ids = self._encode_format(user_end)
|
||||
self._assistant_start_ids = self._encode_format(assistant_start)
|
||||
self._assistant_end_ids = self._encode_format(assistant_end)
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "chatml"
|
||||
|
||||
def assemble_prompt(self, query_tokens: List[int]) -> List[int]:
|
||||
return (
|
||||
self._user_start_ids
|
||||
+ query_tokens
|
||||
+ self._user_end_ids
|
||||
+ self._assistant_start_ids
|
||||
)
|
||||
|
||||
def assemble_response(self, response_tokens: List[int]) -> List[int]:
|
||||
return response_tokens + self._assistant_end_ids
|
||||
@@ -0,0 +1,70 @@
|
||||
"""Factory for creating and registering prompt strategies."""
|
||||
from typing import Dict, List, Type
|
||||
|
||||
from pipeline.tokenizer import BpeTokenizer
|
||||
from pipeline.strategies.base import PromptStrategy
|
||||
|
||||
|
||||
class StrategyFactory:
|
||||
"""Registry and factory for PromptStrategy implementations.
|
||||
|
||||
Supports decorator-based registration for extensible strategy types.
|
||||
|
||||
Example usage::
|
||||
|
||||
@StrategyFactory.register("custom")
|
||||
class CustomStrategy(PromptStrategy):
|
||||
...
|
||||
|
||||
strategy = StrategyFactory.create("custom", tokenizer, **kwargs)
|
||||
"""
|
||||
|
||||
STRATEGY_MAP: Dict[str, Type[PromptStrategy]] = {}
|
||||
|
||||
@classmethod
|
||||
def register(cls, name: str):
|
||||
"""Decorator to register a new strategy class.
|
||||
|
||||
Args:
|
||||
name: Registration name for the strategy.
|
||||
|
||||
Returns:
|
||||
Decorator function that registers the strategy class.
|
||||
"""
|
||||
|
||||
def decorator(strategy_cls: Type[PromptStrategy]) -> Type[PromptStrategy]:
|
||||
if not issubclass(strategy_cls, PromptStrategy):
|
||||
raise TypeError(
|
||||
f"{strategy_cls.__name__} must inherit from PromptStrategy"
|
||||
)
|
||||
cls.STRATEGY_MAP[name] = strategy_cls
|
||||
return strategy_cls
|
||||
|
||||
return decorator
|
||||
|
||||
@classmethod
|
||||
def create(cls, name: str, tokenizer: BpeTokenizer, **kwargs) -> PromptStrategy:
|
||||
"""Create a strategy by name.
|
||||
|
||||
Args:
|
||||
name: Registered strategy name (e.g. ``"chatml"``, ``"alpaca"``).
|
||||
tokenizer: Tokenizer instance (required by all strategies).
|
||||
**kwargs: Forwarded to the strategy constructor.
|
||||
|
||||
Returns:
|
||||
Strategy instance.
|
||||
|
||||
Raises:
|
||||
ValueError: If name is not registered.
|
||||
"""
|
||||
if name not in cls.STRATEGY_MAP:
|
||||
raise ValueError(
|
||||
f"Unknown strategy: '{name}'. "
|
||||
f"Supported types: {sorted(cls.STRATEGY_MAP.keys())}"
|
||||
)
|
||||
return cls.STRATEGY_MAP[name](tokenizer=tokenizer, **kwargs)
|
||||
|
||||
@classmethod
|
||||
def available_types(cls) -> List[str]:
|
||||
"""Return list of registered strategy type names."""
|
||||
return list(cls.STRATEGY_MAP.keys())
|
||||
@@ -15,7 +15,6 @@ dependencies = [
|
||||
"tqdm>=4.64.0",
|
||||
"torch>=2.1.0",
|
||||
"h5py>=3.7.0",
|
||||
"regex>=2022.03.15"
|
||||
]
|
||||
keywords = ["nlp", "datasets", "language-models", "machine-learning"]
|
||||
[project.urls]
|
||||
|
||||
+25
-34
@@ -1,46 +1,36 @@
|
||||
"""JSONL → H5 缓存脚本
|
||||
"""JSONL to H5 caching script.
|
||||
|
||||
将 dataset/ 下的 JSONL 文件 tokenize 并打包为 HDF5 格式。
|
||||
Tokenize JSONL files and pack them into HDF5 format.
|
||||
|
||||
用法:
|
||||
Usage:
|
||||
python scripts/cache_h5.py pt ./dataset/chinese-c4-pretrain
|
||||
python scripts/cache_h5.py sft ./dataset/belle-sft --pack-size 4096
|
||||
python scripts/cache_h5.py sft ./dataset/belle-sft --pack-size 4096 --strategy alpaca
|
||||
python scripts/cache_h5.py sft ./dataset/Ling-Coder-sft --tokenizer ./my_tokenizer.json
|
||||
"""
|
||||
import argparse
|
||||
import os
|
||||
|
||||
from pipeline import BpeTokenizer, ProcessorFactory, cache_jsonl, IOHandler
|
||||
|
||||
|
||||
def collect_jsonl_files(input_dir: str) -> list[str]:
|
||||
"""收集目录下所有 JSONL 文件"""
|
||||
files = [
|
||||
os.path.join(root, f)
|
||||
for root, _, filenames in os.walk(input_dir)
|
||||
for f in filenames
|
||||
if f.endswith(".jsonl")
|
||||
]
|
||||
files.sort()
|
||||
return files
|
||||
from pipeline import BpeTokenizer, ProcessorFactory, cache_jsonl
|
||||
from pipeline.io import IOHandler
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="JSONL → H5 缓存")
|
||||
parser.add_argument("type", choices=["pt", "sft", "dpo"], help="处理器类型")
|
||||
parser.add_argument("input_dir", help="JSONL 文件所在目录")
|
||||
parser = argparse.ArgumentParser(description="JSONL -> H5 cache")
|
||||
parser.add_argument("type", choices=["pt", "sft", "dpo"], help="Processor type")
|
||||
parser.add_argument("input_dir", help="Directory containing JSONL files")
|
||||
parser.add_argument("-o", "--output-dir", default=None,
|
||||
help="H5 输出目录 (默认: <input_dir>/cached)")
|
||||
help="H5 output dir (default: <input_dir>/cached)")
|
||||
parser.add_argument("-t", "--tokenizer", default="./tokenizer.json",
|
||||
help="Tokenizer 路径 (默认: ./tokenizer.json)")
|
||||
help="Tokenizer path (default: ./tokenizer.json)")
|
||||
parser.add_argument("-s", "--strategy", default=None,
|
||||
help="Prompt strategy: chatml, alpaca (default: chatml)")
|
||||
parser.add_argument("-p", "--pack-size", type=int, default=-1,
|
||||
help="序列打包长度，<=0 不打包 (默认: -1)")
|
||||
help="Pack size, <=0 to disable (default: -1)")
|
||||
parser.add_argument("--pad-value", type=int, default=1,
|
||||
help="打包填充值 (默认: 1 即 <eos>)")
|
||||
help="Padding value (default: 1)")
|
||||
args = parser.parse_args()
|
||||
|
||||
# 收集 JSONL 文件
|
||||
jsonl_files = collect_jsonl_files(args.input_dir)
|
||||
jsonl_files = IOHandler.fetch_files(args.input_dir, suffix=".jsonl")
|
||||
if not jsonl_files:
|
||||
print(f"[ERROR] No JSONL files found in {args.input_dir}")
|
||||
return
|
||||
@@ -49,37 +39,38 @@ def main():
|
||||
for f in jsonl_files:
|
||||
print(f" - {f}")
|
||||
|
||||
# 加载 tokenizer
|
||||
if not os.path.exists(args.tokenizer):
|
||||
print(f"[ERROR] Tokenizer not found: {args.tokenizer}")
|
||||
return
|
||||
tokenizer = BpeTokenizer(args.tokenizer)
|
||||
print(f"Tokenizer loaded: vocab_size={len(tokenizer)}")
|
||||
|
||||
# 创建处理器
|
||||
processor = ProcessorFactory.create(args.type, tokenizer)
|
||||
if args.strategy:
|
||||
processor = ProcessorFactory.create_with_strategy_name(
|
||||
args.type, tokenizer, args.strategy
|
||||
)
|
||||
else:
|
||||
processor = ProcessorFactory.create(args.type, tokenizer)
|
||||
|
||||
print(f"Processor: {args.type} ({processor.__class__.__name__})")
|
||||
print(f"Output keys: {processor.output_keys}")
|
||||
|
||||
# 输出目录
|
||||
output_dir = args.output_dir or os.path.join(args.input_dir, "cached")
|
||||
|
||||
# 执行缓存
|
||||
print(f"\nStart caching...")
|
||||
if args.pack_size > 0:
|
||||
print(f" pack_size={args.pack_size}, pad_value={args.pad_value}")
|
||||
else:
|
||||
print(f" no packing")
|
||||
|
||||
output_files = cache_jsonl(
|
||||
cache_jsonl(
|
||||
files=jsonl_files,
|
||||
output_dir=output_dir,
|
||||
processor=processor,
|
||||
pack_size=args.pack_size,
|
||||
pad_value=args.pad_value,
|
||||
)
|
||||
|
||||
print(f"\nDone! {len(output_files)} H5 files saved to {output_dir}")
|
||||
print(f"\nDone! Output saved to {output_dir}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -4,7 +4,7 @@ from pipeline import export_dataset
|
||||
|
||||
def process_func(input_dict: dict):
|
||||
return {
|
||||
"prompt": input_dict["prompt"],
|
||||
"query": input_dict["prompt"],
|
||||
"chosen": input_dict["chosen"],
|
||||
"rejected": input_dict["rejected"],
|
||||
}
|
||||
|
||||
@@ -17,12 +17,12 @@ if __name__ == "__main__":
|
||||
'psycho-10k-dpsk-r1', 'sof-c-zh', 'industryinstruction', 'Chinese-QA-AFAF',
|
||||
]
|
||||
|
||||
datasets = []
|
||||
dataset_list = []
|
||||
for subset in all_data:
|
||||
ds = load_dataset("Mxode/Chinese-Instruct", name=subset)
|
||||
datasets.append(ds["train"])
|
||||
dataset_list.append(ds["train"])
|
||||
|
||||
combined_dataset = concatenate_datasets(datasets)
|
||||
combined_dataset = concatenate_datasets(dataset_list)
|
||||
export_dataset(
|
||||
dataset=combined_dataset,
|
||||
output_dir="./dataset",
|
||||
|
||||
+17
-73
@@ -1,4 +1,4 @@
|
||||
"""单元测试：pipeline.cache 模块中的 cache_jsonl 函数"""
|
||||
"""Tests for pipeline.cache module."""
|
||||
|
||||
import json
|
||||
import os
|
||||
@@ -6,12 +6,12 @@ import tempfile
|
||||
import torch
|
||||
from pathlib import Path
|
||||
|
||||
from pipeline.cache import cache_jsonl
|
||||
from pipeline.io import cache_jsonl
|
||||
from pipeline.processors import BaseProcessor
|
||||
|
||||
|
||||
class DummyProcessor(BaseProcessor):
|
||||
"""用于测试的虚拟处理器"""
|
||||
"""Dummy processor for testing."""
|
||||
|
||||
def __init__(self):
|
||||
self._output_keys = ["sequence", "loss_mask"]
|
||||
@@ -22,8 +22,7 @@ class DummyProcessor(BaseProcessor):
|
||||
|
||||
def process(self, item):
|
||||
text = item.get("text", "")
|
||||
tokens = [ord(c) for c in text[:10]] # 简单模拟tokenize
|
||||
|
||||
tokens = [ord(c) for c in text[:10]]
|
||||
return {
|
||||
"sequence": torch.tensor(tokens, dtype=torch.int32),
|
||||
"loss_mask": torch.ones(len(tokens), dtype=torch.int32),
|
||||
@@ -31,99 +30,57 @@ class DummyProcessor(BaseProcessor):
|
||||
|
||||
|
||||
class TestCacheJsonl:
|
||||
"""cache_jsonl 函数的测试套件"""
|
||||
|
||||
def test_basic_cache_functionality(self):
|
||||
"""测试基本缓存功能：处理简单JSONL文件并生成HDF5"""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
# 创建测试JSONL文件
|
||||
jsonl_path = os.path.join(tmpdir, "test.jsonl")
|
||||
test_data = [
|
||||
{"text": "hello"},
|
||||
{"text": "world"},
|
||||
{"text": "test"},
|
||||
]
|
||||
test_data = [{"text": "hello"}, {"text": "world"}, {"text": "test"}]
|
||||
with open(jsonl_path, "w", encoding="utf-8") as f:
|
||||
for item in test_data:
|
||||
f.write(json.dumps(item) + "\n")
|
||||
|
||||
# 创建处理器
|
||||
processor = DummyProcessor()
|
||||
|
||||
# 调用 cache_jsonl
|
||||
output_files = cache_jsonl(
|
||||
files=[jsonl_path],
|
||||
output_dir=tmpdir,
|
||||
processor=processor,
|
||||
pack_size=-1, # 不打包模式
|
||||
pad_value=0,
|
||||
files=[jsonl_path], output_dir=tmpdir,
|
||||
processor=processor, pack_size=-1, pad_value=0,
|
||||
)
|
||||
|
||||
# 验证输出
|
||||
assert len(output_files) == 1
|
||||
assert os.path.exists(output_files[0])
|
||||
|
||||
def test_packer_state_independence(self):
|
||||
"""测试打包器状态独立性：验证不同 output_key 的打包结果是否独立"""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
# 创建测试JSONL文件，包含不同长度的文本
|
||||
jsonl_path = os.path.join(tmpdir, "test.jsonl")
|
||||
test_data = [
|
||||
{"text": "ab"}, # 2 chars
|
||||
{"text": "abcde"}, # 5 chars
|
||||
{"text": "abc"}, # 3 chars
|
||||
]
|
||||
test_data = [{"text": "ab"}, {"text": "abcde"}, {"text": "abc"}]
|
||||
with open(jsonl_path, "w", encoding="utf-8") as f:
|
||||
for item in test_data:
|
||||
f.write(json.dumps(item) + "\n")
|
||||
|
||||
# 创建处理器
|
||||
processor = DummyProcessor()
|
||||
|
||||
# 调用 cache_jsonl，使用打包模式
|
||||
output_files = cache_jsonl(
|
||||
files=[jsonl_path],
|
||||
output_dir=tmpdir,
|
||||
processor=processor,
|
||||
pack_size=10, # 打包模式
|
||||
pad_value=0,
|
||||
files=[jsonl_path], output_dir=tmpdir,
|
||||
processor=processor, pack_size=10, pad_value=0,
|
||||
)
|
||||
|
||||
# 验证输出文件存在
|
||||
assert len(output_files) == 1
|
||||
assert os.path.exists(output_files[0])
|
||||
|
||||
def test_no_packing_mode(self):
|
||||
"""测试无打包模式（pack_size <= 0）"""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
# 创建测试JSONL文件
|
||||
jsonl_path = os.path.join(tmpdir, "test.jsonl")
|
||||
test_data = [
|
||||
{"text": "hello"},
|
||||
{"text": "world"},
|
||||
]
|
||||
test_data = [{"text": "hello"}, {"text": "world"}]
|
||||
with open(jsonl_path, "w", encoding="utf-8") as f:
|
||||
for item in test_data:
|
||||
f.write(json.dumps(item) + "\n")
|
||||
|
||||
processor = DummyProcessor()
|
||||
|
||||
# 打包大小设为0表示不打包
|
||||
output_files = cache_jsonl(
|
||||
files=[jsonl_path],
|
||||
output_dir=tmpdir,
|
||||
processor=processor,
|
||||
pack_size=0,
|
||||
pad_value=-1,
|
||||
files=[jsonl_path], output_dir=tmpdir,
|
||||
processor=processor, pack_size=0, pad_value=-1,
|
||||
)
|
||||
|
||||
assert len(output_files) == 1
|
||||
assert os.path.exists(output_files[0])
|
||||
|
||||
def test_multiple_files(self):
|
||||
"""测试处理多个文件"""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
# 创建两个测试JSONL文件
|
||||
files = []
|
||||
for i in range(2):
|
||||
jsonl_path = os.path.join(tmpdir, f"test{i}.jsonl")
|
||||
@@ -133,33 +90,20 @@ class TestCacheJsonl:
|
||||
files.append(jsonl_path)
|
||||
|
||||
processor = DummyProcessor()
|
||||
|
||||
output_files = cache_jsonl(
|
||||
files=files,
|
||||
output_dir=tmpdir,
|
||||
processor=processor,
|
||||
pack_size=-1,
|
||||
pad_value=0,
|
||||
files=files, output_dir=tmpdir,
|
||||
processor=processor, pack_size=-1, pad_value=0,
|
||||
)
|
||||
|
||||
assert len(output_files) == 2
|
||||
|
||||
def test_empty_file_handling(self):
|
||||
"""测试处理空文件"""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
# 创建空JSONL文件
|
||||
jsonl_path = os.path.join(tmpdir, "empty.jsonl")
|
||||
Path(jsonl_path).touch()
|
||||
|
||||
processor = DummyProcessor()
|
||||
|
||||
# 不应该抛出异常
|
||||
output_files = cache_jsonl(
|
||||
files=[jsonl_path],
|
||||
output_dir=tmpdir,
|
||||
processor=processor,
|
||||
pack_size=-1,
|
||||
pad_value=0,
|
||||
files=[jsonl_path], output_dir=tmpdir,
|
||||
processor=processor, pack_size=-1, pad_value=0,
|
||||
)
|
||||
|
||||
assert len(output_files) == 1
|
||||
+19
-106
@@ -1,4 +1,4 @@
|
||||
"""单元测试：pipeline.io 模块中的 IOHandler 类"""
|
||||
"""Tests for pipeline.io module."""
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
@@ -11,177 +11,90 @@ from pipeline.io import IOHandler
|
||||
|
||||
|
||||
class TestIOHandler:
|
||||
"""IOHandler 类的测试套件"""
|
||||
|
||||
def test_fetch_files_in_directory(self):
|
||||
"""测试 fetch_files 方法能正确获取目录中的文件"""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
# 创建测试文件
|
||||
test_file1 = os.path.join(tmpdir, "file1.txt")
|
||||
test_file2 = os.path.join(tmpdir, "file2.txt")
|
||||
Path(test_file1).touch()
|
||||
Path(test_file2).touch()
|
||||
|
||||
# 创建子目录和文件
|
||||
Path(tmpdir, "file1.txt").touch()
|
||||
Path(tmpdir, "file2.txt").touch()
|
||||
subdir = os.path.join(tmpdir, "subdir")
|
||||
os.makedirs(subdir)
|
||||
test_file3 = os.path.join(subdir, "file3.txt")
|
||||
Path(test_file3).touch()
|
||||
Path(subdir, "file3.txt").touch()
|
||||
|
||||
# 获取文件列表
|
||||
files = IOHandler.fetch_files(tmpdir)
|
||||
|
||||
# 验证
|
||||
assert len(files) == 3
|
||||
assert any("file1.txt" in f for f in files)
|
||||
assert any("file2.txt" in f for f in files)
|
||||
assert any("file3.txt" in f for f in files)
|
||||
|
||||
def test_fetch_files_empty_directory(self):
|
||||
"""测试空目录返回空列表"""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
files = IOHandler.fetch_files(tmpdir)
|
||||
assert files == []
|
||||
assert IOHandler.fetch_files(tmpdir) == []
|
||||
|
||||
def test_fetch_folders_in_directory(self):
|
||||
"""测试 fetch_folders 方法能正确获取子目录"""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
# 创建子目录
|
||||
subdir1 = os.path.join(tmpdir, "folder1")
|
||||
subdir2 = os.path.join(tmpdir, "folder2")
|
||||
os.makedirs(subdir1)
|
||||
os.makedirs(subdir2)
|
||||
os.makedirs(os.path.join(tmpdir, "folder1"))
|
||||
os.makedirs(os.path.join(tmpdir, "folder2"))
|
||||
os.makedirs(os.path.join(tmpdir, "folder1", "nested"))
|
||||
|
||||
# 创建嵌套子目录
|
||||
nested = os.path.join(subdir1, "nested")
|
||||
os.makedirs(nested)
|
||||
|
||||
# 获取文件夹列表
|
||||
folders = IOHandler.fetch_folders(tmpdir)
|
||||
|
||||
# 验证
|
||||
assert len(folders) == 3
|
||||
assert any("folder1" in f for f in folders)
|
||||
assert any("folder2" in f for f in folders)
|
||||
assert any("nested" in f for f in folders)
|
||||
|
||||
def test_fetch_folders_with_filter(self):
|
||||
"""测试 fetch_folders 方法的过滤功能"""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
# 创建子目录
|
||||
subdir1 = os.path.join(tmpdir, "folder1")
|
||||
subdir2 = os.path.join(tmpdir, "folder2")
|
||||
os.makedirs(subdir1)
|
||||
os.makedirs(subdir2)
|
||||
os.makedirs(os.path.join(tmpdir, "folder1"))
|
||||
os.makedirs(os.path.join(tmpdir, "folder2"))
|
||||
|
||||
# 使用过滤器只获取 folder1
|
||||
folders = IOHandler.fetch_folders(
|
||||
tmpdir,
|
||||
filter_func=lambda x: "folder1" in x
|
||||
)
|
||||
|
||||
# 验证
|
||||
folders = IOHandler.fetch_folders(tmpdir, filter_func=lambda x: "folder1" in x)
|
||||
assert len(folders) == 1
|
||||
assert "folder1" in folders[0]
|
||||
|
||||
def test_save_and_load_h5(self):
|
||||
"""测试 save_h5 和 load_h5 方法的读写功能"""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
# 创建测试数据
|
||||
tensor_group = {
|
||||
"sequence": [torch.tensor([1, 2, 3], dtype=torch.int32)],
|
||||
"labels": [torch.tensor([4, 5], dtype=torch.int32)],
|
||||
}
|
||||
|
||||
# 保存
|
||||
IOHandler.save_h5(tmpdir, "test", tensor_group)
|
||||
|
||||
# 验证文件已创建
|
||||
h5_path = os.path.join(tmpdir, "test.h5")
|
||||
assert os.path.exists(h5_path)
|
||||
|
||||
# 加载 - 传入目录而不是单个文件
|
||||
assert os.path.exists(os.path.join(tmpdir, "test.h5"))
|
||||
loaded = IOHandler.load_h5(tmpdir, share_memory=False)
|
||||
|
||||
# 验证数据
|
||||
assert "sequence" in loaded
|
||||
assert "labels" in loaded
|
||||
assert len(loaded["sequence"]) == 1
|
||||
assert len(loaded["labels"]) == 1
|
||||
assert torch.equal(loaded["sequence"][0], torch.tensor([1, 2, 3], dtype=torch.int32))
|
||||
assert torch.equal(loaded["labels"][0], torch.tensor([4, 5], dtype=torch.int32))
|
||||
|
||||
def test_save_h5_creates_directory(self):
|
||||
"""测试 save_h5 自动创建输出目录"""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
output_dir = os.path.join(tmpdir, "nested", "output")
|
||||
|
||||
tensor_group = {
|
||||
"data": [torch.tensor([1, 2, 3])],
|
||||
}
|
||||
|
||||
# 保存到不存在的目录
|
||||
IOHandler.save_h5(output_dir, "test", tensor_group)
|
||||
|
||||
# 验证目录已创建
|
||||
IOHandler.save_h5(output_dir, "test", {"data": [torch.tensor([1, 2, 3])]})
|
||||
assert os.path.exists(output_dir)
|
||||
assert os.path.exists(os.path.join(output_dir, "test.h5"))
|
||||
|
||||
def test_load_h5_multiple_files(self):
|
||||
"""测试 load_h5 方法能处理多个 H5 文件"""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
# 创建第一个 H5 文件
|
||||
h5_path1 = os.path.join(tmpdir, "file1.h5")
|
||||
with h5py.File(h5_path1, 'w') as f:
|
||||
grp = f.create_group("data")
|
||||
grp.create_dataset('data_0', data=[1, 2, 3])
|
||||
for i, data in enumerate([[1, 2, 3], [4, 5, 6]]):
|
||||
h5_path = os.path.join(tmpdir, f"file{i}.h5")
|
||||
with h5py.File(h5_path, 'w') as f:
|
||||
grp = f.create_group("data")
|
||||
grp.create_dataset('data_0', data=data)
|
||||
|
||||
# 创建第二个 H5 文件
|
||||
h5_path2 = os.path.join(tmpdir, "file2.h5")
|
||||
with h5py.File(h5_path2, 'w') as f:
|
||||
grp = f.create_group("data")
|
||||
grp.create_dataset('data_0', data=[4, 5, 6])
|
||||
|
||||
# 加载目录
|
||||
loaded = IOHandler.load_h5(tmpdir, share_memory=False)
|
||||
|
||||
# 验证
|
||||
assert "data" in loaded
|
||||
assert len(loaded["data"]) == 2
|
||||
|
||||
def test_load_h5_with_rglob(self):
|
||||
"""测试 load_h5 能递归查找 H5 文件"""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
# 在子目录中创建 H5 文件
|
||||
subdir = os.path.join(tmpdir, "subdir")
|
||||
os.makedirs(subdir)
|
||||
h5_path = os.path.join(subdir, "nested.h5")
|
||||
|
||||
with h5py.File(h5_path, 'w') as f:
|
||||
grp = f.create_group("test")
|
||||
grp.create_dataset('data_0', data=[1, 2])
|
||||
|
||||
# 加载根目录
|
||||
loaded = IOHandler.load_h5(tmpdir, share_memory=False)
|
||||
|
||||
# 验证能找到子目录中的文件
|
||||
assert "test" in loaded
|
||||
assert len(loaded["test"]) == 1
|
||||
|
||||
def test_save_h5_multiple_tensors_per_key(self):
|
||||
"""测试 save_h5 能保存多个张量到同一键"""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
tensor_group = {
|
||||
"batch": [
|
||||
torch.tensor([1, 2]),
|
||||
torch.tensor([3, 4, 5]),
|
||||
torch.tensor([6]),
|
||||
],
|
||||
"batch": [torch.tensor([1, 2]), torch.tensor([3, 4, 5]), torch.tensor([6])],
|
||||
}
|
||||
|
||||
IOHandler.save_h5(tmpdir, "multi", tensor_group)
|
||||
|
||||
# 加载目录而不是单个文件
|
||||
loaded = IOHandler.load_h5(tmpdir, share_memory=False)
|
||||
|
||||
assert len(loaded["batch"]) == 3
|
||||
+59
-189
@@ -1,4 +1,4 @@
|
||||
"""单元测试：pipeline.packing 模块中的 SequencePacker 类"""
|
||||
"""Tests for pipeline.packing module."""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
@@ -6,252 +6,122 @@ from pipeline.packing import SequencePacker
|
||||
|
||||
|
||||
class TestSequencePacker:
|
||||
"""SequencePacker 类的测试套件"""
|
||||
|
||||
def test_normal_packing(self):
|
||||
"""测试正常打包场景：多个序列正确打包成固定长度的包"""
|
||||
packer = SequencePacker(pack_size=10, pad_value=0)
|
||||
|
||||
sequences = [
|
||||
torch.tensor([1, 2, 3], dtype=torch.int32),
|
||||
torch.tensor([4, 5], dtype=torch.int32),
|
||||
torch.tensor([6, 7, 8, 9], dtype=torch.int32),
|
||||
]
|
||||
|
||||
packages = packer.pack(sequences)
|
||||
|
||||
# 验证至少有包输出
|
||||
assert len(packages) >= 1
|
||||
|
||||
# 验证每个包的长度是正确的
|
||||
for pkg in packages:
|
||||
assert pkg.shape == (10,)
|
||||
# 验证填充值
|
||||
# 检查所有非零元素都在前几个位置，或者包是满的
|
||||
non_zero_count = (pkg != 0).sum().item()
|
||||
# 非零元素的数量应该等于原始序列元素的总和
|
||||
total_elements = sum(s.numel() for s in sequences)
|
||||
# 由于打包，第一个包包含3+2=5个元素，第二个包包含4个元素
|
||||
# 第一个包应该包含前两个序列
|
||||
pkg1 = packages[0]
|
||||
# 序列[1,2,3]和[4,5]按长度降序排序后是[1,2,3]在前，然后[4,5]
|
||||
# 但排序是原地修改...等等，我们已经修复了使用sorted()
|
||||
# 所以排序后的顺序是[6,7,8,9], [1,2,3], [4,5]
|
||||
# 第一个包包含[6,7,8,9]和部分[1,2,3] = 4+3=7，剩余3个位置放[4,5]
|
||||
# 所以第一个包应该是[6,7,8,9,1,2,3,4,5,0]
|
||||
|
||||
# 简化测试：验证打包后的张量包含所有原始数据
|
||||
# Verify all original values are present
|
||||
all_values = []
|
||||
for pkg in packages:
|
||||
non_zero = pkg[pkg != 0].tolist()
|
||||
all_values.extend(non_zero)
|
||||
|
||||
# 检查所有原始数据是否都被包含
|
||||
original_values = [1, 2, 3, 4, 5, 6, 7, 8, 9]
|
||||
for val in original_values:
|
||||
assert val in all_values, f"Value {val} not found in packages"
|
||||
all_values.extend(pkg[pkg != 0].tolist())
|
||||
for val in [1, 2, 3, 4, 5, 6, 7, 8, 9]:
|
||||
assert val in all_values
|
||||
|
||||
def test_empty_list_input(self):
|
||||
"""测试空列表输入"""
|
||||
packer = SequencePacker(pack_size=10)
|
||||
|
||||
packages = packer.pack([])
|
||||
|
||||
assert packages == []
|
||||
|
||||
# 验证内部状态已正确初始化
|
||||
assert packer._current_pack is not None
|
||||
assert packer._current_pos == 0
|
||||
assert packer.pack([]) == []
|
||||
|
||||
def test_single_sequence_input(self):
|
||||
"""测试单个序列输入"""
|
||||
packer = SequencePacker(pack_size=10, pad_value=-1)
|
||||
|
||||
sequences = [torch.tensor([1, 2, 3], dtype=torch.int32)]
|
||||
|
||||
packages = packer.pack(sequences)
|
||||
|
||||
packages = packer.pack([torch.tensor([1, 2, 3], dtype=torch.int32)])
|
||||
assert len(packages) == 1
|
||||
pkg = packages[0]
|
||||
assert pkg.shape == (10,)
|
||||
assert pkg[:3].tolist() == [1, 2, 3]
|
||||
assert pkg[3:].tolist() == [-1] * 7
|
||||
assert packages[0][:3].tolist() == [1, 2, 3]
|
||||
assert packages[0][3:].tolist() == [-1] * 7
|
||||
|
||||
def test_truncate_long_sequence(self, caplog):
|
||||
"""测试超长序列截断，验证警告日志是否触发"""
|
||||
packer = SequencePacker(pack_size=5, pad_value=0)
|
||||
|
||||
sequences = [
|
||||
torch.tensor([1, 2, 3, 4, 5, 6, 7, 8], dtype=torch.int32), # 长度8，超过pack_size=5
|
||||
]
|
||||
|
||||
packages = packer.pack(sequences)
|
||||
|
||||
packages = packer.pack([torch.tensor([1, 2, 3, 4, 5, 6, 7, 8], dtype=torch.int32)])
|
||||
assert len(packages) == 1
|
||||
pkg = packages[0]
|
||||
assert pkg.shape == (5,)
|
||||
assert pkg.tolist() == [1, 2, 3, 4, 5] # 只保留前5个元素
|
||||
|
||||
# 验证警告日志已触发
|
||||
assert packages[0].tolist() == [1, 2, 3, 4, 5]
|
||||
assert "truncating" in caplog.text.lower() or "exceeds" in caplog.text.lower()
|
||||
|
||||
def test_padding_value(self):
|
||||
"""测试填充值正确应用"""
|
||||
packer = SequencePacker(pack_size=8, pad_value=99)
|
||||
|
||||
sequences = [
|
||||
packages = packer.pack([
|
||||
torch.tensor([1, 2], dtype=torch.int32),
|
||||
torch.tensor([3], dtype=torch.int32),
|
||||
]
|
||||
|
||||
packages = packer.pack(sequences)
|
||||
|
||||
assert len(packages) == 1
|
||||
pkg = packages[0]
|
||||
|
||||
# 前3个元素是数据
|
||||
assert pkg[:3].tolist() == [1, 2, 3]
|
||||
# 后5个元素是填充值
|
||||
assert pkg[3:].tolist() == [99] * 5
|
||||
])
|
||||
assert packages[0][:3].tolist() == [1, 2, 3]
|
||||
assert packages[0][3:].tolist() == [99] * 5
|
||||
|
||||
def test_different_dtypes(self):
|
||||
"""测试支持不同 dtype (int32, int64, float32)"""
|
||||
# int32
|
||||
packer_int32 = SequencePacker(pack_size=10, pad_value=0, dtype=torch.int32)
|
||||
sequences_int32 = [torch.tensor([1, 2, 3], dtype=torch.int32)]
|
||||
packages_int32 = packer_int32.pack(sequences_int32)
|
||||
assert packages_int32[0].dtype == torch.int32
|
||||
for dtype in [torch.int32, torch.int64, torch.float32]:
|
||||
packer = SequencePacker(pack_size=10, dtype=dtype)
|
||||
val = 1.0 if dtype == torch.float32 else 1
|
||||
packages = packer.pack([torch.tensor([val, 2, 3], dtype=dtype)])
|
||||
assert packages[0].dtype == dtype
|
||||
|
||||
# int64
|
||||
packer_int64 = SequencePacker(pack_size=10, pad_value=0, dtype=torch.int64)
|
||||
sequences_int64 = [torch.tensor([1, 2, 3], dtype=torch.int64)]
|
||||
packages_int64 = packer_int64.pack(sequences_int64)
|
||||
assert packages_int64[0].dtype == torch.int64
|
||||
|
||||
# float32
|
||||
packer_float32 = SequencePacker(pack_size=10, pad_value=0.0, dtype=torch.float32)
|
||||
sequences_float32 = [torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32)]
|
||||
packages_float32 = packer_float32.pack(sequences_float32)
|
||||
assert packages_float32[0].dtype == torch.float32
|
||||
def test_dtype_conversion_on_mismatch(self, caplog):
|
||||
"""Tensors with mismatched dtype are silently converted."""
|
||||
packer = SequencePacker(pack_size=10, dtype=torch.int32)
|
||||
packages = packer.pack([torch.tensor([1, 2, 3], dtype=torch.int64)])
|
||||
assert packages[0].dtype == torch.int32
|
||||
assert packages[0][:3].tolist() == [1, 2, 3]
|
||||
|
||||
def test_non_1d_tensor_raises_error(self):
|
||||
"""测试非1D张量是否抛出异常"""
|
||||
packer = SequencePacker(pack_size=10)
|
||||
|
||||
# 2D 张量应该抛出异常
|
||||
sequences_2d = [torch.tensor([[1, 2], [3, 4]])] # shape: (2, 2)
|
||||
with pytest.raises(ValueError, match="Expected 1D tensor"):
|
||||
packer.pack(sequences_2d)
|
||||
|
||||
# 0D 张量 (标量) 应该抛出异常
|
||||
sequences_0d = [torch.tensor(5)] # shape: ()
|
||||
packer.pack([torch.tensor([[1, 2], [3, 4]])])
|
||||
with pytest.raises(ValueError, match="Expected 1D tensor"):
|
||||
packer.pack(sequences_0d)
|
||||
|
||||
# 3D 张量应该抛出异常
|
||||
sequences_3d = [torch.tensor([[[1, 2]]])] # shape: (1, 1, 2)
|
||||
with pytest.raises(ValueError, match="Expected 1D tensor"):
|
||||
packer.pack(sequences_3d)
|
||||
|
||||
def test_reset_method(self):
|
||||
"""测试 reset() 方法是否正确重置内部状态"""
|
||||
packer = SequencePacker(pack_size=10, pad_value=0)
|
||||
|
||||
# 第一次打包
|
||||
sequences1 = [torch.tensor([1, 2, 3], dtype=torch.int32)]
|
||||
packer.pack(sequences1)
|
||||
|
||||
# 验证内部状态已更新
|
||||
assert packer._current_pos == 0
|
||||
assert packer._current_pack is None # 最后一个包已发送，设置为None
|
||||
|
||||
# 重置
|
||||
packer.reset()
|
||||
|
||||
# 验证重置后的状态
|
||||
assert packer._current_pos == 0
|
||||
assert packer._current_pack is not None
|
||||
assert packer._current_pack.shape == (10,)
|
||||
assert packer._current_pack.tolist() == [0] * 10
|
||||
|
||||
# 验证重置后可以继续正常使用
|
||||
sequences2 = [torch.tensor([4, 5, 6], dtype=torch.int32)]
|
||||
packages = packer.pack(sequences2)
|
||||
|
||||
assert len(packages) == 1
|
||||
assert packages[0][:3].tolist() == [4, 5, 6]
|
||||
packer.pack([torch.tensor(5)])
|
||||
|
||||
def test_input_list_not_modified(self):
|
||||
"""测试输入列表是否未被修改（使用 sorted 而非 sort）"""
|
||||
packer = SequencePacker(pack_size=10)
|
||||
|
||||
# 创建原始序列列表（故意不按长度排序）
|
||||
original_sequences = [
|
||||
torch.tensor([3], dtype=torch.int32), # 长度1
|
||||
torch.tensor([1, 2], dtype=torch.int32), # 长度2
|
||||
torch.tensor([4, 5, 6, 7], dtype=torch.int32), # 长度4
|
||||
original = [
|
||||
torch.tensor([3], dtype=torch.int32),
|
||||
torch.tensor([1, 2], dtype=torch.int32),
|
||||
torch.tensor([4, 5, 6, 7], dtype=torch.int32),
|
||||
]
|
||||
|
||||
# 保存原始顺序的字符串表示
|
||||
original_repr = [seq.tolist() for seq in original_sequences]
|
||||
|
||||
# 打包
|
||||
packer.pack(original_sequences)
|
||||
|
||||
# 验证输入列表未被修改
|
||||
current_repr = [seq.tolist() for seq in original_sequences]
|
||||
assert current_repr == original_repr, "输入列表被修改了，应该使用 sorted() 而非 sort()"
|
||||
original_repr = [seq.tolist() for seq in original]
|
||||
packer.pack(original)
|
||||
assert [seq.tolist() for seq in original] == original_repr
|
||||
|
||||
def test_exact_pack_size_fit(self):
|
||||
"""测试序列长度恰好等于 pack_size 的情况"""
|
||||
packer = SequencePacker(pack_size=5, pad_value=0)
|
||||
|
||||
sequences = [
|
||||
packages = packer.pack([
|
||||
torch.tensor([1, 2, 3, 4, 5], dtype=torch.int32),
|
||||
torch.tensor([6, 7, 8, 9, 10], dtype=torch.int32),
|
||||
]
|
||||
|
||||
packages = packer.pack(sequences)
|
||||
|
||||
# 每个序列恰好占满一个包
|
||||
])
|
||||
assert len(packages) == 2
|
||||
|
||||
assert packages[0].tolist() == [1, 2, 3, 4, 5]
|
||||
assert packages[1].tolist() == [6, 7, 8, 9, 10]
|
||||
|
||||
def test_multiple_packs_full_utilization(self):
|
||||
"""测试多个包的高效利用"""
|
||||
packer = SequencePacker(pack_size=10, pad_value=-1)
|
||||
|
||||
# 创建多个小序列，确保高效打包
|
||||
sequences = [
|
||||
torch.tensor([1], dtype=torch.int32),
|
||||
torch.tensor([2], dtype=torch.int32),
|
||||
torch.tensor([3], dtype=torch.int32),
|
||||
torch.tensor([4], dtype=torch.int32),
|
||||
torch.tensor([5], dtype=torch.int32),
|
||||
torch.tensor([6], dtype=torch.int32),
|
||||
torch.tensor([7], dtype=torch.int32),
|
||||
torch.tensor([8], dtype=torch.int32),
|
||||
torch.tensor([9], dtype=torch.int32),
|
||||
torch.tensor([10], dtype=torch.int32),
|
||||
torch.tensor([11], dtype=torch.int32),
|
||||
]
|
||||
|
||||
sequences = [torch.tensor([i], dtype=torch.int32) for i in range(1, 12)]
|
||||
packages = packer.pack(sequences)
|
||||
|
||||
# 前10个序列打包成一个包，最后一个序列单独一个包
|
||||
assert len(packages) == 2
|
||||
assert packages[0].tolist() == list(range(1, 11))
|
||||
assert packages[1].tolist() == [11] + [-1] * 9
|
||||
|
||||
def test_dtype_mismatch_warning(self, caplog):
|
||||
"""测试 dtype 不匹配时的警告"""
|
||||
packer = SequencePacker(pack_size=10, pad_value=0, dtype=torch.int32)
|
||||
def test_cross_group_ordering(self):
|
||||
"""Tensor groups with identical per-item lengths are sorted identically."""
|
||||
packer = SequencePacker(pack_size=10, pad_value=0)
|
||||
# sequences: lengths [3, 1, 4] -> after sort desc: [4, 3, 1]
|
||||
seqs = [
|
||||
torch.tensor([1, 2, 3], dtype=torch.int32),
|
||||
torch.tensor([10], dtype=torch.int32),
|
||||
torch.tensor([4, 5, 6, 7], dtype=torch.int32),
|
||||
]
|
||||
masks = [
|
||||
torch.tensor([True, True, True], dtype=torch.bool),
|
||||
torch.tensor([True], dtype=torch.bool),
|
||||
torch.tensor([True, True, True, True], dtype=torch.bool),
|
||||
]
|
||||
packed_seqs = packer.pack(seqs)
|
||||
packer.reset()
|
||||
packed_masks = packer.pack(masks)
|
||||
|
||||
sequences = [torch.tensor([1, 2, 3], dtype=torch.int64)]
|
||||
|
||||
packages = packer.pack(sequences)
|
||||
|
||||
# 应该触发 dtype 不匹配警告
|
||||
assert "dtype" in caplog.text.lower() or "converted" in caplog.text.lower()
|
||||
# Verify mask packer uses bool dtype
|
||||
assert packed_masks[0].dtype == torch.bool
|
||||
# Both groups should produce the same number of packages
|
||||
assert len(packed_seqs) == len(packed_masks)
|
||||
|
||||
+32
-128
@@ -1,4 +1,4 @@
|
||||
"""单元测试：pipeline.processors 模块中的处理器类"""
|
||||
"""Tests for pipeline.processors module."""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
@@ -13,191 +13,97 @@ from pipeline.processors import (
|
||||
|
||||
|
||||
class DummyTokenizer:
|
||||
"""用于测试的虚拟分词器"""
|
||||
|
||||
def encode(self, text: str):
|
||||
# 简单模拟：返回文本字符的ASCII码列表
|
||||
def encode(self, text: str, add_special_tokens: bool = False):
|
||||
return [ord(c) for c in text]
|
||||
|
||||
|
||||
class TestBaseProcessor:
|
||||
"""BaseProcessor 抽象基类的测试"""
|
||||
|
||||
def test_abstract_class_cannot_be_instantiated(self):
|
||||
"""测试 BaseProcessor 不能直接实例化"""
|
||||
with pytest.raises(TypeError):
|
||||
BaseProcessor()
|
||||
|
||||
|
||||
class TestPreTrainProcessor:
|
||||
"""PreTrainProcessor 类的测试套件"""
|
||||
|
||||
def test_output_keys(self):
|
||||
"""测试 output_keys 属性"""
|
||||
tokenizer = DummyTokenizer()
|
||||
processor = PreTrainProcessor(tokenizer)
|
||||
assert processor.output_keys == ["sequence"]
|
||||
assert PreTrainProcessor(DummyTokenizer()).output_keys == ["sequence"]
|
||||
|
||||
def test_process_returns_tensor(self):
|
||||
"""测试 process 方法返回正确的张量"""
|
||||
tokenizer = DummyTokenizer()
|
||||
processor = PreTrainProcessor(tokenizer)
|
||||
|
||||
processor = PreTrainProcessor(DummyTokenizer())
|
||||
result = processor.process({"text": "hello world"})
|
||||
|
||||
assert "sequence" in result
|
||||
assert isinstance(result["sequence"], torch.Tensor)
|
||||
assert result["sequence"].dtype == torch.int32
|
||||
|
||||
def test_process_adds_eos(self):
|
||||
"""测试 process 方法添加 EOS 标记"""
|
||||
tokenizer = DummyTokenizer()
|
||||
processor = PreTrainProcessor(tokenizer)
|
||||
|
||||
# 文本 "a" 的 ASCII 码是 97
|
||||
result = processor.process({"text": "a"})
|
||||
|
||||
# 应该包含文本的ASCII码 + <eos> (假设是 4)
|
||||
seq = result["sequence"]
|
||||
# 基本验证：返回的张量长度应该大于0
|
||||
assert len(seq) > 0
|
||||
result = PreTrainProcessor(DummyTokenizer()).process({"text": "a"})
|
||||
assert len(result["sequence"]) > 0
|
||||
|
||||
|
||||
class TestSFTProcessor:
|
||||
"""SFTProcessor 类的测试套件"""
|
||||
|
||||
def test_output_keys(self):
|
||||
"""测试 output_keys 属性"""
|
||||
tokenizer = DummyTokenizer()
|
||||
processor = SFTProcessor(tokenizer)
|
||||
assert processor.output_keys == ["sequence", "loss_mask"]
|
||||
assert SFTProcessor(DummyTokenizer()).output_keys == ["sequence", "loss_mask"]
|
||||
|
||||
def test_process_returns_both_keys(self):
|
||||
"""测试 process 方法返回所有键"""
|
||||
tokenizer = DummyTokenizer()
|
||||
processor = SFTProcessor(tokenizer)
|
||||
|
||||
result = processor.process({
|
||||
"query": "hello",
|
||||
"response": "world"
|
||||
})
|
||||
|
||||
result = SFTProcessor(DummyTokenizer()).process({"query": "hello", "response": "world"})
|
||||
assert "sequence" in result
|
||||
assert "loss_mask" in result
|
||||
assert isinstance(result["sequence"], torch.Tensor)
|
||||
assert isinstance(result["loss_mask"], torch.Tensor)
|
||||
|
||||
def test_loss_mask_correct_length(self):
|
||||
"""测试 loss_mask 长度与 sequence 一致"""
|
||||
tokenizer = DummyTokenizer()
|
||||
processor = SFTProcessor(tokenizer)
|
||||
|
||||
result = processor.process({
|
||||
"query": "hi",
|
||||
"response": "bye"
|
||||
})
|
||||
|
||||
result = SFTProcessor(DummyTokenizer()).process({"query": "hi", "response": "bye"})
|
||||
assert len(result["sequence"]) == len(result["loss_mask"])
|
||||
|
||||
def test_loss_mask_after_query_is_true(self):
|
||||
"""测试 loss_mask 在响应部分为 True"""
|
||||
tokenizer = DummyTokenizer()
|
||||
processor = SFTProcessor(tokenizer)
|
||||
|
||||
result = processor.process({
|
||||
"query": "ab", # 2 chars
|
||||
"response": "cd", # 2 chars
|
||||
})
|
||||
|
||||
# 验证 loss_mask 是 bool 类型
|
||||
def test_loss_mask_is_bool(self):
|
||||
result = SFTProcessor(DummyTokenizer()).process({"query": "ab", "response": "cd"})
|
||||
assert result["loss_mask"].dtype == torch.bool
|
||||
|
||||
|
||||
class TestDPOProcessor:
|
||||
"""DPOProcessor 类的测试套件"""
|
||||
|
||||
def test_output_keys(self):
|
||||
"""测试 output_keys 属性"""
|
||||
tokenizer = DummyTokenizer()
|
||||
processor = DPOProcessor(tokenizer)
|
||||
assert processor.output_keys == ["chosen", "chosen_mask", "rejected", "rejected_mask"]
|
||||
keys = DPOProcessor(DummyTokenizer()).output_keys
|
||||
assert keys == ["chosen", "chosen_mask", "rejected", "rejected_mask"]
|
||||
|
||||
def test_process_returns_all_keys(self):
|
||||
"""测试 process 方法返回所有键"""
|
||||
tokenizer = DummyTokenizer()
|
||||
processor = DPOProcessor(tokenizer)
|
||||
|
||||
result = processor.process({
|
||||
"query": "hello",
|
||||
"chosen": "response1",
|
||||
"rejected": "response2"
|
||||
})
|
||||
|
||||
expected_keys = ["chosen", "chosen_mask", "rejected", "rejected_mask"]
|
||||
for key in expected_keys:
|
||||
result = DPOProcessor(DummyTokenizer()).process(
|
||||
{"query": "hello", "chosen": "r1", "rejected": "r2"}
|
||||
)
|
||||
for key in ["chosen", "chosen_mask", "rejected", "rejected_mask"]:
|
||||
assert key in result
|
||||
assert isinstance(result[key], torch.Tensor)
|
||||
|
||||
def test_chosen_and_rejected_same_length_as_mask(self):
|
||||
"""测试 chosen/rejected 长度与 mask 一致"""
|
||||
tokenizer = DummyTokenizer()
|
||||
processor = DPOProcessor(tokenizer)
|
||||
|
||||
result = processor.process({
|
||||
"query": "test",
|
||||
"chosen": "yes",
|
||||
"rejected": "no"
|
||||
})
|
||||
|
||||
def test_masks_match_lengths(self):
|
||||
result = DPOProcessor(DummyTokenizer()).process(
|
||||
{"query": "test", "chosen": "yes", "rejected": "no"}
|
||||
)
|
||||
assert len(result["chosen"]) == len(result["chosen_mask"])
|
||||
assert len(result["rejected"]) == len(result["rejected_mask"])
|
||||
|
||||
def test_masks_are_bool(self):
|
||||
"""测试 mask 张量是 bool 类型"""
|
||||
tokenizer = DummyTokenizer()
|
||||
processor = DPOProcessor(tokenizer)
|
||||
|
||||
result = processor.process({
|
||||
"query": "test",
|
||||
"chosen": "yes",
|
||||
"rejected": "no"
|
||||
})
|
||||
|
||||
result = DPOProcessor(DummyTokenizer()).process(
|
||||
{"query": "test", "chosen": "yes", "rejected": "no"}
|
||||
)
|
||||
assert result["chosen_mask"].dtype == torch.bool
|
||||
assert result["rejected_mask"].dtype == torch.bool
|
||||
|
||||
|
||||
class TestProcessorFactory:
|
||||
"""ProcessorFactory 类的测试套件"""
|
||||
|
||||
def test_create_pre_train_processor(self):
|
||||
"""测试创建预训练处理器"""
|
||||
tokenizer = DummyTokenizer()
|
||||
processor = ProcessorFactory.create("pt", tokenizer)
|
||||
assert isinstance(processor, PreTrainProcessor)
|
||||
assert isinstance(ProcessorFactory.create("pt", DummyTokenizer()), PreTrainProcessor)
|
||||
|
||||
def test_create_sft_processor(self):
|
||||
"""测试创建 SFT 处理器"""
|
||||
tokenizer = DummyTokenizer()
|
||||
processor = ProcessorFactory.create("sft", tokenizer)
|
||||
assert isinstance(processor, SFTProcessor)
|
||||
assert isinstance(ProcessorFactory.create("sft", DummyTokenizer()), SFTProcessor)
|
||||
|
||||
def test_create_dpo_processor(self):
|
||||
"""测试创建 DPO 处理器"""
|
||||
tokenizer = DummyTokenizer()
|
||||
processor = ProcessorFactory.create("dpo", tokenizer)
|
||||
assert isinstance(processor, DPOProcessor)
|
||||
assert isinstance(ProcessorFactory.create("dpo", DummyTokenizer()), DPOProcessor)
|
||||
|
||||
def test_create_invalid_processor_raises_error(self):
|
||||
"""测试创建无效处理器类型抛出异常"""
|
||||
tokenizer = DummyTokenizer()
|
||||
with pytest.raises(ValueError, match="Invalid processor type"):
|
||||
ProcessorFactory.create("invalid", tokenizer)
|
||||
with pytest.raises(ValueError, match="Unknown processor type"):
|
||||
ProcessorFactory.create("invalid", DummyTokenizer())
|
||||
|
||||
def test_register_and_create_custom_processor(self):
|
||||
"""测试注册和创建自定义处理器"""
|
||||
class CustomProcessor(BaseProcessor):
|
||||
def __init__(self, tokenizer=None): # 接受 tokenizer 参数
|
||||
def __init__(self, tokenizer=None):
|
||||
self._tokenizer = tokenizer
|
||||
|
||||
@property
|
||||
@@ -207,7 +113,5 @@ class TestProcessorFactory:
|
||||
def process(self, input_dict):
|
||||
return {"custom": torch.tensor([1, 2, 3])}
|
||||
|
||||
tokenizer = DummyTokenizer()
|
||||
ProcessorFactory.register("custom", CustomProcessor)
|
||||
processor = ProcessorFactory.create("custom", tokenizer)
|
||||
assert isinstance(processor, CustomProcessor)
|
||||
ProcessorFactory.register("custom")(CustomProcessor)
|
||||
assert isinstance(ProcessorFactory.create("custom", DummyTokenizer()), CustomProcessor)
|
||||
|
||||
@@ -0,0 +1,118 @@
|
||||
"""Tests for strategy module."""
|
||||
|
||||
import pytest
|
||||
from pipeline.strategies import (
|
||||
PromptStrategy,
|
||||
ChatMLStrategy,
|
||||
AlpacaStrategy,
|
||||
StrategyFactory,
|
||||
)
|
||||
from pipeline.tokenizer import BpeTokenizer
|
||||
|
||||
|
||||
class DummyTokenizer:
|
||||
def encode(self, text: str, add_special_tokens: bool = False):
|
||||
return [ord(c) for c in text]
|
||||
|
||||
|
||||
class DummyStrategy(PromptStrategy):
|
||||
def __init__(self, tokenizer):
|
||||
super().__init__(tokenizer)
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "dummy"
|
||||
|
||||
def assemble_prompt(self, query_tokens):
|
||||
prefix = self._encode_format("Q: ")
|
||||
return prefix + query_tokens
|
||||
|
||||
def assemble_response(self, response_tokens):
|
||||
suffix = self._encode_format("<eos>")
|
||||
return response_tokens + suffix
|
||||
|
||||
|
||||
def _decode(tokens):
|
||||
return "".join(chr(t) for t in tokens)
|
||||
|
||||
|
||||
class TestChatMLStrategy:
|
||||
def test_name(self):
|
||||
assert ChatMLStrategy(DummyTokenizer()).name == "chatml"
|
||||
|
||||
def test_assemble_prompt(self):
|
||||
tk = DummyTokenizer()
|
||||
strategy = ChatMLStrategy(tk)
|
||||
query_tokens = tk.encode("hello")
|
||||
prompt = strategy.assemble_prompt(query_tokens)
|
||||
text = _decode(prompt)
|
||||
assert "<|im_start|>user" in text
|
||||
assert "hello" in text
|
||||
assert "<|im_start|>assistant" in text
|
||||
|
||||
def test_assemble_response(self):
|
||||
tk = DummyTokenizer()
|
||||
strategy = ChatMLStrategy(tk)
|
||||
response_tokens = tk.encode("world")
|
||||
response = strategy.assemble_response(response_tokens)
|
||||
text = _decode(response)
|
||||
assert "world" in text
|
||||
assert "<|im_end|>" in text
|
||||
assert "<eos>" in text
|
||||
|
||||
def test_prompt_ends_with_assistant_start(self):
|
||||
tk = DummyTokenizer()
|
||||
strategy = ChatMLStrategy(tk)
|
||||
prompt = strategy.assemble_prompt(tk.encode("hi"))
|
||||
# prompt 末尾应该是 assistant_start 的 token ids
|
||||
assert prompt[-len(strategy._assistant_start_ids):] == strategy._assistant_start_ids
|
||||
|
||||
|
||||
class TestAlpacaStrategy:
|
||||
def test_name(self):
|
||||
assert AlpacaStrategy(DummyTokenizer()).name == "alpaca"
|
||||
|
||||
def test_assemble_prompt(self):
|
||||
tk = DummyTokenizer()
|
||||
strategy = AlpacaStrategy(tk)
|
||||
query_tokens = tk.encode("hello")
|
||||
prompt = strategy.assemble_prompt(query_tokens)
|
||||
text = _decode(prompt)
|
||||
assert "### Instruction:" in text
|
||||
assert "hello" in text
|
||||
assert "### Response:" in text
|
||||
|
||||
def test_assemble_response(self):
|
||||
tk = DummyTokenizer()
|
||||
strategy = AlpacaStrategy(tk)
|
||||
response_tokens = tk.encode("world")
|
||||
response = strategy.assemble_response(response_tokens)
|
||||
text = _decode(response)
|
||||
assert "world" in text
|
||||
assert "<eos>" in text
|
||||
|
||||
|
||||
class TestStrategyFactory:
|
||||
def test_create_chatml(self):
|
||||
tk = DummyTokenizer()
|
||||
assert isinstance(StrategyFactory.create("chatml", tk), ChatMLStrategy)
|
||||
|
||||
def test_create_alpaca(self):
|
||||
tk = DummyTokenizer()
|
||||
assert isinstance(StrategyFactory.create("alpaca", tk), AlpacaStrategy)
|
||||
|
||||
def test_create_invalid_raises_error(self):
|
||||
with pytest.raises(ValueError, match="Unknown strategy"):
|
||||
StrategyFactory.create("invalid_strategy", DummyTokenizer())
|
||||
|
||||
def test_register_and_create(self):
|
||||
StrategyFactory.register("dummy")(DummyStrategy)
|
||||
tk = DummyTokenizer()
|
||||
strategy = StrategyFactory.create("dummy", tk)
|
||||
assert isinstance(strategy, DummyStrategy)
|
||||
assert strategy.name == "dummy"
|
||||
|
||||
def test_available_strategies(self):
|
||||
strategies = StrategyFactory.available_types()
|
||||
assert "chatml" in strategies
|
||||
assert "alpaca" in strategies
|
||||
@@ -1,229 +0,0 @@
|
||||
"""单元测试：pipeline.tokenizer 模块中的 BpeTokenizer 类"""
|
||||
|
||||
import pytest
|
||||
from pathlib import Path
|
||||
import tempfile
|
||||
|
||||
from pipeline.tokenizer import BpeTokenizer
|
||||
|
||||
|
||||
class TestBpeTokenizer:
|
||||
"""BpeTokenizer 类的测试套件"""
|
||||
|
||||
def test_initialization_without_path(self):
|
||||
"""测试不加载外部文件初始化"""
|
||||
tokenizer = BpeTokenizer()
|
||||
assert tokenizer is not None
|
||||
assert hasattr(tokenizer, '_tokenizer')
|
||||
|
||||
def test_initialization_with_path(self):
|
||||
"""测试加载外部文件初始化"""
|
||||
# 这个测试假设没有预训练的分词器文件，所以只测试不抛出异常
|
||||
# 实际使用中需要提供有效的分词器文件路径
|
||||
try:
|
||||
tokenizer = BpeTokenizer(path="nonexistent.json")
|
||||
except Exception:
|
||||
# 预期会抛出异常，因为文件不存在
|
||||
pass
|
||||
|
||||
def test_vocab_size(self):
|
||||
"""测试获取词汇表大小"""
|
||||
tokenizer = BpeTokenizer()
|
||||
vocab_size = len(tokenizer)
|
||||
assert isinstance(vocab_size, int)
|
||||
assert vocab_size >= 0
|
||||
|
||||
def test_special_tokens_exist(self):
|
||||
"""测试特殊token是否存在"""
|
||||
tokenizer = BpeTokenizer()
|
||||
|
||||
# 检查控制token
|
||||
assert hasattr(tokenizer, '_control_tokens')
|
||||
assert '<bos>' in tokenizer._control_tokens
|
||||
assert '<eos>' in tokenizer._control_tokens
|
||||
assert '<pad>' in tokenizer._control_tokens
|
||||
|
||||
# 检查特殊token
|
||||
assert hasattr(tokenizer, '_special_tokens')
|
||||
assert '<|im_start|>' in tokenizer._special_tokens
|
||||
assert '<|im_end|>' in tokenizer._special_tokens
|
||||
|
||||
def test_encode_string(self):
|
||||
"""测试编码单个字符串"""
|
||||
tokenizer = BpeTokenizer()
|
||||
|
||||
# 使用简单的ASCII字符测试
|
||||
result = tokenizer.encode("hello")
|
||||
|
||||
# 返回应该是 token IDs 列表
|
||||
assert isinstance(result, list)
|
||||
|
||||
def test_encode_list(self):
|
||||
"""测试编码字符串列表"""
|
||||
tokenizer = BpeTokenizer()
|
||||
|
||||
texts = ["hello", "world", "test"]
|
||||
result = tokenizer.encode(texts)
|
||||
|
||||
# 返回应该是列表的列表
|
||||
assert isinstance(result, list)
|
||||
assert len(result) == len(texts)
|
||||
for item in result:
|
||||
assert isinstance(item, list)
|
||||
|
||||
def test_encode_with_output_tokens(self):
|
||||
"""测试编码返回tokens而非ids"""
|
||||
tokenizer = BpeTokenizer()
|
||||
|
||||
result = tokenizer.encode("hello", out_ids=False)
|
||||
|
||||
# 应该返回 token 字符串列表
|
||||
assert isinstance(result, list)
|
||||
|
||||
def test_encode_with_special_tokens(self):
|
||||
"""测试编码添加特殊token"""
|
||||
tokenizer = BpeTokenizer()
|
||||
|
||||
result = tokenizer.encode("hello", add_special_tokens=True)
|
||||
|
||||
assert isinstance(result, list)
|
||||
|
||||
def test_decode(self):
|
||||
"""测试解码token IDs"""
|
||||
tokenizer = BpeTokenizer()
|
||||
|
||||
# 解码空列表
|
||||
result = tokenizer.decode([])
|
||||
assert isinstance(result, str)
|
||||
|
||||
# 解码包含一些ID的列表（假设有 vocab）
|
||||
# 如果分词器未训练，可能无法正确解码
|
||||
result = tokenizer.decode([104, 101, 108, 108, 111]) # "hello" 的 ASCII
|
||||
assert isinstance(result, str)
|
||||
|
||||
def test_decode_with_special_tokens(self):
|
||||
"""测试解码保留特殊token"""
|
||||
tokenizer = BpeTokenizer()
|
||||
|
||||
# 解码空列表
|
||||
result = tokenizer.decode([], skip_special_tokens=False)
|
||||
assert isinstance(result, str)
|
||||
|
||||
def test_stop_ids_property(self):
|
||||
"""测试 stop_ids 属性"""
|
||||
tokenizer = BpeTokenizer()
|
||||
|
||||
stop_ids = tokenizer.stop_ids
|
||||
assert isinstance(stop_ids, list)
|
||||
|
||||
def test_special_token_properties(self):
|
||||
"""测试特殊token ID属性"""
|
||||
tokenizer = BpeTokenizer()
|
||||
|
||||
# 这些属性可能返回 None 如果分词器未训练
|
||||
bos_id = tokenizer.bos_id
|
||||
eos_id = tokenizer.eos_id
|
||||
pad_id = tokenizer.pad_id
|
||||
|
||||
# 只验证属性存在且为 int 或 None
|
||||
assert isinstance(bos_id, (int, type(None)))
|
||||
assert isinstance(eos_id, (int, type(None)))
|
||||
assert isinstance(pad_id, (int, type(None)))
|
||||
|
||||
def test_save_method_exists(self):
|
||||
"""测试 save 方法存在"""
|
||||
tokenizer = BpeTokenizer()
|
||||
assert hasattr(tokenizer, 'save')
|
||||
assert callable(tokenizer.save)
|
||||
|
||||
def test_load_method_exists(self):
|
||||
"""测试 load 方法存在"""
|
||||
tokenizer = BpeTokenizer()
|
||||
assert hasattr(tokenizer, 'load')
|
||||
assert callable(tokenizer.load)
|
||||
|
||||
def test_train_method_exists(self):
|
||||
"""测试 train 方法存在"""
|
||||
tokenizer = BpeTokenizer()
|
||||
assert hasattr(tokenizer, 'train')
|
||||
assert callable(tokenizer.train)
|
||||
|
||||
def test_train_from_iterator_method_exists(self):
|
||||
"""测试 train_from_iterator 方法存在"""
|
||||
tokenizer = BpeTokenizer()
|
||||
assert hasattr(tokenizer, 'train_from_iterator')
|
||||
assert callable(tokenizer.train_from_iterator)
|
||||
|
||||
|
||||
class TestBpeTokenizerIntegration:
|
||||
"""BpeTokenizer 集成测试"""
|
||||
|
||||
def test_encode_decode_roundtrip(self):
|
||||
"""测试编码解码往返"""
|
||||
tokenizer = BpeTokenizer()
|
||||
|
||||
original = "hello world"
|
||||
encoded = tokenizer.encode(original)
|
||||
decoded = tokenizer.decode(encoded)
|
||||
|
||||
# 往返后应该得到类似的结果
|
||||
# 注意：由于分词器可能未训练，结果可能不完全一致
|
||||
assert isinstance(encoded, list)
|
||||
assert isinstance(decoded, str)
|
||||
|
||||
def test_train_from_iterator_small_corpus(self, tmp_path):
|
||||
"""测试使用小语料库训练"""
|
||||
tokenizer = BpeTokenizer()
|
||||
|
||||
# 创建临时训练文件
|
||||
train_file = tmp_path / "train.txt"
|
||||
train_content = "hello world\nthis is a test\nmachine learning\n"
|
||||
train_file.write_text(train_content)
|
||||
|
||||
# 训练分词器（使用较小的 vocab size 加快测试）
|
||||
try:
|
||||
tokenizer.train(
|
||||
files=[str(train_file)],
|
||||
vocab_size=100,
|
||||
min_freq=1,
|
||||
reserved_token_size=10
|
||||
)
|
||||
|
||||
# 验证训练后分词器可用
|
||||
result = tokenizer.encode("hello")
|
||||
assert isinstance(result, list)
|
||||
assert len(result) > 0
|
||||
except Exception as e:
|
||||
pytest.skip(f"Training failed: {e}")
|
||||
|
||||
def test_save_and_load_tokenizer(self, tmp_path):
|
||||
"""测试保存和加载分词器"""
|
||||
tokenizer = BpeTokenizer()
|
||||
|
||||
# 创建临时训练文件并训练
|
||||
train_file = tmp_path / "train.txt"
|
||||
train_content = "hello world\ntest data\n"
|
||||
train_file.write_text(train_content)
|
||||
|
||||
try:
|
||||
tokenizer.train(
|
||||
files=[str(train_file)],
|
||||
vocab_size=50,
|
||||
min_freq=1,
|
||||
reserved_token_size=5
|
||||
)
|
||||
|
||||
# 保存
|
||||
save_path = tmp_path / "tokenizer.json"
|
||||
tokenizer.save(str(save_path))
|
||||
|
||||
# 加载到新实例
|
||||
new_tokenizer = BpeTokenizer()
|
||||
new_tokenizer.load(str(save_path))
|
||||
|
||||
# 验证加载后分词器可用
|
||||
result = new_tokenizer.encode("hello")
|
||||
assert isinstance(result, list)
|
||||
|
||||
except Exception as e:
|
||||
pytest.skip(f"Save/load test failed: {e}")
|
||||
Reference in New Issue
Block a user