fix: cache_h5 移除重复 --batch-size 参数，默认输出格式改为 bin

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

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

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

pipeline/io/export.py

@@ -4,14 +4,18 @@ import json
 import logging
 import os
 from pathlib import Path
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, 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.io.writers import create_writer, BaseWriter
 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 +79,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,60 +113,150 @@ def cache_jsonl(
     *,
     pack_size: int = -1,
     pad_value: int = 0,
+    group_size: int = 1_000,
+    pack_algo: Optional[str] = None,
+    output_format: str = "h5",
+    batch_size: int = 1000,
 ) -> 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 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.
+        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.
+        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 H5 file paths.
+        List of generated file paths.
     """
     os.makedirs(output_dir, exist_ok=True)
     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
 
+        buf: List[Tuple[int, str]] = []
+
+        def flush_buf():
+            nonlocal batch_tokens
+            if not buf:
+                return
+            samples = []
+            for line_num, line in buf:
+                try:
+                    samples.append((line_num, json.loads(line)))
+                except json.JSONDecodeError as e:
+                    logger.warning(
+                        f"JSON decode error in {file_path} line {line_num}: "
+                        f"{e}. Skipping line."
+                    )
+            buf.clear()
+            if not samples:
+                return
+            items = [item for _, item in samples]
+            try:
+                results = (
+                    processor.process_batch(items)
+                    if hasattr(processor, "process_batch")
+                    else [processor.process(s) for s in items]
+                )
+                if len(results) != len(items):
+                    raise RuntimeError(
+                        "Batch processor returned a different number of results"
+                    )
+            except Exception:
+                results = []
+                for line_num, item in samples:
+                    try:
+                        results.append(processor.process(item))
+                    except Exception as e:
+                        logger.warning(
+                            f"Unexpected error processing line {line_num} "
+                            f"in {file_path}: {e}. Skipping line."
+                        )
+                        results.append(None)
+            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])
+
+        batch_size = max(1, batch_size)
         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:
+                buf.append((line_num, line))
-                    result = processor.process(json.loads(line))
+                if len(buf) >= batch_size:
-                    if result is not None:
+                    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[key].append(result[key])
+                            all_packed[key].extend(packed[key])
-                except json.JSONDecodeError as e:
+                            arrows_batch[key] = []
-                    logger.warning(
+                        batch_tokens = 0
-                        f"JSON decode error in {file_path} line {line_num}: {e}. Skipping line."
-                    )
-                    continue
-                except Exception as e:
-                    logger.warning(
-                        f"Unexpected error processing line {line_num} in {file_path}: {e}. Skipping line."
-                    )
-                    continue
 
-        if pack_size > 0:
+        flush_buf()
-            dtypes = (
+
-                dict(processor.schema.output_fields)
+        if arrows_batch[output_keys[0]]:
-                if processor.schema is not None
+            if pack_size > 0:
-                else None
+                packed = pack_tensors(arrows_batch, pack_size, pad_value, dtypes, pad_values=pad_values, algo=pack_algo)
-            )
+                for key in output_keys:
-            output = pack_tensors(arrows, pack_size, pad_value, dtypes)
+                    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:
+            output = all_packed
+        elif group_size > 0 and all_packed[output_keys[0]]:
+            output = {
+                key: merge_tensors(tensors, group_size)
+                for key, tensors in all_packed.items()
+            }
         else:
-            output = arrows
+            output = all_packed
 
-        h5_path = HDF5Handler.save(output_dir, file_name, output)
+        writer: BaseWriter = create_writer(output_format)
-        output_files.append(h5_path)
+        saved = writer.save(output_dir, file_name, output)
-        logger.info(f"Saved {h5_path}")
+        output_files.append(saved)
+        logger.info(f"Saved {saved}")
 
     return output_files

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)

pipeline/packing.py

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

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

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

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

pipeline/processors/base.py

@@ -76,12 +76,33 @@ class BaseProcessor(ABC):
         """
         pass
 
+    def process_batch(
+        self, input_dicts: List[Dict[str, Any]]
+    ) -> List[Dict[str, Tensor]]:
+        """Process a batch, falling back to the single-record implementation."""
+        return [self.process(input_dict) for input_dict in input_dicts]
+
     @property
     @abstractmethod
     def output_keys(self) -> List[str]:
         """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.
 

pipeline/processors/dpo.py

@@ -74,6 +74,33 @@ class DPOProcessor(BaseProcessor):
             "rejected_mask": rejected_m,
         }
 
+    def process_batch(self, input_dicts: List[Dict[str, Any]]) -> List[Dict[str, Tensor]]:
+        query_batch = self.tokenizer.encode([item["query"] for item in input_dicts])
+        chosen_batch = self.tokenizer.encode([item["chosen"] for item in input_dicts])
+        rejected_batch = self.tokenizer.encode(
+            [item["rejected"] for item in input_dicts]
+        )
+        results = []
+        for query_tokens, chosen_tokens, rejected_tokens in zip(
+            query_batch, chosen_batch, rejected_batch
+        ):
+            prompt = self.strategy.assemble_prompt(query_tokens)
+            chosen_t, chosen_m = encode_with_mask(
+                prompt, self.strategy.assemble_response(chosen_tokens)
+            )
+            rejected_t, rejected_m = encode_with_mask(
+                prompt, self.strategy.assemble_response(rejected_tokens)
+            )
+            results.append(
+                {
+                    "chosen": chosen_t,
+                    "chosen_mask": chosen_m,
+                    "rejected": rejected_t,
+                    "rejected_mask": rejected_m,
+                }
+            )
+        return results
+
     @property
     def output_keys(self) -> List[str]:
         return ["chosen", "chosen_mask", "rejected", "rejected_mask"]

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

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,29 +15,34 @@ 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.
+         Multi-turn and system prompts are supported.  Each assistant
-         The tokenizer's ``apply_chat_template`` is used for rendering.
+         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
+
+    Only the final assistant message is trained (mask_history behavior).
+    All earlier turns are context/prompt and masked from loss.
     """
 
     def __init__(
         self,
         tokenizer: AutoTokenizer,
         strategy: Optional[PromptStrategy] = None,
+        max_seq_len: Optional[int] = None,
     ):
         self.tokenizer = tokenizer
         self.strategy = strategy
+        self.max_seq_len = max_seq_len
 
     @property
     def schema(self) -> ProcessorSchema:
@@ -58,49 +63,97 @@ 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"
         )
 
+    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")
         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"])
+        prompt, resp = strategy.format_messages(messages)
-        response_tokens = self.tokenizer.encode(input_dict["response"])
 
-        prompt = strategy.assemble_prompt(query_tokens)
+        sequence = torch.tensor(prompt + resp, dtype=torch.int32)
-        response = strategy.assemble_response(response_tokens)
+        loss_mask = torch.zeros(len(sequence), dtype=torch.bool)
+        loss_mask[len(prompt) :] = True
+        if self.max_seq_len and len(sequence) > self.max_seq_len:
+            sequence = sequence[: self.max_seq_len]
+            loss_mask = loss_mask[: self.max_seq_len]
+        position_ids = torch.arange(len(sequence), dtype=torch.int32)
+        return {
+            "sequence": sequence,
+            "loss_mask": loss_mask,
+            "position_ids": position_ids,
+        }
 
-        tokens, loss_mask = encode_with_mask(prompt, response)
+    def process_batch(self, input_dicts: List[Dict[str, Any]]) -> List[Optional[Dict[str, Tensor]]]:
-        position_ids = torch.arange(len(tokens), dtype=torch.int32)
+        strategy = self.strategy or ChatMLStrategy(self.tokenizer)
-        return {"sequence": tokens, "loss_mask": loss_mask, "position_ids": position_ids}
+
+        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
+            if self.max_seq_len and len(sequence) > self.max_seq_len:
+                sequence = sequence[: self.max_seq_len]
+                loss_mask = loss_mask[: self.max_seq_len]
+            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]:

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__(
+    Renders messages using the tokenizer's jinja chat_template from
-        self,
+    ``tokenizer_config.json``.  Falls back to DEFAULT_CHATML_TEMPLATE
-        tokenizer: AutoTokenizer,
+    when no template is configured.
-        user_start: str = "<｜im▁start｜>user",
+
-        user_end: str = "<｜im▁end｜>",
+    The strategy does **not** hard-code any special tokens – all
-        assistant_start: str = "<｜im▁start｜>assistant",
+    formatting is driven by the jinja template.
-        assistant_end: str = "<｜im▁end｜>",
+    """
-    ):
+
+    def __init__(self, tokenizer: AutoTokenizer):
         super().__init__(tokenizer)
-        nl_id = tokenizer.encode("a\nb", add_special_tokens=False)[1]
+        if tokenizer._chat_template is None:
-
+            tokenizer.set_chat_template(DEFAULT_CHATML_TEMPLATE)
-        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:
         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 (
+        text = self.tokenizer.decode(query_tokens)
-            self._user_start_ids
+        return self.tokenizer.apply_chat_template(
-            + query_tokens
+            [{"role": "user", "content": text}],
-            + self._user_end_ids
+            add_generation_prompt=True,
-            + self._assistant_start_ids
+            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) :]

pipeline/tokenize/chat_template.py

@@ -3,6 +3,7 @@ Chat template module with Jinja2 rendering support.
 """
 
 from dataclasses import dataclass, field
+from functools import cached_property
 from typing import Any, Dict, List, Optional
 
 from jinja2 import Template
@@ -32,6 +33,10 @@ class ChatTemplate:
     default_variables: Dict[str, Any] = field(default_factory=dict)
     special_tokens: Dict[str, str] = field(default_factory=dict)
 
+    @cached_property
+    def _compiled(self) -> Template:
+        return Template(self.template_str)
+
     @classmethod
     def from_string(
         cls,
@@ -79,8 +84,7 @@ class ChatTemplate:
         if system_prompt is not None:
             variables["system_prompt"] = system_prompt
 
-        jinja_template = Template(self.template_str)
+        return self._compiled.render(**variables)
-        return jinja_template.render(**variables)
 
 
 # Default ChatML template

pipeline/tokenize/tokenizer.py

@@ -3,6 +3,7 @@ Tokenizer module with BPE implementation and auto-loading support.
 """
 
 from dataclasses import dataclass
+from functools import cached_property
 import json
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Union
@@ -102,6 +103,10 @@ class ChatTemplate:
         if self.special_tokens is None:
             self.special_tokens = {}
 
+    @cached_property
+    def _compiled(self) -> Template:
+        return Template(self.template_str)
+
     @classmethod
     def from_string(
         cls,
@@ -142,8 +147,7 @@ class ChatTemplate:
         if system_prompt is not None:
             variables["system_prompt"] = system_prompt
 
-        jinja_template = Template(self.template_str)
+        return self._compiled.render(**variables)
-        return jinja_template.render(**variables)
 
 
 
@@ -240,22 +244,18 @@ class AutoTokenizer:
                 "Tokenizer not initialized. Load or create a tokenizer first."
             )
 
-        if isinstance(tokens, str):
+        single = isinstance(tokens, str)
-            encoded = self._tokenizer.encode(
+        if single:
-                tokens,
+            tokens = [tokens]
-                is_pretokenized=is_pretokenized,
+        encoded_list = self._tokenizer.encode_batch(
-                add_special_tokens=add_special_tokens,
+            tokens,
-            )
+            is_pretokenized=is_pretokenized,
-            return encoded.ids if out_ids else encoded.tokens
+            add_special_tokens=add_special_tokens,
-        else:
+        )
-            encoded_list = self._tokenizer.encode_batch(
+        result = [
-                tokens,
+            encoded.ids if out_ids else encoded.tokens for encoded in encoded_list
-                is_pretokenized=is_pretokenized,
+        ]
-                add_special_tokens=add_special_tokens,
+        return result[0] if single else result
-            )
-            return [
-                encoded.ids if out_ids else encoded.tokens for encoded in encoded_list
-            ]
 
     def decode(self, tokens: List[int], skip_special_tokens: bool = True) -> str:
         """Decode token IDs to text."""
@@ -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
@@ -326,7 +332,9 @@ class AutoTokenizer:
             KeyError: If template name is not registered.
         """
         if isinstance(template, str):
-            self._chat_template = ChatTemplate.from_string(template)
+            self._chat_template = ChatTemplate.from_string(
+                template, special_tokens=self._special_token_map
+            )
         elif isinstance(template, ChatTemplate):
             self._chat_template = template
         else:

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"

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/belle-sft --pack-size 4096 --output-format bin
-    python scripts/cache_h5.py sft ./dataset/Ling-Coder-sft --tokenizer ./my_tokenizer.json
 """
 
 import argparse
@@ -29,13 +28,13 @@ def main():
         "-o",
         "--output-dir",
         default=None,
-        help="H5 output dir (default: <input_dir>/cached)",
+        help="Output dir (default: <input_dir>/cached)",
     )
     parser.add_argument(
         "-t",
         "--tokenizer",
-        default="./tokenizer.json",
+        default="./tokenizer",
-        help="Tokenizer path (default: ./tokenizer.json)",
+        help="Tokenizer dir (default: ./tokenizer)",
     )
     parser.add_argument(
         "-s",
@@ -43,6 +42,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 +57,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",
@@ -59,6 +72,19 @@ 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",
+        default="bin",
+        choices=["h5", "bin"],
+        help="Output format: h5 or bin (default: bin)",
+    )
     args = parser.parse_args()
 
     # Initialize logging explicitly (not automatic anymore)
@@ -95,9 +121,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 +136,10 @@ def main():
         processor=processor,
         pack_size=args.pack_size,
         pad_value=args.pad_value,
+        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}")
 

scripts/dedup_pretrain.py

@@ -0,0 +1,36 @@
+"""MinHash + LSH deduplication CLI.
+
+Usage:
+    python scripts/dedup_pretrain.py --input-dir <data_dir> --output-dir <out_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()

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

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

scripts/pre_train/english-dolma.py

@@ -0,0 +1,10 @@
+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",
+    )

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

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

scripts/pre_train/ultra-fineweb-l3-qa-synthetic-sample-shards.py

@@ -0,0 +1,67 @@
+import os
+import random
+
+from datasets import load_dataset
+from huggingface_hub import HfApi
+
+from pipeline import export_dataset
+
+REPO = "openbmb/Ultra-FineWeb-L3"
+FRACTION = 0.1
+SEED = 42
+SAVE_ARROW = False
+
+CONFIGS = {
+    "Ultra-FineWeb-L3-en-QA-Synthetic": "data/ultrafineweb_en_l3/qa/",
+    "Ultra-FineWeb-L3-zh-QA-Synthetic": "data/ultrafineweb_zh_l3/qa/",
+}
+
+HF_CACHE_DIR = "./cached_pt/ultra-fineweb-l3-qa-synthetic"
+OUTPUT_DIR = "./dataset"
+
+
+def process_func(input_dict: dict):
+    return {"text": input_dict["content"]}
+
+
+def main():
+    api = HfApi()
+    for config, prefix in CONFIGS.items():
+        lang = "en" if "-en-" in config else "zh"
+
+        shards = [
+            f.path
+            for f in api.list_repo_tree(
+                REPO, path_in_repo=prefix, recursive=True, repo_type="dataset"
+            )
+            if f.path.endswith(".parquet")
+        ]
+        k = max(1, int(len(shards) * FRACTION))
+        selected = random.Random(SEED).sample(shards, k)
+        print(f"[{config}] total shards={len(shards)}, selected={k}", flush=True)
+
+        dataset = load_dataset(
+            REPO,
+            data_files=selected,
+            split="train",
+            cache_dir=HF_CACHE_DIR,
+        )
+        print(f"[{config}] loaded {len(dataset)} rows", flush=True)
+
+        if SAVE_ARROW:
+            arrow_dir = os.path.join(
+                HF_CACHE_DIR, f"arrow-{lang}"
+            )
+            dataset.save_to_disk(arrow_dir)
+            print(f"[{config}] cached arrow to {arrow_dir}", flush=True)
+
+        export_dataset(
+            dataset=dataset,
+            output_dir=OUTPUT_DIR,
+            output_prefix=f"ultra-fineweb-l3-{lang}-qa-synthetic-10pct-pretrain",
+            process_func=process_func,
+        )
+
+
+if __name__ == "__main__":
+    main()

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

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
+        content = instruction
-    return {"query": query, "response": input_dict["output"]}
+    return {"messages": [
+        {"role": "user", "content": content},
+        {"role": "assistant", "content": input_dict["output"]},
+    ]}
 
 
 if __name__ == "__main__":

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

scripts/supervised_finetuning/sft_belle.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("BelleGroup/train_2M_CN")
+    export_dataset(
+        dataset=dataset["train"],
+        output_dir="./dataset",
+        output_prefix="belle-sft",
+        process_func=process_func,
+    )

scripts/supervised_finetuning/sft_firefly-1.1m-rephrased.py

@@ -1,16 +0,0 @@
-from datasets import load_dataset
-from pipeline import export_dataset
-
-
-def process_func(input_dict: dict):
-    return {"query": input_dict["instruction"], "response": input_dict["output"]}
-
-
-if __name__ == "__main__":
-    dataset = load_dataset("Mxode/Firefly-1.1M-Rephrased")
-    export_dataset(
-        dataset=dataset["train"],
-        output_dir="./dataset",
-        output_prefix="Firefly-1.1M-Rephrased",
-        process_func=process_func,
-    )

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__":

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__":

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):
+    for i in range(idx, len(conversations) - 1, 2):
-        user_msg = conversations[i]["value"]
+        user_msg = conversations[i]
-        assistant_msg = conversations[i + 1]["value"]
+        assistant_msg = conversations[i + 1]
-        examples.append({"query": user_msg, "response": assistant_msg})
+        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
 
 

tests/test_cache.py

@@ -29,6 +29,16 @@ class DummyProcessor(BaseProcessor):
         }
 
 
+class BatchTrackingProcessor(DummyProcessor):
+    def __init__(self):
+        super().__init__()
+        self.batch_sizes = []
+
+    def process_batch(self, items):
+        self.batch_sizes.append(len(items))
+        return super().process_batch(items)
+
+
 class TestCacheJsonl:
     def test_basic_cache_functionality(self):
         with tempfile.TemporaryDirectory() as tmpdir:
@@ -120,4 +130,20 @@ class TestCacheJsonl:
                 pack_size=-1,
                 pad_value=0,
             )
-            assert len(output_files) == 1
+            assert len(output_files) == 0
+
+    def test_uses_configured_batch_size(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            jsonl_path = os.path.join(tmpdir, "test.jsonl")
+            with open(jsonl_path, "w", encoding="utf-8") as f:
+                for i in range(5):
+                    f.write(json.dumps({"text": str(i)}) + "\n")
+
+            processor = BatchTrackingProcessor()
+            cache_jsonl(
+                files=[jsonl_path],
+                output_dir=tmpdir,
+                processor=processor,
+                batch_size=2,
+            )
+            assert processor.batch_sizes == [2, 2, 1]

tests/test_io.py

@@ -140,19 +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 apply_chat_template(
+    def decode(self, tokens, skip_special_tokens=True):
-        self, messages, add_generation_prompt=True, tokenize=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"
+            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
 
 
@@ -168,7 +177,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)

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:
+class TestBfdPacker:
-    def test_normal_packing(self):
+    def test_best_fit_tight(self):
-        packer = SequencePacker(pack_size=10, pad_value=0)
+        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 packages[0].shape == (10,)
-            assert pkg.shape == (10,)
-
-        # Verify all original values are present in order
         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):
+    def test_overlong_sample_truncated(self):
-        packer = SequencePacker(pack_size=10)
+        """Overlong sample is truncated to pack_size."""
-        assert packer.pack([]) == []
+        packer = GreedyPacker(pack_size=5, pad_value=0)
-
-    def test_single_sequence_input(self):
-        packer = SequencePacker(pack_size=10, pad_value=-1)
-        packages = packer.pack([torch.tensor([1, 2, 3], dtype=torch.int32)])
-        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)
         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):
+    def test_multiple_fill(self):
-        packer = SequencePacker(pack_size=8, pad_value=99)
+        packer = GreedyPacker(pack_size=6, pad_value=0)
-        packages = packer.pack(
+        sequences = [
-            [
-                torch.tensor([1, 2], dtype=torch.int32),
-                torch.tensor([3], dtype=torch.int32),
-            ]
-        )
-        assert packages[0][:3].tolist() == [1, 2, 3]
-        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),
             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))
+        for pkg in packages:
-        assert packages[1].tolist() == [11] + [-1] * 9
+            assert pkg.shape == (6,)
 
-    def test_cross_group_ordering(self):
+    def test_reset(self):
-        """Separate packers for different dtypes produce identical chunk boundaries."""
+        packer = GreedyPacker(pack_size=10)
-        seq_packer = SequencePacker(pack_size=10, pad_value=0, dtype=torch.int32)
+        packer.pack([torch.tensor([1, 2, 3], dtype=torch.int32)])
-        mask_packer = SequencePacker(pack_size=10, pad_value=False, dtype=torch.bool)
+        assert len(packer._bins) == 1
-        # 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
         packer.reset()
-        assert len(packer._packages) == 0
+        assert len(packer._bins) == 0
-        assert packer._pos == 0
-        assert packer._buffer == []
 
-    def test_no_sorting_needed(self):
+    def test_empty_input(self):
-        """Streaming concat preserves input order, no sorting."""
+        packer = GreedyPacker(pack_size=10)
-        packer = SequencePacker(pack_size=4, pad_value=-1)
+        assert packer.pack([]) == []
-        # short then long (fits in 2 chunks)
+
-        packages = packer.pack(
+
-            [
+class TestPackTensors:
-                torch.tensor([1], dtype=torch.int32),
+    def test_default_is_bfd(self):
-                torch.tensor([2, 3, 4, 5, 6, 7], dtype=torch.int32),
+        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 result["input_ids"][0].tolist() == [1, 2, 3, 4, 5]
-        assert packages[1].tolist() == [5, 6, 7, -1]
+
+    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

tests/test_processors.py

@@ -13,19 +13,30 @@ 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):
+    def encode(self, text, add_special_tokens: bool = False):
+        if isinstance(text, list):
+            return [[ord(c) for c in item] for item in text]
         return [ord(c) for c in text]
 
-    def apply_chat_template(
+    def decode(self, tokens, skip_special_tokens=True):
-        self, messages, add_generation_prompt=True, tokenize=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"
+            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
 
 
@@ -50,6 +61,16 @@ class TestPreTrainProcessor:
         result = PreTrainProcessor(DummyTokenizer()).process({"text": "a"})
         assert len(result["sequence"]) > 0
 
+    def test_process_batch_matches_single(self):
+        processor = PreTrainProcessor(DummyTokenizer())
+        items = [{"text": "hello"}, {"text": "world"}]
+        batch = processor.process_batch(items)
+        single = [processor.process(item) for item in items]
+        assert all(
+            torch.equal(batch_item["sequence"], single_item["sequence"])
+            for batch_item, single_item in zip(batch, single)
+        )
+
 
 class TestSFTProcessor:
     def test_output_keys(self):
@@ -121,6 +142,23 @@ class TestSFTProcessor:
         })
         assert "sequence" in result
 
+    def test_process_batch_matches_single(self):
+        processor = SFTProcessor(DummyTokenizer())
+        items = [
+            {
+                "messages": [
+                    {"role": "user", "content": "q1"},
+                    {"role": "assistant", "content": "a1"},
+                ]
+            },
+            {"query": "q2", "response": "a2"},
+        ]
+        batch = processor.process_batch(items)
+        single = [processor.process(item) for item in items]
+        for batch_item, single_item in zip(batch, single):
+            for key in processor.output_keys:
+                assert torch.equal(batch_item[key], single_item[key])
+
     def test_messages_empty_raises(self):
         with pytest.raises(ValueError, match="Messages list is empty"):
             SFTProcessor(DummyTokenizer()).process({"messages": []})
@@ -173,6 +211,18 @@ class TestDPOProcessor:
         assert result["chosen_mask"].dtype == torch.bool
         assert result["rejected_mask"].dtype == torch.bool
 
+    def test_process_batch_matches_single(self):
+        processor = DPOProcessor(DummyTokenizer())
+        items = [
+            {"query": "q1", "chosen": "yes", "rejected": "no"},
+            {"query": "q2", "chosen": "good", "rejected": "bad"},
+        ]
+        batch = processor.process_batch(items)
+        single = [processor.process(item) for item in items]
+        for batch_item, single_item in zip(batch, single):
+            for key in processor.output_keys:
+                assert torch.equal(batch_item[key], single_item[key])
+
 
 class TestProcessorFactory:
     def test_create_pre_train_processor(self):

tests/test_strategies.py

@@ -10,9 +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):
@@ -62,11 +83,8 @@ class TestChatMLStrategy:
         tk = DummyTokenizer()
         strategy = ChatMLStrategy(tk)
         prompt = strategy.assemble_prompt(tk.encode("hi"))
-        # prompt 末尾应该是 assistant_start 的 token ids
+        assistant_start = tk.encode("<｜im▁start｜>assistant\n")
-        assert (
+        assert prompt[-len(assistant_start):] == assistant_start
-            prompt[-len(strategy._assistant_start_ids) :]
-            == strategy._assistant_start_ids
-        )
 
 
 class TestAlpacaStrategy: