From 35963bcb08b7bae7052d6f32d7932bd517d028a0 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Mon, 30 Mar 2026 20:58:51 +0800 Subject: [PATCH] =?UTF-8?q?reafactor:=20=E9=87=8D=E6=9E=84=E9=A1=B9?= =?UTF-8?q?=E7=9B=AE?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- README.md | 205 +++++++++----- docs/Design.md | 185 +++++++++++++ pipeline/__init__.py | 23 +- pipeline/cache.py | 82 ------ pipeline/export.py | 63 ----- pipeline/io.py | 153 ++++++++++- pipeline/packing.py | 167 +++++++----- pipeline/processors.py | 122 --------- pipeline/processors/__init__.py | 18 ++ pipeline/processors/base.py | 31 +++ pipeline/processors/dpo.py | 51 ++++ pipeline/processors/factory.py | 122 +++++++++ pipeline/processors/pretrain.py | 26 ++ pipeline/processors/sft.py | 40 +++ pipeline/strategies/__init__.py | 14 + pipeline/strategies/alpaca.py | 43 +++ pipeline/strategies/base.py | 41 +++ pipeline/strategies/chatml.py | 41 +++ pipeline/strategies/factory.py | 70 +++++ pyproject.toml | 1 - scripts/cache_h5.py | 59 ++-- .../reforce_learning/dpp_chinese_dpo_pairs.py | 2 +- .../sft_chinese_instruct.py | 6 +- tests/test_cache.py | 118 +++----- tests/test_io.py | 157 +++-------- tests/test_packing.py | 252 +++++------------- tests/test_processors.py | 190 ++++--------- tests/test_strategies.py | 118 ++++++++ tests/test_tokenizer.py | 229 ---------------- 29 files changed, 1395 insertions(+), 1234 deletions(-) create mode 100644 docs/Design.md delete mode 100644 pipeline/cache.py delete mode 100644 pipeline/export.py delete mode 100644 pipeline/processors.py create mode 100644 pipeline/processors/__init__.py create mode 100644 pipeline/processors/base.py create mode 100644 pipeline/processors/dpo.py create mode 100644 pipeline/processors/factory.py create mode 100644 pipeline/processors/pretrain.py create mode 100644 pipeline/processors/sft.py create mode 100644 pipeline/strategies/__init__.py create mode 100644 pipeline/strategies/alpaca.py create mode 100644 pipeline/strategies/base.py create mode 100644 pipeline/strategies/chatml.py create mode 100644 pipeline/strategies/factory.py create mode 100644 tests/test_strategies.py delete mode 100644 tests/test_tokenizer.py diff --git a/README.md b/README.md index 276616f..52c9ee4 100644 --- a/README.md +++ b/README.md @@ -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 + "") +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]}) - -# 加载 -data = IOHandler.load_h5("./out") # → {"key": [tensor1, ...]} +# 自定义 ChatML 的特殊 token +strategy = StrategyFactory.create("chatml", + user_start="user\n", + user_end="\n", + assistant_start="assistant\n", + assistant_end="\n", +) +processor = ProcessorFactory.create_with_strategy("sft", tokenizer, strategy) ``` + +## 策略格式 + +| Strategy | Key | Default Tokens | +|-----------|------------|--------------------------------------------------------------------------------------------------| +| ChatML | `"chatml"` | `<\|im_start\|>user`, `<\|im_end\|>`, `<\|im_start\|>assistant`, `` | +| Alpaca | `"alpaca"` | `### Instruction:`, `### Response:`, `` | + +所有策略的 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 +``` \ No newline at end of file diff --git a/docs/Design.md b/docs/Design.md new file mode 100644 index 0000000..0d7609a --- /dev/null +++ b/docs/Design.md @@ -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 + "") +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 后缀（含 ``） +- `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` | + +Alpaca: + +| Parameter | Default Token | +|------------------|------------------------| +| instruction_start | `### Instruction:\n` | +| response_start | `### Response:\n` | +| response_suffix | `\n` | + +**自定义示例**: +```python +strategy = StrategyFactory.create("chatml", + user_start="user\n", + user_end="\n", + assistant_start="assistant\n", + assistant_end="\n", +) +``` + +**注册新格式**: +```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 分词器，支持从文件加载、训练、保存。内置 ``/``/`` 控制符和 `<|im_start|>`/`<|im_end|>` 特殊 token。 + +## API 参考 + +各模块的详细 API 文档请参阅对应的源文件。 diff --git a/pipeline/__init__.py b/pipeline/__init__.py index 8ab4c2e..a5bc42a 100644 --- a/pipeline/__init__.py +++ b/pipeline/__init__.py @@ -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', ] diff --git a/pipeline/cache.py b/pipeline/cache.py deleted file mode 100644 index 9efa337..0000000 --- a/pipeline/cache.py +++ /dev/null @@ -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 diff --git a/pipeline/export.py b/pipeline/export.py deleted file mode 100644 index 6150f67..0000000 --- a/pipeline/export.py +++ /dev/null @@ -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 diff --git a/pipeline/io.py b/pipeline/io.py index be34e4a..e5231f1 100644 --- a/pipeline/io.py +++ b/pipeline/io.py @@ -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]: @@ -34,7 +47,7 @@ class IOHandler: def save_h5(output_dir: str, file_name: str, tensor_group: Dict[str, List[Tensor]]) -> None: os.makedirs(output_dir, exist_ok=True) full_path = os.path.join(output_dir, f"{file_name}.h5") - + with h5py.File(full_path, 'w') as f: for key, tensors in tensor_group.items(): grp = f.create_group(key) @@ -43,12 +56,12 @@ 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) h5_files = list(root_path.rglob("*.h5")) + list(root_path.rglob("*.hdf5")) - + for h5_file in h5_files: with h5py.File(h5_file, 'r') as f: for key in f.keys(): @@ -60,9 +73,129 @@ class IOHandler: if share_memory: tensor = tensor.share_memory_() dsets.append(tensor) - + if tensor_group.get(key) is None: tensor_group[key] = [] tensor_group[key].extend(dsets) - return tensor_group \ No newline at end of file + 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 diff --git a/pipeline/packing.py b/pipeline/packing.py index fabbee2..0fdc8fa 100644 --- a/pipeline/packing.py +++ b/pipeline/packing.py @@ -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) - 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] + # --- sort by length descending (FFD heuristic) --- + normalized.sort(key=lambda x: x[1], reverse=True) - # 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 + # --- greedy fill --- + buf = self._buffer + pos = self._pos + packages = self._packages + pack_size = self.pack_size + pad_value = self.pad_value - # Place tensor in current package - self._buffer[self._current_pos : self._current_pos + tensor_size] = tensor - self._current_pos += tensor_size + 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 - # Handle the last package (pad to pack_size) - if self._current_pos > 0: - package = self._buffer.clone() - self._packages.append(package) + # flush the last (possibly partial) package + if pos > 0: + packages.append(buf.clone()) - # Clear buffer and reset state for backward compatibility - self._buffer = None - self._current_pack = None - self._current_pos = 0 - + # 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() \ No newline at end of file diff --git a/pipeline/processors.py b/pipeline/processors.py deleted file mode 100644 index 9ff49e3..0000000 --- a/pipeline/processors.py +++ /dev/null @@ -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}") - 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") - - 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") - 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") - 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 diff --git a/pipeline/processors/__init__.py b/pipeline/processors/__init__.py new file mode 100644 index 0000000..0d3542e --- /dev/null +++ b/pipeline/processors/__init__.py @@ -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", +] diff --git a/pipeline/processors/base.py b/pipeline/processors/base.py new file mode 100644 index 0000000..77afb8f --- /dev/null +++ b/pipeline/processors/base.py @@ -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 diff --git a/pipeline/processors/dpo.py b/pipeline/processors/dpo.py new file mode 100644 index 0000000..635169f --- /dev/null +++ b/pipeline/processors/dpo.py @@ -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"] diff --git a/pipeline/processors/factory.py b/pipeline/processors/factory.py new file mode 100644 index 0000000..e7e4116 --- /dev/null +++ b/pipeline/processors/factory.py @@ -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()) diff --git a/pipeline/processors/pretrain.py b/pipeline/processors/pretrain.py new file mode 100644 index 0000000..b0207be --- /dev/null +++ b/pipeline/processors/pretrain.py @@ -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}") + return {"sequence": torch.tensor(tokens, dtype=torch.int32)} + + @property + def output_keys(self) -> List[str]: + return ["sequence"] diff --git a/pipeline/processors/sft.py b/pipeline/processors/sft.py new file mode 100644 index 0000000..0ac71e2 --- /dev/null +++ b/pipeline/processors/sft.py @@ -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"] diff --git a/pipeline/strategies/__init__.py b/pipeline/strategies/__init__.py new file mode 100644 index 0000000..ca481d7 --- /dev/null +++ b/pipeline/strategies/__init__.py @@ -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", +] diff --git a/pipeline/strategies/alpaca.py b/pipeline/strategies/alpaca.py new file mode 100644 index 0000000..20e97ad --- /dev/null +++ b/pipeline/strategies/alpaca.py @@ -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: ... ``""" + + def __init__( + self, + tokenizer: BpeTokenizer, + instruction_start: str = "### Instruction:\n", + response_start: str = "### Response:\n", + response_suffix: str = "\n", + ): + 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 diff --git a/pipeline/strategies/base.py b/pipeline/strategies/base.py new file mode 100644 index 0000000..1f7e02e --- /dev/null +++ b/pipeline/strategies/base.py @@ -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}')" diff --git a/pipeline/strategies/chatml.py b/pipeline/strategies/chatml.py new file mode 100644 index 0000000..38f5941 --- /dev/null +++ b/pipeline/strategies/chatml.py @@ -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|> ``""" + + 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", + ): + 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 diff --git a/pipeline/strategies/factory.py b/pipeline/strategies/factory.py new file mode 100644 index 0000000..7585e09 --- /dev/null +++ b/pipeline/strategies/factory.py @@ -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()) diff --git a/pyproject.toml b/pyproject.toml index 59aa694..fdfa49a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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] diff --git a/scripts/cache_h5.py b/scripts/cache_h5.py index e1d6d1b..4e97a84 100644 --- a/scripts/cache_h5.py +++ b/scripts/cache_h5.py @@ -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 输出目录 (默认: /cached)") + help="H5 output dir (default: /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 即 )") + 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__": diff --git a/scripts/reforce_learning/dpp_chinese_dpo_pairs.py b/scripts/reforce_learning/dpp_chinese_dpo_pairs.py index 159d0a8..3cb4070 100644 --- a/scripts/reforce_learning/dpp_chinese_dpo_pairs.py +++ b/scripts/reforce_learning/dpp_chinese_dpo_pairs.py @@ -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"], } diff --git a/scripts/supervised_finetuning/sft_chinese_instruct.py b/scripts/supervised_finetuning/sft_chinese_instruct.py index b491466..5ecb015 100644 --- a/scripts/supervised_finetuning/sft_chinese_instruct.py +++ b/scripts/supervised_finetuning/sft_chinese_instruct.py @@ -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", diff --git a/tests/test_cache.py b/tests/test_cache.py index cd6dddf..d3b08be 100644 --- a/tests/test_cache.py +++ b/tests/test_cache.py @@ -1,4 +1,4 @@ -"""单元测试：pipeline.cache 模块中的 cache_jsonl 函数""" +"""Tests for pipeline.cache module.""" import json import os @@ -6,24 +6,23 @@ 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"] - + @property def output_keys(self): return self._output_keys - + 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") @@ -131,35 +88,22 @@ class TestCacheJsonl: with open(jsonl_path, "w", encoding="utf-8") as f: f.write(json.dumps(test_data[0]) + "\n") 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 \ No newline at end of file + assert len(output_files) == 1 diff --git a/tests/test_io.py b/tests/test_io.py index 245bc5e..c000b6d 100644 --- a/tests/test_io.py +++ b/tests/test_io.py @@ -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) - - # 创建嵌套子目录 - nested = os.path.join(subdir1, "nested") - os.makedirs(nested) - - # 获取文件夹列表 + os.makedirs(os.path.join(tmpdir, "folder1")) + os.makedirs(os.path.join(tmpdir, "folder2")) + os.makedirs(os.path.join(tmpdir, "folder1", "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) - - # 使用过滤器只获取 folder1 - folders = IOHandler.fetch_folders( - tmpdir, - filter_func=lambda x: "folder1" in x - ) - - # 验证 + os.makedirs(os.path.join(tmpdir, "folder1")) + os.makedirs(os.path.join(tmpdir, "folder2")) + + 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]) - - # 创建第二个 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]) - - # 加载目录 + 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) + 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 \ No newline at end of file + assert len(loaded["batch"]) == 3 diff --git a/tests/test_packing.py b/tests/test_packing.py index 55e32db..28130bd 100644 --- a/tests/test_packing.py +++ b/tests/test_packing.py @@ -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) - - 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() \ No newline at end of file + 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) + + # 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) diff --git a/tests/test_processors.py b/tests/test_processors.py index 9ab51fd..fc63a22 100644 --- a/tests/test_processors.py +++ b/tests/test_processors.py @@ -1,4 +1,4 @@ -"""单元测试：pipeline.processors 模块中的处理器类""" +"""Tests for pipeline.processors module.""" import pytest import torch @@ -13,201 +13,105 @@ 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码 + (假设是 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 def output_keys(self): return ["custom"] - + 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) \ No newline at end of file + + ProcessorFactory.register("custom")(CustomProcessor) + assert isinstance(ProcessorFactory.create("custom", DummyTokenizer()), CustomProcessor) diff --git a/tests/test_strategies.py b/tests/test_strategies.py new file mode 100644 index 0000000..99f117d --- /dev/null +++ b/tests/test_strategies.py @@ -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("") + 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 "" 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 "" 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 diff --git a/tests/test_tokenizer.py b/tests/test_tokenizer.py deleted file mode 100644 index cfc50f8..0000000 --- a/tests/test_tokenizer.py +++ /dev/null @@ -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 '' in tokenizer._control_tokens - assert '' in tokenizer._control_tokens - assert '' 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}") \ No newline at end of file