fix: 修复特殊token 问题

README.md

@@ -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>`                       |
+| ChatML    | `"chatml"` | `<\|im_start\|>user`, `<\|im_end\|>`, `<\|im_start\|>assistant`, `<｜end▁of▁sentence｜>`                       |
-| Alpaca    | `"alpaca"` | `### Instruction:`, `### Response:`, `<eos>`                                                     |
+| Alpaca    | `"alpaca"` | `### Instruction:`, `### Response:`, `<｜end▁of▁sentence｜>`                                                     |
 
 所有策略的 token 均可通过构造函数参数自定义，同时支持通过 `StrategyFactory.register()` 注册新格式。
 

docs/Design.md

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

pipeline/__init__.py

@@ -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',
+    "BpeTokenizer",
-    'TextNormalizer',
+    "TextNormalizer",
-    'SequencePacker',
+    "SequencePacker",
-    'IOHandler',
+    "IOHandler",
-    'ProcessorFactory',
+    "ProcessorFactory",
-    'BaseProcessor',
+    "BaseProcessor",
-    'export_dataset',
+    "export_dataset",
-    'cache_jsonl',
+    "cache_jsonl",
     # Strategy pattern
-    'PromptStrategy',
+    "PromptStrategy",
-    'ChatMLStrategy',
+    "ChatMLStrategy",
-    'AlpacaStrategy',
+    "AlpacaStrategy",
-    'StrategyFactory',
+    "StrategyFactory",
 ]

pipeline/io.py

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

pipeline/packing.py

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

pipeline/processors/__init__.py

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

pipeline/processors/base.py

@@ -1,4 +1,5 @@
 """Processor base class and shared utilities."""
+
 from abc import ABC, abstractmethod
 from typing import Dict, List, Any, Tuple
 

pipeline/processors/dpo.py

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

pipeline/processors/factory.py

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

pipeline/processors/pretrain.py

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

pipeline/processors/sft.py

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

pipeline/strategies/__init__.py

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

pipeline/strategies/alpaca.py

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

pipeline/strategies/base.py

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

pipeline/strategies/chatml.py

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

pipeline/strategies/factory.py

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

pipeline/text.py

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

pipeline/tokenizer.py

@@ -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._control_tokens = [
-        self._special_tokens = ["<|im_start|>", "<|im_end|>"]
+            "<｜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([
+        self._tokenizer.normalizer = normalizers.Sequence(
-            normalizers.NFC(),
+            [normalizers.NFC(), normalizers.Strip()]
-            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.UnicodeScripts(),
-        ])
+                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))]
+        reserved_tokens = [
-        detail_vocab_size = vocab_size - (len(reserved_tokens) + len(self._special_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,15 +105,26 @@ 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)
+            encoded_list: List[Encoding] = self._tokenizer.encode_batch(
-            return [encoded.ids if out_ids else encoded.tokens for encoded in encoded_list]
+                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:
+    def decode(self, tokens: List[int], skip_special_tokens: bool = True) -> str:
         return self._tokenizer.decode(tokens, skip_special_tokens=skip_special_tokens)
 
     def __len__(self) -> int:
@@ -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▁｜>")

pipeline/utils.py

@@ -30,7 +30,9 @@ def error_handler(
                 if reraise:
                     raise
                 return None
+
         return wrapper
+
     return decorator
 
 

scripts/cache_h5.py

@@ -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,
+    parser.add_argument(
-                        help="H5 output dir (default: <input_dir>/cached)")
+        "-o",
-    parser.add_argument("-t", "--tokenizer", default="./tokenizer.json",
+        "--output-dir",
-                        help="Tokenizer path (default: ./tokenizer.json)")
+        default=None,
-    parser.add_argument("-s", "--strategy", default=None,
+        help="H5 output dir (default: <input_dir>/cached)",
-                        help="Prompt strategy: chatml, alpaca (default: chatml)")
+    )
-    parser.add_argument("-p", "--pack-size", type=int, default=-1,
+    parser.add_argument(
-                        help="Pack size, <=0 to disable (default: -1)")
+        "-t",
-    parser.add_argument("--pad-value", type=int, default=1,
+        "--tokenizer",
-                        help="Padding value (default: 1)")
+        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")

scripts/pre_train/chinese-cosmopedia.py

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

scripts/supervised_finetuning/sft_chinese_instruct.py

@@ -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',
+        "stem_zh",
-        'coig-cqia', 'disc-law', 'neo_sft_phase2', 'chinese-medical', 'chinese-reasoning-distil',
+        "infinity-instruct",
-        'psycho-10k-dpsk-r1', 'sof-c-zh', 'industryinstruction', 'Chinese-QA-AFAF',
+        "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 = []

tests/test_cache.py

@@ -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,
+                files=[jsonl_path],
-                processor=processor, pack_size=-1, pad_value=0,
+                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,
+                files=[jsonl_path],
-                processor=processor, pack_size=10, pad_value=0,
+                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,
+                files=[jsonl_path],
-                processor=processor, pack_size=0, pad_value=-1,
+                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,
+                files=files,
-                processor=processor, pack_size=-1, pad_value=0,
+                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,
+                files=[jsonl_path],
-                processor=processor, pack_size=-1, pad_value=0,
+                output_dir=tmpdir,
+                processor=processor,
+                pack_size=-1,
+                pad_value=0,
             )
             assert len(output_files) == 1

tests/test_io.py

@@ -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(
-            assert torch.equal(loaded["labels"][0], torch.tensor([4, 5], dtype=torch.int32))
+                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)

tests/test_packing.py

@@ -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),
+                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),
+                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),
+                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),
+                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]

tests/test_processors.py

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

tests/test_strategies.py

@@ -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 "<｜im▁end｜>" in text
-        assert "<eos>" 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: