fix: 修复特殊token 问题

2026-04-02 16:16:02 +08:00
parent e01ec081b3
commit f44ad6912e
28 changed files with 334 additions and 163 deletions
@@ -59,7 +59,7 @@ Stage 1: Export Dataset                     Stage 2: Tokenize & Cache
 **PT (Pre-training)**
 ```
 Input:  {"text": "Hello world"}
-Action: tokenizer.encode(text + "<eos>")
+Action: tokenizer.encode(text + "<｜end▁of▁sentence｜>")
 Output: {"sequence": Tensor[int32]}
 ```

@@ -67,7 +67,7 @@ Output: {"sequence": Tensor[int32]}
 ```
 Input:  {"query": "...", "response": "..."}
 Action:
-  1. strategy.build_prompt(input_dict)  ->  "<|im_start|>user\n...\n<|im_start|>assistant\n"
+  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
@@ -145,7 +145,7 @@ strategy = StrategyFactory.create("chatml",
    user_start="<s>user\n",
    user_end="</s>\n",
    assistant_start="<s>assistant\n",
-    assistant_end="</s>\n<eos>",
+    assistant_end="</s>\n<｜end▁of▁sentence｜>",
 )
 processor = ProcessorFactory.create_with_strategy("sft", tokenizer, strategy)
 ```
@@ -154,8 +154,8 @@ 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>`                                                     |
+| ChatML    | `"chatml"` | `<\|im_start\|>user`, `<\|im_end\|>`, `<\|im_start\|>assistant`, `<｜end▁of▁sentence｜>`                       |
+| Alpaca    | `"alpaca"` | `### Instruction:`, `### Response:`, `<｜end▁of▁sentence｜>`                                                     |

 所有策略的 token 均可通过构造函数参数自定义，同时支持通过 `StrategyFactory.register()` 注册新格式。

@@ -70,7 +70,7 @@
 **PreTrainProcessor** (`"pt"`)
 ```
 Input:  {"text": "Hello world"}
-Action: tokenizer.encode(text + "<eos>")
+Action: tokenizer.encode(text + "<｜end▁of▁sentence｜>")
 Output: {"sequence": Tensor[int32]}
 ```

@@ -106,7 +106,7 @@ Output: {"chosen": Tensor, "chosen_mask": Tensor[bool],
 **接口**:
 - `build_prompt(input_dict)` — 构建包含 query 的完整 prompt
 - `build_response_prefix()` — response 前缀（当前均返回空串）
- `build_response_suffix()` — response 后缀（含 `<eos>`）
+- `build_response_suffix()` — response 后缀（含 `<｜end▁of▁sentence｜>`）
 - `response_start_token` — response 起始 token（用于 DPO 的 loss mask 定位）
 - `eos_tokens` — 结束 token

@@ -119,7 +119,7 @@ ChatML:
 | user_start      | `<\|im_start\|>user\n`                    |
 | user_end        | `<\|im_end\|>\n`                          |
 | assistant_start | `<\|im_start\|>assistant\n`               |
-| assistant_end   | `<\|im_end\|>\n<eos>`                     |
+| assistant_end   | `<\|im_end\|>\n<｜end▁of▁sentence｜>`                     |

 Alpaca:

@@ -127,7 +127,7 @@ Alpaca:
 |------------------|------------------------|
 | instruction_start | `### Instruction:\n`  |
 | response_start    | `### Response:\n`     |
-| response_suffix   | `\n<eos>`             |
+| response_suffix   | `\n<｜end▁of▁sentence｜>`             |

 **自定义示例**:
 ```python
@@ -135,7 +135,7 @@ strategy = StrategyFactory.create("chatml",
    user_start="<s>user\n",
    user_end="</s>\n",
    assistant_start="<s>assistant\n",
-    assistant_end="</s>\n<eos>",
+    assistant_end="</s>\n<｜end▁of▁sentence｜>",
 )
 ```

@@ -178,7 +178,7 @@ StrategyFactory.register("my_format", MyStrategy)

 ### BpeTokenizer (`pipeline/tokenizer.py`)

-基于 HuggingFace `tokenizers` 库的 BPE 分词器，支持从文件加载、训练、保存。内置 `<bos>`/`<eos>`/`<pad>` 控制符和 `<|im_start|>`/`<|im_end|>` 特殊 token。
+基于 HuggingFace `tokenizers` 库的 BPE 分词器，支持从文件加载、训练、保存。内置 `<｜begin▁of▁sentence｜>`/`<｜end▁of▁sentence｜>`/`<｜▁pad▁｜>` 控制符和 `<｜im▁start｜>`/`<｜im▁end｜>` 特殊 token。

 ## API 参考

@@ -5,24 +5,29 @@ from pipeline.packing import SequencePacker
 from pipeline.io import IOHandler, export_dataset, cache_jsonl
 from pipeline.processors import ProcessorFactory, BaseProcessor
 from pipeline.utils import setup_logging
-from pipeline.strategies import PromptStrategy, ChatMLStrategy, AlpacaStrategy, StrategyFactory
+from pipeline.strategies import (
+    PromptStrategy,
+    ChatMLStrategy,
+    AlpacaStrategy,
+    StrategyFactory,
+)

 # Configure project-level logging
 setup_logging()

 __all__ = [
    # Core modules
-    'BpeTokenizer',
-    'TextNormalizer',
-    'SequencePacker',
-    'IOHandler',
-    'ProcessorFactory',
-    'BaseProcessor',
-    'export_dataset',
-    'cache_jsonl',
+    "BpeTokenizer",
+    "TextNormalizer",
+    "SequencePacker",
+    "IOHandler",
+    "ProcessorFactory",
+    "BaseProcessor",
+    "export_dataset",
+    "cache_jsonl",
    # Strategy pattern
-    'PromptStrategy',
-    'ChatMLStrategy',
-    'AlpacaStrategy',
-    'StrategyFactory',
+    "PromptStrategy",
+    "ChatMLStrategy",
+    "AlpacaStrategy",
+    "StrategyFactory",
 ]
@@ -1,4 +1,5 @@
 """File, HDF5, JSONL I/O operations."""
+
 import json
 import os
 import logging
@@ -33,7 +34,9 @@ class IOHandler:
        return sorted(files)

    @staticmethod
-    def fetch_folders(root_dir: str, filter_func: Optional[Callable[[str], bool]] = None) -> List[str]:
+    def fetch_folders(
+        root_dir: str, filter_func: Optional[Callable[[str], bool]] = None
+    ) -> List[str]:
        folders = []
        for root, dirs, _ in os.walk(root_dir):
            for dir_name in dirs:
@@ -44,15 +47,17 @@ class IOHandler:

    @staticmethod
    @error_handler()
-    def save_h5(output_dir: str, file_name: str, tensor_group: Dict[str, List[Tensor]]) -> None:
+    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:
+        with h5py.File(full_path, "w") as f:
            for key, tensors in tensor_group.items():
                grp = f.create_group(key)
                for idx, tensor in enumerate(tensors):
-                    grp.create_dataset(f'data_{idx}', data=tensor.cpu().numpy())
+                    grp.create_dataset(f"data_{idx}", data=tensor.cpu().numpy())

    @staticmethod
    @error_handler()
@@ -63,7 +68,7 @@ class IOHandler:
        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:
+            with h5py.File(h5_file, "r") as f:
                for key in f.keys():
                    grp = f[key]
                    dsets = []
@@ -92,7 +97,9 @@ def export_dataset(
    *,
    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,
+    process_func: Optional[
+        Callable[[Dict[str, Any]], Union[Dict[str, Any], List[Dict[str, Any]]]]
+    ] = None,
    column: str = "text",
 ) -> List[str]:
    """
@@ -125,7 +132,11 @@ def export_dataset(
        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]}
+                    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")
@@ -172,17 +183,23 @@ def cache_jsonl(
        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):
+            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.")
+                    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.")
+                    logger.warning(
+                        f"Unexpected error processing line {line_num} in {file_path}: {e}. Skipping line."
+                    )
                    continue

        if pack_size > 0:
@@ -37,7 +37,9 @@ class SequencePacker:
        identical chunk boundaries. Element-level correspondence is preserved.
    """

-    def __init__(self, pack_size: int, pad_value: int = 0, dtype: torch.dtype = torch.int32):
+    def __init__(
+        self, pack_size: int, pad_value: int = 0, dtype: torch.dtype = torch.int32
+    ):
        self.pack_size = pack_size
        self.pad_value = pad_value
        self.dtype = dtype
@@ -3,6 +3,7 @@
 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
@@ -1,4 +1,5 @@
 """Processor base class and shared utilities."""
+
 from abc import ABC, abstractmethod
 from typing import Dict, List, Any, Tuple

@@ -1,4 +1,5 @@
 """DPO preference learning data processor."""
+
 from typing import Dict, List, Any, Optional

 import torch
@@ -1,4 +1,5 @@
 """Factory for creating and registering processors."""
+
 from typing import Dict, List, Any, Optional, Type

 from pipeline.processors.base import BaseProcessor
@@ -1,4 +1,5 @@
 """Pre-training data processor."""
+
 from typing import Dict, List, Any

 import torch
@@ -18,7 +19,7 @@ class PreTrainProcessor(BaseProcessor):

    def process(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]:
        segment = input_dict["text"]
-        tokens = self.tokenizer.encode(f"{segment}<eos>")
+        tokens = self.tokenizer.encode(f"{segment}<｜end▁of▁sentence｜>")
        return {"sequence": torch.tensor(tokens, dtype=torch.int32)}

    @property
@@ -1,4 +1,5 @@
 """Supervised fine-tuning data processor."""
+
 from typing import Dict, List, Any, Optional

 import torch
@@ -1,4 +1,5 @@
 """Strategy pattern for prompt/response format abstraction."""
+
 from pipeline.strategies.base import PromptStrategy
 from pipeline.strategies.factory import StrategyFactory

@@ -1,4 +1,5 @@
 """Alpaca format strategy."""
+
 from typing import List

 from pipeline.tokenizer import BpeTokenizer
@@ -8,14 +9,14 @@ from pipeline.strategies.factory import StrategyFactory

@StrategyFactory.register("alpaca")
 class AlpacaStrategy(PromptStrategy):
-    """Alpaca format: ``### Instruction: ... \\n\\n### Response: ... <eos>``"""
+    """Alpaca format:"""

    def __init__(
        self,
        tokenizer: BpeTokenizer,
        instruction_start: str = "### Instruction:\n",
        response_start: str = "### Response:\n",
-        response_suffix: str = "\n<eos>",
+        response_suffix: str = "\n<｜end▁of▁sentence｜>",
    ):
        super().__init__(tokenizer)
        self.instruction_start = instruction_start
@@ -1,4 +1,5 @@
 """Abstract base class for prompt construction strategies."""
+
 from abc import ABC, abstractmethod
 from typing import List

@@ -30,7 +31,7 @@ class PromptStrategy(ABC):
        """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``.
+        start marker, e.g. ``<｜im▁start｜>assistant\n``.
        """

    @abstractmethod
@@ -1,4 +1,5 @@
 """ChatML format strategy."""
+
 from typing import List

 from pipeline.tokenizer import BpeTokenizer
@@ -8,15 +9,15 @@ 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>``"""
+    """ChatML format strategy."""

    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>",
+        user_start: str = "<｜im▁start｜>user\n",
+        user_end: str = "<｜im▁end｜>\n",
+        assistant_start: str = "<｜im▁start｜>assistant\n",
+        assistant_end: str = "<｜im▁end｜>\n",
    ):
        super().__init__(tokenizer)

@@ -1,4 +1,5 @@
 """Factory for creating and registering prompt strategies."""
+
 from typing import Dict, List, Type

 from pipeline.tokenizer import BpeTokenizer
@@ -6,16 +6,25 @@ class TextNormalizer:
    """Text normalization."""

    DEFAULT_REPLACEMENTS = {
-        "\\[": "$$", "\\]": "$$", "\\(": "$", "\\)": "$",
-        '\u2018': "'", '\u2019': "'", '\u0060': "'",
-        '\u201C': '"', '\u201D': '"',
-        '\u2013': '-', '\u2014': '--', '\u2212': '-',
-        '\u00A0': ' ', '\u2026': '...'
+        "\\[": "$$",
+        "\\]": "$$",
+        "\\(": "$",
+        "\\)": "$",
+        "\u2018": "'",
+        "\u2019": "'",
+        "\u0060": "'",
+        "\u201c": '"',
+        "\u201d": '"',
+        "\u2013": "-",
+        "\u2014": "--",
+        "\u2212": "-",
+        "\u00a0": " ",
+        "\u2026": "...",
    }

    def __init__(self, custom_rules: Optional[Dict[str, str]] = None):
        self.replacements = {**self.DEFAULT_REPLACEMENTS, **(custom_rules or {})}
-        self._pattern = re.compile('|'.join(re.escape(k) for k in self.replacements))
+        self._pattern = re.compile("|".join(re.escape(k) for k in self.replacements))

    def normalize(self, text: str) -> str:
        return self._pattern.sub(lambda m: self.replacements[m.group()], text)
@@ -7,20 +7,25 @@ from typing import List, Union, Optional, Tuple, Iterator

 class BpeTokenizer:
    def __init__(self, path: Optional[str] = None):
-        self._control_tokens = ["<bos>", "<eos>", "<pad>"]
-        self._special_tokens = ["<|im_start|>", "<|im_end|>"]
+        self._control_tokens = [
+            "<｜begin▁of▁sentence｜>",
+            "<｜end▁of▁sentence｜>",
+            "<｜▁pad▁｜>",
+        ]
+        self._special_tokens = ["<｜im▁start｜>", "<｜im▁end｜>"]

        model = BPE()
        self._tokenizer = Tokenizer(model)
-        self._tokenizer.normalizer = normalizers.Sequence([
-            normalizers.NFC(),
-            normalizers.Strip() 
-        ])
+        self._tokenizer.normalizer = normalizers.Sequence(
+            [normalizers.NFC(), normalizers.Strip()]
+        )

-        self._tokenizer.pre_tokenizer = pre_tokenizers.Sequence([
+        self._tokenizer.pre_tokenizer = pre_tokenizers.Sequence(
+            [
                pre_tokenizers.UnicodeScripts(),
-            pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=True)
-        ])
+                pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=True),
+            ]
+        )

        self._tokenizer.decoder = decoders.ByteLevel()
        self._tokenizer.post_processor = processors.ByteLevel(trim_offsets=True)
@@ -28,10 +33,21 @@ class BpeTokenizer:
        if path is not None:
            self._tokenizer = Tokenizer.from_file(path)

-    def _prepare_trainer(self, vocab_size: int, min_freq: int, reserved_token_size: int, max_token_length: int = 18) -> Tuple[BpeTrainer, int, List[str]]:
+    def _prepare_trainer(
+        self,
+        vocab_size: int,
+        min_freq: int,
+        reserved_token_size: int,
+        max_token_length: int = 18,
+    ) -> Tuple[BpeTrainer, int, List[str]]:
        assert reserved_token_size > len(self._special_tokens)
-        reserved_tokens = [f"<|reserve{i:02d}|>" for i in range(reserved_token_size - len(self._special_tokens))]
-        detail_vocab_size = vocab_size - (len(reserved_tokens) + len(self._special_tokens))
+        reserved_tokens = [
+            f"<｜reserve{i:02d}｜>"
+            for i in range(reserved_token_size - len(self._special_tokens))
+        ]
+        detail_vocab_size = vocab_size - (
+            len(reserved_tokens) + len(self._special_tokens)
+        )

        alphabet = pre_tokenizers.ByteLevel.alphabet()
        min_size = len(alphabet) + len(self._control_tokens)
@@ -42,30 +58,46 @@ class BpeTokenizer:
            min_frequency=min_freq,
            limit_alphabet=detail_vocab_size // 6,
            max_token_length=max_token_length,
-            special_tokens=self._control_tokens,
+            special_tokens=self._control_tokens + self._special_tokens,
            initial_alphabet=alphabet,
            show_progress=True,
        )

        return trainer, detail_vocab_size, reserved_tokens

-    def train(self, files: List[str], vocab_size: int, min_freq: int, reserved_token_size: int = 100) -> None:
+    def train(
+        self,
+        files: List[str],
+        vocab_size: int,
+        min_freq: int,
+        reserved_token_size: int = 100,
+    ) -> None:
        trainer, _, reserved_tokens = self._prepare_trainer(
            vocab_size=vocab_size,
            min_freq=min_freq,
-            reserved_token_size=reserved_token_size
+            reserved_token_size=reserved_token_size,
        )
        self._tokenizer.train(files=files, trainer=trainer)
-        self._tokenizer.add_special_tokens(self._special_tokens + reserved_tokens)
+        self._tokenizer.add_special_tokens(
+            self._control_tokens + self._special_tokens + reserved_tokens
+        )

-    def train_from_iterator(self, iterator: Iterator[str], vocab_size: int, min_freq: int, reserved_token_size: int = 100) -> None:
+    def train_from_iterator(
+        self,
+        iterator: Iterator[str],
+        vocab_size: int,
+        min_freq: int,
+        reserved_token_size: int = 100,
+    ) -> None:
        trainer, _, reserved_tokens = self._prepare_trainer(
            vocab_size=vocab_size,
            min_freq=min_freq,
-            reserved_token_size=reserved_token_size
+            reserved_token_size=reserved_token_size,
        )
        self._tokenizer.train_from_iterator(iterator=iterator, trainer=trainer)
-        self._tokenizer.add_special_tokens(self._special_tokens + reserved_tokens)
+        self._tokenizer.add_special_tokens(
+            self._control_tokens + self._special_tokens + reserved_tokens
+        )

    def save(self, path: str) -> None:
        self._tokenizer.save(path)
@@ -73,13 +105,24 @@ class BpeTokenizer:
    def load(self, path: str) -> None:
        self._tokenizer = Tokenizer.from_file(path)

-    def encode(self, tokens: Union[str, List[str]], out_ids: bool = True, add_special_tokens: bool = False) -> Union[List[int], List[str], List[List[int]], List[List[str]]]:
+    def encode(
+        self,
+        tokens: Union[str, List[str]],
+        out_ids: bool = True,
+        add_special_tokens: bool = False,
+    ) -> Union[List[int], List[str], List[List[int]], List[List[str]]]:
        if isinstance(tokens, str):
-            encoded: Encoding = self._tokenizer.encode(tokens, add_special_tokens=add_special_tokens)
+            encoded: Encoding = self._tokenizer.encode(
+                tokens, add_special_tokens=add_special_tokens
+            )
            return encoded.ids if out_ids else encoded.tokens
        elif isinstance(tokens, list):
-            encoded_list: List[Encoding] = self._tokenizer.encode_batch(tokens, add_special_tokens=add_special_tokens)
-            return [encoded.ids if out_ids else encoded.tokens for encoded in encoded_list]
+            encoded_list: List[Encoding] = self._tokenizer.encode_batch(
+                tokens, add_special_tokens=add_special_tokens
+            )
+            return [
+                encoded.ids if out_ids else encoded.tokens for encoded in encoded_list
+            ]

    def decode(self, tokens: List[int], skip_special_tokens: bool = True) -> str:
        return self._tokenizer.decode(tokens, skip_special_tokens=skip_special_tokens)
@@ -95,12 +138,12 @@ class BpeTokenizer:

    @property
    def bos_id(self) -> int:
-        return self._tokenizer.token_to_id("<bos>")
+        return self._tokenizer.token_to_id("<｜begin▁of▁sentence｜>")

    @property
    def eos_id(self) -> int:
-        return self._tokenizer.token_to_id("<eos>")
+        return self._tokenizer.token_to_id("<｜end▁of▁sentence｜>")

    @property
    def pad_id(self) -> int:
-        return self._tokenizer.token_to_id("<pad>")
+        return self._tokenizer.token_to_id("<｜▁pad▁｜>")
@@ -30,7 +30,9 @@ def error_handler(
                if reraise:
                    raise
                return None
+
        return wrapper
+
    return decorator


@@ -7,6 +7,7 @@ Usage:
    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

@@ -18,16 +19,34 @@ def main():
    parser = argparse.ArgumentParser(description="JSONL -> H5 cache")
    parser.add_argument("type", choices=["pt", "sft", "dpo"], help="Processor type")
    parser.add_argument("input_dir", help="Directory containing JSONL files")
-    parser.add_argument("-o", "--output-dir", default=None,
-                        help="H5 output dir (default: <input_dir>/cached)")
-    parser.add_argument("-t", "--tokenizer", default="./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="Pack size, <=0 to disable (default: -1)")
-    parser.add_argument("--pad-value", type=int, default=1,
-                        help="Padding value (default: 1)")
+    parser.add_argument(
+        "-o",
+        "--output-dir",
+        default=None,
+        help="H5 output dir (default: <input_dir>/cached)",
+    )
+    parser.add_argument(
+        "-t",
+        "--tokenizer",
+        default="./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="Pack size, <=0 to disable (default: -1)",
+    )
+    parser.add_argument(
+        "--pad-value", type=int, default=1, help="Padding value (default: 1)"
+    )
    args = parser.parse_args()

    jsonl_files = IOHandler.fetch_files(args.input_dir, suffix=".jsonl")
@@ -4,7 +4,7 @@ from pipeline import export_dataset
 if __name__ == "__main__":
    dataset = load_dataset(
        "opencsg/chinese-cosmopedia",
-        data_files={"train": [f"data/000{i:02d}.parquet" for i in range(25)]}
+        data_files={"train": [f"data/000{i:02d}.parquet" for i in range(25)]},
    )
    export_dataset(
        dataset=dataset["train"],
@@ -7,14 +7,28 @@ normalizer = TextNormalizer()
 def process_func(input_dict: dict):
    query = input_dict["prompt"] if input_dict["prompt"] else ""
    resp = input_dict["response"] if input_dict["response"] else ""
-    return {"query": normalizer.normalize(query), "response": normalizer.normalize(resp)}
+    return {
+        "query": normalizer.normalize(query),
+        "response": normalizer.normalize(resp),
+    }


 if __name__ == "__main__":
    all_data = [
-        'stem_zh', 'infinity-instruct', 'firefly', 'magpie', 'dpsk-r1-distil',
-        'coig-cqia', 'disc-law', 'neo_sft_phase2', 'chinese-medical', 'chinese-reasoning-distil',
-        'psycho-10k-dpsk-r1', 'sof-c-zh', 'industryinstruction', 'Chinese-QA-AFAF',
+        "stem_zh",
+        "infinity-instruct",
+        "firefly",
+        "magpie",
+        "dpsk-r1-distil",
+        "coig-cqia",
+        "disc-law",
+        "neo_sft_phase2",
+        "chinese-medical",
+        "chinese-reasoning-distil",
+        "psycho-10k-dpsk-r1",
+        "sof-c-zh",
+        "industryinstruction",
+        "Chinese-QA-AFAF",
    ]

    dataset_list = []
@@ -30,7 +30,6 @@ class DummyProcessor(BaseProcessor):


 class TestCacheJsonl:
-
    def test_basic_cache_functionality(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            jsonl_path = os.path.join(tmpdir, "test.jsonl")
@@ -41,8 +40,11 @@ class TestCacheJsonl:

            processor = DummyProcessor()
            output_files = cache_jsonl(
-                files=[jsonl_path], output_dir=tmpdir,
-                processor=processor, pack_size=-1, pad_value=0,
+                files=[jsonl_path],
+                output_dir=tmpdir,
+                processor=processor,
+                pack_size=-1,
+                pad_value=0,
            )
            assert len(output_files) == 1
            assert os.path.exists(output_files[0])
@@ -57,8 +59,11 @@ class TestCacheJsonl:

            processor = DummyProcessor()
            output_files = cache_jsonl(
-                files=[jsonl_path], output_dir=tmpdir,
-                processor=processor, pack_size=10, pad_value=0,
+                files=[jsonl_path],
+                output_dir=tmpdir,
+                processor=processor,
+                pack_size=10,
+                pad_value=0,
            )
            assert len(output_files) == 1
            assert os.path.exists(output_files[0])
@@ -73,8 +78,11 @@ class TestCacheJsonl:

            processor = DummyProcessor()
            output_files = cache_jsonl(
-                files=[jsonl_path], output_dir=tmpdir,
-                processor=processor, pack_size=0, pad_value=-1,
+                files=[jsonl_path],
+                output_dir=tmpdir,
+                processor=processor,
+                pack_size=0,
+                pad_value=-1,
            )
            assert len(output_files) == 1
            assert os.path.exists(output_files[0])
@@ -91,8 +99,11 @@ class TestCacheJsonl:

            processor = DummyProcessor()
            output_files = cache_jsonl(
-                files=files, output_dir=tmpdir,
-                processor=processor, pack_size=-1, pad_value=0,
+                files=files,
+                output_dir=tmpdir,
+                processor=processor,
+                pack_size=-1,
+                pad_value=0,
            )
            assert len(output_files) == 2

@@ -103,7 +114,10 @@ class TestCacheJsonl:

            processor = DummyProcessor()
            output_files = cache_jsonl(
-                files=[jsonl_path], output_dir=tmpdir,
-                processor=processor, pack_size=-1, pad_value=0,
+                files=[jsonl_path],
+                output_dir=tmpdir,
+                processor=processor,
+                pack_size=-1,
+                pad_value=0,
            )
            assert len(output_files) == 1
@@ -11,7 +11,6 @@ from pipeline.io import IOHandler


 class TestIOHandler:
-
    def test_fetch_files_in_directory(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            Path(tmpdir, "file1.txt").touch()
@@ -41,7 +40,9 @@ class TestIOHandler:
            os.makedirs(os.path.join(tmpdir, "folder1"))
            os.makedirs(os.path.join(tmpdir, "folder2"))

-            folders = IOHandler.fetch_folders(tmpdir, filter_func=lambda x: "folder1" in x)
+            folders = IOHandler.fetch_folders(
+                tmpdir, filter_func=lambda x: "folder1" in x
+            )
            assert len(folders) == 1

    def test_save_and_load_h5(self):
@@ -56,8 +57,12 @@ class TestIOHandler:
            loaded = IOHandler.load_h5(tmpdir, share_memory=False)
            assert "sequence" in loaded
            assert "labels" in loaded
-            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))
+            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):
        with tempfile.TemporaryDirectory() as tmpdir:
@@ -70,9 +75,9 @@ class TestIOHandler:
        with tempfile.TemporaryDirectory() as tmpdir:
            for i, data in enumerate([[1, 2, 3], [4, 5, 6]]):
                h5_path = os.path.join(tmpdir, f"file{i}.h5")
-                with h5py.File(h5_path, 'w') as f:
+                with h5py.File(h5_path, "w") as f:
                    grp = f.create_group("data")
-                    grp.create_dataset('data_0', data=data)
+                    grp.create_dataset("data_0", data=data)

            loaded = IOHandler.load_h5(tmpdir, share_memory=False)
            assert len(loaded["data"]) == 2
@@ -82,9 +87,9 @@ class TestIOHandler:
            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:
+            with h5py.File(h5_path, "w") as f:
                grp = f.create_group("test")
-                grp.create_dataset('data_0', data=[1, 2])
+                grp.create_dataset("data_0", data=[1, 2])

            loaded = IOHandler.load_h5(tmpdir, share_memory=False)
            assert "test" in loaded
@@ -93,7 +98,11 @@ class TestIOHandler:
    def test_save_h5_multiple_tensors_per_key(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            tensor_group = {
-                "batch": [torch.tensor([1, 2]), torch.tensor([3, 4, 5]), torch.tensor([6])],
+                "batch": [
+                    torch.tensor([1, 2]),
+                    torch.tensor([3, 4, 5]),
+                    torch.tensor([6]),
+                ],
            }
            IOHandler.save_h5(tmpdir, "multi", tensor_group)
            loaded = IOHandler.load_h5(tmpdir, share_memory=False)
@@ -6,7 +6,6 @@ from pipeline.packing import SequencePacker


 class TestSequencePacker:
-
    def test_normal_packing(self):
        packer = SequencePacker(pack_size=10, pad_value=0)
        sequences = [
@@ -37,17 +36,21 @@ class TestSequencePacker:
    def test_long_sequence_split_across_chunks(self):
        """Sequences longer than pack_size are split across multiple chunks."""
        packer = SequencePacker(pack_size=5, pad_value=0)
-        packages = packer.pack([torch.tensor([1, 2, 3, 4, 5, 6, 7, 8], dtype=torch.int32)])
+        packages = packer.pack(
+            [torch.tensor([1, 2, 3, 4, 5, 6, 7, 8], dtype=torch.int32)]
+        )
        assert len(packages) == 2
        assert packages[0].tolist() == [1, 2, 3, 4, 5]
        assert packages[1].tolist() == [6, 7, 8, 0, 0]

    def test_padding_value(self):
        packer = SequencePacker(pack_size=8, pad_value=99)
-        packages = packer.pack([
+        packages = packer.pack(
+            [
                torch.tensor([1, 2], dtype=torch.int32),
                torch.tensor([3], dtype=torch.int32),
-        ])
+            ]
+        )
        assert packages[0][:3].tolist() == [1, 2, 3]
        assert packages[0][3:].tolist() == [99] * 5

@@ -85,10 +88,12 @@ class TestSequencePacker:

    def test_exact_pack_size_fit(self):
        packer = SequencePacker(pack_size=5, pad_value=0)
-        packages = packer.pack([
+        packages = packer.pack(
+            [
                torch.tensor([1, 2, 3, 4, 5], dtype=torch.int32),
                torch.tensor([6, 7, 8, 9, 10], dtype=torch.int32),
-        ])
+            ]
+        )
        assert len(packages) == 2
        assert packages[0].tolist() == [1, 2, 3, 4, 5]
        assert packages[1].tolist() == [6, 7, 8, 9, 10]
@@ -127,10 +132,12 @@ class TestSequencePacker:
    def test_stream_split_across_chunks(self):
        """Sequences are split across chunks in streaming mode."""
        packer = SequencePacker(pack_size=5, pad_value=0)
-        packages = packer.pack([
+        packages = packer.pack(
+            [
                torch.tensor([1, 2, 3], dtype=torch.int32),
                torch.tensor([4, 5, 6, 7, 8], dtype=torch.int32),
-        ])
+            ]
+        )
        assert len(packages) == 2
        # First chunk: [1, 2, 3, 4, 5] — first seq + part of second
        assert packages[0].tolist() == [1, 2, 3, 4, 5]
@@ -151,9 +158,11 @@ class TestSequencePacker:
        """Streaming concat preserves input order, no sorting."""
        packer = SequencePacker(pack_size=4, pad_value=-1)
        # short then long (fits in 2 chunks)
-        packages = packer.pack([
+        packages = packer.pack(
+            [
                torch.tensor([1], dtype=torch.int32),
                torch.tensor([2, 3, 4, 5, 6, 7], dtype=torch.int32),
-        ])
+            ]
+        )
        assert packages[0].tolist() == [1, 2, 3, 4]
        assert packages[1].tolist() == [5, 6, 7, -1]
@@ -44,18 +44,24 @@ class TestSFTProcessor:
        assert SFTProcessor(DummyTokenizer()).output_keys == ["sequence", "loss_mask"]

    def test_process_returns_both_keys(self):
-        result = SFTProcessor(DummyTokenizer()).process({"query": "hello", "response": "world"})
+        result = SFTProcessor(DummyTokenizer()).process(
+            {"query": "hello", "response": "world"}
+        )
        assert "sequence" in result
        assert "loss_mask" in result
        assert isinstance(result["sequence"], torch.Tensor)
        assert isinstance(result["loss_mask"], torch.Tensor)

    def test_loss_mask_correct_length(self):
-        result = SFTProcessor(DummyTokenizer()).process({"query": "hi", "response": "bye"})
+        result = SFTProcessor(DummyTokenizer()).process(
+            {"query": "hi", "response": "bye"}
+        )
        assert len(result["sequence"]) == len(result["loss_mask"])

    def test_loss_mask_is_bool(self):
-        result = SFTProcessor(DummyTokenizer()).process({"query": "ab", "response": "cd"})
+        result = SFTProcessor(DummyTokenizer()).process(
+            {"query": "ab", "response": "cd"}
+        )
        assert result["loss_mask"].dtype == torch.bool


@@ -89,13 +95,19 @@ class TestDPOProcessor:

 class TestProcessorFactory:
    def test_create_pre_train_processor(self):
-        assert isinstance(ProcessorFactory.create("pt", DummyTokenizer()), PreTrainProcessor)
+        assert isinstance(
+            ProcessorFactory.create("pt", DummyTokenizer()), PreTrainProcessor
+        )

    def test_create_sft_processor(self):
-        assert isinstance(ProcessorFactory.create("sft", DummyTokenizer()), SFTProcessor)
+        assert isinstance(
+            ProcessorFactory.create("sft", DummyTokenizer()), SFTProcessor
+        )

    def test_create_dpo_processor(self):
-        assert isinstance(ProcessorFactory.create("dpo", DummyTokenizer()), DPOProcessor)
+        assert isinstance(
+            ProcessorFactory.create("dpo", DummyTokenizer()), DPOProcessor
+        )

    def test_create_invalid_processor_raises_error(self):
        with pytest.raises(ValueError, match="Unknown processor type"):
@@ -114,4 +126,6 @@ class TestProcessorFactory:
                return {"custom": torch.tensor([1, 2, 3])}

        ProcessorFactory.register("custom")(CustomProcessor)
-        assert isinstance(ProcessorFactory.create("custom", DummyTokenizer()), CustomProcessor)
+        assert isinstance(
+            ProcessorFactory.create("custom", DummyTokenizer()), CustomProcessor
+        )
@@ -28,7 +28,7 @@ class DummyStrategy(PromptStrategy):
        return prefix + query_tokens

    def assemble_response(self, response_tokens):
-        suffix = self._encode_format("<eos>")
+        suffix = self._encode_format("<｜end▁of▁sentence｜>")
        return response_tokens + suffix


@@ -46,9 +46,9 @@ class TestChatMLStrategy:
        query_tokens = tk.encode("hello")
        prompt = strategy.assemble_prompt(query_tokens)
        text = _decode(prompt)
-        assert "<|im_start|>user" in text
+        assert "<｜im▁start｜>user" in text
        assert "hello" in text
-        assert "<|im_start|>assistant" in text
+        assert "<｜im▁start｜>assistant" in text

    def test_assemble_response(self):
        tk = DummyTokenizer()
@@ -57,15 +57,18 @@ class TestChatMLStrategy:
        response = strategy.assemble_response(response_tokens)
        text = _decode(response)
        assert "world" in text
-        assert "<|im_end|>" in text
-        assert "<eos>" in text
+        assert "<｜im▁end｜>" in text
+        assert "<｜end▁of▁sentence｜>" 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
+        assert (
+            prompt[-len(strategy._assistant_start_ids) :]
+            == strategy._assistant_start_ids
+        )


 class TestAlpacaStrategy:
@@ -89,7 +92,7 @@ class TestAlpacaStrategy:
        response = strategy.assemble_response(response_tokens)
        text = _decode(response)
        assert "world" in text
-        assert "<eos>" in text
+        assert "<｜end▁of▁sentence｜>" in text


 class TestStrategyFactory: