reafactor: 重构项目

README.md

@@ -1,111 +1,188 @@
 # DataPipeline
 
-数据集处理工具，支持预训练 / SFT / DPO 三种训练范式。
+语言模型训练数据集处理工具包。支持 PT / SFT / DPO 范式。
 
 ## 项目结构
 
 ```
 pipeline/
-├── tokenizer.py      # BPE 分词器
+├── tokenizer.py           # BPE tokenizer
-├── text.py           # 文本规范化
+├── text.py                # Text normalization
-├── packing.py        # 序列打包
+├── packing.py             # Sequence packing
-├── io.py             # 文件/HDF5 读写
+├── io.py                  # File / HDF5 I/O
-├── processors.py     # PT / SFT / DPO 处理器
+├── processors.py          # PT / SFT / DPO processors
-├── export.py         # Dataset → JSONL
+├── export.py              # Dataset -> JSONL
-└── cache.py          # JSONL → Tokenize → H5
+├── 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
+Stage 1: Export Dataset                     Stage 2: Tokenize & Cache
-                              ↑
+┌──────────────────────────────────────┐    ┌──────────────────────────────────────┐
-                        processors.py
+│                                      │    │                                      │
-                        packing.py
+│  HuggingFace Dataset                 │    │  JSONL files (one dict per line)     │
-                        io.py
+│       │                              │    │       │                              │
+│       ▼  process_func (optional)     │    │       ▼  json.loads()                │
+│  export_dataset()                    │    │  Processor.process(input_dict)       │
+│       │  chunk by chunk_size         │    │       │                              │
+│       ▼                              │    │       │  ┌───────────────────────┐   │
+│  ./dataset/prefix_chunk_0.jsonl ────────────────|>   │ PT / SFT / DPO        │   │
+│  ./dataset/prefix_chunk_1.jsonl      │    │       │  │ processor details     │   │
+│  ...                                 │    │       │  └───────────────────────┘   │
+│                                      │    │       ▼                              │
+│  modules: export.py                  │    │  List[Tensor]                        │
+│                                      │    │       │                              │
+│                                      │    │       ▼  SequencePacker (optional)   │
+│                                      │    │  Fixed-length packed tensors         │
+│                                      │    │       │                              │
+│                                      │    │       ▼  IOHandler.save_h5()         │
+│                                      │    │  ./cached/chunk_0.h5                 │
+│                                      │    │  ./cached/chunk_1.h5                 │
+│                                      │    │  ...                                 │
+│                                      │    │                                      │
+│                                      │    │  modules: io.py, processors.py,      │
+│                                      │    │           packing.py, tokenizer.py   │
+└──────────────────────────────────────┘    └──────────────────────────────────────┘
 ```
 
-## 使用方法
+### Processor 转换规则
 
-### 1. 导出数据集
+**PT (Pre-training)**
+```
+Input:  {"text": "Hello world"}
+Action: tokenizer.encode(text + "<eos>")
+Output: {"sequence": Tensor[int32]}
+```
+
+**SFT (Supervised Fine-tuning)** — uses PromptStrategy
+```
+Input:  {"query": "...", "response": "..."}
+Action:
+  1. strategy.build_prompt(input_dict)  ->  "<|im_start|>user\n...\n<|im_start|>assistant\n"
+  2. concat response + response_suffix
+  3. tokenizer.encode full string
+  4. build loss_mask: query part=False, response part=True
+Output: {"sequence": Tensor[int32], "loss_mask": Tensor[bool]}
+```
+
+**DPO (Direct Preference Optimization)** — uses PromptStrategy
+```
+Input:  {"query": "...", "chosen": "...", "rejected": "..."}
+Action:
+  1. strategy.build_prompt(input_dict)  ->  shared query prompt
+  2. encode chosen  = query + response_start + chosen  + response_suffix
+  3. encode rejected = query + response_start + rejected + response_suffix
+  4. build masks: query part=False, response part=True
+Output: {"chosen": Tensor, "chosen_mask": Tensor[bool],
+         "rejected": Tensor, "rejected_mask": Tensor[bool]}
+```
+
+### 序列打包
+
+当 `pack_size > 0` 时启用，使用 First-Fit Decreasing 算法将变长序列填充到固定长度：
+
+- 按序列长度降序排列
+- 逐个放入当前包，超长则截断
+- 当前包放不下时开新包
+- 不足部分用 `pad_value` 填充
+
+## 快速开始
+
+### 1. 导出数据集为 JSONL
 
 ```python
 from datasets import load_dataset
-from pipeline.export import export_dataset
+from pipeline import export_dataset
 
 dataset = load_dataset("your-dataset")
 export_dataset(
     dataset=dataset["train"],
     output_dir="./data",
     output_prefix="train",
-    process_func=lambda x: {"text": x["content"]},  # 可选
+    process_func=lambda x: {"text": x["content"]},
 )
 ```
 
-### 2. Tokenize 并缓存
+### 2. 分词并缓存为 HDF5
 
 ```python
 from pipeline import BpeTokenizer, ProcessorFactory, cache_jsonl
 
 tokenizer = BpeTokenizer("tokenizer.json")
-processor = ProcessorFactory.create("pt", tokenizer)  # "pt" | "sft" | "dpo"
+processor = ProcessorFactory.create("pt", tokenizer)
 
 cache_jsonl(
     files=["./data/train.jsonl"],
     output_dir="./cached",
     processor=processor,
-    pack_size=4096,  # <=0 不打包
+    pack_size=4096,
-    pad_value=1,
 )
 ```
 
-### 3. 处理器类型
+### 3. 使用策略模式（默认 token）
-
-| 类型 | 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
 
 ```python
-packer = SequencePacker(pack_size=4096, pad_value=0)
+from pipeline import StrategyFactory, ProcessorFactory
-packed = packer.pack([tensor1, tensor2, ...])  # → List[Tensor]
+
+strategy = StrategyFactory.create("alpaca")
+processor = ProcessorFactory.create_with_strategy("sft", tokenizer, strategy)
 ```
 
-### IOHandler
+### 4. 使用策略模式（自定义 token）
 
 ```python
-# 保存
+# 自定义 ChatML 的特殊 token
-IOHandler.save_h5("./out", "name", {"key": [tensor1, tensor2]})
+strategy = StrategyFactory.create("chatml",
-
+    user_start="<s>user\n",
-# 加载
+    user_end="</s>\n",
-data = IOHandler.load_h5("./out")  # → {"key": [tensor1, ...]}
+    assistant_start="<s>assistant\n",
+    assistant_end="</s>\n<eos>",
+)
+processor = ProcessorFactory.create_with_strategy("sft", tokenizer, strategy)
 ```
+
+## 策略格式
+
+| Strategy  | Key        | Default Tokens                                                                                   |
+|-----------|------------|--------------------------------------------------------------------------------------------------|
+| ChatML    | `"chatml"` | `<\|im_start\|>user`, `<\|im_end\|>`, `<\|im_start\|>assistant`, `<eos>`                       |
+| Alpaca    | `"alpaca"` | `### Instruction:`, `### Response:`, `<eos>`                                                     |
+
+所有策略的 token 均可通过构造函数参数自定义，同时支持通过 `StrategyFactory.register()` 注册新格式。
+
+## 命令行工具
+
+```bash
+# 缓存 JSONL 到 H5
+python scripts/cache_h5.py pt ./dataset/chinese-c4-pretrain
+python scripts/cache_h5.py sft ./dataset/belle-sft --pack-size 4096 --strategy alpaca
+```
+
+## 脚本示例
+
+```
+scripts/
+├── cache_h5.py                              # Stage 2: JSONL -> H5 (CLI tool)
+├── pre_train/
+│   ├── chinese-c4.py                        # Chinese pretrain data export
+│   ├── chinese-cosmopedia.py                # Chinese pretrain data export
+│   ├── english-fineweb.py                   # English pretrain data export
+│   └── english-wiki.py                      # English pretrain data export
+├── supervised_finetuning/
+│   ├── sft_belle.py                         # Belle Chinese SFT data export
+│   ├── sft_chinese_instruct.py              # Chinese instruct SFT (with TextNormalizer)
+│   ├── sft_coder.py                         # Code SFT data export
+│   ├── sft_firefly-1.1m-rephrased.py        # Firefly SFT data export
+│   └── sft_magpie-pro-300k.py               # Magpie Pro SFT data export
+└── reforce_learning/
+    └── dpp_chinese_dpo_pairs.py              # Chinese DPO preference pairs export
+```

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 + "<eos>")
+Output: {"sequence": Tensor[int32]}
+```
+
+**SFTProcessor** (`"sft"`) — 注入 PromptStrategy
+```
+Input:  {"query": "...", "response": "..."}
+Action:
+  1. strategy.build_prompt(input_dict) -> query prompt with special tokens
+  2. concat: prompt + response + response_suffix
+  3. tokenizer.encode -> token ids
+  4. build loss_mask: query=False, response=True
+Output: {"sequence": Tensor[int32], "loss_mask": Tensor[bool]}
+```
+
+**DPOProcessor** (`"dpo"`) — 注入 PromptStrategy
+```
+Input:  {"query": "...", "chosen": "...", "rejected": "..."}
+Action:
+  1. strategy.build_prompt(input_dict) -> shared query prompt
+  2. encode chosen  = query + response_start + chosen  + suffix
+  3. encode rejected = query + response_start + rejected + suffix
+  4. build masks: query=False, response=True
+Output: {"chosen": Tensor, "chosen_mask": Tensor[bool],
+         "rejected": Tensor, "rejected_mask": Tensor[bool]}
+```
+
+## 策略模式
+
+### PromptStrategy
+
+抽象 prompt/response 格式，所有特殊 token 可通过构造函数自定义配置。
+
+**接口**:
+- `build_prompt(input_dict)` — 构建包含 query 的完整 prompt
+- `build_response_prefix()` — response 前缀（当前均返回空串）
+- `build_response_suffix()` — response 后缀（含 `<eos>`）
+- `response_start_token` — response 起始 token（用于 DPO 的 loss mask 定位）
+- `eos_tokens` — 结束 token
+
+**内置实现**:
+
+ChatML:
+
+| Parameter       | Default Token                              |
+|-----------------|--------------------------------------------|
+| user_start      | `<\|im_start\|>user\n`                    |
+| user_end        | `<\|im_end\|>\n`                          |
+| assistant_start | `<\|im_start\|>assistant\n`               |
+| assistant_end   | `<\|im_end\|>\n<eos>`                     |
+
+Alpaca:
+
+| Parameter        | Default Token          |
+|------------------|------------------------|
+| instruction_start | `### Instruction:\n`  |
+| response_start    | `### Response:\n`     |
+| response_suffix   | `\n<eos>`             |
+
+**自定义示例**:
+```python
+strategy = StrategyFactory.create("chatml",
+    user_start="<s>user\n",
+    user_end="</s>\n",
+    assistant_start="<s>assistant\n",
+    assistant_end="</s>\n<eos>",
+)
+```
+
+**注册新格式**:
+```python
+from pipeline import StrategyFactory
+from pipeline.strategies import PromptStrategy
+
+class MyStrategy(PromptStrategy):
+    @property
+    def name(self) -> str:
+        return "my_format"
+    # ... 实现抽象方法
+
+StrategyFactory.register("my_format", MyStrategy)
+```
+
+## 序列打包
+
+`SequencePacker` 将变长序列打包为固定长度的 tensor，使用 First-Fit Decreasing 算法：
+
+1. 按序列长度降序排列
+2. 逐个放入当前包，超长则截断并警告
+3. 当前包放不下时保存并开新包
+4. 不足部分用 `pad_value` 填充
+
+打包在 `cache_jsonl()` 中按 `pack_size` 启用（`> 0` 时生效，`<= 0` 跳过打包）。
+
+## 辅助模块
+
+### TextNormalizer (`pipeline/text.py`)
+
+文本标准化：统一标点符号（弯引号→直引号、破折号、省略号等），通过 `DEFAULT_REPLACEMENTS` 配置，支持 `custom_rules` 扩展。
+
+### IOHandler (`pipeline/io.py`)
+
+- `save_h5(output_dir, file_name, tensor_group)` — 按 key 分组存储 tensor 到 HDF5
+- `load_h5(file_path, share_memory)` — 加载 HDF5，支持共享内存（用于 DataLoader 多进程）
+- `fetch_files(directory)` / `fetch_folders(root_dir)` — 文件/目录遍历
+
+### BpeTokenizer (`pipeline/tokenizer.py`)
+
+基于 HuggingFace `tokenizers` 库的 BPE 分词器，支持从文件加载、训练、保存。内置 `<bos>`/`<eos>`/`<pad>` 控制符和 `<|im_start|>`/`<|im_end|>` 特殊 token。
+
+## API 参考
+
+各模块的详细 API 文档请参阅对应的源文件。

pipeline/__init__.py

@@ -1,17 +1,17 @@
 import logging
-from .tokenizer import BpeTokenizer
+from pipeline.tokenizer import BpeTokenizer
-from .text import TextNormalizer
+from pipeline.text import TextNormalizer
-from .packing import SequencePacker
+from pipeline.packing import SequencePacker
-from .io import IOHandler
+from pipeline.io import IOHandler, export_dataset, cache_jsonl
-from .processors import ProcessorFactory, BaseProcessor
+from pipeline.processors import ProcessorFactory, BaseProcessor
-from .export import export_dataset
+from pipeline.utils import setup_logging
-from .cache import cache_jsonl
+from pipeline.strategies import PromptStrategy, ChatMLStrategy, AlpacaStrategy, StrategyFactory
-from .utils import setup_logging
 
-# 配置项目级日志记录
+# 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',
 ]

pipeline/cache.py

@@ -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

pipeline/export.py

@@ -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

pipeline/io.py

@@ -1,23 +1,36 @@
-from pathlib import Path
+"""File, HDF5, JSONL I/O operations."""
-from typing import Dict, List, Optional, Callable
+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]:
+    def fetch_files(directory: str, suffix: Optional[str] = None) -> List[str]:
-        return [
+        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
+        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

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 pipeline.utils import error_handler
-from .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
+        self.dtype = dtype  # None = follow input dtype
-        # Pre-allocate buffer for better performance
+        self._buffer: Tensor | None = None
-        self._buffer: Optional[Tensor] = None
+        self._pos = 0
-        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._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
+        # --- validate & normalize in a single pass ---
-        tensor_sizes = []
+        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 != target_dtype:
-            if seq.dtype != self.dtype:
+                seq = seq.to(target_dtype)
-                logger.warning(
+            length = seq.numel()
-                    f"Input tensor dtype {seq.dtype} does not match packer dtype {self.dtype}, "
+            if length > self.pack_size:
-                    f"will be converted. This may affect packing efficiency."
+                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:
+        # --- sort by length descending (FFD heuristic) ---
-            # Truncate sequences that exceed pack_size
+        normalized.sort(key=lambda x: x[1], reverse=True)
-            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]
 
-            # Current package is full, create a new one
+        # --- greedy fill ---
-            if self._current_pos + tensor_size > self.pack_size:
+        buf = self._buffer
-                # Finish current package (pad to pack_size)
+        pos = self._pos
-                package = self._buffer.clone()
+        packages = self._packages
-                self._packages.append(package)
+        pack_size = self.pack_size
-                # Reset buffer for reuse
+        pad_value = self.pad_value
-                self._buffer.fill_(self.pad_value)
-                self._current_pos = 0
 
-            # Place tensor in current package
+        for tensor, length in normalized:
-            self._buffer[self._current_pos : self._current_pos + tensor_size] = tensor
+            if pos + length > pack_size:
-            self._current_pos += tensor_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)
+        # flush the last (possibly partial) package
-        if self._current_pos > 0:
+        if pos > 0:
-            package = self._buffer.clone()
+            packages.append(buf.clone())
-            self._packages.append(package)
 
-        # Clear buffer and reset state for backward compatibility
+        # write back state
-        self._buffer = None
+        self._pos = pos
-        self._current_pack = None
-        self._current_pos = 0
-        
         return self._packages
-
-    def reset(self) -> None:
-        """Reset packer state for reuse. More efficient than creating a new instance."""
-        self._reset()

pipeline/processors.py

@@ -1,122 +0,0 @@
-from abc import ABC, abstractmethod
-from typing import Dict, List, Any
-import torch
-from torch import Tensor
-from .tokenizer import BpeTokenizer
-
-
-class BaseProcessor(ABC):
-    """Abstract base class for processors."""
-
-    @abstractmethod
-    def process(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]:
-        pass
-
-    @property
-    @abstractmethod
-    def output_keys(self) -> List[str]:
-        pass
-
-
-class PreTrainProcessor(BaseProcessor):
-    """Pre-training data processor."""
-
-    def __init__(self, tokenizer: BpeTokenizer):
-        self.tokenizer = tokenizer
-
-    def process(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]:
-        segment = input_dict["text"]
-        tokens = self.tokenizer.encode(f"{segment}<eos>")
-        return {'sequence': torch.tensor(tokens, dtype=torch.int32)}
-
-    @property
-    def output_keys(self) -> List[str]:
-        return ["sequence"]
-
-
-class SFTProcessor(BaseProcessor):
-    """Supervised fine-tuning data processor."""
-
-    def __init__(self, tokenizer: BpeTokenizer):
-        self.tokenizer = tokenizer
-
-    def process(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]:
-        query = input_dict["query"]
-        response = input_dict["response"]
-        
-        q = self.tokenizer.encode(
-            f"<|im_start|>user\n{query}<|im_end|>\n<|im_start|>assistant\n"
-        )
-        a = self.tokenizer.encode(f"{response}<|im_end|>\n<eos>")
-        
-        q_len = len(q)
-        tokens = torch.tensor(q + a, dtype=torch.int32)
-        loss_mask = torch.zeros(q_len + len(a), dtype=torch.bool)
-        loss_mask[q_len:] = True
-        return {"sequence": tokens, "loss_mask": loss_mask}
-
-    @property
-    def output_keys(self) -> List[str]:
-        return ["sequence", "loss_mask"]
-
-
-class DPOProcessor(BaseProcessor):
-    """DPO preference learning data processor."""
-
-    def __init__(self, tokenizer: BpeTokenizer):
-        self.tokenizer = tokenizer
-
-    def process(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]:
-        query = input_dict["query"]
-        chosen_response = input_dict["chosen"]
-        rejected_response = input_dict["rejected"]
-
-        q = self.tokenizer.encode(
-            f"<|im_start|>user\n{query}<|im_end|>\n<|im_start|>assistant\n"
-        )
-
-        chosen = self.tokenizer.encode(f"{chosen_response}<|im_end|>\n<eos>")
-        q_len = len(q)
-        chosen_len = len(chosen)
-        
-        chosen_tokens = torch.tensor(q + chosen, dtype=torch.int32)
-        chosen_mask = torch.zeros(q_len + chosen_len, dtype=torch.bool)
-        chosen_mask[q_len:] = True
-
-        rejected = self.tokenizer.encode(f"{rejected_response}<|im_end|>\n<eos>")
-        rejected_len = len(rejected)
-        
-        rejected_tokens = torch.tensor(q + rejected, dtype=torch.int32)
-        rejected_mask = torch.zeros(q_len + rejected_len, dtype=torch.bool)
-        rejected_mask[q_len:] = True
-
-        return {
-            "chosen": chosen_tokens,
-            "chosen_mask": chosen_mask,
-            "rejected": rejected_tokens,
-            "rejected_mask": rejected_mask,
-        }
-
-    @property
-    def output_keys(self) -> List[str]:
-        return ["chosen", "chosen_mask", "rejected", "rejected_mask"]
-
-
-class ProcessorFactory:
-    """Processor factory."""
-
-    _processors = {
-        "pt": PreTrainProcessor,
-        "sft": SFTProcessor,
-        "dpo": DPOProcessor,
-    }
-
-    @classmethod
-    def create(cls, processor_type: str, tokenizer: BpeTokenizer) -> BaseProcessor:
-        if processor_type not in cls._processors:
-            raise ValueError(f"Invalid processor type: {processor_type}")
-        return cls._processors[processor_type](tokenizer)
-
-    @classmethod
-    def register(cls, processor_type: str, processor_class: type):
-        cls._processors[processor_type] = processor_class

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",
+]

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

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"]

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())

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}<eos>")
+        return {"sequence": torch.tensor(tokens, dtype=torch.int32)}
+
+    @property
+    def output_keys(self) -> List[str]:
+        return ["sequence"]

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"]

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",
+]

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: ... <eos>``"""
+
+    def __init__(
+        self,
+        tokenizer: BpeTokenizer,
+        instruction_start: str = "### Instruction:\n",
+        response_start: str = "### Response:\n",
+        response_suffix: str = "\n<eos>",
+    ):
+        super().__init__(tokenizer)
+        self.instruction_start = instruction_start
+        self.response_start = response_start
+        self.response_suffix = response_suffix
+
+        self._instruction_start_ids = self._encode_format(instruction_start)
+        self._separator_ids = self._encode_format("\n\n")
+        self._response_start_ids = self._encode_format(response_start)
+        self._response_suffix_ids = self._encode_format(response_suffix)
+
+    @property
+    def name(self) -> str:
+        return "alpaca"
+
+    def assemble_prompt(self, query_tokens: List[int]) -> List[int]:
+        return (
+            self._instruction_start_ids
+            + query_tokens
+            + self._separator_ids
+            + self._response_start_ids
+        )
+
+    def assemble_response(self, response_tokens: List[int]) -> List[int]:
+        return response_tokens + self._response_suffix_ids

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}')"

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|> <eos>``"""
+
+    def __init__(
+        self,
+        tokenizer: BpeTokenizer,
+        user_start: str = "<|im_start|>user\n",
+        user_end: str = "<|im_end|>\n",
+        assistant_start: str = "<|im_start|>assistant\n",
+        assistant_end: str = "<|im_end|>\n<eos>",
+    ):
+        super().__init__(tokenizer)
+
+        self._user_start_ids = self._encode_format(user_start)
+        self._user_end_ids = self._encode_format(user_end)
+        self._assistant_start_ids = self._encode_format(assistant_start)
+        self._assistant_end_ids = self._encode_format(assistant_end)
+
+    @property
+    def name(self) -> str:
+        return "chatml"
+
+    def assemble_prompt(self, query_tokens: List[int]) -> List[int]:
+        return (
+            self._user_start_ids
+            + query_tokens
+            + self._user_end_ids
+            + self._assistant_start_ids
+        )
+
+    def assemble_response(self, response_tokens: List[int]) -> List[int]:
+        return response_tokens + self._assistant_end_ids

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())

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]

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
+from pipeline import BpeTokenizer, ProcessorFactory, cache_jsonl
-
+from pipeline.io import 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
 
 
 def main():
-    parser = argparse.ArgumentParser(description="JSONL → H5 缓存")
+    parser = argparse.ArgumentParser(description="JSONL -> H5 cache")
-    parser.add_argument("type", choices=["pt", "sft", "dpo"], help="处理器类型")
+    parser.add_argument("type", choices=["pt", "sft", "dpo"], help="Processor type")
-    parser.add_argument("input_dir", help="JSONL 文件所在目录")
+    parser.add_argument("input_dir", help="Directory containing JSONL files")
     parser.add_argument("-o", "--output-dir", default=None,
-                        help="H5 输出目录 (默认: <input_dir>/cached)")
+                        help="H5 output dir (default: <input_dir>/cached)")
     parser.add_argument("-t", "--tokenizer", default="./tokenizer.json",
-                        help="Tokenizer 路径 (默认: ./tokenizer.json)")
+                        help="Tokenizer path (default: ./tokenizer.json)")
+    parser.add_argument("-s", "--strategy", default=None,
+                        help="Prompt strategy: chatml, alpaca (default: chatml)")
     parser.add_argument("-p", "--pack-size", type=int, default=-1,
-                        help="序列打包长度，<=0 不打包 (默认: -1)")
+                        help="Pack size, <=0 to disable (default: -1)")
     parser.add_argument("--pad-value", type=int, default=1,
-                        help="打包填充值 (默认: 1 即 <eos>)")
+                        help="Padding value (default: 1)")
     args = parser.parse_args()
 
-    # 收集 JSONL 文件
+    jsonl_files = IOHandler.fetch_files(args.input_dir, suffix=".jsonl")
-    jsonl_files = collect_jsonl_files(args.input_dir)
     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)}")
 
-    # 创建处理器
+    if args.strategy:
-    processor = ProcessorFactory.create(args.type, tokenizer)
+        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! Output saved to {output_dir}")
-    print(f"\nDone! {len(output_files)} H5 files saved to {output_dir}")
 
 
 if __name__ == "__main__":

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"],
     }

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",

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 = [
+            test_data = [{"text": "hello"}, {"text": "world"}, {"text": "test"}]
-                {"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],
+                files=[jsonl_path], output_dir=tmpdir,
-                output_dir=tmpdir,
+                processor=processor, pack_size=-1, pad_value=0,
-                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 = [
+            test_data = [{"text": "ab"}, {"text": "abcde"}, {"text": "abc"}]
-                {"text": "ab"},  # 2 chars
-                {"text": "abcde"},  # 5 chars
-                {"text": "abc"},  # 3 chars
-            ]
             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],
+                files=[jsonl_path], output_dir=tmpdir,
-                output_dir=tmpdir,
+                processor=processor, pack_size=10, pad_value=0,
-                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 = [
+            test_data = [{"text": "hello"}, {"text": "world"}]
-                {"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],
+                files=[jsonl_path], output_dir=tmpdir,
-                output_dir=tmpdir,
+                processor=processor, pack_size=0, pad_value=-1,
-                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,
+                files=files, output_dir=tmpdir,
-                output_dir=tmpdir,
+                processor=processor, pack_size=-1, pad_value=0,
-                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],
+                files=[jsonl_path], output_dir=tmpdir,
-                output_dir=tmpdir,
+                processor=processor, pack_size=-1, pad_value=0,
-                processor=processor,
-                pack_size=-1,
-                pad_value=0,
             )
-            
+            assert len(output_files) == 1
-            assert len(output_files) == 1

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:
-            # 创建测试文件
+            Path(tmpdir, "file1.txt").touch()
-            test_file1 = os.path.join(tmpdir, "file1.txt")
+            Path(tmpdir, "file2.txt").touch()
-            test_file2 = os.path.join(tmpdir, "file2.txt")
-            Path(test_file1).touch()
-            Path(test_file2).touch()
-            
-            # 创建子目录和文件
             subdir = os.path.join(tmpdir, "subdir")
             os.makedirs(subdir)
-            test_file3 = os.path.join(subdir, "file3.txt")
+            Path(subdir, "file3.txt").touch()
-            Path(test_file3).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 IOHandler.fetch_files(tmpdir) == []
-            assert files == []
+
-    
     def test_fetch_folders_in_directory(self):
-        """测试 fetch_folders 方法能正确获取子目录"""
         with tempfile.TemporaryDirectory() as tmpdir:
-            # 创建子目录
+            os.makedirs(os.path.join(tmpdir, "folder1"))
-            subdir1 = os.path.join(tmpdir, "folder1")
+            os.makedirs(os.path.join(tmpdir, "folder2"))
-            subdir2 = os.path.join(tmpdir, "folder2")
+            os.makedirs(os.path.join(tmpdir, "folder1", "nested"))
-            os.makedirs(subdir1)
+
-            os.makedirs(subdir2)
-            
-            # 创建嵌套子目录
-            nested = os.path.join(subdir1, "nested")
-            os.makedirs(nested)
-            
-            # 获取文件夹列表
             folders = IOHandler.fetch_folders(tmpdir)
-            
-            # 验证
             assert len(folders) == 3
-            assert any("folder1" in f for f in folders)
+
-            assert any("folder2" in f for f in folders)
-            assert any("nested" in f for f in folders)
-    
     def test_fetch_folders_with_filter(self):
-        """测试 fetch_folders 方法的过滤功能"""
         with tempfile.TemporaryDirectory() as tmpdir:
-            # 创建子目录
+            os.makedirs(os.path.join(tmpdir, "folder1"))
-            subdir1 = os.path.join(tmpdir, "folder1")
+            os.makedirs(os.path.join(tmpdir, "folder2"))
-            subdir2 = os.path.join(tmpdir, "folder2")
+
-            os.makedirs(subdir1)
+            folders = IOHandler.fetch_folders(tmpdir, filter_func=lambda x: "folder1" in x)
-            os.makedirs(subdir2)
-            
-            # 使用过滤器只获取 folder1
-            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)
-            
+
-            # 验证文件已创建
+            assert os.path.exists(os.path.join(tmpdir, "test.h5"))
-            h5_path = os.path.join(tmpdir, "test.h5")
-            assert os.path.exists(h5_path)
-            
-            # 加载 - 传入目录而不是单个文件
             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")
-            
+            IOHandler.save_h5(output_dir, "test", {"data": [torch.tensor([1, 2, 3])]})
-            tensor_group = {
-                "data": [torch.tensor([1, 2, 3])],
-            }
-            
-            # 保存到不存在的目录
-            IOHandler.save_h5(output_dir, "test", tensor_group)
-            
-            # 验证目录已创建
             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 文件
+            for i, data in enumerate([[1, 2, 3], [4, 5, 6]]):
-            h5_path1 = os.path.join(tmpdir, "file1.h5")
+                h5_path = os.path.join(tmpdir, f"file{i}.h5")
-            with h5py.File(h5_path1, 'w') as f:
+                with h5py.File(h5_path, 'w') as f:
-                grp = f.create_group("data")
+                    grp = f.create_group("data")
-                grp.create_dataset('data_0', data=[1, 2, 3])
+                    grp.create_dataset('data_0', data=data)
-            
+
-            # 创建第二个 H5 文件
-            h5_path2 = os.path.join(tmpdir, "file2.h5")
-            with h5py.File(h5_path2, 'w') as f:
-                grp = f.create_group("data")
-                grp.create_dataset('data_0', data=[4, 5, 6])
-            
-            # 加载目录
             loaded = IOHandler.load_h5(tmpdir, share_memory=False)
-            
-            # 验证
-            assert "data" in loaded
             assert len(loaded["data"]) == 2
-    
+
     def test_load_h5_with_rglob(self):
-        """测试 load_h5 能递归查找 H5 文件"""
         with tempfile.TemporaryDirectory() as tmpdir:
-            # 在子目录中创建 H5 文件
             subdir = os.path.join(tmpdir, "subdir")
             os.makedirs(subdir)
             h5_path = os.path.join(subdir, "nested.h5")
-            
             with h5py.File(h5_path, 'w') as f:
                 grp = f.create_group("test")
                 grp.create_dataset('data_0', data=[1, 2])
-            
+
-            # 加载根目录
             loaded = IOHandler.load_h5(tmpdir, share_memory=False)
-            
-            # 验证能找到子目录中的文件
             assert "test" in loaded
             assert len(loaded["test"]) == 1
-    
+
     def test_save_h5_multiple_tensors_per_key(self):
-        """测试 save_h5 能保存多个张量到同一键"""
         with tempfile.TemporaryDirectory() as tmpdir:
             tensor_group = {
-                "batch": [
+                "batch": [torch.tensor([1, 2]), torch.tensor([3, 4, 5]), torch.tensor([6])],
-                    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
-            assert len(loaded["batch"]) == 3

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,)
-            # 验证填充值
+
-            # 检查所有非零元素都在前几个位置，或者包是满的
+        # Verify all original values are present
-            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]
-        
-        # 简化测试：验证打包后的张量包含所有原始数据
         all_values = []
         for pkg in packages:
-            non_zero = pkg[pkg != 0].tolist()
+            all_values.extend(pkg[pkg != 0].tolist())
-            all_values.extend(non_zero)
+        for val in [1, 2, 3, 4, 5, 6, 7, 8, 9]:
-        
+            assert val in all_values
-        # 检查所有原始数据是否都被包含
-        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"
 
     def test_empty_list_input(self):
-        """测试空列表输入"""
         packer = SequencePacker(pack_size=10)
-        
+        assert packer.pack([]) == []
-        packages = packer.pack([])
-        
-        assert packages == []
-        
-        # 验证内部状态已正确初始化
-        assert packer._current_pack is not None
-        assert packer._current_pos == 0
 
     def test_single_sequence_input(self):
-        """测试单个序列输入"""
         packer = SequencePacker(pack_size=10, pad_value=-1)
-        
+        packages = packer.pack([torch.tensor([1, 2, 3], dtype=torch.int32)])
-        sequences = [torch.tensor([1, 2, 3], dtype=torch.int32)]
-        
-        packages = packer.pack(sequences)
-        
         assert len(packages) == 1
-        pkg = packages[0]
+        assert packages[0][:3].tolist() == [1, 2, 3]
-        assert pkg.shape == (10,)
+        assert packages[0][3:].tolist() == [-1] * 7
-        assert pkg[:3].tolist() == [1, 2, 3]
-        assert pkg[3:].tolist() == [-1] * 7
 
     def test_truncate_long_sequence(self, caplog):
-        """测试超长序列截断，验证警告日志是否触发"""
         packer = SequencePacker(pack_size=5, pad_value=0)
-        
+        packages = packer.pack([torch.tensor([1, 2, 3, 4, 5, 6, 7, 8], dtype=torch.int32)])
-        sequences = [
-            torch.tensor([1, 2, 3, 4, 5, 6, 7, 8], dtype=torch.int32),  # 长度8，超过pack_size=5
-        ]
-        
-        packages = packer.pack(sequences)
-        
         assert len(packages) == 1
-        pkg = packages[0]
+        assert packages[0].tolist() == [1, 2, 3, 4, 5]
-        assert pkg.shape == (5,)
-        assert pkg.tolist() == [1, 2, 3, 4, 5]  # 只保留前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)
-        
+        packages = packer.pack([
-        sequences = [
             torch.tensor([1, 2], dtype=torch.int32),
             torch.tensor([3], dtype=torch.int32),
-        ]
+        ])
-        
+        assert packages[0][:3].tolist() == [1, 2, 3]
-        packages = packer.pack(sequences)
+        assert packages[0][3:].tolist() == [99] * 5
-        
-        assert len(packages) == 1
-        pkg = packages[0]
-        
-        # 前3个元素是数据
-        assert pkg[:3].tolist() == [1, 2, 3]
-        # 后5个元素是填充值
-        assert pkg[3:].tolist() == [99] * 5
 
     def test_different_dtypes(self):
-        """测试支持不同 dtype (int32, int64, float32)"""
+        for dtype in [torch.int32, torch.int64, torch.float32]:
-        # int32
+            packer = SequencePacker(pack_size=10, dtype=dtype)
-        packer_int32 = SequencePacker(pack_size=10, pad_value=0, dtype=torch.int32)
+            val = 1.0 if dtype == torch.float32 else 1
-        sequences_int32 = [torch.tensor([1, 2, 3], dtype=torch.int32)]
+            packages = packer.pack([torch.tensor([val, 2, 3], dtype=dtype)])
-        packages_int32 = packer_int32.pack(sequences_int32)
+            assert packages[0].dtype == dtype
-        assert packages_int32[0].dtype == torch.int32
 
-        # int64
+    def test_dtype_conversion_on_mismatch(self, caplog):
-        packer_int64 = SequencePacker(pack_size=10, pad_value=0, dtype=torch.int64)
+        """Tensors with mismatched dtype are silently converted."""
-        sequences_int64 = [torch.tensor([1, 2, 3], dtype=torch.int64)]
+        packer = SequencePacker(pack_size=10, dtype=torch.int32)
-        packages_int64 = packer_int64.pack(sequences_int64)
+        packages = packer.pack([torch.tensor([1, 2, 3], dtype=torch.int64)])
-        assert packages_int64[0].dtype == torch.int64
+        assert packages[0].dtype == torch.int32
-
+        assert packages[0][:3].tolist() == [1, 2, 3]
-        # 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_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)
+            packer.pack([torch.tensor([[1, 2], [3, 4]])])
-        
-        # 0D 张量 (标量) 应该抛出异常
-        sequences_0d = [torch.tensor(5)]  # shape: ()
         with pytest.raises(ValueError, match="Expected 1D tensor"):
-            packer.pack(sequences_0d)
+            packer.pack([torch.tensor(5)])
-        
-        # 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]
 
     def test_input_list_not_modified(self):
-        """测试输入列表是否未被修改（使用 sorted 而非 sort）"""
         packer = SequencePacker(pack_size=10)
-        
+        original = [
-        # 创建原始序列列表（故意不按长度排序）
+            torch.tensor([3], dtype=torch.int32),
-        original_sequences = [
+            torch.tensor([1, 2], dtype=torch.int32),
-            torch.tensor([3], dtype=torch.int32),   # 长度1
+            torch.tensor([4, 5, 6, 7], dtype=torch.int32),
-            torch.tensor([1, 2], dtype=torch.int32), # 长度2
-            torch.tensor([4, 5, 6, 7], dtype=torch.int32), # 长度4
         ]
-        
+        original_repr = [seq.tolist() for seq in original]
-        # 保存原始顺序的字符串表示
+        packer.pack(original)
-        original_repr = [seq.tolist() for seq in original_sequences]
+        assert [seq.tolist() for seq in original] == original_repr
-        
-        # 打包
-        packer.pack(original_sequences)
-        
-        # 验证输入列表未被修改
-        current_repr = [seq.tolist() for seq in original_sequences]
-        assert current_repr == original_repr, "输入列表被修改了，应该使用 sorted() 而非 sort()"
 
     def test_exact_pack_size_fit(self):
-        """测试序列长度恰好等于 pack_size 的情况"""
         packer = SequencePacker(pack_size=5, pad_value=0)
-        
+        packages = packer.pack([
-        sequences = [
             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([i], dtype=torch.int32) for i in range(1, 12)]
-        # 创建多个小序列，确保高效打包
-        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),
-        ]
-        
         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):
+    def test_cross_group_ordering(self):
-        """测试 dtype 不匹配时的警告"""
+        """Tensor groups with identical per-item lengths are sorted identically."""
-        packer = SequencePacker(pack_size=10, pad_value=0, dtype=torch.int32)
+        packer = SequencePacker(pack_size=10, pad_value=0)
-        
+        # sequences: lengths [3, 1, 4] -> after sort desc: [4, 3, 1]
-        sequences = [torch.tensor([1, 2, 3], dtype=torch.int64)]
+        seqs = [
-        
+            torch.tensor([1, 2, 3], dtype=torch.int32),
-        packages = packer.pack(sequences)
+            torch.tensor([10], dtype=torch.int32),
-        
+            torch.tensor([4, 5, 6, 7], dtype=torch.int32),
-        # 应该触发 dtype 不匹配警告
+        ]
-        assert "dtype" in caplog.text.lower() or "converted" in caplog.text.lower()
+        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)

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, add_special_tokens: bool = False):
-    
-    def encode(self, text: str):
-        # 简单模拟：返回文本字符的ASCII码列表
         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 属性"""
+        assert PreTrainProcessor(DummyTokenizer()).output_keys == ["sequence"]
-        tokenizer = DummyTokenizer()
+
-        processor = PreTrainProcessor(tokenizer)
-        assert processor.output_keys == ["sequence"]
-    
     def test_process_returns_tensor(self):
-        """测试 process 方法返回正确的张量"""
+        processor = PreTrainProcessor(DummyTokenizer())
-        tokenizer = DummyTokenizer()
-        processor = PreTrainProcessor(tokenizer)
-        
         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 标记"""
+        result = PreTrainProcessor(DummyTokenizer()).process({"text": "a"})
-        tokenizer = DummyTokenizer()
+        assert len(result["sequence"]) > 0
-        processor = PreTrainProcessor(tokenizer)
-        
-        # 文本 "a" 的 ASCII 码是 97
-        result = processor.process({"text": "a"})
-        
-        # 应该包含文本的ASCII码 + <eos> (假设是 4)
-        seq = result["sequence"]
-        # 基本验证：返回的张量长度应该大于0
-        assert len(seq) > 0
 
 
 class TestSFTProcessor:
-    """SFTProcessor 类的测试套件"""
-    
     def test_output_keys(self):
-        """测试 output_keys 属性"""
+        assert SFTProcessor(DummyTokenizer()).output_keys == ["sequence", "loss_mask"]
-        tokenizer = DummyTokenizer()
+
-        processor = SFTProcessor(tokenizer)
-        assert processor.output_keys == ["sequence", "loss_mask"]
-    
     def test_process_returns_both_keys(self):
-        """测试 process 方法返回所有键"""
+        result = SFTProcessor(DummyTokenizer()).process({"query": "hello", "response": "world"})
-        tokenizer = DummyTokenizer()
-        processor = SFTProcessor(tokenizer)
-        
-        result = processor.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 一致"""
+        result = SFTProcessor(DummyTokenizer()).process({"query": "hi", "response": "bye"})
-        tokenizer = DummyTokenizer()
-        processor = SFTProcessor(tokenizer)
-        
-        result = processor.process({
-            "query": "hi",
-            "response": "bye"
-        })
-        
         assert len(result["sequence"]) == len(result["loss_mask"])
-    
+
-    def test_loss_mask_after_query_is_true(self):
+    def test_loss_mask_is_bool(self):
-        """测试 loss_mask 在响应部分为 True"""
+        result = SFTProcessor(DummyTokenizer()).process({"query": "ab", "response": "cd"})
-        tokenizer = DummyTokenizer()
-        processor = SFTProcessor(tokenizer)
-        
-        result = processor.process({
-            "query": "ab",  # 2 chars
-            "response": "cd",  # 2 chars
-        })
-        
-        # 验证 loss_mask 是 bool 类型
         assert result["loss_mask"].dtype == torch.bool
 
 
 class TestDPOProcessor:
-    """DPOProcessor 类的测试套件"""
-    
     def test_output_keys(self):
-        """测试 output_keys 属性"""
+        keys = DPOProcessor(DummyTokenizer()).output_keys
-        tokenizer = DummyTokenizer()
+        assert keys == ["chosen", "chosen_mask", "rejected", "rejected_mask"]
-        processor = DPOProcessor(tokenizer)
+
-        assert processor.output_keys == ["chosen", "chosen_mask", "rejected", "rejected_mask"]
-    
     def test_process_returns_all_keys(self):
-        """测试 process 方法返回所有键"""
+        result = DPOProcessor(DummyTokenizer()).process(
-        tokenizer = DummyTokenizer()
+            {"query": "hello", "chosen": "r1", "rejected": "r2"}
-        processor = DPOProcessor(tokenizer)
+        )
-        
+        for key in ["chosen", "chosen_mask", "rejected", "rejected_mask"]:
-        result = processor.process({
-            "query": "hello",
-            "chosen": "response1",
-            "rejected": "response2"
-        })
-        
-        expected_keys = ["chosen", "chosen_mask", "rejected", "rejected_mask"]
-        for key in expected_keys:
             assert key in result
             assert isinstance(result[key], torch.Tensor)
-    
+
-    def test_chosen_and_rejected_same_length_as_mask(self):
+    def test_masks_match_lengths(self):
-        """测试 chosen/rejected 长度与 mask 一致"""
+        result = DPOProcessor(DummyTokenizer()).process(
-        tokenizer = DummyTokenizer()
+            {"query": "test", "chosen": "yes", "rejected": "no"}
-        processor = DPOProcessor(tokenizer)
+        )
-        
-        result = processor.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 类型"""
+        result = DPOProcessor(DummyTokenizer()).process(
-        tokenizer = DummyTokenizer()
+            {"query": "test", "chosen": "yes", "rejected": "no"}
-        processor = DPOProcessor(tokenizer)
+        )
-        
-        result = processor.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):
-        """测试创建预训练处理器"""
+        assert isinstance(ProcessorFactory.create("pt", DummyTokenizer()), PreTrainProcessor)
-        tokenizer = DummyTokenizer()
+
-        processor = ProcessorFactory.create("pt", tokenizer)
-        assert isinstance(processor, PreTrainProcessor)
-    
     def test_create_sft_processor(self):
-        """测试创建 SFT 处理器"""
+        assert isinstance(ProcessorFactory.create("sft", DummyTokenizer()), SFTProcessor)
-        tokenizer = DummyTokenizer()
+
-        processor = ProcessorFactory.create("sft", tokenizer)
-        assert isinstance(processor, SFTProcessor)
-    
     def test_create_dpo_processor(self):
-        """测试创建 DPO 处理器"""
+        assert isinstance(ProcessorFactory.create("dpo", DummyTokenizer()), DPOProcessor)
-        tokenizer = DummyTokenizer()
+
-        processor = ProcessorFactory.create("dpo", tokenizer)
-        assert isinstance(processor, DPOProcessor)
-    
     def test_create_invalid_processor_raises_error(self):
-        """测试创建无效处理器类型抛出异常"""
+        with pytest.raises(ValueError, match="Unknown processor type"):
-        tokenizer = DummyTokenizer()
+            ProcessorFactory.create("invalid", DummyTokenizer())
-        with pytest.raises(ValueError, match="Invalid processor type"):
+
-            ProcessorFactory.create("invalid", tokenizer)
-    
     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)
-        ProcessorFactory.register("custom", CustomProcessor)
+        assert isinstance(ProcessorFactory.create("custom", DummyTokenizer()), CustomProcessor)
-        processor = ProcessorFactory.create("custom", tokenizer)
-        assert isinstance(processor, CustomProcessor)

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("<eos>")
+        return response_tokens + suffix
+
+
+def _decode(tokens):
+    return "".join(chr(t) for t in tokens)
+
+
+class TestChatMLStrategy:
+    def test_name(self):
+        assert ChatMLStrategy(DummyTokenizer()).name == "chatml"
+
+    def test_assemble_prompt(self):
+        tk = DummyTokenizer()
+        strategy = ChatMLStrategy(tk)
+        query_tokens = tk.encode("hello")
+        prompt = strategy.assemble_prompt(query_tokens)
+        text = _decode(prompt)
+        assert "<|im_start|>user" in text
+        assert "hello" in text
+        assert "<|im_start|>assistant" in text
+
+    def test_assemble_response(self):
+        tk = DummyTokenizer()
+        strategy = ChatMLStrategy(tk)
+        response_tokens = tk.encode("world")
+        response = strategy.assemble_response(response_tokens)
+        text = _decode(response)
+        assert "world" in text
+        assert "<|im_end|>" in text
+        assert "<eos>" in text
+
+    def test_prompt_ends_with_assistant_start(self):
+        tk = DummyTokenizer()
+        strategy = ChatMLStrategy(tk)
+        prompt = strategy.assemble_prompt(tk.encode("hi"))
+        # prompt 末尾应该是 assistant_start 的 token ids
+        assert prompt[-len(strategy._assistant_start_ids):] == strategy._assistant_start_ids
+
+
+class TestAlpacaStrategy:
+    def test_name(self):
+        assert AlpacaStrategy(DummyTokenizer()).name == "alpaca"
+
+    def test_assemble_prompt(self):
+        tk = DummyTokenizer()
+        strategy = AlpacaStrategy(tk)
+        query_tokens = tk.encode("hello")
+        prompt = strategy.assemble_prompt(query_tokens)
+        text = _decode(prompt)
+        assert "### Instruction:" in text
+        assert "hello" in text
+        assert "### Response:" in text
+
+    def test_assemble_response(self):
+        tk = DummyTokenizer()
+        strategy = AlpacaStrategy(tk)
+        response_tokens = tk.encode("world")
+        response = strategy.assemble_response(response_tokens)
+        text = _decode(response)
+        assert "world" in text
+        assert "<eos>" in text
+
+
+class TestStrategyFactory:
+    def test_create_chatml(self):
+        tk = DummyTokenizer()
+        assert isinstance(StrategyFactory.create("chatml", tk), ChatMLStrategy)
+
+    def test_create_alpaca(self):
+        tk = DummyTokenizer()
+        assert isinstance(StrategyFactory.create("alpaca", tk), AlpacaStrategy)
+
+    def test_create_invalid_raises_error(self):
+        with pytest.raises(ValueError, match="Unknown strategy"):
+            StrategyFactory.create("invalid_strategy", DummyTokenizer())
+
+    def test_register_and_create(self):
+        StrategyFactory.register("dummy")(DummyStrategy)
+        tk = DummyTokenizer()
+        strategy = StrategyFactory.create("dummy", tk)
+        assert isinstance(strategy, DummyStrategy)
+        assert strategy.name == "dummy"
+
+    def test_available_strategies(self):
+        strategies = StrategyFactory.available_types()
+        assert "chatml" in strategies
+        assert "alpaca" in strategies

tests/test_tokenizer.py

@@ -1,229 +0,0 @@
-"""单元测试：pipeline.tokenizer 模块中的 BpeTokenizer 类"""
-
-import pytest
-from pathlib import Path
-import tempfile
-
-from pipeline.tokenizer import BpeTokenizer
-
-
-class TestBpeTokenizer:
-    """BpeTokenizer 类的测试套件"""
-    
-    def test_initialization_without_path(self):
-        """测试不加载外部文件初始化"""
-        tokenizer = BpeTokenizer()
-        assert tokenizer is not None
-        assert hasattr(tokenizer, '_tokenizer')
-    
-    def test_initialization_with_path(self):
-        """测试加载外部文件初始化"""
-        # 这个测试假设没有预训练的分词器文件，所以只测试不抛出异常
-        # 实际使用中需要提供有效的分词器文件路径
-        try:
-            tokenizer = BpeTokenizer(path="nonexistent.json")
-        except Exception:
-            # 预期会抛出异常，因为文件不存在
-            pass
-    
-    def test_vocab_size(self):
-        """测试获取词汇表大小"""
-        tokenizer = BpeTokenizer()
-        vocab_size = len(tokenizer)
-        assert isinstance(vocab_size, int)
-        assert vocab_size >= 0
-    
-    def test_special_tokens_exist(self):
-        """测试特殊token是否存在"""
-        tokenizer = BpeTokenizer()
-        
-        # 检查控制token
-        assert hasattr(tokenizer, '_control_tokens')
-        assert '<bos>' in tokenizer._control_tokens
-        assert '<eos>' in tokenizer._control_tokens
-        assert '<pad>' in tokenizer._control_tokens
-        
-        # 检查特殊token
-        assert hasattr(tokenizer, '_special_tokens')
-        assert '<|im_start|>' in tokenizer._special_tokens
-        assert '<|im_end|>' in tokenizer._special_tokens
-    
-    def test_encode_string(self):
-        """测试编码单个字符串"""
-        tokenizer = BpeTokenizer()
-        
-        # 使用简单的ASCII字符测试
-        result = tokenizer.encode("hello")
-        
-        # 返回应该是 token IDs 列表
-        assert isinstance(result, list)
-    
-    def test_encode_list(self):
-        """测试编码字符串列表"""
-        tokenizer = BpeTokenizer()
-        
-        texts = ["hello", "world", "test"]
-        result = tokenizer.encode(texts)
-        
-        # 返回应该是列表的列表
-        assert isinstance(result, list)
-        assert len(result) == len(texts)
-        for item in result:
-            assert isinstance(item, list)
-    
-    def test_encode_with_output_tokens(self):
-        """测试编码返回tokens而非ids"""
-        tokenizer = BpeTokenizer()
-        
-        result = tokenizer.encode("hello", out_ids=False)
-        
-        # 应该返回 token 字符串列表
-        assert isinstance(result, list)
-    
-    def test_encode_with_special_tokens(self):
-        """测试编码添加特殊token"""
-        tokenizer = BpeTokenizer()
-        
-        result = tokenizer.encode("hello", add_special_tokens=True)
-        
-        assert isinstance(result, list)
-    
-    def test_decode(self):
-        """测试解码token IDs"""
-        tokenizer = BpeTokenizer()
-        
-        # 解码空列表
-        result = tokenizer.decode([])
-        assert isinstance(result, str)
-        
-        # 解码包含一些ID的列表（假设有 vocab）
-        # 如果分词器未训练，可能无法正确解码
-        result = tokenizer.decode([104, 101, 108, 108, 111])  # "hello" 的 ASCII
-        assert isinstance(result, str)
-    
-    def test_decode_with_special_tokens(self):
-        """测试解码保留特殊token"""
-        tokenizer = BpeTokenizer()
-        
-        # 解码空列表
-        result = tokenizer.decode([], skip_special_tokens=False)
-        assert isinstance(result, str)
-    
-    def test_stop_ids_property(self):
-        """测试 stop_ids 属性"""
-        tokenizer = BpeTokenizer()
-        
-        stop_ids = tokenizer.stop_ids
-        assert isinstance(stop_ids, list)
-    
-    def test_special_token_properties(self):
-        """测试特殊token ID属性"""
-        tokenizer = BpeTokenizer()
-        
-        # 这些属性可能返回 None 如果分词器未训练
-        bos_id = tokenizer.bos_id
-        eos_id = tokenizer.eos_id
-        pad_id = tokenizer.pad_id
-        
-        # 只验证属性存在且为 int 或 None
-        assert isinstance(bos_id, (int, type(None)))
-        assert isinstance(eos_id, (int, type(None)))
-        assert isinstance(pad_id, (int, type(None)))
-    
-    def test_save_method_exists(self):
-        """测试 save 方法存在"""
-        tokenizer = BpeTokenizer()
-        assert hasattr(tokenizer, 'save')
-        assert callable(tokenizer.save)
-    
-    def test_load_method_exists(self):
-        """测试 load 方法存在"""
-        tokenizer = BpeTokenizer()
-        assert hasattr(tokenizer, 'load')
-        assert callable(tokenizer.load)
-    
-    def test_train_method_exists(self):
-        """测试 train 方法存在"""
-        tokenizer = BpeTokenizer()
-        assert hasattr(tokenizer, 'train')
-        assert callable(tokenizer.train)
-    
-    def test_train_from_iterator_method_exists(self):
-        """测试 train_from_iterator 方法存在"""
-        tokenizer = BpeTokenizer()
-        assert hasattr(tokenizer, 'train_from_iterator')
-        assert callable(tokenizer.train_from_iterator)
-
-
-class TestBpeTokenizerIntegration:
-    """BpeTokenizer 集成测试"""
-    
-    def test_encode_decode_roundtrip(self):
-        """测试编码解码往返"""
-        tokenizer = BpeTokenizer()
-        
-        original = "hello world"
-        encoded = tokenizer.encode(original)
-        decoded = tokenizer.decode(encoded)
-        
-        # 往返后应该得到类似的结果
-        # 注意：由于分词器可能未训练，结果可能不完全一致
-        assert isinstance(encoded, list)
-        assert isinstance(decoded, str)
-    
-    def test_train_from_iterator_small_corpus(self, tmp_path):
-        """测试使用小语料库训练"""
-        tokenizer = BpeTokenizer()
-        
-        # 创建临时训练文件
-        train_file = tmp_path / "train.txt"
-        train_content = "hello world\nthis is a test\nmachine learning\n"
-        train_file.write_text(train_content)
-        
-        # 训练分词器（使用较小的 vocab size 加快测试）
-        try:
-            tokenizer.train(
-                files=[str(train_file)],
-                vocab_size=100,
-                min_freq=1,
-                reserved_token_size=10
-            )
-            
-            # 验证训练后分词器可用
-            result = tokenizer.encode("hello")
-            assert isinstance(result, list)
-            assert len(result) > 0
-        except Exception as e:
-            pytest.skip(f"Training failed: {e}")
-    
-    def test_save_and_load_tokenizer(self, tmp_path):
-        """测试保存和加载分词器"""
-        tokenizer = BpeTokenizer()
-        
-        # 创建临时训练文件并训练
-        train_file = tmp_path / "train.txt"
-        train_content = "hello world\ntest data\n"
-        train_file.write_text(train_content)
-        
-        try:
-            tokenizer.train(
-                files=[str(train_file)],
-                vocab_size=50,
-                min_freq=1,
-                reserved_token_size=5
-            )
-            
-            # 保存
-            save_path = tmp_path / "tokenizer.json"
-            tokenizer.save(str(save_path))
-            
-            # 加载到新实例
-            new_tokenizer = BpeTokenizer()
-            new_tokenizer.load(str(save_path))
-            
-            # 验证加载后分词器可用
-            result = new_tokenizer.encode("hello")
-            assert isinstance(result, list)
-            
-        except Exception as e:
-            pytest.skip(f"Save/load test failed: {e}")