From fd99b9d852a19f57d4f8951e6f0a424c0edf9694 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Thu, 4 Jun 2026 14:27:04 +0800 Subject: [PATCH 01/13] feat: add OpenHermes-2.5 SFT export script --- .../supervised_finetuning/sft_openhermes.py | 22 +++++++++++++++++++ 1 file changed, 22 insertions(+) create mode 100644 scripts/supervised_finetuning/sft_openhermes.py diff --git a/scripts/supervised_finetuning/sft_openhermes.py b/scripts/supervised_finetuning/sft_openhermes.py new file mode 100644 index 0000000..b84d79c --- /dev/null +++ b/scripts/supervised_finetuning/sft_openhermes.py @@ -0,0 +1,22 @@ +from datasets import load_dataset +from pipeline import export_dataset + + +def process_func(input_dict: dict): + conversations = input_dict["conversations"] + examples = [] + for i in range(0, len(conversations) - 1, 2): + user_msg = conversations[i]["value"] + assistant_msg = conversations[i + 1]["value"] + examples.append({"query": user_msg, "response": assistant_msg}) + return examples + + +if __name__ == "__main__": + dataset = load_dataset("teknium/OpenHermes-2.5") + export_dataset( + dataset=dataset["train"], + output_dir="./dataset", + output_prefix="OpenHermes-2.5", + process_func=process_func, + ) From 2f919e9243cb07d0449799bb0beea9e008d76928 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Thu, 2 Jul 2026 17:40:03 +0800 Subject: [PATCH 02/13] feat: add MetaMathQA SFT export script and fix ChatML newline encoding --- pipeline/strategies/chatml.py | 17 +++++++++-------- scripts/supervised_finetuning/sft_metamathqa.py | 17 +++++++++++++++++ 2 files changed, 26 insertions(+), 8 deletions(-) create mode 100644 scripts/supervised_finetuning/sft_metamathqa.py diff --git a/pipeline/strategies/chatml.py b/pipeline/strategies/chatml.py index cd8411c..b4f1493 100644 --- a/pipeline/strategies/chatml.py +++ b/pipeline/strategies/chatml.py @@ -14,17 +14,18 @@ class ChatMLStrategy(PromptStrategy): def __init__( self, tokenizer: AutoTokenizer, - 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", + user_end: str = "<｜im▁end｜>", + assistant_start: str = "<｜im▁start｜>assistant", + assistant_end: str = "<｜im▁end｜>", ): super().__init__(tokenizer) + nl_id = tokenizer.encode("a\nb", add_special_tokens=False)[1] - self._user_start_ids = self._encode_format(user_start) - self._user_end_ids = self._encode_format(user_end) - self._assistant_start_ids = self._encode_format(assistant_start) - self._assistant_end_ids = self._encode_format(assistant_end) + self._user_start_ids = self._encode_format(user_start) + [nl_id] + self._user_end_ids = self._encode_format(user_end) + [nl_id] + self._assistant_start_ids = self._encode_format(assistant_start) + [nl_id] + self._assistant_end_ids = self._encode_format(assistant_end) + [nl_id] @property def name(self) -> str: diff --git a/scripts/supervised_finetuning/sft_metamathqa.py b/scripts/supervised_finetuning/sft_metamathqa.py new file mode 100644 index 0000000..a26ff59 --- /dev/null +++ b/scripts/supervised_finetuning/sft_metamathqa.py @@ -0,0 +1,17 @@ +from datasets import load_dataset +from pipeline import export_dataset + + +def process_func(sample: dict) -> dict: + return {"query": sample["query"], "response": sample["response"]} + + +if __name__ == "__main__": + dataset = load_dataset("meta-math/MetaMathQA", split="train") + export_dataset( + dataset=dataset, + output_dir="./dataset", + output_prefix="MetaMathQA", + process_func=process_func, + chunk_size=1_000_000, + ) From 598e1ce4ae7e9d2395a67c5028c2002c3b082f33 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Fri, 3 Jul 2026 16:17:27 +0800 Subject: [PATCH 03/13] =?UTF-8?q?refactor:=20=E9=87=8D=E6=9E=84=E6=89=93?= =?UTF-8?q?=E5=8C=85=E6=A8=A1=E5=9D=97=EF=BC=8C=E6=96=B0=E5=A2=9E=20BFD/FF?= =?UTF-8?q?D/Greedy=20=E4=B8=89=E7=A7=8D=20bin-packing=20=E7=AE=97?= =?UTF-8?q?=E6=B3=95=EF=BC=8C=E9=BB=98=E8=AE=A4=20BFD?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 将 pipeline/packing.py 拆分为 packing/ 子包 (base/stream/binpack) - 新增 BfdPacker(默认)/FfDPacker/GreedyPacker，移除 StreamingPacker - 超长序列直接截断至 pack_size - group_size 语义改为"每 N 个 chunk 合并为一块"，默认 1000 - 新增 AutoTokenizer.token_to_id()，修复 ChatML 中 hacky 的 nl_id 获取 - pad_value 默认改为 2(pad_token_id)，position_ids pad=0, loss_mask pad=False - 新增 position_ids 打包后归零一致性测试 - scripts/cache_h5.py 新增 --pack-algo 参数 --- pipeline/__init__.py | 14 +- pipeline/io/export.py | 50 +++- pipeline/packing.py | 146 ----------- pipeline/packing/__init__.py | 84 +++++++ pipeline/packing/base.py | 49 ++++ pipeline/packing/binpack.py | 174 +++++++++++++ pipeline/strategies/chatml.py | 2 +- pipeline/tokenize/tokenizer.py | 6 + scripts/cache_h5.py | 25 +- tests/test_cache.py | 2 +- tests/test_io.py | 5 +- tests/test_packing.py | 447 ++++++++++++++++++++++----------- tests/test_processors.py | 3 + tests/test_strategies.py | 3 + 14 files changed, 715 insertions(+), 295 deletions(-) delete mode 100644 pipeline/packing.py create mode 100644 pipeline/packing/__init__.py create mode 100644 pipeline/packing/base.py create mode 100644 pipeline/packing/binpack.py diff --git a/pipeline/__init__.py b/pipeline/__init__.py index a1790ad..a0e9135 100644 --- a/pipeline/__init__.py +++ b/pipeline/__init__.py @@ -28,7 +28,13 @@ Usage:: from pipeline.pipeline import Pipeline, PipelineConfig, Stage, TransformStage from pipeline.tokenize import AutoTokenizer, ChatTemplate, train_bpe_tokenizer from pipeline.text import TextNormalizer -from pipeline.packing import SequencePacker +from pipeline.packing import ( + GreedyPacker, + FfDPacker, + BfdPacker, + BasePacker, + pack_tensors, +) # I/O module from pipeline.io import FileScanner, HDF5Handler, export_dataset, cache_jsonl @@ -70,7 +76,11 @@ __all__ = [ "train_bpe_tokenizer", # Text processing "TextNormalizer", - "SequencePacker", + "GreedyPacker", + "FfDPacker", + "BfdPacker", + "BasePacker", + "pack_tensors", # I/O "FileScanner", "HDF5Handler", diff --git a/pipeline/io/export.py b/pipeline/io/export.py index 2061f42..7576693 100644 --- a/pipeline/io/export.py +++ b/pipeline/io/export.py @@ -5,13 +5,16 @@ import logging import os from pathlib import Path from typing import Any, Callable, Dict, List, Optional, Union + +import torch from datasets import Dataset +from torch import Tensor from tqdm import tqdm from pipeline.io.file_scanner import FileScanner from pipeline.io.hdf5_handler import HDF5Handler from pipeline.processors import BaseProcessor -from pipeline.packing import pack_tensors +from pipeline.packing import pack_tensors, BasePacker from pipeline.utils import error_handler logger = logging.getLogger(__name__) @@ -75,6 +78,32 @@ def export_dataset( return output_files +def merge_tensors( + tensors: List[Tensor], + group_size: int, +) -> List[Tensor]: + """Merge a list of tensors into fewer larger tensors. + + Concatenates every group_size consecutive tensors into one merged + tensor. This reduces the number of shm blocks when loading. + + Args: + tensors: List of 1D tensors. + group_size: Number of tensors to merge into each group. + + Returns: + List of merged tensors. + """ + if not tensors: + return [] + + merged: List[Tensor] = [] + for i in range(0, len(tensors), group_size): + merged.append(torch.cat(tensors[i : i + group_size])) + + return merged + + @error_handler() def cache_jsonl( files: List[str], @@ -83,6 +112,8 @@ def cache_jsonl( *, pack_size: int = -1, pad_value: int = 0, + group_size: int = 1_000, + pack_algo: Optional[str] = None, ) -> List[str]: """Tokenize JSONL files and pack them into HDF5 storage. @@ -92,6 +123,10 @@ def cache_jsonl( processor: Initialized Processor instance. pack_size: Packing length, <=0 means no packing. pad_value: Padding value. + group_size: Merge every this many packed chunks into one tensor, + <=0 means no merging. + pack_algo: Packing algorithm: 'bfd' (default), 'ffd', + 'greedy'. Only used when pack_size > 0. Returns: List of generated H5 file paths. @@ -125,16 +160,27 @@ def cache_jsonl( ) continue + if not arrows[output_keys[0]]: + logger.warning(f"No valid samples in {file_path}, skipping") + continue + if pack_size > 0: dtypes = ( dict(processor.schema.output_fields) if processor.schema is not None else None ) - output = pack_tensors(arrows, pack_size, pad_value, dtypes) + pad_values = {k: (0 if k == "position_ids" else (False if k.endswith("_mask") else pad_value)) for k in output_keys} + output = pack_tensors(arrows, pack_size, pad_value, dtypes, pad_values=pad_values, algo=pack_algo) else: output = arrows + if group_size > 0 and output[output_keys[0]]: + output = { + key: merge_tensors(tensors, group_size) + for key, tensors in output.items() + } + h5_path = HDF5Handler.save(output_dir, file_name, output) output_files.append(h5_path) logger.info(f"Saved {h5_path}") diff --git a/pipeline/packing.py b/pipeline/packing.py deleted file mode 100644 index 42531ea..0000000 --- a/pipeline/packing.py +++ /dev/null @@ -1,146 +0,0 @@ -import logging -from typing import Any, Dict, List, Optional, Tuple, Union - -import torch -from torch import Tensor - -from pipeline.utils import error_handler - -logger = logging.getLogger(__name__) - - -class SequencePacker: - """ - Stream-concatenation packer for LLM training sequences. - - Algorithm (streaming concat): - - Input: sequences = [A(len=3), B(len=5), C(len=2)], pack_size = 6 - - 1. Validate & Normalize - - Check 1D dimension, unify dtype, warn on overlong sequences - - Result: [A, B, C] - - 2. Stream into buffer, slice off full chunks - - buffer += A(3) -> [a1 a2 a3], pos=3 - - buffer += B(5) -> [a1 a2 a3 b1 b2 b3 b4 b5], pos=8 - pos >= 6 -> flush [a1 a2 a3 b1 b2 b3], buffer=[b4 b5], pos=2 - - buffer += C(2) -> [b4 b5 c1 c2], pos=4 - loop ends -> flush tail [b4 b5 c1 c2 PAD PAD] - - Output: [[a1 a2 a3 b1 b2 b3], [b4 b5 c1 c2 PAD PAD]] - - Samples may be split across chunks — this is intentional and standard - practice in LLM training (TRL, Megatron-LM, etc.). - - Cross-group consistency: - Different tensor groups (e.g. input_ids, loss_masks) packed with - separate packer instances on samples with matching lengths produce - identical chunk boundaries. Element-level correspondence is preserved. - """ - - def __init__( - self, - pack_size: int, - pad_value: Union[int, bool] = 0, - dtype: Optional[torch.dtype] = None, - ): - self.pack_size = pack_size - self.pad_value = pad_value - self.dtype = dtype - self._buffer: List = [] - self._pos: int = 0 - self._packages: List[Tensor] = [] - - def reset(self) -> None: - """Reset packer state for instance reuse.""" - self._buffer = [] - self._pos = 0 - self._packages = [] - - @error_handler() - def pack(self, sequences: List[Tensor]) -> List[Tensor]: - """ - Pack sequences via streaming concatenation into fixed-size chunks. - - Sequences are concatenated in order and sliced at pack_size boundaries. - The final chunk is padded with pad_value. - - When dtype is not set at init, it is inferred from the first input tensor. - - Args: - sequences: List of 1D input tensors. - - Returns: - List of packed tensors, each with length equal to pack_size. - """ - if not sequences: - return [] - - # --- auto-infer dtype from first sequence --- - if self.dtype is None: - self.dtype = sequences[0].dtype - - # --- validate & normalize --- - normalized: List[Tensor] = [] - for i, seq in enumerate(sequences): - if seq.dim() != 1: - raise ValueError( - f"Expected 1D tensor at index {i}, got {seq.dim()}D tensor with shape {seq.shape}" - ) - if seq.dtype != self.dtype: - seq = seq.to(self.dtype) - normalized.append(seq) - - # --- stream into buffer, slice off full chunks --- - self._buffer = [] - self._packages = [] - pack_size = self.pack_size - buf = self._buffer - - for seq in normalized: - buf.extend(seq.tolist()) - while len(buf) >= pack_size: - self._packages.append(torch.tensor(buf[:pack_size], dtype=self.dtype)) - buf = buf[pack_size:] - - # flush tail with padding - if buf: - padded = buf + [self.pad_value] * (pack_size - len(buf)) - self._packages.append(torch.tensor(padded, dtype=self.dtype)) - - self._pos = len(buf) - return self._packages - - -def pack_tensors( - tensors: Dict[str, List[Tensor]], - pack_size: int, - pad_value: Union[int, bool] = 0, - dtypes: Optional[Dict[str, torch.dtype]] = None, -) -> Dict[str, List[Tensor]]: - """ - Pack multiple named tensor groups in parallel. - - Each group is packed independently with its own SequencePacker instance. - When dtypes is provided, packers use the declared dtype per key; - otherwise dtype is auto-inferred from the first tensor in each group. - - Args: - tensors: Dict mapping key names to lists of 1D tensors. - pack_size: Fixed chunk length. - pad_value: Padding value for non-bool tensors. - dtypes: Optional per-key dtype declarations. - - Returns: - Dict mapping key names to lists of packed tensors. - """ - if dtypes is None: - dtypes = {} - - output: Dict[str, List[Tensor]] = {} - for key, seqs in tensors.items(): - dtype = dtypes.get(key) - packer = SequencePacker(pack_size, pad_value, dtype=dtype) - output[key] = packer.pack(seqs) - return output diff --git a/pipeline/packing/__init__.py b/pipeline/packing/__init__.py new file mode 100644 index 0000000..b3a0a32 --- /dev/null +++ b/pipeline/packing/__init__.py @@ -0,0 +1,84 @@ +"""Sequence packing algorithms for LLM training data. + +Available packers: + - BfdPacker: Best-Fit Decreasing, samples never split (default) + - FfDPacker: First-Fit Decreasing, samples never split + - GreedyPacker: First-fit in input order, samples never split +""" + +from typing import Dict, List, Optional, Union + +import torch + +from pipeline.packing.base import BasePacker +from pipeline.packing.binpack import GreedyPacker, FfDPacker, BfdPacker + + +def pack_tensors( + tensors: Dict[str, List[torch.Tensor]], + pack_size: int, + pad_value: Union[int, bool] = 0, + dtypes: Optional[Dict[str, torch.dtype]] = None, + pad_values: Optional[Dict[str, Union[int, bool]]] = None, + algo: Optional[Union[str, BasePacker]] = None, +) -> Dict[str, List[torch.Tensor]]: + """Pack multiple named tensor groups in parallel. + + Each group is packed independently with its own packer instance. + + Args: + tensors: Dict mapping key names to lists of 1D tensors. + pack_size: Fixed chunk length. + pad_value: Default padding value, used for keys not in pad_values. + dtypes: Optional per-key dtype declarations. + pad_values: Optional per-key padding values (e.g. pad_token_id for + 'sequence', False for 'loss_mask', 0 for 'position_ids'). + algo: Packing algorithm to use. Can be 'bfd' (default), + 'ffd', 'greedy', or a BasePacker instance. + + Returns: + Dict mapping key names to lists of packed tensors. + """ + if dtypes is None: + dtypes = {} + if pad_values is None: + pad_values = {} + + output: Dict[str, List[torch.Tensor]] = {} + for key, seqs in tensors.items(): + key_pad = pad_values.get(key, pad_value) + actual_packer = _resolve_algo(algo, pack_size, key_pad) + dtype = dtypes.get(key) + if dtype is not None: + actual_packer.dtype = dtype + output[key] = actual_packer.pack(seqs) + return output + + +def _resolve_algo( + algo: Optional[Union[str, BasePacker]], + pack_size: int, + pad_value: Union[int, bool], +) -> BasePacker: + if algo is None or algo == "bfd": + return BfdPacker(pack_size, pad_value) + if isinstance(algo, BasePacker): + cls = type(algo) + return cls(pack_size, pad_value) + if algo == "ffd": + return FfDPacker(pack_size, pad_value) + if algo == "greedy": + return GreedyPacker(pack_size, pad_value) + raise ValueError( + f"Unknown packing algorithm: {algo}. " + f"Choose from: bfd, ffd, greedy" + ) + + +__all__ = [ + "BasePacker", + "BfdPacker", + "FfDPacker", + "GreedyPacker", + "pack_tensors", +] diff --git a/pipeline/packing/base.py b/pipeline/packing/base.py new file mode 100644 index 0000000..27ed628 --- /dev/null +++ b/pipeline/packing/base.py @@ -0,0 +1,49 @@ +from abc import ABC, abstractmethod +from typing import List, Optional, Union + +import torch +from torch import Tensor + + +class BasePacker(ABC): + """Abstract base class for sequence packing algorithms. + + All packers must implement pack() and reset(). + pack() takes a list of 1D tensors and returns a list of packed fixed-size tensors. + """ + + def __init__( + self, + pack_size: int, + pad_value: Union[int, bool] = 0, + dtype: Optional[torch.dtype] = None, + ): + self.pack_size = pack_size + self.pad_value = pad_value + self.dtype = dtype + + @abstractmethod + def pack(self, sequences: List[Tensor]) -> List[Tensor]: + """Pack sequences into fixed-size chunks.""" + ... + + @abstractmethod + def reset(self) -> None: + """Reset packer state for instance reuse.""" + ... + + def _validate_and_normalize(self, sequences: List[Tensor]) -> List[Tensor]: + """Validate 1D tensors and unify dtype.""" + if self.dtype is None and sequences: + self.dtype = sequences[0].dtype + + normalized: List[Tensor] = [] + for i, seq in enumerate(sequences): + if seq.dim() != 1: + raise ValueError( + f"Expected 1D tensor at index {i}, got {seq.dim()}D tensor with shape {seq.shape}" + ) + if seq.dtype != self.dtype: + seq = seq.to(self.dtype) + normalized.append(seq) + return normalized diff --git a/pipeline/packing/binpack.py b/pipeline/packing/binpack.py new file mode 100644 index 0000000..20b7d12 --- /dev/null +++ b/pipeline/packing/binpack.py @@ -0,0 +1,174 @@ +from typing import List, Optional, Union + +import torch +from torch import Tensor + +from pipeline.packing.base import BasePacker +from pipeline.utils import error_handler + + +def _truncate(tokens: List, max_len: int) -> List: + return tokens[:max_len] + + +def _pad_bin(bin_list: List, target_len: int, pad_value: Union[int, bool], dtype: torch.dtype) -> Tensor: + bin_list.extend([pad_value] * (target_len - len(bin_list))) + return torch.tensor(bin_list, dtype=dtype) + + +class GreedyPacker(BasePacker): + """Greedy first-fit packer (no sorting). + + Sequences are packed in input order into the first bin with enough space. + Overlong sequences (> pack_size) are truncated to pack_size. + """ + + def __init__( + self, + pack_size: int, + pad_value: Union[int, bool] = 0, + dtype: Optional[torch.dtype] = None, + ): + super().__init__(pack_size, pad_value, dtype) + self._bins: List[List] = [] + + def reset(self) -> None: + self._bins = [] + + @error_handler() + def pack(self, sequences: List[Tensor]) -> List[Tensor]: + if not sequences: + return [] + + normalized = self._validate_and_normalize(sequences) + self._bins = [] + pack_size = self.pack_size + pad_value = self.pad_value + + for seq in normalized: + seq_len = int(seq.shape[0]) + if seq_len > pack_size: + self._bins.append(_truncate(seq.tolist(), pack_size)) + continue + placed = False + for bin_list in self._bins: + if len(bin_list) + seq_len <= pack_size: + bin_list.extend(seq.tolist()) + placed = True + break + if not placed: + self._bins.append(list(seq.tolist())) + + packages: List[Tensor] = [] + for bin_list in self._bins: + packages.append(_pad_bin(bin_list, pack_size, pad_value, self.dtype)) + + return packages + + +class FfDPacker(BasePacker): + """First-Fit Decreasing (FFD) bin-packing packer. + + Sequences are sorted by descending length, then packed into the first + bin with enough space. Overlong sequences are truncated to pack_size. + """ + + def __init__( + self, + pack_size: int, + pad_value: Union[int, bool] = 0, + dtype: Optional[torch.dtype] = None, + ): + super().__init__(pack_size, pad_value, dtype) + self._bins: List[List] = [] + + def reset(self) -> None: + self._bins = [] + + @error_handler() + def pack(self, sequences: List[Tensor]) -> List[Tensor]: + if not sequences: + return [] + + normalized = self._validate_and_normalize(sequences) + self._bins = [] + pack_size = self.pack_size + pad_value = self.pad_value + + indexed = [(int(s.shape[0]), s) for s in normalized] + indexed.sort(key=lambda x: x[0], reverse=True) + + for seq_len, seq in indexed: + if seq_len > pack_size: + self._bins.append(_truncate(seq.tolist(), pack_size)) + continue + placed = False + for bin_list in self._bins: + if len(bin_list) + seq_len <= pack_size: + bin_list.extend(seq.tolist()) + placed = True + break + if not placed: + self._bins.append(list(seq.tolist())) + + packages: List[Tensor] = [] + for bin_list in self._bins: + packages.append(_pad_bin(bin_list, pack_size, pad_value, self.dtype)) + + return packages + + +class BfdPacker(BasePacker): + """Best-Fit Decreasing (BFD) bin-packing packer. + + Sequences are sorted by descending length, then packed into the bin + that minimizes remaining space (tightest fit). + Overlong sequences are truncated to pack_size. + """ + + def __init__( + self, + pack_size: int, + pad_value: Union[int, bool] = 0, + dtype: Optional[torch.dtype] = None, + ): + super().__init__(pack_size, pad_value, dtype) + self._bins: List[List] = [] + + def reset(self) -> None: + self._bins = [] + + @error_handler() + def pack(self, sequences: List[Tensor]) -> List[Tensor]: + if not sequences: + return [] + + normalized = self._validate_and_normalize(sequences) + self._bins = [] + pack_size = self.pack_size + pad_value = self.pad_value + + indexed = [(int(s.shape[0]), s) for s in normalized] + indexed.sort(key=lambda x: x[0], reverse=True) + + for seq_len, seq in indexed: + if seq_len > pack_size: + self._bins.append(_truncate(seq.tolist(), pack_size)) + continue + best_idx = -1 + best_remain = pack_size + 1 + for i, bin_list in enumerate(self._bins): + remain = pack_size - len(bin_list) + if seq_len <= remain < best_remain: + best_remain = remain + best_idx = i + if best_idx >= 0: + self._bins[best_idx].extend(seq.tolist()) + else: + self._bins.append(list(seq.tolist())) + + packages: List[Tensor] = [] + for bin_list in self._bins: + packages.append(_pad_bin(bin_list, pack_size, pad_value, self.dtype)) + + return packages diff --git a/pipeline/strategies/chatml.py b/pipeline/strategies/chatml.py index b4f1493..b1cbb3f 100644 --- a/pipeline/strategies/chatml.py +++ b/pipeline/strategies/chatml.py @@ -20,7 +20,7 @@ class ChatMLStrategy(PromptStrategy): assistant_end: str = "<｜im▁end｜>", ): super().__init__(tokenizer) - nl_id = tokenizer.encode("a\nb", add_special_tokens=False)[1] + nl_id = tokenizer.token_to_id("\n") self._user_start_ids = self._encode_format(user_start) + [nl_id] self._user_end_ids = self._encode_format(user_end) + [nl_id] diff --git a/pipeline/tokenize/tokenizer.py b/pipeline/tokenize/tokenizer.py index 83f0f63..060adbe 100644 --- a/pipeline/tokenize/tokenizer.py +++ b/pipeline/tokenize/tokenizer.py @@ -266,6 +266,12 @@ class AutoTokenizer: return self._tokenizer.decode(tokens, skip_special_tokens=skip_special_tokens) + def token_to_id(self, token: str) -> Optional[int]: + """Convert a token string to its integer ID.""" + if self._tokenizer is None: + raise RuntimeError("Tokenizer not initialized.") + return self._tokenizer.token_to_id(token) + def __len__(self) -> int: if self._tokenizer is None: return 0 diff --git a/scripts/cache_h5.py b/scripts/cache_h5.py index 8dc8899..47b7ef9 100644 --- a/scripts/cache_h5.py +++ b/scripts/cache_h5.py @@ -43,6 +43,13 @@ def main(): default=None, help="Prompt strategy: chatml, alpaca (default: chatml)", ) + parser.add_argument( + "-a", + "--pack-algo", + default=None, + choices=[None, "bfd", "ffd", "greedy"], + help="Packing algorithm: bfd (default), ffd, greedy", + ) parser.add_argument( "-p", "--pack-size", @@ -51,7 +58,14 @@ def main(): help="Pack size, <=0 to disable (default: -1)", ) parser.add_argument( - "--pad-value", type=int, default=0, help="Padding value (default: 0)" + "--pad-value", type=int, default=2, help="Padding token ID (default: 2 = <|pad|>)" + ) + parser.add_argument( + "-g", + "--group-size", + type=int, + default=1_000, + help="Merge every N packed chunks into one tensor, <=0 to disable (default: 1000)", ) parser.add_argument( "--log-level", @@ -95,9 +109,14 @@ def main(): print(f"\nStart caching...") if args.pack_size > 0: - print(f" pack_size={args.pack_size}, pad_value={args.pad_value}") + algo = args.pack_algo or "bfd" + print(f" pack_size={args.pack_size}, pad_value={args.pad_value}, algo={algo}") else: print(f" no packing") + if args.group_size > 0: + print(f" group_size={args.group_size} chunks per tensor") + else: + print(f" no grouping") cache_jsonl( files=jsonl_files, @@ -105,6 +124,8 @@ def main(): processor=processor, pack_size=args.pack_size, pad_value=args.pad_value, + group_size=args.group_size, + pack_algo=args.pack_algo, ) print(f"\nDone! Output saved to {output_dir}") diff --git a/tests/test_cache.py b/tests/test_cache.py index c2c8585..c11e23b 100644 --- a/tests/test_cache.py +++ b/tests/test_cache.py @@ -120,4 +120,4 @@ class TestCacheJsonl: pack_size=-1, pad_value=0, ) - assert len(output_files) == 1 + assert len(output_files) == 0 diff --git a/tests/test_io.py b/tests/test_io.py index a496722..b07099d 100644 --- a/tests/test_io.py +++ b/tests/test_io.py @@ -145,6 +145,9 @@ class DummyTokenizer: def encode(self, text: str, add_special_tokens: bool = False): return [ord(c) for c in text] + def token_to_id(self, token: str): + return ord(token) + def apply_chat_template( self, messages, add_generation_prompt=True, tokenize=True ): @@ -168,7 +171,7 @@ class TestPositionIds: processor = SFTProcessor(DummyTokenizer()) out_dir = os.path.join(tmpdir, "cached") - cache_jsonl([jsonl_path], out_dir, processor, pack_size=-1) + cache_jsonl([jsonl_path], out_dir, processor, pack_size=-1, group_size=0) h5_path = os.path.join(out_dir, "data.h5") loaded = HDF5Handler.load(h5_path, share_memory=False) diff --git a/tests/test_packing.py b/tests/test_packing.py index 5390d2d..89810fc 100644 --- a/tests/test_packing.py +++ b/tests/test_packing.py @@ -2,12 +2,147 @@ import pytest import torch -from pipeline.packing import SequencePacker +from pipeline.packing import ( + GreedyPacker, + FfDPacker, + BfdPacker, + pack_tensors, +) -class TestSequencePacker: - def test_normal_packing(self): - packer = SequencePacker(pack_size=10, pad_value=0) +class TestBfdPacker: + def test_best_fit_tight(self): + packer = BfdPacker(pack_size=10, pad_value=-1) + sequences = [ + torch.tensor([5, 6], dtype=torch.int32), + torch.tensor([1, 2, 3, 4], dtype=torch.int32), + torch.tensor([5, 6, 7, 8], dtype=torch.int32), + ] + packages = packer.pack(sequences) + assert len(packages) == 1 + assert packages[0].tolist() == [1, 2, 3, 4, 5, 6, 7, 8, 5, 6] + + def test_different_dtypes(self): + for dtype in [torch.int32, torch.int64, torch.float32]: + packer = BfdPacker(pack_size=10, dtype=dtype) + val = 1.0 if dtype == torch.float32 else 1 + packages = packer.pack([torch.tensor([val, 2, 3], dtype=dtype)]) + assert packages[0].dtype == dtype + + def test_dtype_conversion_on_mismatch(self): + packer = BfdPacker(pack_size=10, dtype=torch.int32) + packages = packer.pack([torch.tensor([1, 2, 3], dtype=torch.int64)]) + assert packages[0].dtype == torch.int32 + assert packages[0][:3].tolist() == [1, 2, 3] + + def test_non_1d_tensor_raises_error(self): + packer = BfdPacker(pack_size=10) + with pytest.raises(ValueError, match="Expected 1D tensor"): + packer.pack([torch.tensor([[1, 2], [3, 4]])]) + with pytest.raises(ValueError, match="Expected 1D tensor"): + packer.pack([torch.tensor(5)]) + + def test_empty_input(self): + packer = BfdPacker(pack_size=10) + assert packer.pack([]) == [] + + def test_reset(self): + packer = BfdPacker(pack_size=10) + packer.pack([torch.tensor([1, 2, 3], dtype=torch.int32)]) + assert len(packer._bins) == 1 + packer.reset() + assert len(packer._bins) == 0 + + def test_overlong_sample_truncated(self): + """Overlong sample is truncated to pack_size.""" + packer = BfdPacker(pack_size=6, pad_value=-1) + packages = packer.pack( + [ + torch.tensor([1, 2, 3, 4, 5, 6, 7], dtype=torch.int32), + torch.tensor([8, 9], dtype=torch.int32), + ] + ) + assert len(packages) == 2 + assert packages[0].tolist() == [1, 2, 3, 4, 5, 6] + assert packages[1].tolist() == [8, 9, -1, -1, -1, -1] + + def test_uses_two_bins_when_needed(self): + packer = BfdPacker(pack_size=10, pad_value=0) + sequences = [ + torch.tensor([1, 2, 3], dtype=torch.int32), + torch.tensor([4, 5, 6, 7], dtype=torch.int32), + torch.tensor([8, 9, 10], dtype=torch.int32), + torch.tensor([11, 12, 13, 14, 15, 16], dtype=torch.int32), + ] + packages = packer.pack(sequences) + assert len(packages) == 2 + for pkg in packages: + assert pkg.shape == (10,) + + def test_minimizes_waste_vs_ffd(self): + sequences = [ + torch.tensor([6] * i, dtype=torch.int32) + for i in [3, 5, 5, 7, 2, 4, 1, 4, 6, 2] + ] + bfd = BfdPacker(pack_size=10, pad_value=0) + ffd = FfDPacker(pack_size=10, pad_value=0) + assert len(bfd.pack(sequences)) <= len(ffd.pack(sequences)) + + +class TestFfDPacker: + def test_fills_tightly(self): + packer = FfDPacker(pack_size=10, pad_value=0) + sequences = [ + torch.tensor([7, 8], dtype=torch.int32), + torch.tensor([1, 2, 3, 4, 5, 6], dtype=torch.int32), + torch.tensor([9, 10], dtype=torch.int32), + ] + packages = packer.pack(sequences) + assert len(packages) == 1 + + def test_overlong_sample_truncated(self): + packer = FfDPacker(pack_size=5, pad_value=0) + packages = packer.pack( + [torch.tensor([1, 2, 3, 4, 5, 6], dtype=torch.int32)] + ) + assert len(packages) == 1 + assert packages[0].tolist() == [1, 2, 3, 4, 5] + + def test_sort_descending_order(self): + packer = FfDPacker(pack_size=10, pad_value=-1) + sequences = [ + torch.tensor([1, 2], dtype=torch.int32), + torch.tensor([3, 4, 5, 6, 7, 8], dtype=torch.int32), + torch.tensor([9, 10], dtype=torch.int32), + ] + packages = packer.pack(sequences) + assert len(packages) == 1 + assert packages[0].tolist() == [3, 4, 5, 6, 7, 8, 1, 2, 9, 10] + + def test_reduces_bins_vs_greedy(self): + sequences = [ + torch.tensor([6] * i, dtype=torch.int32) + for i in [3, 8, 2, 7, 1, 4, 5, 3, 2, 6] + ] + greedy = GreedyPacker(pack_size=10, pad_value=0) + ffd = FfDPacker(pack_size=10, pad_value=0) + assert len(ffd.pack(sequences)) <= len(greedy.pack(sequences)) + + def test_reset(self): + packer = FfDPacker(pack_size=10) + packer.pack([torch.tensor([1, 2, 3], dtype=torch.int32)]) + assert len(packer._bins) == 1 + packer.reset() + assert len(packer._bins) == 0 + + def test_empty_input(self): + packer = FfDPacker(pack_size=10) + assert packer.pack([]) == [] + + +class TestGreedyPacker: + def test_basic_packing(self): + packer = GreedyPacker(pack_size=10, pad_value=0) sequences = [ torch.tensor([1, 2, 3], dtype=torch.int32), torch.tensor([4, 5], dtype=torch.int32), @@ -15,154 +150,186 @@ class TestSequencePacker: ] packages = packer.pack(sequences) assert len(packages) == 1 - for pkg in packages: - assert pkg.shape == (10,) - - # Verify all original values are present in order + assert packages[0].shape == (10,) assert packages[0][:9].tolist() == [1, 2, 3, 4, 5, 6, 7, 8, 9] - assert packages[0][9] == 0 # padding + assert packages[0][9] == 0 - def test_empty_list_input(self): - packer = SequencePacker(pack_size=10) - assert packer.pack([]) == [] - - def test_single_sequence_input(self): - packer = SequencePacker(pack_size=10, pad_value=-1) - packages = packer.pack([torch.tensor([1, 2, 3], dtype=torch.int32)]) - assert len(packages) == 1 - assert packages[0][:3].tolist() == [1, 2, 3] - assert packages[0][3:].tolist() == [-1] * 7 - - 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) + def test_overlong_sample_truncated(self): + """Overlong sample is truncated to pack_size.""" + packer = GreedyPacker(pack_size=5, pad_value=0) packages = packer.pack( [torch.tensor([1, 2, 3, 4, 5, 6, 7, 8], dtype=torch.int32)] ) - assert len(packages) == 2 + assert len(packages) == 1 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), - ] - ) - assert packages[0][:3].tolist() == [1, 2, 3] - assert packages[0][3:].tolist() == [99] * 5 - - def test_different_dtypes(self): - for dtype in [torch.int32, torch.int64, torch.float32]: - packer = SequencePacker(pack_size=10, dtype=dtype) - val = 1.0 if dtype == torch.float32 else 1 - packages = packer.pack([torch.tensor([val, 2, 3], dtype=dtype)]) - assert packages[0].dtype == dtype - - def test_dtype_conversion_on_mismatch(self, caplog): - """Tensors with mismatched dtype are silently converted.""" - packer = SequencePacker(pack_size=10, dtype=torch.int32) - packages = packer.pack([torch.tensor([1, 2, 3], dtype=torch.int64)]) - assert packages[0].dtype == torch.int32 - assert packages[0][:3].tolist() == [1, 2, 3] - - def test_non_1d_tensor_raises_error(self): - packer = SequencePacker(pack_size=10) - with pytest.raises(ValueError, match="Expected 1D tensor"): - packer.pack([torch.tensor([[1, 2], [3, 4]])]) - with pytest.raises(ValueError, match="Expected 1D tensor"): - packer.pack([torch.tensor(5)]) - - def test_input_list_not_modified(self): - packer = SequencePacker(pack_size=10) - original = [ - torch.tensor([3], dtype=torch.int32), + def test_multiple_fill(self): + packer = GreedyPacker(pack_size=6, pad_value=0) + sequences = [ torch.tensor([1, 2], dtype=torch.int32), - torch.tensor([4, 5, 6, 7], dtype=torch.int32), + torch.tensor([3, 4], dtype=torch.int32), + torch.tensor([5, 6], dtype=torch.int32), + torch.tensor([7], dtype=torch.int32), ] - original_repr = [seq.tolist() for seq in original] - packer.pack(original) - assert [seq.tolist() for seq in original] == original_repr - - def test_exact_pack_size_fit(self): - 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), - ] - ) - assert len(packages) == 2 - assert packages[0].tolist() == [1, 2, 3, 4, 5] - assert packages[1].tolist() == [6, 7, 8, 9, 10] - - def test_multiple_packs_full_utilization(self): - packer = SequencePacker(pack_size=10, pad_value=-1) - sequences = [torch.tensor([i], dtype=torch.int32) for i in range(1, 12)] packages = packer.pack(sequences) assert len(packages) == 2 - assert packages[0].tolist() == list(range(1, 11)) - assert packages[1].tolist() == [11] + [-1] * 9 + for pkg in packages: + assert pkg.shape == (6,) - def test_cross_group_ordering(self): - """Separate packers for different dtypes produce identical chunk boundaries.""" - seq_packer = SequencePacker(pack_size=10, pad_value=0, dtype=torch.int32) - mask_packer = SequencePacker(pack_size=10, pad_value=False, dtype=torch.bool) - # sequences: lengths [3, 1, 4] - seqs = [ - torch.tensor([1, 2, 3], dtype=torch.int32), - torch.tensor([10], dtype=torch.int32), - torch.tensor([4, 5, 6, 7], dtype=torch.int32), - ] - masks = [ - torch.tensor([False, False, True], dtype=torch.bool), - torch.tensor([False], dtype=torch.bool), - torch.tensor([False, False, False, True], dtype=torch.bool), - ] - packed_seqs = seq_packer.pack(seqs) - packed_masks = mask_packer.pack(masks) - - # Verify mask packer uses bool dtype - assert packed_masks[0].dtype == torch.bool - # Both groups should produce the same number of packages - assert len(packed_seqs) == len(packed_masks) - - 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), - ] - ) - 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] - # Second chunk: [6, 7, 8, 0, 0] — rest of second + padding - assert packages[1].tolist() == [6, 7, 8, 0, 0] - - def test_reset_method(self): - packer = SequencePacker(pack_size=10, pad_value=0) - seqs = [torch.tensor([1, 2, 3], dtype=torch.int32)] - packer.pack(seqs) - assert len(packer._packages) == 1 + def test_reset(self): + packer = GreedyPacker(pack_size=10) + packer.pack([torch.tensor([1, 2, 3], dtype=torch.int32)]) + assert len(packer._bins) == 1 packer.reset() - assert len(packer._packages) == 0 - assert packer._pos == 0 - assert packer._buffer == [] + assert len(packer._bins) == 0 - def test_no_sorting_needed(self): - """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), - ] + def test_empty_input(self): + packer = GreedyPacker(pack_size=10) + assert packer.pack([]) == [] + + +class TestPackTensors: + def test_default_is_bfd(self): + result = pack_tensors( + tensors={ + "input_ids": [ + torch.tensor([1, 2], dtype=torch.int32), + torch.tensor([3, 4], dtype=torch.int32), + torch.tensor([5], dtype=torch.int32), + ], + }, + pack_size=5, + pad_value=0, ) - assert packages[0].tolist() == [1, 2, 3, 4] - assert packages[1].tolist() == [5, 6, 7, -1] + assert result["input_ids"][0].tolist() == [1, 2, 3, 4, 5] + + def test_greedy(self): + result = pack_tensors( + tensors={ + "input_ids": [ + torch.tensor([1, 2, 3], dtype=torch.int32), + torch.tensor([4, 5], dtype=torch.int32), + ], + }, + pack_size=5, + pad_value=0, + algo="greedy", + ) + assert result["input_ids"][0].tolist() == [1, 2, 3, 4, 5] + + def test_ffd(self): + result = pack_tensors( + tensors={ + "input_ids": [ + torch.tensor([1], dtype=torch.int32), + torch.tensor([2, 3, 4], dtype=torch.int32), + torch.tensor([5], dtype=torch.int32), + ], + }, + pack_size=5, + pad_value=0, + algo="ffd", + ) + assert result["input_ids"][0].tolist() == [2, 3, 4, 1, 5] + + def test_bfd_explicit(self): + result = pack_tensors( + tensors={ + "input_ids": [ + torch.tensor([1, 2], dtype=torch.int32), + torch.tensor([3, 4], dtype=torch.int32), + torch.tensor([5], dtype=torch.int32), + ], + }, + pack_size=5, + pad_value=0, + algo="bfd", + ) + assert result["input_ids"][0].tolist() == [1, 2, 3, 4, 5] + + def test_unknown_algo_raises(self): + with pytest.raises(ValueError, match="Unknown packing algorithm"): + pack_tensors( + tensors={"input_ids": [torch.tensor([1, 2, 3])]}, + pack_size=10, + pad_value=0, + algo="unknown_algo", + ) + + +class TestPositionIdsPacking: + """Verify position_ids reset to zero at sample boundaries after packing.""" + + def test_position_ids_reset_in_packed_chunk(self): + """After packing multiple SFT samples, position_ids restart from 0 at each boundary.""" + seqs = [ + torch.tensor([0, 1, 2, 3, 4], dtype=torch.int32), # len=5 + torch.tensor([0, 1, 2], dtype=torch.int32), # len=3 + torch.tensor([0, 1, 2, 3, 4, 5, 6], dtype=torch.int32), # len=7 + ] + result = pack_tensors( + tensors={"position_ids": seqs}, + pack_size=16, + pad_value=-1, + algo="greedy", + ) + packed = result["position_ids"][0].tolist() + assert packed == [0, 1, 2, 3, 4, 0, 1, 2, 0, 1, 2, 3, 4, 5, 6, -1] + + def test_position_ids_reset_with_bfd(self): + """BFD may reorder, but each sample's position_ids still start from 0.""" + seqs = [ + torch.tensor([0, 1, 2], dtype=torch.int32), + torch.tensor([0, 1, 2, 3, 4, 5], dtype=torch.int32), + torch.tensor([0, 1, 2, 3], dtype=torch.int32), + ] + result = pack_tensors( + tensors={"position_ids": seqs}, + pack_size=16, + pad_value=-1, + algo="bfd", + ) + packed = result["position_ids"][0].tolist() + assert packed[0] == 0 + zeros = [i for i, v in enumerate(packed) if v == 0 and (i == 0 or packed[i - 1] != 0)] + assert len(zeros) == 3 + + def test_multiple_keys_share_same_boundaries(self): + """sequence, loss_mask, position_ids share identical chunk boundaries after packing.""" + seq_a = torch.tensor([101, 102, 103, 104], dtype=torch.int32) + seq_b = torch.tensor([201, 202, 203, 204, 205, 206, 207], dtype=torch.int32) + seq_c = torch.tensor([301, 302, 303, 304, 305], dtype=torch.int32) + + mask_a = torch.tensor([False, False, True, True], dtype=torch.bool) + mask_b = torch.tensor([False, False, False, False, True, True, True], dtype=torch.bool) + mask_c = torch.tensor([False, False, False, True, True], dtype=torch.bool) + + pos_a = torch.tensor([0, 1, 2, 3], dtype=torch.int32) + pos_b = torch.tensor([0, 1, 2, 3, 4, 5, 6], dtype=torch.int32) + pos_c = torch.tensor([0, 1, 2, 3, 4], dtype=torch.int32) + + result = pack_tensors( + tensors={ + "sequence": [seq_a, seq_b, seq_c], + "loss_mask": [mask_a, mask_b, mask_c], + "position_ids": [pos_a, pos_b, pos_c], + }, + pack_size=16, + pad_value=-1, + algo="greedy", + ) + + seq_chunk = result["sequence"][0] + mask_chunk = result["loss_mask"][0] + pos_chunk = result["position_ids"][0] + + assert len(seq_chunk) == len(mask_chunk) == len(pos_chunk) == 16 + + for i in range(16): + if seq_chunk[i] == -1: + assert mask_chunk[i] == -1 + assert pos_chunk[i] == -1 + + pos_ids = pos_chunk.tolist() + zeros = [i for i, v in enumerate(pos_ids) if v == 0] + assert len(zeros) == 3 diff --git a/tests/test_processors.py b/tests/test_processors.py index e10ba3d..bec2400 100644 --- a/tests/test_processors.py +++ b/tests/test_processors.py @@ -18,6 +18,9 @@ class DummyTokenizer: def encode(self, text: str, add_special_tokens: bool = False): return [ord(c) for c in text] + def token_to_id(self, token: str): + return ord(token) + def apply_chat_template( self, messages, add_generation_prompt=True, tokenize=True ): diff --git a/tests/test_strategies.py b/tests/test_strategies.py index d4600b5..2447024 100644 --- a/tests/test_strategies.py +++ b/tests/test_strategies.py @@ -13,6 +13,9 @@ class DummyTokenizer: def encode(self, text: str, add_special_tokens: bool = False): return [ord(c) for c in text] + def token_to_id(self, token: str): + return ord(token) + class DummyStrategy(PromptStrategy): def __init__(self, tokenizer): From 06735b9cb32d4265013af41564cb04bd11736a74 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Fri, 3 Jul 2026 17:07:26 +0800 Subject: [PATCH 04/13] =?UTF-8?q?fix:=20=E4=BF=AE=E5=A4=8D=20nl=5Fid=20+?= =?UTF-8?q?=20BFD=20=E6=8C=89=20group=5Fsize=20=E5=88=86=E6=89=B9=E6=89=93?= =?UTF-8?q?=E5=8C=85=EF=BC=88=E5=AD=98=E7=9B=98=E4=B8=8D=E6=8B=86=E5=88=86?= =?UTF-8?q?=E6=96=87=E4=BB=B6=EF=BC=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- pipeline/io/export.py | 61 ++++++++++++++++++++++++----------- pipeline/strategies/chatml.py | 2 +- 2 files changed, 44 insertions(+), 19 deletions(-) diff --git a/pipeline/io/export.py b/pipeline/io/export.py index 7576693..c955c27 100644 --- a/pipeline/io/export.py +++ b/pipeline/io/export.py @@ -117,14 +117,17 @@ def cache_jsonl( ) -> List[str]: """Tokenize JSONL files and pack them into HDF5 storage. + BFD packs in group_size-bounded batches to avoid O(N²), then all + packed chunks are merged and saved as one HDF5 file per input file. + Args: files: List of JSONL file paths. output_dir: H5 output directory. processor: Initialized Processor instance. pack_size: Packing length, <=0 means no packing. pad_value: Padding value. - group_size: Merge every this many packed chunks into one tensor, - <=0 means no merging. + group_size: BFD batch granularity (token count threshold for each + packing batch) and merge granularity, <=0 means no merging. pack_algo: Packing algorithm: 'bfd' (default), 'ffd', 'greedy'. Only used when pack_size > 0. @@ -135,10 +138,21 @@ def cache_jsonl( output_files: List[str] = [] output_keys = processor.output_keys + dtypes = ( + dict(processor.schema.output_fields) + if processor.schema is not None + else None + ) + pad_values = {k: (0 if k == "position_ids" else (False if k.endswith("_mask") else pad_value)) for k in output_keys} + + target_tokens = group_size * pack_size if group_size > 0 and pack_size > 0 else 0 + for file_path in files: file_name = Path(file_path).stem - arrows: Dict[str, List] = {key: [] for key in output_keys} + all_packed: Dict[str, List[Tensor]] = {key: [] for key in output_keys} + arrows_batch: Dict[str, List] = {key: [] for key in output_keys} + batch_tokens: int = 0 with open(file_path, "r", encoding="utf-8") as f: for line_num, line in enumerate( @@ -148,7 +162,9 @@ def cache_jsonl( result = processor.process(json.loads(line)) if result is not None: for key in output_keys: - arrows[key].append(result[key]) + arrows_batch[key].append(result[key]) + if target_tokens > 0: + batch_tokens += int(result[output_keys[0]].shape[0]) except json.JSONDecodeError as e: logger.warning( f"JSON decode error in {file_path} line {line_num}: {e}. Skipping line." @@ -160,26 +176,35 @@ def cache_jsonl( ) continue - if not arrows[output_keys[0]]: + if target_tokens > 0 and batch_tokens >= target_tokens: + packed = pack_tensors(arrows_batch, pack_size, pad_value, dtypes, pad_values=pad_values, algo=pack_algo) + for key in output_keys: + all_packed[key].extend(packed[key]) + arrows_batch[key] = [] + batch_tokens = 0 + + if arrows_batch[output_keys[0]]: + if pack_size > 0: + packed = pack_tensors(arrows_batch, pack_size, pad_value, dtypes, pad_values=pad_values, algo=pack_algo) + for key in output_keys: + all_packed[key].extend(packed[key]) + else: + for key in output_keys: + all_packed[key].extend(arrows_batch[key]) + + if not all_packed[output_keys[0]]: logger.warning(f"No valid samples in {file_path}, skipping") continue - if pack_size > 0: - dtypes = ( - dict(processor.schema.output_fields) - if processor.schema is not None - else None - ) - pad_values = {k: (0 if k == "position_ids" else (False if k.endswith("_mask") else pad_value)) for k in output_keys} - output = pack_tensors(arrows, pack_size, pad_value, dtypes, pad_values=pad_values, algo=pack_algo) - else: - output = arrows - - if group_size > 0 and output[output_keys[0]]: + if pack_size <= 0: + output = all_packed + elif group_size > 0 and all_packed[output_keys[0]]: output = { key: merge_tensors(tensors, group_size) - for key, tensors in output.items() + for key, tensors in all_packed.items() } + else: + output = all_packed h5_path = HDF5Handler.save(output_dir, file_name, output) output_files.append(h5_path) diff --git a/pipeline/strategies/chatml.py b/pipeline/strategies/chatml.py index b1cbb3f..b4f1493 100644 --- a/pipeline/strategies/chatml.py +++ b/pipeline/strategies/chatml.py @@ -20,7 +20,7 @@ class ChatMLStrategy(PromptStrategy): assistant_end: str = "<｜im▁end｜>", ): super().__init__(tokenizer) - nl_id = tokenizer.token_to_id("\n") + nl_id = tokenizer.encode("a\nb", add_special_tokens=False)[1] self._user_start_ids = self._encode_format(user_start) + [nl_id] self._user_end_ids = self._encode_format(user_end) + [nl_id] From 816c02dab0d940c06bb4e620af59ad4286287429 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Sat, 4 Jul 2026 14:32:35 +0800 Subject: [PATCH 05/13] =?UTF-8?q?refactor:=20SFT=20=E7=BB=9F=E4=B8=80=20me?= =?UTF-8?q?ssages=20=E6=A0=BC=E5=BC=8F=20+=20ChatML=20=E7=BA=AF=20jinja=20?= =?UTF-8?q?=E6=B8=B2=E6=9F=93?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- pipeline/processors/sft.py | 60 +++++------- pipeline/strategies/chatml.py | 92 ++++++++++++++----- scripts/cache_h5.py | 4 +- .../supervised_finetuning/sft_alpaca_gpt4.py | 9 +- .../sft_firefly-1.1m-rephrased.py | 5 +- .../supervised_finetuning/sft_magicoder.py | 5 +- .../supervised_finetuning/sft_metamathqa.py | 5 +- .../supervised_finetuning/sft_openhermes.py | 25 ++++- tests/test_io.py | 18 ++-- tests/test_processors.py | 18 ++-- tests/test_strategies.py | 25 ++++- 11 files changed, 177 insertions(+), 89 deletions(-) diff --git a/pipeline/processors/sft.py b/pipeline/processors/sft.py index c11464c..acad467 100644 --- a/pipeline/processors/sft.py +++ b/pipeline/processors/sft.py @@ -7,7 +7,7 @@ from torch import Tensor from pipeline.tokenize import AutoTokenizer from pipeline.strategies import PromptStrategy, ChatMLStrategy -from pipeline.processors.base import BaseProcessor, ProcessorSchema, encode_with_mask +from pipeline.processors.base import BaseProcessor, ProcessorSchema from pipeline.processors.factory import ProcessorFactory @@ -15,20 +15,20 @@ from pipeline.processors.factory import ProcessorFactory class SFTProcessor(BaseProcessor): """Supervised fine-tuning data processor. - Supports two input formats: + Input formats: 1. messages (recommended): ``{"messages": [{"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]}`` - Multi-turn and system prompts are supported. - The tokenizer's ``apply_chat_template`` is used for rendering. + Multi-turn and system prompts are supported. Each assistant + turn gets ``loss_mask = 1``; all other roles get 0. 2. legacy query/response: ``{"query": "...", "response": "..."}`` - Falls back to the configured PromptStrategy (ChatML by default). + Internally converted to messages. Output schema: - - sequence: int32 tensor - Combined token IDs (prompt + response) - - loss_mask: bool tensor - True for response tokens (compute loss) - - position_ids: int32 tensor - Per-sample position IDs starting from 0 + - sequence: int32 tensor - Combined token IDs + - loss_mask: bool tensor - True for assistant response tokens + - position_ids: int32 tensor - Per-sample position IDs, start from 0 """ def __init__( @@ -58,7 +58,10 @@ class SFTProcessor(BaseProcessor): if "messages" in input_dict: return self._process_messages(input_dict["messages"]) if "query" in input_dict and "response" in input_dict: - return self._process_legacy(input_dict) + return self._process_messages([ + {"role": "user", "content": input_dict["query"]}, + {"role": "assistant", "content": input_dict["response"]}, + ]) raise KeyError( "Input must contain 'messages' or 'query'/'response' pair" ) @@ -69,38 +72,19 @@ class SFTProcessor(BaseProcessor): if messages[-1]["role"] != "assistant": raise ValueError("Last message must have role 'assistant'") - last_asst_idx = max( - i for i, m in enumerate(messages) if m["role"] == "assistant" - ) - - prompt_tokens = self.tokenizer.apply_chat_template( - messages[:last_asst_idx], - add_generation_prompt=True, - tokenize=True, - ) - - resp_content = messages[last_asst_idx]["content"] - im_end = getattr(self.tokenizer, "im_end", "<|im_end|>") - resp_tokens = self.tokenizer.encode( - f"{resp_content}{im_end}\n", add_special_tokens=False - ) - - tokens, loss_mask = encode_with_mask(prompt_tokens, resp_tokens) - position_ids = torch.arange(len(tokens), dtype=torch.int32) - return {"sequence": tokens, "loss_mask": loss_mask, "position_ids": position_ids} - - def _process_legacy(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]: strategy = self.strategy or ChatMLStrategy(self.tokenizer) - query_tokens = self.tokenizer.encode(input_dict["query"]) - response_tokens = self.tokenizer.encode(input_dict["response"]) + prompt, resp = strategy.format_messages(messages) - prompt = strategy.assemble_prompt(query_tokens) - response = strategy.assemble_response(response_tokens) - - tokens, loss_mask = encode_with_mask(prompt, response) - position_ids = torch.arange(len(tokens), dtype=torch.int32) - return {"sequence": tokens, "loss_mask": loss_mask, "position_ids": position_ids} + sequence = torch.tensor(prompt + resp, dtype=torch.int32) + loss_mask = torch.zeros(len(sequence), dtype=torch.bool) + loss_mask[len(prompt) :] = True + position_ids = torch.arange(len(sequence), dtype=torch.int32) + return { + "sequence": sequence, + "loss_mask": loss_mask, + "position_ids": position_ids, + } @property def output_keys(self) -> List[str]: diff --git a/pipeline/strategies/chatml.py b/pipeline/strategies/chatml.py index b4f1493..947e1f6 100644 --- a/pipeline/strategies/chatml.py +++ b/pipeline/strategies/chatml.py @@ -1,43 +1,91 @@ """ChatML format strategy.""" -from typing import List +from typing import Dict, List, Tuple from pipeline.tokenize import AutoTokenizer from pipeline.strategies.base import PromptStrategy from pipeline.strategies.factory import StrategyFactory +DEFAULT_CHATML_TEMPLATE = ( + "{% for message in messages %}" + "{% if message['role'] == 'system' %}" + "{{ '<|im_start|>system\n' + message['content'] + '<|im_end|>\n' }}" + "{% elif message['role'] == 'user' %}" + "{{ '<|im_start|>user\n' + message['content'] + '<|im_end|>\n' }}" + "{% elif message['role'] == 'assistant' %}" + "{{ '<|im_start|>assistant\n' + message['content'] + '<|im_end|>\n' }}" + "{% endif %}" + "{% endfor %}" + "{% if add_generation_prompt %}" + "{{ '<|im_start|>assistant\n' }}" + "{% endif %}" +) + @StrategyFactory.register("chatml") class ChatMLStrategy(PromptStrategy): - """ChatML format strategy.""" + """ChatML format strategy. - def __init__( - self, - tokenizer: AutoTokenizer, - user_start: str = "<｜im▁start｜>user", - user_end: str = "<｜im▁end｜>", - assistant_start: str = "<｜im▁start｜>assistant", - assistant_end: str = "<｜im▁end｜>", - ): + Renders messages using the tokenizer's jinja chat_template from + ``tokenizer_config.json``. Falls back to DEFAULT_CHATML_TEMPLATE + when no template is configured. + + The strategy does **not** hard-code any special tokens – all + formatting is driven by the jinja template. + """ + + def __init__(self, tokenizer: AutoTokenizer): super().__init__(tokenizer) - nl_id = tokenizer.encode("a\nb", add_special_tokens=False)[1] - - self._user_start_ids = self._encode_format(user_start) + [nl_id] - self._user_end_ids = self._encode_format(user_end) + [nl_id] - self._assistant_start_ids = self._encode_format(assistant_start) + [nl_id] - self._assistant_end_ids = self._encode_format(assistant_end) + [nl_id] + if tokenizer._chat_template is None: + tokenizer.set_chat_template(DEFAULT_CHATML_TEMPLATE) @property def name(self) -> str: return "chatml" + def format_messages( + self, + messages: List[Dict[str, str]], + ) -> Tuple[List[int], List[int]]: + """Render a single-turn messages conversation. + + Returns ``(prompt_tokens, response_tokens)`` where + *prompt_tokens* contains everything up to (and including) the + last assistant start marker, and *response_tokens* is the + assistant content plus the closing markers. + """ + last_asst = max( + i for i, m in enumerate(messages) if m["role"] == "assistant" + ) + + prompt = self.tokenizer.apply_chat_template( + messages[:last_asst], + add_generation_prompt=True, + tokenize=True, + ) + full = self.tokenizer.apply_chat_template( + messages[: last_asst + 1], + add_generation_prompt=False, + tokenize=True, + ) + return prompt, full[len(prompt) :] + def assemble_prompt(self, query_tokens: List[int]) -> List[int]: - return ( - self._user_start_ids - + query_tokens - + self._user_end_ids - + self._assistant_start_ids + text = self.tokenizer.decode(query_tokens) + return self.tokenizer.apply_chat_template( + [{"role": "user", "content": text}], + add_generation_prompt=True, + tokenize=True, ) def assemble_response(self, response_tokens: List[int]) -> List[int]: - return response_tokens + self._assistant_end_ids + text = self.tokenizer.decode(response_tokens) + full = self.tokenizer.apply_chat_template( + [{"role": "assistant", "content": text}], + add_generation_prompt=False, + tokenize=True, + ) + opening = self.tokenizer.apply_chat_template( + [], add_generation_prompt=True, tokenize=True + ) + return full[len(opening) :] diff --git a/scripts/cache_h5.py b/scripts/cache_h5.py index 47b7ef9..7b32bb1 100644 --- a/scripts/cache_h5.py +++ b/scripts/cache_h5.py @@ -34,8 +34,8 @@ def main(): parser.add_argument( "-t", "--tokenizer", - default="./tokenizer.json", - help="Tokenizer path (default: ./tokenizer.json)", + default="./tokenizer", + help="Tokenizer dir (default: ./tokenizer)", ) parser.add_argument( "-s", diff --git a/scripts/supervised_finetuning/sft_alpaca_gpt4.py b/scripts/supervised_finetuning/sft_alpaca_gpt4.py index a4ad21c..0ae2def 100644 --- a/scripts/supervised_finetuning/sft_alpaca_gpt4.py +++ b/scripts/supervised_finetuning/sft_alpaca_gpt4.py @@ -6,10 +6,13 @@ def process_func(input_dict: dict): instruction = input_dict["instruction"] inp = input_dict.get("input", "") if inp: - query = instruction + "\n" + inp + content = instruction + "\n" + inp else: - query = instruction - return {"query": query, "response": input_dict["output"]} + content = instruction + return {"messages": [ + {"role": "user", "content": content}, + {"role": "assistant", "content": input_dict["output"]}, + ]} if __name__ == "__main__": diff --git a/scripts/supervised_finetuning/sft_firefly-1.1m-rephrased.py b/scripts/supervised_finetuning/sft_firefly-1.1m-rephrased.py index 69b7c20..76d6e90 100644 --- a/scripts/supervised_finetuning/sft_firefly-1.1m-rephrased.py +++ b/scripts/supervised_finetuning/sft_firefly-1.1m-rephrased.py @@ -3,7 +3,10 @@ from pipeline import export_dataset def process_func(input_dict: dict): - return {"query": input_dict["instruction"], "response": input_dict["output"]} + return {"messages": [ + {"role": "user", "content": input_dict["instruction"]}, + {"role": "assistant", "content": input_dict["output"]}, + ]} if __name__ == "__main__": diff --git a/scripts/supervised_finetuning/sft_magicoder.py b/scripts/supervised_finetuning/sft_magicoder.py index be029e5..f6c59fe 100644 --- a/scripts/supervised_finetuning/sft_magicoder.py +++ b/scripts/supervised_finetuning/sft_magicoder.py @@ -3,7 +3,10 @@ from pipeline import export_dataset def process_func(input_dict: dict): - return {"query": input_dict["instruction"], "response": input_dict["response"]} + return {"messages": [ + {"role": "user", "content": input_dict["instruction"]}, + {"role": "assistant", "content": input_dict["response"]}, + ]} if __name__ == "__main__": diff --git a/scripts/supervised_finetuning/sft_metamathqa.py b/scripts/supervised_finetuning/sft_metamathqa.py index a26ff59..0e8aec1 100644 --- a/scripts/supervised_finetuning/sft_metamathqa.py +++ b/scripts/supervised_finetuning/sft_metamathqa.py @@ -3,7 +3,10 @@ from pipeline import export_dataset def process_func(sample: dict) -> dict: - return {"query": sample["query"], "response": sample["response"]} + return {"messages": [ + {"role": "user", "content": sample["query"]}, + {"role": "assistant", "content": sample["response"]}, + ]} if __name__ == "__main__": diff --git a/scripts/supervised_finetuning/sft_openhermes.py b/scripts/supervised_finetuning/sft_openhermes.py index b84d79c..c1f40ae 100644 --- a/scripts/supervised_finetuning/sft_openhermes.py +++ b/scripts/supervised_finetuning/sft_openhermes.py @@ -2,13 +2,30 @@ from datasets import load_dataset from pipeline import export_dataset +ROLE_MAP = {"system": "system", "human": "user", "gpt": "assistant"} + + def process_func(input_dict: dict): conversations = input_dict["conversations"] + + system_msgs = [] + idx = 0 + if conversations and conversations[0]["from"] == "system": + system_msgs.append({ + "role": "system", + "content": conversations[0]["value"], + }) + idx = 1 + examples = [] - for i in range(0, len(conversations) - 1, 2): - user_msg = conversations[i]["value"] - assistant_msg = conversations[i + 1]["value"] - examples.append({"query": user_msg, "response": assistant_msg}) + for i in range(idx, len(conversations) - 1, 2): + user_msg = conversations[i] + assistant_msg = conversations[i + 1] + messages = system_msgs + [ + {"role": ROLE_MAP[user_msg["from"]], "content": user_msg["value"]}, + {"role": ROLE_MAP[assistant_msg["from"]], "content": assistant_msg["value"]}, + ] + examples.append({"messages": messages}) return examples diff --git a/tests/test_io.py b/tests/test_io.py index b07099d..9368d23 100644 --- a/tests/test_io.py +++ b/tests/test_io.py @@ -140,22 +140,28 @@ class TestHDF5Handler: class DummyTokenizer: - im_end = "<|im_end|>" + def __init__(self): + self._special_token_map = {} + self._chat_template = None def encode(self, text: str, add_special_tokens: bool = False): return [ord(c) for c in text] + def decode(self, tokens, skip_special_tokens=True): + return "".join(chr(t) for t in tokens) + def token_to_id(self, token: str): return ord(token) - def apply_chat_template( - self, messages, add_generation_prompt=True, tokenize=True - ): + def set_chat_template(self, template): + self._chat_template = template + + def apply_chat_template(self, messages, add_generation_prompt=True, tokenize=True): text = "" for m in messages: - text += f"<|im_start|>{m['role']}\n{m['content']}<|im_end|>\n" + text += f"<｜im▁start｜>{m['role']}\n{m['content']}<｜im▁end｜>\n" if add_generation_prompt: - text += "<|im_start|>assistant\n" + text += "<｜im▁start｜>assistant\n" return self.encode(text) if tokenize else text diff --git a/tests/test_processors.py b/tests/test_processors.py index bec2400..41c1ff3 100644 --- a/tests/test_processors.py +++ b/tests/test_processors.py @@ -13,22 +13,28 @@ from pipeline.processors import ( class DummyTokenizer: - im_end = "<|im_end|>" + def __init__(self): + self._special_token_map = {} + self._chat_template = None def encode(self, text: str, add_special_tokens: bool = False): return [ord(c) for c in text] + def decode(self, tokens, skip_special_tokens=True): + return "".join(chr(t) for t in tokens) + def token_to_id(self, token: str): return ord(token) - def apply_chat_template( - self, messages, add_generation_prompt=True, tokenize=True - ): + def set_chat_template(self, template): + self._chat_template = template + + def apply_chat_template(self, messages, add_generation_prompt=True, tokenize=True): text = "" for m in messages: - text += f"<|im_start|>{m['role']}\n{m['content']}<|im_end|>\n" + text += f"<｜im▁start｜>{m['role']}\n{m['content']}<｜im▁end｜>\n" if add_generation_prompt: - text += "<|im_start|>assistant\n" + text += "<｜im▁start｜>assistant\n" return self.encode(text) if tokenize else text diff --git a/tests/test_strategies.py b/tests/test_strategies.py index 2447024..809a7ee 100644 --- a/tests/test_strategies.py +++ b/tests/test_strategies.py @@ -10,12 +10,30 @@ from pipeline.strategies import ( class DummyTokenizer: + def __init__(self): + self._special_token_map = {} + self._chat_template = None + def encode(self, text: str, add_special_tokens: bool = False): return [ord(c) for c in text] + def decode(self, tokens, skip_special_tokens=True): + return "".join(chr(t) for t in tokens) + def token_to_id(self, token: str): return ord(token) + def set_chat_template(self, template): + self._chat_template = template + + def apply_chat_template(self, messages, add_generation_prompt=True, tokenize=True): + text = "" + for m in messages: + text += f"<｜im▁start｜>{m['role']}\n{m['content']}<｜im▁end｜>\n" + if add_generation_prompt: + text += "<｜im▁start｜>assistant\n" + return self.encode(text) if tokenize else text + class DummyStrategy(PromptStrategy): def __init__(self, tokenizer): @@ -65,11 +83,8 @@ class TestChatMLStrategy: 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 - ) + assistant_start = tk.encode("<｜im▁start｜>assistant\n") + assert prompt[-len(assistant_start):] == assistant_start class TestAlpacaStrategy: From 900cd91798e1898360209943dba691722c178ade Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Sat, 4 Jul 2026 14:47:38 +0800 Subject: [PATCH 06/13] =?UTF-8?q?feat:=20MinHash+LSH=20=E5=8E=BB=E9=87=8D?= =?UTF-8?q?=20+=20Strategy/Factory=20=E5=AD=98=E5=82=A8=E5=90=8E=E7=AB=AF?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- pipeline/io/__init__.py | 12 ++- pipeline/io/dedup.py | 167 ++++++++++++++++++++++++++++++++++++++ pipeline/io/export.py | 18 ++-- pipeline/io/writers.py | 107 ++++++++++++++++++++++++ scripts/cache_h5.py | 17 ++-- scripts/dedup_pretrain.py | 36 ++++++++ 6 files changed, 344 insertions(+), 13 deletions(-) create mode 100644 pipeline/io/dedup.py create mode 100644 pipeline/io/writers.py create mode 100644 scripts/dedup_pretrain.py diff --git a/pipeline/io/__init__.py b/pipeline/io/__init__.py index f475f8e..9caabc4 100644 --- a/pipeline/io/__init__.py +++ b/pipeline/io/__init__.py @@ -4,16 +4,26 @@ This module provides: - FileScanner: File and directory scanning utilities - HDF5Handler: Tensor data persistence - export_dataset: HuggingFace Dataset to JSONL export -- cache_jsonl: JSONL to HDF5 tokenization and caching +- cache_jsonl: JSONL to HDF5/binary tokenization and caching +- dedup_jsonl: MinHash+LSH deduplication for pretraining text +- writers: BaseWriter / H5Writer / BinWriter / TextWriter (Strategy + Factory) """ from pipeline.io.file_scanner import FileScanner from pipeline.io.hdf5_handler import HDF5Handler from pipeline.io.export import export_dataset, cache_jsonl +from pipeline.io.dedup import dedup_jsonl +from pipeline.io.writers import BaseWriter, H5Writer, BinWriter, TextWriter, create_writer __all__ = [ "FileScanner", "HDF5Handler", "export_dataset", "cache_jsonl", + "dedup_jsonl", + "BaseWriter", + "H5Writer", + "BinWriter", + "TextWriter", + "create_writer", ] diff --git a/pipeline/io/dedup.py b/pipeline/io/dedup.py new file mode 100644 index 0000000..7e8f6a2 --- /dev/null +++ b/pipeline/io/dedup.py @@ -0,0 +1,167 @@ +"""MinHash + LSH deduplication for pretraining text data.""" + +import json +import logging +import os +from pathlib import Path +from typing import Iterator, List, Set, Tuple + +from tqdm import tqdm + +from pipeline.io.writers import TextWriter +from pipeline.utils import error_handler + +logger = logging.getLogger(__name__) + + +def _tokenize(text: str, ngram: int = 3) -> Set[str]: + return {text[i : i + ngram] for i in range(len(text) - ngram + 1)} + + +def _iter_docs(input_dir: Path) -> Iterator[Tuple[str, dict]]: + for fpath in sorted(input_dir.glob("*.jsonl")): + with open(fpath, encoding="utf-8") as f: + for line in f: + line = line.strip() + if not line: + continue + record = json.loads(line) + text = record.get("text", "") + if text: + yield text, record + + +def _write_h5(records: List[dict], output_dir: str, chunk_idx: int): + import h5py + + fname = os.path.join(output_dir, f"chunk_{chunk_idx}.h5") + texts = [rec.get("text", "") for rec in records] + with h5py.File(fname, "w") as f: + dt = h5py.special_dtype(vlen=str) + ds = f.create_dataset("text", (len(texts),), dtype=dt) + for i, t in enumerate(texts): + ds[i] = t + + +def _write_bin(records: List[dict], output_dir: Path, chunk_idx: int): + output_dir.mkdir(parents=True, exist_ok=True) + texts = [rec.get("text", "") + "\n" for rec in records] + + meta = {"chunk": chunk_idx, "count": len(texts), "format": "text", "encoding": "utf-8"} + meta_path = output_dir / "meta.json" + existing = json.loads(meta_path.read_text()) if meta_path.exists() else {} + existing[str(chunk_idx)] = meta + meta_path.write_text(json.dumps(existing, indent=2)) + + (output_dir / f"text_{chunk_idx}.bin").write_bytes("".join(texts).encode("utf-8")) + + +_WRITERS = { + "jsonl": TextWriter, + "h5": lambda: None, # handled inline below + "bin": lambda: None, +} + + +@error_handler() +def dedup_jsonl( + input_dir: str, + output_dir: str, + *, + threshold: float = 0.8, + num_perm: int = 128, + ngram: int = 3, + output_format: str = "jsonl", + chunk_size: int = 1_000_000, +) -> Tuple[int, int]: + """Deduplicate JSONL text files using MinHash + LSH. + + Args: + input_dir: Directory with source ``*.jsonl`` files. + output_dir: Directory for deduplicated output. + threshold: Jaccard similarity threshold (0–1). + num_perm: Number of MinHash permutations. + ngram: Character n-gram size. + output_format: ``"jsonl"``, ``"h5"``, or ``"bin"``. + chunk_size: Records per output chunk file. + + Returns: + ``(kept, removed)`` counts. + """ + from datasketch import MinHash, MinHashLSH + + input_path = Path(input_dir) + output_path = Path(output_dir) + output_path.mkdir(parents=True, exist_ok=True) + + logger.info( + f"Deduplicating {input_dir} -> {output_dir} " + f"(threshold={threshold}, perm={num_perm}, fmt={output_format})" + ) + + lsh = MinHashLSH(threshold=threshold, num_perm=num_perm) + + kept = 0 + removed = 0 + + dup_doc_ids: Set[int] = set() + for doc_id, (text, _record) in enumerate(tqdm(_iter_docs(input_path), desc="indexing", unit="docs")): + shingles = _tokenize(text, ngram=ngram) + if len(shingles) < ngram * 2: + dup_doc_ids.add(doc_id) + continue + + m = MinHash(num_perm=num_perm) + for s in shingles: + m.update(s.encode("utf-8")) + + if lsh.query(m): + dup_doc_ids.add(doc_id) + else: + lsh.insert(doc_id, m) + + logger.info(f"Found {len(dup_doc_ids)} duplicates, writing deduplicated data") + + buffer: List[dict] = [] + chunk_idx = 0 + writer = TextWriter(chunk_size) if output_format == "jsonl" else None + + for doc_id, (_text, record) in enumerate(tqdm(_iter_docs(input_path), desc="writing", unit="docs")): + if doc_id in dup_doc_ids: + removed += 1 + continue + + kept += 1 + buffer.append(record) + + if len(buffer) >= chunk_size: + _flush_chunk(buffer, output_path, chunk_idx, output_format, writer) + chunk_idx += 1 + buffer = [] + + if buffer: + _flush_chunk(buffer, output_path, chunk_idx, output_format, writer) + + if writer: + writer.flush(output_path) + + logger.info(f"Done. kept={kept}, removed={removed}") + return kept, removed + + +def _flush_chunk( + records: List[dict], + output_dir: Path, + chunk_idx: int, + output_format: str, + writer=None, +): + if output_format == "jsonl": + for rec in records: + writer.write_record(rec, output_dir) + elif output_format == "h5": + _write_h5(records, str(output_dir), chunk_idx) + elif output_format == "bin": + _write_bin(records, output_dir, chunk_idx) + else: + raise ValueError(f"Unknown output format: {output_format}") diff --git a/pipeline/io/export.py b/pipeline/io/export.py index c955c27..ca978e1 100644 --- a/pipeline/io/export.py +++ b/pipeline/io/export.py @@ -13,6 +13,7 @@ from tqdm import tqdm from pipeline.io.file_scanner import FileScanner from pipeline.io.hdf5_handler import HDF5Handler +from pipeline.io.writers import create_writer, BaseWriter from pipeline.processors import BaseProcessor from pipeline.packing import pack_tensors, BasePacker from pipeline.utils import error_handler @@ -114,15 +115,16 @@ def cache_jsonl( pad_value: int = 0, group_size: int = 1_000, pack_algo: Optional[str] = None, + output_format: str = "h5", ) -> List[str]: - """Tokenize JSONL files and pack them into HDF5 storage. + """Tokenize JSONL files and save as HDF5 or binary. BFD packs in group_size-bounded batches to avoid O(N²), then all - packed chunks are merged and saved as one HDF5 file per input file. + packed chunks are merged and saved as one file per input file. Args: files: List of JSONL file paths. - output_dir: H5 output directory. + output_dir: Output directory. processor: Initialized Processor instance. pack_size: Packing length, <=0 means no packing. pad_value: Padding value. @@ -130,9 +132,10 @@ def cache_jsonl( packing batch) and merge granularity, <=0 means no merging. pack_algo: Packing algorithm: 'bfd' (default), 'ffd', 'greedy'. Only used when pack_size > 0. + output_format: ``"h5"`` or ``"bin"``. Returns: - List of generated H5 file paths. + List of generated file paths. """ os.makedirs(output_dir, exist_ok=True) output_files: List[str] = [] @@ -206,8 +209,9 @@ def cache_jsonl( else: output = all_packed - h5_path = HDF5Handler.save(output_dir, file_name, output) - output_files.append(h5_path) - logger.info(f"Saved {h5_path}") + writer: BaseWriter = create_writer(output_format) + saved = writer.save(output_dir, file_name, output) + output_files.append(saved) + logger.info(f"Saved {saved}") return output_files diff --git a/pipeline/io/writers.py b/pipeline/io/writers.py new file mode 100644 index 0000000..6ccf6d4 --- /dev/null +++ b/pipeline/io/writers.py @@ -0,0 +1,107 @@ +"""Storage backends for tensor / text output (Strategy + Factory). + +Each backend implements a common ``save()`` interface so callers use +polymorphism instead of ``if fmt == "h5" ... elif fmt == "bin" ...``. + +Supports: + - **H5Writer**: HDF5 format (via HDF5Handler) + - **BinWriter**: binary format – meta.json + {key}.bin (memmap-compatible) + - **TextWriter**: raw JSONL text (for dedup output) +""" + +import json +import os +from abc import ABC, abstractmethod +from pathlib import Path +from typing import Dict, List + +import torch +from torch import Tensor + + +class BaseWriter(ABC): + """Abstract writer – call ``save(dir, name, data)`` without caring + about the underlying format.""" + + @abstractmethod + def save(self, output_dir: str, file_name: str, data: Dict[str, List[Tensor]]) -> str: + ... + + +class H5Writer(BaseWriter): + def save(self, output_dir: str, file_name: str, data: Dict[str, List[Tensor]]) -> str: + from pipeline.io.hdf5_handler import HDF5Handler + return HDF5Handler.save(output_dir, file_name, data) + + +class BinWriter(BaseWriter): + def save(self, output_dir: str, file_name: str, data: Dict[str, List[Tensor]]) -> str: + import numpy as np + + os.makedirs(output_dir, exist_ok=True) + sub_dir = os.path.join(output_dir, file_name) + os.makedirs(sub_dir, exist_ok=True) + + meta: Dict[str, Dict] = {} + for key, tensors in data.items(): + cat = torch.cat(tensors, dim=0) + meta[key] = {"shape": list(cat.shape), "dtype": str(cat.dtype).split(".")[-1]} + np.asarray(cat.cpu().numpy()).tofile(os.path.join(sub_dir, f"{key}.bin")) + + with open(os.path.join(sub_dir, "meta.json"), "w") as f: + json.dump(meta, f, indent=2) + + return sub_dir + + +class TextWriter(BaseWriter): + """Write raw text records as JSONL (used by dedup output).""" + + def __init__(self, chunk_size: int = 1_000_000): + self._chunk_size = chunk_size + self._buffer: List[dict] = [] + self._chunk_idx = 0 + + def save(self, output_dir: str, file_name: str, data: Dict[str, List[Tensor]]) -> str: + raise NotImplementedError("TextWriter.save_one is for tensor data; use write_record()") + + def write_record(self, record: dict, output_dir: Path): + self._buffer.append(record) + if len(self._buffer) >= self._chunk_size: + self._flush(output_dir) + + def flush(self, output_dir: Path): + if self._buffer: + self._flush(output_dir) + + def _flush(self, output_dir: Path): + output_dir.mkdir(parents=True, exist_ok=True) + fpath = output_dir / f"chunk_{self._chunk_idx}.jsonl" + with open(fpath, "w", encoding="utf-8") as f: + for rec in self._buffer: + f.write(json.dumps(rec, ensure_ascii=False) + "\n") + self._chunk_idx += 1 + self._buffer = [] + + +_WRITER_REGISTRY: Dict[str, type] = {} + + +def register_writer(name: str): + def decorator(cls): + _WRITER_REGISTRY[name] = cls + return cls + return decorator + + +def create_writer(name: str, **kwargs) -> BaseWriter: + cls = _WRITER_REGISTRY.get(name) + if cls is None: + raise ValueError(f"Unknown writer: {name}. Available: {list(_WRITER_REGISTRY)}") + return cls(**kwargs) + + +# Register built-in writers +register_writer("h5")(H5Writer) +register_writer("bin")(BinWriter) +register_writer("jsonl")(TextWriter) diff --git a/scripts/cache_h5.py b/scripts/cache_h5.py index 7b32bb1..fcae8c7 100644 --- a/scripts/cache_h5.py +++ b/scripts/cache_h5.py @@ -1,11 +1,10 @@ -"""JSONL to H5 caching script. +"""JSONL tokenization and caching script. -Tokenize JSONL files and pack them into HDF5 format. +Tokenize JSONL files and save as HDF5 or binary format. Usage: python scripts/cache_h5.py pt ./dataset/chinese-c4-pretrain - 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 + python scripts/cache_h5.py sft ./dataset/belle-sft --pack-size 4096 --output-format bin """ import argparse @@ -29,7 +28,7 @@ def main(): "-o", "--output-dir", default=None, - help="H5 output dir (default: /cached)", + help="Output dir (default: /cached)", ) parser.add_argument( "-t", @@ -73,6 +72,13 @@ def main(): choices=["DEBUG", "INFO", "WARNING", "ERROR"], help="Logging level (default: INFO)", ) + parser.add_argument( + "-f", + "--output-format", + default="h5", + choices=["h5", "bin"], + help="Output format: h5 or bin (default: h5)", + ) args = parser.parse_args() # Initialize logging explicitly (not automatic anymore) @@ -126,6 +132,7 @@ def main(): pad_value=args.pad_value, group_size=args.group_size, pack_algo=args.pack_algo, + output_format=args.output_format, ) print(f"\nDone! Output saved to {output_dir}") diff --git a/scripts/dedup_pretrain.py b/scripts/dedup_pretrain.py new file mode 100644 index 0000000..44cd318 --- /dev/null +++ b/scripts/dedup_pretrain.py @@ -0,0 +1,36 @@ +"""MinHash + LSH deduplication CLI. + +Usage: + python scripts/dedup_pretrain.py --input-dir --output-dir --threshold 0.8 --num-perm 128 --output-format jsonl +""" + +import argparse + +from pipeline.io import dedup_jsonl + + +def main(): + parser = argparse.ArgumentParser(description="MinHash + LSH deduplication") + parser.add_argument("--input-dir", required=True) + parser.add_argument("--output-dir", required=True) + parser.add_argument("--threshold", type=float, default=0.8) + parser.add_argument("--num-perm", type=int, default=128) + parser.add_argument("--ngram", type=int, default=3) + parser.add_argument("--output-format", default="jsonl", choices=["jsonl", "h5", "bin"]) + args = parser.parse_args() + + kept, removed = dedup_jsonl( + input_dir=args.input_dir, + output_dir=args.output_dir, + threshold=args.threshold, + num_perm=args.num_perm, + ngram=args.ngram, + output_format=args.output_format, + ) + + total = kept + removed + print(f"kept={kept}, removed={removed} ({removed/max(total,1)*100:.1f}%)") + + +if __name__ == "__main__": + main() From fec37545e2d211f7eb2293bb40677cade2de6957 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Mon, 6 Jul 2026 08:53:18 +0800 Subject: [PATCH 07/13] Add max_chunks=18 to dolma-30b export, remove max_chunks=5 from english-wiki, rename project to datapipline --- pyproject.toml | 2 +- scripts/pre_train/english-dolma.py | 11 +++++++++++ scripts/pre_train/english-wiki.py | 1 - 3 files changed, 12 insertions(+), 2 deletions(-) create mode 100644 scripts/pre_train/english-dolma.py diff --git a/pyproject.toml b/pyproject.toml index fdfa49a..e7a159e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,7 +3,7 @@ requires = ["setuptools>=64", "wheel"] build-backend = "setuptools.build_meta" [project] -name = "khaosz_dataset" +name = "datapipline" version = "0.1.0" description = "A dataset processing toolkit for language model training" readme = "README.md" diff --git a/scripts/pre_train/english-dolma.py b/scripts/pre_train/english-dolma.py new file mode 100644 index 0000000..07ef4c7 --- /dev/null +++ b/scripts/pre_train/english-dolma.py @@ -0,0 +1,11 @@ +from datasets import load_dataset +from pipeline import export_dataset + +if __name__ == "__main__": + dataset = load_dataset("emozilla/dolma-v1_7-30B") + export_dataset( + dataset=dataset["train"], + output_dir="./dataset", + output_prefix="english-dolma-30b-pretrain", + max_chunks=18, + ) diff --git a/scripts/pre_train/english-wiki.py b/scripts/pre_train/english-wiki.py index fd54ed4..fe4cca2 100644 --- a/scripts/pre_train/english-wiki.py +++ b/scripts/pre_train/english-wiki.py @@ -7,5 +7,4 @@ if __name__ == "__main__": dataset=dataset["train"], output_dir="./dataset", output_prefix="english-wiki-pretrain", - max_chunks=5, ) From e6787a2036e936dc1c09b909bfbb3fa121628151 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Fri, 17 Jul 2026 23:15:18 +0800 Subject: [PATCH 08/13] feat: add BELLE SFT export script, remove Firefly-1.1M-Rephrased --- .../{sft_firefly-1.1m-rephrased.py => sft_belle.py} | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) rename scripts/supervised_finetuning/{sft_firefly-1.1m-rephrased.py => sft_belle.py} (53%) diff --git a/scripts/supervised_finetuning/sft_firefly-1.1m-rephrased.py b/scripts/supervised_finetuning/sft_belle.py similarity index 53% rename from scripts/supervised_finetuning/sft_firefly-1.1m-rephrased.py rename to scripts/supervised_finetuning/sft_belle.py index 76d6e90..f92dac5 100644 --- a/scripts/supervised_finetuning/sft_firefly-1.1m-rephrased.py +++ b/scripts/supervised_finetuning/sft_belle.py @@ -3,17 +3,23 @@ from pipeline import export_dataset def process_func(input_dict: dict): + instruction = input_dict["instruction"] + inp = input_dict.get("input", "") + if inp: + content = instruction + "\n" + inp + else: + content = instruction return {"messages": [ - {"role": "user", "content": input_dict["instruction"]}, + {"role": "user", "content": content}, {"role": "assistant", "content": input_dict["output"]}, ]} if __name__ == "__main__": - dataset = load_dataset("Mxode/Firefly-1.1M-Rephrased") + dataset = load_dataset("BelleGroup/train_2M_CN") export_dataset( dataset=dataset["train"], output_dir="./dataset", - output_prefix="Firefly-1.1M-Rephrased", + output_prefix="belle-sft", process_func=process_func, ) From 33c8720d693680cf99aa14d01ce1dbdec676f9ba Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Sat, 25 Jul 2026 12:19:45 +0800 Subject: [PATCH 09/13] =?UTF-8?q?feat:=20=E5=B9=B6=E8=A1=8C=20batch=20toke?= =?UTF-8?q?nization=20+=20cache=5Fjsonl=20=E6=89=B9=E5=A4=84=E7=90=86?= =?UTF-8?q?=E6=94=AF=E6=8C=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - pipeline/tokenize/tokenizer.py: encode() 全部走 encode_batch(支持单条/批量) - pipeline/processors/base.py: BaseProcessor 新增 process_batch() - pipeline/processors/pretrain.py: PreTrainProcessor 覆盖 process_batch() 批量编码 - pipeline/io/export.py: cache_jsonl 新增 batch_size 参数默认 1000, 批量处理 - scripts/cache_h5.py: 新增 --batch-size 参数, 默认 tokenizer 路径改为 ../AstrAI/params --- pipeline/io/export.py | 58 ++++++++++++++++++++------------- pipeline/processors/base.py | 15 +++++++++ pipeline/processors/pretrain.py | 8 +++++ pipeline/tokenize/tokenizer.py | 28 +++++++--------- scripts/cache_h5.py | 7 ++++ 5 files changed, 78 insertions(+), 38 deletions(-) diff --git a/pipeline/io/export.py b/pipeline/io/export.py index ca978e1..d077099 100644 --- a/pipeline/io/export.py +++ b/pipeline/io/export.py @@ -116,6 +116,7 @@ def cache_jsonl( group_size: int = 1_000, pack_algo: Optional[str] = None, output_format: str = "h5", + batch_size: int = 1000, ) -> List[str]: """Tokenize JSONL files and save as HDF5 or binary. @@ -133,6 +134,8 @@ def cache_jsonl( pack_algo: Packing algorithm: 'bfd' (default), 'ffd', 'greedy'. Only used when pack_size > 0. output_format: ``"h5"`` or ``"bin"``. + batch_size: Number of lines to batch-process together for parallel + tokenization via encode_batch (default: 1000). Returns: List of generated file paths. @@ -157,34 +160,45 @@ def cache_jsonl( arrows_batch: Dict[str, List] = {key: [] for key in output_keys} batch_tokens: int = 0 + buf: List[str] = [] + buf_num: int = 0 + + def flush_buf(): + nonlocal batch_tokens + if not buf: + return + samples = [] + for line in buf: + try: + samples.append(json.loads(line)) + except json.JSONDecodeError as e: + logger.warning(f"JSON decode error, skipping: {e}") + buf.clear() + if not samples: + return + results = processor.process_batch(samples) if hasattr(processor, "process_batch") else [processor.process(s) for s in samples] + for result in results: + if result is not None: + for key in output_keys: + arrows_batch[key].append(result[key]) + if target_tokens > 0: + batch_tokens += int(result[output_keys[0]].shape[0]) + with open(file_path, "r", encoding="utf-8") as f: for line_num, line in enumerate( tqdm(f, desc=f"Processing {file_name}", leave=False), start=1 ): - try: - result = processor.process(json.loads(line)) - if result is not None: + buf.append(line) + if len(buf) >= batch_size: + flush_buf() + if target_tokens > 0 and batch_tokens >= target_tokens: + packed = pack_tensors(arrows_batch, pack_size, pad_value, dtypes, pad_values=pad_values, algo=pack_algo) for key in output_keys: - arrows_batch[key].append(result[key]) - if target_tokens > 0: - batch_tokens += int(result[output_keys[0]].shape[0]) - except json.JSONDecodeError as e: - logger.warning( - f"JSON decode error in {file_path} line {line_num}: {e}. Skipping line." - ) - continue - except Exception as e: - logger.warning( - f"Unexpected error processing line {line_num} in {file_path}: {e}. Skipping line." - ) - continue + all_packed[key].extend(packed[key]) + arrows_batch[key] = [] + batch_tokens = 0 - if target_tokens > 0 and batch_tokens >= target_tokens: - packed = pack_tensors(arrows_batch, pack_size, pad_value, dtypes, pad_values=pad_values, algo=pack_algo) - for key in output_keys: - all_packed[key].extend(packed[key]) - arrows_batch[key] = [] - batch_tokens = 0 + flush_buf() if arrows_batch[output_keys[0]]: if pack_size > 0: diff --git a/pipeline/processors/base.py b/pipeline/processors/base.py index fb0de2f..4e2e859 100644 --- a/pipeline/processors/base.py +++ b/pipeline/processors/base.py @@ -82,6 +82,21 @@ class BaseProcessor(ABC): """Return list of output tensor key names.""" pass + def process_batch(self, input_dicts: List[Dict[str, Any]]) -> List[Dict[str, Tensor]]: + """Process a batch of input samples. + + Default implementation calls process() for each sample. + Subclasses should override for efficient batch processing + (e.g., using tokenizer.encode_batch). + + Args: + input_dicts: List of input dictionaries. + + Returns: + List of output dictionaries mapping output key names to tensors. + """ + return [self.process(d) for d in input_dicts] + def validate_input(self, input_dict: Dict[str, Any]) -> None: """Validate input against schema before processing. diff --git a/pipeline/processors/pretrain.py b/pipeline/processors/pretrain.py index f3be6f6..4d3aae7 100644 --- a/pipeline/processors/pretrain.py +++ b/pipeline/processors/pretrain.py @@ -43,6 +43,14 @@ class PreTrainProcessor(BaseProcessor): tokens = self.tokenizer.encode(f"{segment}{self._eos_token}") return {"sequence": torch.tensor(tokens, dtype=torch.int32)} + def process_batch(self, input_dicts: List[Dict[str, Any]]) -> List[Dict[str, Tensor]]: + texts = [f"{d['text']}{self._eos_token}" for d in input_dicts] + batch_tokens = self.tokenizer.encode(texts) + return [ + {"sequence": torch.tensor(tokens, dtype=torch.int32)} + for tokens in batch_tokens + ] + @property def output_keys(self) -> List[str]: return ["sequence"] diff --git a/pipeline/tokenize/tokenizer.py b/pipeline/tokenize/tokenizer.py index 060adbe..50d91df 100644 --- a/pipeline/tokenize/tokenizer.py +++ b/pipeline/tokenize/tokenizer.py @@ -240,22 +240,18 @@ class AutoTokenizer: "Tokenizer not initialized. Load or create a tokenizer first." ) - if isinstance(tokens, str): - encoded = self._tokenizer.encode( - tokens, - is_pretokenized=is_pretokenized, - add_special_tokens=add_special_tokens, - ) - return encoded.ids if out_ids else encoded.tokens - else: - encoded_list = self._tokenizer.encode_batch( - tokens, - is_pretokenized=is_pretokenized, - add_special_tokens=add_special_tokens, - ) - return [ - encoded.ids if out_ids else encoded.tokens for encoded in encoded_list - ] + single = isinstance(tokens, str) + if single: + tokens = [tokens] + encoded_list = self._tokenizer.encode_batch( + tokens, + is_pretokenized=is_pretokenized, + add_special_tokens=add_special_tokens, + ) + result = [ + encoded.ids if out_ids else encoded.tokens for encoded in encoded_list + ] + return result[0] if single else result def decode(self, tokens: List[int], skip_special_tokens: bool = True) -> str: """Decode token IDs to text.""" diff --git a/scripts/cache_h5.py b/scripts/cache_h5.py index fcae8c7..a75db88 100644 --- a/scripts/cache_h5.py +++ b/scripts/cache_h5.py @@ -72,6 +72,12 @@ def main(): choices=["DEBUG", "INFO", "WARNING", "ERROR"], help="Logging level (default: INFO)", ) + parser.add_argument( + "--batch-size", + type=int, + default=1000, + help="Lines per batch for parallel tokenization via encode_batch (default: 1000)", + ) parser.add_argument( "-f", "--output-format", @@ -133,6 +139,7 @@ def main(): group_size=args.group_size, pack_algo=args.pack_algo, output_format=args.output_format, + batch_size=args.batch_size, ) print(f"\nDone! Output saved to {output_dir}") From e9996290989ae6f4a3453323ea691ad7cc69e34a Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Sun, 26 Jul 2026 10:57:30 +0800 Subject: [PATCH 10/13] =?UTF-8?q?=E7=A7=BB=E9=99=A4=E7=BC=93=E5=AD=98?= =?UTF-8?q?=E8=84=9A=E6=9C=AC=E4=B8=AD=E7=9A=84=20data=5Ffiles=20=E5=92=8C?= =?UTF-8?q?=20max=5Fchunks=20=E9=99=90=E5=88=B6=E5=8F=82=E6=95=B0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- scripts/pre_train/chinese-cosmopedia.py | 1 - scripts/pre_train/english-dolma.py | 1 - 2 files changed, 2 deletions(-) diff --git a/scripts/pre_train/chinese-cosmopedia.py b/scripts/pre_train/chinese-cosmopedia.py index 8a1e769..01f64b3 100644 --- a/scripts/pre_train/chinese-cosmopedia.py +++ b/scripts/pre_train/chinese-cosmopedia.py @@ -4,7 +4,6 @@ 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)]}, ) export_dataset( dataset=dataset["train"], diff --git a/scripts/pre_train/english-dolma.py b/scripts/pre_train/english-dolma.py index 07ef4c7..7e82b3c 100644 --- a/scripts/pre_train/english-dolma.py +++ b/scripts/pre_train/english-dolma.py @@ -7,5 +7,4 @@ if __name__ == "__main__": dataset=dataset["train"], output_dir="./dataset", output_prefix="english-dolma-30b-pretrain", - max_chunks=18, ) From 545104ba705b6d7ea5bf1c9bdb955a3508141338 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Wed, 29 Jul 2026 21:23:20 +0800 Subject: [PATCH 11/13] feat: add Alpaca-GPT4-ZH SFT export script --- .../sft_alpaca_gpt4_zh.py | 25 +++++++++++++++++++ 1 file changed, 25 insertions(+) create mode 100644 scripts/supervised_finetuning/sft_alpaca_gpt4_zh.py diff --git a/scripts/supervised_finetuning/sft_alpaca_gpt4_zh.py b/scripts/supervised_finetuning/sft_alpaca_gpt4_zh.py new file mode 100644 index 0000000..a04205b --- /dev/null +++ b/scripts/supervised_finetuning/sft_alpaca_gpt4_zh.py @@ -0,0 +1,25 @@ +from datasets import load_dataset +from pipeline import export_dataset + + +def process_func(input_dict: dict): + instruction = input_dict["instruction"] + inp = input_dict.get("input", "") + if inp: + content = instruction + "\n" + inp + else: + content = instruction + return {"messages": [ + {"role": "user", "content": content}, + {"role": "assistant", "content": input_dict["output"]}, + ]} + + +if __name__ == "__main__": + dataset = load_dataset("llm-wizard/alpaca-gpt4-data-zh") + export_dataset( + dataset=dataset["train"], + output_dir="./dataset", + output_prefix="alpaca-gpt4-data-zh", + process_func=process_func, + ) From 2b3bf442e907b88d630db4694c50d0f462ff2ae7 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Wed, 29 Jul 2026 22:01:09 +0800 Subject: [PATCH 12/13] feat: add SFT process_batch for parallel tokenization + short QA filter script --- pipeline/processors/sft.py | 58 ++++++++++++++++++++++++ scripts/filter_short.py | 92 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 150 insertions(+) create mode 100644 scripts/filter_short.py diff --git a/pipeline/processors/sft.py b/pipeline/processors/sft.py index acad467..3f70002 100644 --- a/pipeline/processors/sft.py +++ b/pipeline/processors/sft.py @@ -66,6 +66,16 @@ class SFTProcessor(BaseProcessor): "Input must contain 'messages' or 'query'/'response' pair" ) + def _extract_messages(self, input_dict: Dict[str, Any]) -> Optional[List[Dict[str, str]]]: + if "messages" in input_dict: + return input_dict["messages"] + if "query" in input_dict and "response" in input_dict: + return [ + {"role": "user", "content": input_dict["query"]}, + {"role": "assistant", "content": input_dict["response"]}, + ] + return None + def _process_messages(self, messages: List[Dict[str, str]]) -> Dict[str, Tensor]: if not messages: raise ValueError("Messages list is empty") @@ -86,6 +96,54 @@ class SFTProcessor(BaseProcessor): "position_ids": position_ids, } + def process_batch(self, input_dicts: List[Dict[str, Any]]) -> List[Optional[Dict[str, Tensor]]]: + strategy = self.strategy or ChatMLStrategy(self.tokenizer) + + prompts_text: List[str] = [] + fulls_text: List[str] = [] + indices: List[int] = [] + results: List[Optional[Dict[str, Tensor]]] = [None] * len(input_dicts) + + for i, d in enumerate(input_dicts): + try: + messages = self._extract_messages(d) + if not messages or messages[-1]["role"] != "assistant": + continue + last_asst = max(j for j, m in enumerate(messages) if m["role"] == "assistant") + prompt_text = self.tokenizer.apply_chat_template( + messages[:last_asst], add_generation_prompt=True, tokenize=False + ) + full_text = self.tokenizer.apply_chat_template( + messages[: last_asst + 1], add_generation_prompt=False, tokenize=False + ) + prompts_text.append(prompt_text) + fulls_text.append(full_text) + indices.append(i) + except Exception: + continue + + if not prompts_text: + return results + + prompt_tokens_list = self.tokenizer.encode(prompts_text) + full_tokens_list = self.tokenizer.encode(fulls_text) + + for j, idx in enumerate(indices): + prompt_tokens = prompt_tokens_list[j] + full_tokens = full_tokens_list[j] + resp_tokens = full_tokens[len(prompt_tokens):] + sequence = torch.tensor(prompt_tokens + resp_tokens, dtype=torch.int32) + loss_mask = torch.zeros(len(sequence), dtype=torch.bool) + loss_mask[len(prompt_tokens):] = True + position_ids = torch.arange(len(sequence), dtype=torch.int32) + results[idx] = { + "sequence": sequence, + "loss_mask": loss_mask, + "position_ids": position_ids, + } + + return results + @property def output_keys(self) -> List[str]: return ["sequence", "loss_mask", "position_ids"] diff --git a/scripts/filter_short.py b/scripts/filter_short.py new file mode 100644 index 0000000..88e70fc --- /dev/null +++ b/scripts/filter_short.py @@ -0,0 +1,92 @@ +import argparse +import json +import os +import tempfile +import shutil + +MIN_LEN = 15 + + +def filter_sft(input_path: str) -> tuple[int, int]: + """Filter SFT JSONL (messages format), remove if any msg content < MIN_LEN chars.""" + kept, total = 0, 0 + tmp_fd, tmp_path = tempfile.mkstemp(dir=os.path.dirname(input_path)) + try: + with open(input_path, encoding="utf-8") as fin, open(tmp_fd, "w", encoding="utf-8") as fout: + for line in fin: + line = line.strip() + if not line: + continue + total += 1 + try: + obj = json.loads(line) + except json.JSONDecodeError: + continue + messages = obj.get("messages", []) + short = any(len(m.get("content", "")) < MIN_LEN for m in messages) + if not short: + fout.write(line + "\n") + kept += 1 + shutil.move(tmp_path, input_path) + except Exception: + if os.path.exists(tmp_path): + os.unlink(tmp_path) + raise + return kept, total + + +def filter_pretrain(input_path: str) -> tuple[int, int]: + """Filter pretrain JSONL (text format), remove if text < MIN_LEN chars.""" + kept, total = 0, 0 + tmp_fd, tmp_path = tempfile.mkstemp(dir=os.path.dirname(input_path)) + try: + with open(input_path, encoding="utf-8") as fin, open(tmp_fd, "w", encoding="utf-8") as fout: + for line in fin: + line = line.strip() + if not line: + continue + total += 1 + try: + obj = json.loads(line) + except json.JSONDecodeError: + continue + text = obj.get("text", "") + if len(text) >= MIN_LEN: + fout.write(line + "\n") + kept += 1 + shutil.move(tmp_path, input_path) + except Exception: + if os.path.exists(tmp_path): + os.unlink(tmp_path) + raise + return kept, total + + +def main(): + parser = argparse.ArgumentParser(description="Filter short samples from JSONL datasets") + parser.add_argument("input_dir", help="Directory containing JSONL files") + parser.add_argument("--type", choices=["sft", "pt"], required=True, help="Dataset type") + args = parser.parse_args() + + from pipeline import FileScanner + + jsonl_files = FileScanner.scan(args.input_dir, suffix=".jsonl") + if not jsonl_files: + print(f"No JSONL files found in {args.input_dir}") + return + + filter_fn = filter_sft if args.type == "sft" else filter_pretrain + + total_kept, total_lines = 0, 0 + for fpath in jsonl_files: + kept, lines = filter_fn(fpath) + total_kept += kept + total_lines += lines + removed = lines - kept + print(f" {os.path.basename(fpath)}: {lines} -> {kept} (removed {removed})") + + print(f"\nTotal: {total_lines} -> {total_kept} (removed {total_lines - total_kept})") + + +if __name__ == "__main__": + main() From cc451e54927c65abc520291c606f6c163e2ecda8 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Fri, 31 Jul 2026 10:03:49 +0800 Subject: [PATCH 13/13] feat: add Ultra-FineWeb-L3 QA Synthetic pre-train export scripts (en/zh) --- .../ultra-fineweb-l3-en-qa-synthetic.py | 20 +++++++++++++++++++ .../ultra-fineweb-l3-zh-qa-synthetic.py | 20 +++++++++++++++++++ 2 files changed, 40 insertions(+) create mode 100644 scripts/pre_train/ultra-fineweb-l3-en-qa-synthetic.py create mode 100644 scripts/pre_train/ultra-fineweb-l3-zh-qa-synthetic.py diff --git a/scripts/pre_train/ultra-fineweb-l3-en-qa-synthetic.py b/scripts/pre_train/ultra-fineweb-l3-en-qa-synthetic.py new file mode 100644 index 0000000..eab46eb --- /dev/null +++ b/scripts/pre_train/ultra-fineweb-l3-en-qa-synthetic.py @@ -0,0 +1,20 @@ +from datasets import load_dataset +from pipeline import export_dataset + + +def process_func(input_dict: dict): + return {"text": input_dict["content"]} + + +if __name__ == "__main__": + dataset = load_dataset( + "openbmb/Ultra-FineWeb-L3", + "Ultra-FineWeb-L3-en-QA-Synthetic", + split="train", + ) + export_dataset( + dataset=dataset, + output_dir="./dataset", + output_prefix="ultra-fineweb-l3-en-qa-synthetic-pretrain", + process_func=process_func, + ) diff --git a/scripts/pre_train/ultra-fineweb-l3-zh-qa-synthetic.py b/scripts/pre_train/ultra-fineweb-l3-zh-qa-synthetic.py new file mode 100644 index 0000000..e7b005f --- /dev/null +++ b/scripts/pre_train/ultra-fineweb-l3-zh-qa-synthetic.py @@ -0,0 +1,20 @@ +from datasets import load_dataset +from pipeline import export_dataset + + +def process_func(input_dict: dict): + return {"text": input_dict["content"]} + + +if __name__ == "__main__": + dataset = load_dataset( + "openbmb/Ultra-FineWeb-L3", + "Ultra-FineWeb-L3-zh-QA-Synthetic", + split="train", + ) + export_dataset( + dataset=dataset, + output_dir="./dataset", + output_prefix="ultra-fineweb-l3-zh-qa-synthetic-pretrain", + process_func=process_func, + )