From f44ad6912e40b3e78498da67ad395a8b6a3b09a3 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Thu, 2 Apr 2026 16:16:02 +0800 Subject: [PATCH] =?UTF-8?q?fix:=20=E4=BF=AE=E5=A4=8D=E7=89=B9=E6=AE=8Atoke?= =?UTF-8?q?n=20=E9=97=AE=E9=A2=98?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- README.md | 10 +- docs/Design.md | 12 +- pipeline/__init__.py | 31 +++-- pipeline/io.py | 37 +++-- pipeline/packing.py | 4 +- pipeline/processors/__init__.py | 1 + pipeline/processors/base.py | 1 + pipeline/processors/dpo.py | 1 + pipeline/processors/factory.py | 1 + pipeline/processors/pretrain.py | 3 +- pipeline/processors/sft.py | 1 + pipeline/strategies/__init__.py | 1 + pipeline/strategies/alpaca.py | 5 +- pipeline/strategies/base.py | 3 +- pipeline/strategies/chatml.py | 11 +- pipeline/strategies/factory.py | 1 + pipeline/text.py | 21 ++- pipeline/tokenizer.py | 131 ++++++++++++------ pipeline/utils.py | 4 +- scripts/cache_h5.py | 39 ++++-- scripts/pre_train/chinese-cosmopedia.py | 2 +- .../sft_chinese_instruct.py | 22 ++- tests/__init__.py | 2 +- tests/test_cache.py | 36 +++-- tests/test_io.py | 27 ++-- tests/test_packing.py | 45 +++--- tests/test_processors.py | 28 +++- tests/test_strategies.py | 17 ++- 28 files changed, 334 insertions(+), 163 deletions(-) diff --git a/README.md b/README.md index 52c9ee4..114a206 100644 --- a/README.md +++ b/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 + "") +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="user\n", user_end="\n", assistant_start="assistant\n", - assistant_end="\n", + assistant_end="\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`, `` | -| Alpaca | `"alpaca"` | `### Instruction:`, `### Response:`, `` | +| ChatML | `"chatml"` | `<\|im_start\|>user`, `<\|im_end\|>`, `<\|im_start\|>assistant`, `<｜end▁of▁sentence｜>` | +| Alpaca | `"alpaca"` | `### Instruction:`, `### Response:`, `<｜end▁of▁sentence｜>` | 所有策略的 token 均可通过构造函数参数自定义，同时支持通过 `StrategyFactory.register()` 注册新格式。 diff --git a/docs/Design.md b/docs/Design.md index 0d7609a..6fb77b1 100644 --- a/docs/Design.md +++ b/docs/Design.md @@ -70,7 +70,7 @@ **PreTrainProcessor** (`"pt"`) ``` Input: {"text": "Hello world"} -Action: tokenizer.encode(text + "") +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 后缀（含 ``） +- `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` | +| 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` | +| response_suffix | `\n<｜end▁of▁sentence｜>` | **自定义示例**: ```python @@ -135,7 +135,7 @@ strategy = StrategyFactory.create("chatml", user_start="user\n", user_end="\n", assistant_start="assistant\n", - assistant_end="\n", + assistant_end="\n<｜end▁of▁sentence｜>", ) ``` @@ -178,7 +178,7 @@ StrategyFactory.register("my_format", MyStrategy) ### BpeTokenizer (`pipeline/tokenizer.py`) -基于 HuggingFace `tokenizers` 库的 BPE 分词器，支持从文件加载、训练、保存。内置 ``/``/`` 控制符和 `<|im_start|>`/`<|im_end|>` 特殊 token。 +基于 HuggingFace `tokenizers` 库的 BPE 分词器，支持从文件加载、训练、保存。内置 `<｜begin▁of▁sentence｜>`/`<｜end▁of▁sentence｜>`/`<｜▁pad▁｜>` 控制符和 `<｜im▁start｜>`/`<｜im▁end｜>` 特殊 token。 ## API 参考 diff --git a/pipeline/__init__.py b/pipeline/__init__.py index a5bc42a..76d7a02 100644 --- a/pipeline/__init__.py +++ b/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', - '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", ] diff --git a/pipeline/io.py b/pipeline/io.py index e5231f1..668ef28 100644 --- a/pipeline/io.py +++ b/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: diff --git a/pipeline/packing.py b/pipeline/packing.py index 67fa6fb..9542250 100644 --- a/pipeline/packing.py +++ b/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 diff --git a/pipeline/processors/__init__.py b/pipeline/processors/__init__.py index 0d3542e..0cfa442 100644 --- a/pipeline/processors/__init__.py +++ b/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 diff --git a/pipeline/processors/base.py b/pipeline/processors/base.py index 77afb8f..3c21f4e 100644 --- a/pipeline/processors/base.py +++ b/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 diff --git a/pipeline/processors/dpo.py b/pipeline/processors/dpo.py index 635169f..5f60a96 100644 --- a/pipeline/processors/dpo.py +++ b/pipeline/processors/dpo.py @@ -1,4 +1,5 @@ """DPO preference learning data processor.""" + from typing import Dict, List, Any, Optional import torch diff --git a/pipeline/processors/factory.py b/pipeline/processors/factory.py index e7e4116..da80eaf 100644 --- a/pipeline/processors/factory.py +++ b/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 diff --git a/pipeline/processors/pretrain.py b/pipeline/processors/pretrain.py index b0207be..2691cee 100644 --- a/pipeline/processors/pretrain.py +++ b/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}") + tokens = self.tokenizer.encode(f"{segment}<｜end▁of▁sentence｜>") return {"sequence": torch.tensor(tokens, dtype=torch.int32)} @property diff --git a/pipeline/processors/sft.py b/pipeline/processors/sft.py index 0ac71e2..a831ad9 100644 --- a/pipeline/processors/sft.py +++ b/pipeline/processors/sft.py @@ -1,4 +1,5 @@ """Supervised fine-tuning data processor.""" + from typing import Dict, List, Any, Optional import torch diff --git a/pipeline/strategies/__init__.py b/pipeline/strategies/__init__.py index ca481d7..0c31c0a 100644 --- a/pipeline/strategies/__init__.py +++ b/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 diff --git a/pipeline/strategies/alpaca.py b/pipeline/strategies/alpaca.py index 20e97ad..153de83 100644 --- a/pipeline/strategies/alpaca.py +++ b/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: ... ``""" + """Alpaca format:""" def __init__( self, tokenizer: BpeTokenizer, instruction_start: str = "### Instruction:\n", response_start: str = "### Response:\n", - response_suffix: str = "\n", + response_suffix: str = "\n<｜end▁of▁sentence｜>", ): super().__init__(tokenizer) self.instruction_start = instruction_start diff --git a/pipeline/strategies/base.py b/pipeline/strategies/base.py index 1f7e02e..f0753d8 100644 --- a/pipeline/strategies/base.py +++ b/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 diff --git a/pipeline/strategies/chatml.py b/pipeline/strategies/chatml.py index 38f5941..139c7c9 100644 --- a/pipeline/strategies/chatml.py +++ b/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|> ``""" + """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", + 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) diff --git a/pipeline/strategies/factory.py b/pipeline/strategies/factory.py index 7585e09..d06c855 100644 --- a/pipeline/strategies/factory.py +++ b/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 diff --git a/pipeline/text.py b/pipeline/text.py index ffdbdaf..150d64a 100644 --- a/pipeline/text.py +++ b/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) diff --git a/pipeline/tokenizer.py b/pipeline/tokenizer.py index e78f17a..66ac2dd 100644 --- a/pipeline/tokenizer.py +++ b/pipeline/tokenizer.py @@ -7,100 +7,143 @@ from typing import List, Union, Optional, Tuple, Iterator class BpeTokenizer: def __init__(self, path: Optional[str] = None): - self._control_tokens = ["", "", ""] - 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.pre_tokenizer = pre_tokenizers.Sequence([ - pre_tokenizers.UnicodeScripts(), - pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=True) - ]) - + self._tokenizer.normalizer = normalizers.Sequence( + [normalizers.NFC(), normalizers.Strip()] + ) + + self._tokenizer.pre_tokenizer = pre_tokenizers.Sequence( + [ + 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) - + 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) assert detail_vocab_size > min_size - + trainer = BpeTrainer( vocab_size=detail_vocab_size, 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) - - def train_from_iterator(self, iterator: Iterator[str], vocab_size: int, min_freq: int, reserved_token_size: int = 100) -> None: + 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: 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) - + 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: + 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: return self._tokenizer.get_vocab_size() - + @property def stop_ids(self) -> List[int]: stop_token = self._control_tokens + self._special_tokens stop_ids = [self._tokenizer.token_to_id(token) for token in stop_token] return stop_ids - + @property def bos_id(self) -> int: - return self._tokenizer.token_to_id("") - + return self._tokenizer.token_to_id("<｜begin▁of▁sentence｜>") + @property def eos_id(self) -> int: - return self._tokenizer.token_to_id("") - + return self._tokenizer.token_to_id("<｜end▁of▁sentence｜>") + @property def pad_id(self) -> int: - return self._tokenizer.token_to_id("") \ No newline at end of file + return self._tokenizer.token_to_id("<｜▁pad▁｜>") diff --git a/pipeline/utils.py b/pipeline/utils.py index 6283d1d..9faef75 100644 --- a/pipeline/utils.py +++ b/pipeline/utils.py @@ -30,7 +30,9 @@ def error_handler( if reraise: raise return None + return wrapper + return decorator @@ -58,4 +60,4 @@ def setup_logging(level: Optional[int] = None) -> None: root_logger.addHandler(console_handler) logging.getLogger("h5py").setLevel(logging.WARNING) - logging.getLogger("torch").setLevel(logging.WARNING) \ No newline at end of file + logging.getLogger("torch").setLevel(logging.WARNING) diff --git a/scripts/cache_h5.py b/scripts/cache_h5.py index 4e97a84..c17b51d 100644 --- a/scripts/cache_h5.py +++ b/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, - help="H5 output dir (default: /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: /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") diff --git a/scripts/pre_train/chinese-cosmopedia.py b/scripts/pre_train/chinese-cosmopedia.py index b27f815..8a1e769 100644 --- a/scripts/pre_train/chinese-cosmopedia.py +++ b/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"], diff --git a/scripts/supervised_finetuning/sft_chinese_instruct.py b/scripts/supervised_finetuning/sft_chinese_instruct.py index 5ecb015..15b243a 100644 --- a/scripts/supervised_finetuning/sft_chinese_instruct.py +++ b/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', - '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 = [] diff --git a/tests/__init__.py b/tests/__init__.py index b6b1342..e2d0b1d 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -1 +1 @@ -# Test suite for DataPipeline \ No newline at end of file +# Test suite for DataPipeline diff --git a/tests/test_cache.py b/tests/test_cache.py index d3b08be..c2c8585 100644 --- a/tests/test_cache.py +++ b/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, - 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 diff --git a/tests/test_io.py b/tests/test_io.py index c000b6d..5ae9cde 100644 --- a/tests/test_io.py +++ b/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(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) diff --git a/tests/test_packing.py b/tests/test_packing.py index a84c510..5390d2d 100644 --- a/tests/test_packing.py +++ b/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([ - torch.tensor([1, 2], dtype=torch.int32), - torch.tensor([3], dtype=torch.int32), - ]) + 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([ - torch.tensor([1, 2, 3, 4, 5], dtype=torch.int32), - torch.tensor([6, 7, 8, 9, 10], dtype=torch.int32), - ]) + 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([ - torch.tensor([1, 2, 3], dtype=torch.int32), - torch.tensor([4, 5, 6, 7, 8], dtype=torch.int32), - ]) + 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([ - torch.tensor([1], dtype=torch.int32), - torch.tensor([2, 3, 4, 5, 6, 7], dtype=torch.int32), - ]) + 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] diff --git a/tests/test_processors.py b/tests/test_processors.py index fc63a22..e9d6dab 100644 --- a/tests/test_processors.py +++ b/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 + ) diff --git a/tests/test_strategies.py b/tests/test_strategies.py index 99f117d..03dd280 100644 --- a/tests/test_strategies.py +++ b/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("") + 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 "" 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 "" in text + assert "<｜end▁of▁sentence｜>" in text class TestStrategyFactory: