diff --git a/pipeline/__init__.py b/pipeline/__init__.py index a1790ad..a0e9135 100644 --- a/pipeline/__init__.py +++ b/pipeline/__init__.py @@ -28,7 +28,13 @@ Usage:: from pipeline.pipeline import Pipeline, PipelineConfig, Stage, TransformStage from pipeline.tokenize import AutoTokenizer, ChatTemplate, train_bpe_tokenizer from pipeline.text import TextNormalizer -from pipeline.packing import SequencePacker +from pipeline.packing import ( + GreedyPacker, + FfDPacker, + BfdPacker, + BasePacker, + pack_tensors, +) # I/O module from pipeline.io import FileScanner, HDF5Handler, export_dataset, cache_jsonl @@ -70,7 +76,11 @@ __all__ = [ "train_bpe_tokenizer", # Text processing "TextNormalizer", - "SequencePacker", + "GreedyPacker", + "FfDPacker", + "BfdPacker", + "BasePacker", + "pack_tensors", # I/O "FileScanner", "HDF5Handler", diff --git a/pipeline/io/__init__.py b/pipeline/io/__init__.py index f475f8e..9caabc4 100644 --- a/pipeline/io/__init__.py +++ b/pipeline/io/__init__.py @@ -4,16 +4,26 @@ This module provides: - FileScanner: File and directory scanning utilities - HDF5Handler: Tensor data persistence - export_dataset: HuggingFace Dataset to JSONL export -- cache_jsonl: JSONL to HDF5 tokenization and caching +- cache_jsonl: JSONL to HDF5/binary tokenization and caching +- dedup_jsonl: MinHash+LSH deduplication for pretraining text +- writers: BaseWriter / H5Writer / BinWriter / TextWriter (Strategy + Factory) """ from pipeline.io.file_scanner import FileScanner from pipeline.io.hdf5_handler import HDF5Handler from pipeline.io.export import export_dataset, cache_jsonl +from pipeline.io.dedup import dedup_jsonl +from pipeline.io.writers import BaseWriter, H5Writer, BinWriter, TextWriter, create_writer __all__ = [ "FileScanner", "HDF5Handler", "export_dataset", "cache_jsonl", + "dedup_jsonl", + "BaseWriter", + "H5Writer", + "BinWriter", + "TextWriter", + "create_writer", ] diff --git a/pipeline/io/dedup.py b/pipeline/io/dedup.py new file mode 100644 index 0000000..7e8f6a2 --- /dev/null +++ b/pipeline/io/dedup.py @@ -0,0 +1,167 @@ +"""MinHash + LSH deduplication for pretraining text data.""" + +import json +import logging +import os +from pathlib import Path +from typing import Iterator, List, Set, Tuple + +from tqdm import tqdm + +from pipeline.io.writers import TextWriter +from pipeline.utils import error_handler + +logger = logging.getLogger(__name__) + + +def _tokenize(text: str, ngram: int = 3) -> Set[str]: + return {text[i : i + ngram] for i in range(len(text) - ngram + 1)} + + +def _iter_docs(input_dir: Path) -> Iterator[Tuple[str, dict]]: + for fpath in sorted(input_dir.glob("*.jsonl")): + with open(fpath, encoding="utf-8") as f: + for line in f: + line = line.strip() + if not line: + continue + record = json.loads(line) + text = record.get("text", "") + if text: + yield text, record + + +def _write_h5(records: List[dict], output_dir: str, chunk_idx: int): + import h5py + + fname = os.path.join(output_dir, f"chunk_{chunk_idx}.h5") + texts = [rec.get("text", "") for rec in records] + with h5py.File(fname, "w") as f: + dt = h5py.special_dtype(vlen=str) + ds = f.create_dataset("text", (len(texts),), dtype=dt) + for i, t in enumerate(texts): + ds[i] = t + + +def _write_bin(records: List[dict], output_dir: Path, chunk_idx: int): + output_dir.mkdir(parents=True, exist_ok=True) + texts = [rec.get("text", "") + "\n" for rec in records] + + meta = {"chunk": chunk_idx, "count": len(texts), "format": "text", "encoding": "utf-8"} + meta_path = output_dir / "meta.json" + existing = json.loads(meta_path.read_text()) if meta_path.exists() else {} + existing[str(chunk_idx)] = meta + meta_path.write_text(json.dumps(existing, indent=2)) + + (output_dir / f"text_{chunk_idx}.bin").write_bytes("".join(texts).encode("utf-8")) + + +_WRITERS = { + "jsonl": TextWriter, + "h5": lambda: None, # handled inline below + "bin": lambda: None, +} + + +@error_handler() +def dedup_jsonl( + input_dir: str, + output_dir: str, + *, + threshold: float = 0.8, + num_perm: int = 128, + ngram: int = 3, + output_format: str = "jsonl", + chunk_size: int = 1_000_000, +) -> Tuple[int, int]: + """Deduplicate JSONL text files using MinHash + LSH. + + Args: + input_dir: Directory with source ``*.jsonl`` files. + output_dir: Directory for deduplicated output. + threshold: Jaccard similarity threshold (0–1). + num_perm: Number of MinHash permutations. + ngram: Character n-gram size. + output_format: ``"jsonl"``, ``"h5"``, or ``"bin"``. + chunk_size: Records per output chunk file. + + Returns: + ``(kept, removed)`` counts. + """ + from datasketch import MinHash, MinHashLSH + + input_path = Path(input_dir) + output_path = Path(output_dir) + output_path.mkdir(parents=True, exist_ok=True) + + logger.info( + f"Deduplicating {input_dir} -> {output_dir} " + f"(threshold={threshold}, perm={num_perm}, fmt={output_format})" + ) + + lsh = MinHashLSH(threshold=threshold, num_perm=num_perm) + + kept = 0 + removed = 0 + + dup_doc_ids: Set[int] = set() + for doc_id, (text, _record) in enumerate(tqdm(_iter_docs(input_path), desc="indexing", unit="docs")): + shingles = _tokenize(text, ngram=ngram) + if len(shingles) < ngram * 2: + dup_doc_ids.add(doc_id) + continue + + m = MinHash(num_perm=num_perm) + for s in shingles: + m.update(s.encode("utf-8")) + + if lsh.query(m): + dup_doc_ids.add(doc_id) + else: + lsh.insert(doc_id, m) + + logger.info(f"Found {len(dup_doc_ids)} duplicates, writing deduplicated data") + + buffer: List[dict] = [] + chunk_idx = 0 + writer = TextWriter(chunk_size) if output_format == "jsonl" else None + + for doc_id, (_text, record) in enumerate(tqdm(_iter_docs(input_path), desc="writing", unit="docs")): + if doc_id in dup_doc_ids: + removed += 1 + continue + + kept += 1 + buffer.append(record) + + if len(buffer) >= chunk_size: + _flush_chunk(buffer, output_path, chunk_idx, output_format, writer) + chunk_idx += 1 + buffer = [] + + if buffer: + _flush_chunk(buffer, output_path, chunk_idx, output_format, writer) + + if writer: + writer.flush(output_path) + + logger.info(f"Done. kept={kept}, removed={removed}") + return kept, removed + + +def _flush_chunk( + records: List[dict], + output_dir: Path, + chunk_idx: int, + output_format: str, + writer=None, +): + if output_format == "jsonl": + for rec in records: + writer.write_record(rec, output_dir) + elif output_format == "h5": + _write_h5(records, str(output_dir), chunk_idx) + elif output_format == "bin": + _write_bin(records, output_dir, chunk_idx) + else: + raise ValueError(f"Unknown output format: {output_format}") diff --git a/pipeline/io/export.py b/pipeline/io/export.py index 8788c6b..7c53410 100644 --- a/pipeline/io/export.py +++ b/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,41 +113,85 @@ def cache_jsonl( *, pack_size: int = -1, pad_value: int = 0, - batch_size: int = 256, + 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. - batch_size: Number of records passed to the processor at once. + 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 - def append_batch(batch): - items = [item for _, item in batch] + 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) + 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 batch: + for line_num, item in samples: try: results.append(processor.process(item)) except Exception as e: @@ -126,42 +200,63 @@ def cache_jsonl( 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[key].append(result[key]) + arrows_batch[key].append(result[key]) + if target_tokens > 0: + batch_tokens += int(result[output_keys[0]].shape[0]) - batch = [] 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: - batch.append((line_num, json.loads(line))) - if len(batch) >= batch_size: - append_batch(batch) - batch = [] - except json.JSONDecodeError as e: - logger.warning( - f"JSON decode error in {file_path} line {line_num}: {e}. Skipping line." - ) - if batch: - append_batch(batch) + buf.append((line_num, line)) + if len(buf) >= batch_size: + flush_buf() + if target_tokens > 0 and batch_tokens >= target_tokens: + packed = pack_tensors( + arrows_batch, + pack_size, + pad_value, + dtypes, + pad_values=pad_values, + algo=pack_algo, + ) + for key in output_keys: + all_packed[key].extend(packed[key]) + arrows_batch[key] = [] + batch_tokens = 0 - if pack_size > 0: - dtypes = ( - dict(processor.schema.output_fields) - if processor.schema is not None - else None - ) - output = pack_tensors(arrows, pack_size, pad_value, dtypes) + flush_buf() + + if arrows_batch[output_keys[0]]: + if pack_size > 0: + packed = pack_tensors(arrows_batch, pack_size, pad_value, dtypes, pad_values=pad_values, algo=pack_algo) + for key in output_keys: + all_packed[key].extend(packed[key]) + else: + for key in output_keys: + all_packed[key].extend(arrows_batch[key]) + + if not all_packed[output_keys[0]]: + logger.warning(f"No valid samples in {file_path}, skipping") + continue + + if pack_size <= 0: + 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) - output_files.append(h5_path) - logger.info(f"Saved {h5_path}") + writer: BaseWriter = create_writer(output_format) + saved = writer.save(output_dir, file_name, output) + output_files.append(saved) + logger.info(f"Saved {saved}") return output_files diff --git a/pipeline/io/writers.py b/pipeline/io/writers.py new file mode 100644 index 0000000..6ccf6d4 --- /dev/null +++ b/pipeline/io/writers.py @@ -0,0 +1,107 @@ +"""Storage backends for tensor / text output (Strategy + Factory). + +Each backend implements a common ``save()`` interface so callers use +polymorphism instead of ``if fmt == "h5" ... elif fmt == "bin" ...``. + +Supports: + - **H5Writer**: HDF5 format (via HDF5Handler) + - **BinWriter**: binary format – meta.json + {key}.bin (memmap-compatible) + - **TextWriter**: raw JSONL text (for dedup output) +""" + +import json +import os +from abc import ABC, abstractmethod +from pathlib import Path +from typing import Dict, List + +import torch +from torch import Tensor + + +class BaseWriter(ABC): + """Abstract writer – call ``save(dir, name, data)`` without caring + about the underlying format.""" + + @abstractmethod + def save(self, output_dir: str, file_name: str, data: Dict[str, List[Tensor]]) -> str: + ... + + +class H5Writer(BaseWriter): + def save(self, output_dir: str, file_name: str, data: Dict[str, List[Tensor]]) -> str: + from pipeline.io.hdf5_handler import HDF5Handler + return HDF5Handler.save(output_dir, file_name, data) + + +class BinWriter(BaseWriter): + def save(self, output_dir: str, file_name: str, data: Dict[str, List[Tensor]]) -> str: + import numpy as np + + os.makedirs(output_dir, exist_ok=True) + sub_dir = os.path.join(output_dir, file_name) + os.makedirs(sub_dir, exist_ok=True) + + meta: Dict[str, Dict] = {} + for key, tensors in data.items(): + cat = torch.cat(tensors, dim=0) + meta[key] = {"shape": list(cat.shape), "dtype": str(cat.dtype).split(".")[-1]} + np.asarray(cat.cpu().numpy()).tofile(os.path.join(sub_dir, f"{key}.bin")) + + with open(os.path.join(sub_dir, "meta.json"), "w") as f: + json.dump(meta, f, indent=2) + + return sub_dir + + +class TextWriter(BaseWriter): + """Write raw text records as JSONL (used by dedup output).""" + + def __init__(self, chunk_size: int = 1_000_000): + self._chunk_size = chunk_size + self._buffer: List[dict] = [] + self._chunk_idx = 0 + + def save(self, output_dir: str, file_name: str, data: Dict[str, List[Tensor]]) -> str: + raise NotImplementedError("TextWriter.save_one is for tensor data; use write_record()") + + def write_record(self, record: dict, output_dir: Path): + self._buffer.append(record) + if len(self._buffer) >= self._chunk_size: + self._flush(output_dir) + + def flush(self, output_dir: Path): + if self._buffer: + self._flush(output_dir) + + def _flush(self, output_dir: Path): + output_dir.mkdir(parents=True, exist_ok=True) + fpath = output_dir / f"chunk_{self._chunk_idx}.jsonl" + with open(fpath, "w", encoding="utf-8") as f: + for rec in self._buffer: + f.write(json.dumps(rec, ensure_ascii=False) + "\n") + self._chunk_idx += 1 + self._buffer = [] + + +_WRITER_REGISTRY: Dict[str, type] = {} + + +def register_writer(name: str): + def decorator(cls): + _WRITER_REGISTRY[name] = cls + return cls + return decorator + + +def create_writer(name: str, **kwargs) -> BaseWriter: + cls = _WRITER_REGISTRY.get(name) + if cls is None: + raise ValueError(f"Unknown writer: {name}. Available: {list(_WRITER_REGISTRY)}") + return cls(**kwargs) + + +# Register built-in writers +register_writer("h5")(H5Writer) +register_writer("bin")(BinWriter) +register_writer("jsonl")(TextWriter) diff --git a/pipeline/packing.py b/pipeline/packing.py deleted file mode 100644 index 42531ea..0000000 --- a/pipeline/packing.py +++ /dev/null @@ -1,146 +0,0 @@ -import logging -from typing import Any, Dict, List, Optional, Tuple, Union - -import torch -from torch import Tensor - -from pipeline.utils import error_handler - -logger = logging.getLogger(__name__) - - -class SequencePacker: - """ - Stream-concatenation packer for LLM training sequences. - - Algorithm (streaming concat): - - Input: sequences = [A(len=3), B(len=5), C(len=2)], pack_size = 6 - - 1. Validate & Normalize - - Check 1D dimension, unify dtype, warn on overlong sequences - - Result: [A, B, C] - - 2. Stream into buffer, slice off full chunks - - buffer += A(3) -> [a1 a2 a3], pos=3 - - buffer += B(5) -> [a1 a2 a3 b1 b2 b3 b4 b5], pos=8 - pos >= 6 -> flush [a1 a2 a3 b1 b2 b3], buffer=[b4 b5], pos=2 - - buffer += C(2) -> [b4 b5 c1 c2], pos=4 - loop ends -> flush tail [b4 b5 c1 c2 PAD PAD] - - Output: [[a1 a2 a3 b1 b2 b3], [b4 b5 c1 c2 PAD PAD]] - - Samples may be split across chunks — this is intentional and standard - practice in LLM training (TRL, Megatron-LM, etc.). - - Cross-group consistency: - Different tensor groups (e.g. input_ids, loss_masks) packed with - separate packer instances on samples with matching lengths produce - identical chunk boundaries. Element-level correspondence is preserved. - """ - - def __init__( - self, - pack_size: int, - pad_value: Union[int, bool] = 0, - dtype: Optional[torch.dtype] = None, - ): - self.pack_size = pack_size - self.pad_value = pad_value - self.dtype = dtype - self._buffer: List = [] - self._pos: int = 0 - self._packages: List[Tensor] = [] - - def reset(self) -> None: - """Reset packer state for instance reuse.""" - self._buffer = [] - self._pos = 0 - self._packages = [] - - @error_handler() - def pack(self, sequences: List[Tensor]) -> List[Tensor]: - """ - Pack sequences via streaming concatenation into fixed-size chunks. - - Sequences are concatenated in order and sliced at pack_size boundaries. - The final chunk is padded with pad_value. - - When dtype is not set at init, it is inferred from the first input tensor. - - Args: - sequences: List of 1D input tensors. - - Returns: - List of packed tensors, each with length equal to pack_size. - """ - if not sequences: - return [] - - # --- auto-infer dtype from first sequence --- - if self.dtype is None: - self.dtype = sequences[0].dtype - - # --- validate & normalize --- - normalized: List[Tensor] = [] - for i, seq in enumerate(sequences): - if seq.dim() != 1: - raise ValueError( - f"Expected 1D tensor at index {i}, got {seq.dim()}D tensor with shape {seq.shape}" - ) - if seq.dtype != self.dtype: - seq = seq.to(self.dtype) - normalized.append(seq) - - # --- stream into buffer, slice off full chunks --- - self._buffer = [] - self._packages = [] - pack_size = self.pack_size - buf = self._buffer - - for seq in normalized: - buf.extend(seq.tolist()) - while len(buf) >= pack_size: - self._packages.append(torch.tensor(buf[:pack_size], dtype=self.dtype)) - buf = buf[pack_size:] - - # flush tail with padding - if buf: - padded = buf + [self.pad_value] * (pack_size - len(buf)) - self._packages.append(torch.tensor(padded, dtype=self.dtype)) - - self._pos = len(buf) - return self._packages - - -def pack_tensors( - tensors: Dict[str, List[Tensor]], - pack_size: int, - pad_value: Union[int, bool] = 0, - dtypes: Optional[Dict[str, torch.dtype]] = None, -) -> Dict[str, List[Tensor]]: - """ - Pack multiple named tensor groups in parallel. - - Each group is packed independently with its own SequencePacker instance. - When dtypes is provided, packers use the declared dtype per key; - otherwise dtype is auto-inferred from the first tensor in each group. - - Args: - tensors: Dict mapping key names to lists of 1D tensors. - pack_size: Fixed chunk length. - pad_value: Padding value for non-bool tensors. - dtypes: Optional per-key dtype declarations. - - Returns: - Dict mapping key names to lists of packed tensors. - """ - if dtypes is None: - dtypes = {} - - output: Dict[str, List[Tensor]] = {} - for key, seqs in tensors.items(): - dtype = dtypes.get(key) - packer = SequencePacker(pack_size, pad_value, dtype=dtype) - output[key] = packer.pack(seqs) - return output diff --git a/pipeline/packing/__init__.py b/pipeline/packing/__init__.py new file mode 100644 index 0000000..b3a0a32 --- /dev/null +++ b/pipeline/packing/__init__.py @@ -0,0 +1,84 @@ +"""Sequence packing algorithms for LLM training data. + +Available packers: + - BfdPacker: Best-Fit Decreasing, samples never split (default) + - FfDPacker: First-Fit Decreasing, samples never split + - GreedyPacker: First-fit in input order, samples never split +""" + +from typing import Dict, List, Optional, Union + +import torch + +from pipeline.packing.base import BasePacker +from pipeline.packing.binpack import GreedyPacker, FfDPacker, BfdPacker + + +def pack_tensors( + tensors: Dict[str, List[torch.Tensor]], + pack_size: int, + pad_value: Union[int, bool] = 0, + dtypes: Optional[Dict[str, torch.dtype]] = None, + pad_values: Optional[Dict[str, Union[int, bool]]] = None, + algo: Optional[Union[str, BasePacker]] = None, +) -> Dict[str, List[torch.Tensor]]: + """Pack multiple named tensor groups in parallel. + + Each group is packed independently with its own packer instance. + + Args: + tensors: Dict mapping key names to lists of 1D tensors. + pack_size: Fixed chunk length. + pad_value: Default padding value, used for keys not in pad_values. + dtypes: Optional per-key dtype declarations. + pad_values: Optional per-key padding values (e.g. pad_token_id for + 'sequence', False for 'loss_mask', 0 for 'position_ids'). + algo: Packing algorithm to use. Can be 'bfd' (default), + 'ffd', 'greedy', or a BasePacker instance. + + Returns: + Dict mapping key names to lists of packed tensors. + """ + if dtypes is None: + dtypes = {} + if pad_values is None: + pad_values = {} + + output: Dict[str, List[torch.Tensor]] = {} + for key, seqs in tensors.items(): + key_pad = pad_values.get(key, pad_value) + actual_packer = _resolve_algo(algo, pack_size, key_pad) + dtype = dtypes.get(key) + if dtype is not None: + actual_packer.dtype = dtype + output[key] = actual_packer.pack(seqs) + return output + + +def _resolve_algo( + algo: Optional[Union[str, BasePacker]], + pack_size: int, + pad_value: Union[int, bool], +) -> BasePacker: + if algo is None or algo == "bfd": + return BfdPacker(pack_size, pad_value) + if isinstance(algo, BasePacker): + cls = type(algo) + return cls(pack_size, pad_value) + if algo == "ffd": + return FfDPacker(pack_size, pad_value) + if algo == "greedy": + return GreedyPacker(pack_size, pad_value) + raise ValueError( + f"Unknown packing algorithm: {algo}. " + f"Choose from: bfd, ffd, greedy" + ) + + +__all__ = [ + "BasePacker", + "BfdPacker", + "FfDPacker", + "GreedyPacker", + "pack_tensors", +] diff --git a/pipeline/packing/base.py b/pipeline/packing/base.py new file mode 100644 index 0000000..27ed628 --- /dev/null +++ b/pipeline/packing/base.py @@ -0,0 +1,49 @@ +from abc import ABC, abstractmethod +from typing import List, Optional, Union + +import torch +from torch import Tensor + + +class BasePacker(ABC): + """Abstract base class for sequence packing algorithms. + + All packers must implement pack() and reset(). + pack() takes a list of 1D tensors and returns a list of packed fixed-size tensors. + """ + + def __init__( + self, + pack_size: int, + pad_value: Union[int, bool] = 0, + dtype: Optional[torch.dtype] = None, + ): + self.pack_size = pack_size + self.pad_value = pad_value + self.dtype = dtype + + @abstractmethod + def pack(self, sequences: List[Tensor]) -> List[Tensor]: + """Pack sequences into fixed-size chunks.""" + ... + + @abstractmethod + def reset(self) -> None: + """Reset packer state for instance reuse.""" + ... + + def _validate_and_normalize(self, sequences: List[Tensor]) -> List[Tensor]: + """Validate 1D tensors and unify dtype.""" + if self.dtype is None and sequences: + self.dtype = sequences[0].dtype + + normalized: List[Tensor] = [] + for i, seq in enumerate(sequences): + if seq.dim() != 1: + raise ValueError( + f"Expected 1D tensor at index {i}, got {seq.dim()}D tensor with shape {seq.shape}" + ) + if seq.dtype != self.dtype: + seq = seq.to(self.dtype) + normalized.append(seq) + return normalized diff --git a/pipeline/packing/binpack.py b/pipeline/packing/binpack.py new file mode 100644 index 0000000..20b7d12 --- /dev/null +++ b/pipeline/packing/binpack.py @@ -0,0 +1,174 @@ +from typing import List, Optional, Union + +import torch +from torch import Tensor + +from pipeline.packing.base import BasePacker +from pipeline.utils import error_handler + + +def _truncate(tokens: List, max_len: int) -> List: + return tokens[:max_len] + + +def _pad_bin(bin_list: List, target_len: int, pad_value: Union[int, bool], dtype: torch.dtype) -> Tensor: + bin_list.extend([pad_value] * (target_len - len(bin_list))) + return torch.tensor(bin_list, dtype=dtype) + + +class GreedyPacker(BasePacker): + """Greedy first-fit packer (no sorting). + + Sequences are packed in input order into the first bin with enough space. + Overlong sequences (> pack_size) are truncated to pack_size. + """ + + def __init__( + self, + pack_size: int, + pad_value: Union[int, bool] = 0, + dtype: Optional[torch.dtype] = None, + ): + super().__init__(pack_size, pad_value, dtype) + self._bins: List[List] = [] + + def reset(self) -> None: + self._bins = [] + + @error_handler() + def pack(self, sequences: List[Tensor]) -> List[Tensor]: + if not sequences: + return [] + + normalized = self._validate_and_normalize(sequences) + self._bins = [] + pack_size = self.pack_size + pad_value = self.pad_value + + for seq in normalized: + seq_len = int(seq.shape[0]) + if seq_len > pack_size: + self._bins.append(_truncate(seq.tolist(), pack_size)) + continue + placed = False + for bin_list in self._bins: + if len(bin_list) + seq_len <= pack_size: + bin_list.extend(seq.tolist()) + placed = True + break + if not placed: + self._bins.append(list(seq.tolist())) + + packages: List[Tensor] = [] + for bin_list in self._bins: + packages.append(_pad_bin(bin_list, pack_size, pad_value, self.dtype)) + + return packages + + +class FfDPacker(BasePacker): + """First-Fit Decreasing (FFD) bin-packing packer. + + Sequences are sorted by descending length, then packed into the first + bin with enough space. Overlong sequences are truncated to pack_size. + """ + + def __init__( + self, + pack_size: int, + pad_value: Union[int, bool] = 0, + dtype: Optional[torch.dtype] = None, + ): + super().__init__(pack_size, pad_value, dtype) + self._bins: List[List] = [] + + def reset(self) -> None: + self._bins = [] + + @error_handler() + def pack(self, sequences: List[Tensor]) -> List[Tensor]: + if not sequences: + return [] + + normalized = self._validate_and_normalize(sequences) + self._bins = [] + pack_size = self.pack_size + pad_value = self.pad_value + + indexed = [(int(s.shape[0]), s) for s in normalized] + indexed.sort(key=lambda x: x[0], reverse=True) + + for seq_len, seq in indexed: + if seq_len > pack_size: + self._bins.append(_truncate(seq.tolist(), pack_size)) + continue + placed = False + for bin_list in self._bins: + if len(bin_list) + seq_len <= pack_size: + bin_list.extend(seq.tolist()) + placed = True + break + if not placed: + self._bins.append(list(seq.tolist())) + + packages: List[Tensor] = [] + for bin_list in self._bins: + packages.append(_pad_bin(bin_list, pack_size, pad_value, self.dtype)) + + return packages + + +class BfdPacker(BasePacker): + """Best-Fit Decreasing (BFD) bin-packing packer. + + Sequences are sorted by descending length, then packed into the bin + that minimizes remaining space (tightest fit). + Overlong sequences are truncated to pack_size. + """ + + def __init__( + self, + pack_size: int, + pad_value: Union[int, bool] = 0, + dtype: Optional[torch.dtype] = None, + ): + super().__init__(pack_size, pad_value, dtype) + self._bins: List[List] = [] + + def reset(self) -> None: + self._bins = [] + + @error_handler() + def pack(self, sequences: List[Tensor]) -> List[Tensor]: + if not sequences: + return [] + + normalized = self._validate_and_normalize(sequences) + self._bins = [] + pack_size = self.pack_size + pad_value = self.pad_value + + indexed = [(int(s.shape[0]), s) for s in normalized] + indexed.sort(key=lambda x: x[0], reverse=True) + + for seq_len, seq in indexed: + if seq_len > pack_size: + self._bins.append(_truncate(seq.tolist(), pack_size)) + continue + best_idx = -1 + best_remain = pack_size + 1 + for i, bin_list in enumerate(self._bins): + remain = pack_size - len(bin_list) + if seq_len <= remain < best_remain: + best_remain = remain + best_idx = i + if best_idx >= 0: + self._bins[best_idx].extend(seq.tolist()) + else: + self._bins.append(list(seq.tolist())) + + packages: List[Tensor] = [] + for bin_list in self._bins: + packages.append(_pad_bin(bin_list, pack_size, pad_value, self.dtype)) + + return packages diff --git a/pipeline/processors/base.py b/pipeline/processors/base.py index 5c52996..1878ef6 100644 --- a/pipeline/processors/base.py +++ b/pipeline/processors/base.py @@ -88,6 +88,21 @@ class BaseProcessor(ABC): """Return list of output tensor key names.""" pass + def process_batch(self, input_dicts: List[Dict[str, Any]]) -> List[Dict[str, Tensor]]: + """Process a batch of input samples. + + Default implementation calls process() for each sample. + Subclasses should override for efficient batch processing + (e.g., using tokenizer.encode_batch). + + Args: + input_dicts: List of input dictionaries. + + Returns: + List of output dictionaries mapping output key names to tensors. + """ + return [self.process(d) for d in input_dicts] + def validate_input(self, input_dict: Dict[str, Any]) -> None: """Validate input against schema before processing. diff --git a/pipeline/processors/pretrain.py b/pipeline/processors/pretrain.py index ceaea5e..4d3aae7 100644 --- a/pipeline/processors/pretrain.py +++ b/pipeline/processors/pretrain.py @@ -44,11 +44,11 @@ class PreTrainProcessor(BaseProcessor): return {"sequence": torch.tensor(tokens, dtype=torch.int32)} def process_batch(self, input_dicts: List[Dict[str, Any]]) -> List[Dict[str, Tensor]]: - texts = [f"{item['text']}{self._eos_token}" for item in input_dicts] - encoded = self.tokenizer.encode(texts) + 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 encoded + for tokens in batch_tokens ] @property diff --git a/pipeline/processors/sft.py b/pipeline/processors/sft.py index c929770..2f54070 100644 --- a/pipeline/processors/sft.py +++ b/pipeline/processors/sft.py @@ -7,7 +7,7 @@ from torch import Tensor from pipeline.tokenize import AutoTokenizer from pipeline.strategies import PromptStrategy, ChatMLStrategy -from pipeline.processors.base import BaseProcessor, ProcessorSchema, encode_with_mask +from pipeline.processors.base import BaseProcessor, ProcessorSchema from pipeline.processors.factory import ProcessorFactory @@ -15,15 +15,15 @@ from pipeline.processors.factory import ProcessorFactory class SFTProcessor(BaseProcessor): """Supervised fine-tuning data processor. - Supports two input formats: + Input formats: 1. messages (recommended): ``{"messages": [{"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]}`` - Multi-turn and system prompts are supported. - The tokenizer's ``apply_chat_template`` is used for rendering. + Multi-turn and system prompts are supported. Each assistant + turn gets ``loss_mask = 1``; all other roles get 0. 2. legacy query/response: ``{"query": "...", "response": "..."}`` - Falls back to the configured PromptStrategy (ChatML by default). + Internally converted to messages. Output schema: - sequence: int32 tensor - Combined token IDs (prompt + response) @@ -63,39 +63,23 @@ 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 process_batch(self, input_dicts: List[Dict[str, Any]]) -> List[Dict[str, Tensor]]: - results: List[Optional[Dict[str, Tensor]]] = [None] * len(input_dicts) - message_indices = [i for i, item in enumerate(input_dicts) if "messages" in item] - legacy_indices = [ - i - for i, item in enumerate(input_dicts) - if "messages" not in item and "query" in item and "response" in item - ] - if len(message_indices) + len(legacy_indices) != len(input_dicts): - raise KeyError("Input must contain 'messages' or 'query'/'response' pair") - - if message_indices: - items = [input_dicts[i] for i in message_indices] - batch_results = self._process_messages_batch( - [item["messages"] for item in items] - ) - for index, result in zip(message_indices, batch_results): - results[index] = result - - if legacy_indices: - items = [input_dicts[i] for i in legacy_indices] - batch_results = self._process_legacy_batch(items) - for index, result in zip(legacy_indices, batch_results): - results[index] = result - - if any(result is None for result in results): - raise RuntimeError("Batch processing did not produce all results") - return results + 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: @@ -103,118 +87,72 @@ class SFTProcessor(BaseProcessor): if messages[-1]["role"] != "assistant": raise ValueError("Last message must have role 'assistant'") - last_asst_idx = max( - i for i, m in enumerate(messages) if m["role"] == "assistant" - ) + strategy = self.strategy or ChatMLStrategy(self.tokenizer) - full_text = self.tokenizer.apply_chat_template( - messages, tokenize=False, add_generation_prompt=False - ) - full_ids = self.tokenizer.encode(full_text, add_special_tokens=False) + prompt, resp = strategy.format_messages(messages) - prompt_text = self.tokenizer.apply_chat_template( - messages[:last_asst_idx], - tokenize=False, - add_generation_prompt=True, - ) - prompt_ids = self.tokenizer.encode(prompt_text, add_special_tokens=False) - - resp_ids = full_ids[len(prompt_ids) :] - if not resp_ids: - raise ValueError("Empty assistant response") - - tokens, loss_mask = encode_with_mask(prompt_ids, list(resp_ids)) - - if self.max_seq_len and len(tokens) > self.max_seq_len: - tokens = tokens[: self.max_seq_len] + sequence = torch.tensor(prompt + resp, dtype=torch.int32) + 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(tokens), dtype=torch.int32) + position_ids = torch.arange(len(sequence), dtype=torch.int32) return { - "sequence": tokens, + "sequence": sequence, "loss_mask": loss_mask, "position_ids": position_ids, } - def _process_messages_batch( - self, conversations: List[List[Dict[str, str]]] - ) -> List[Dict[str, Tensor]]: - for messages in conversations: - if not messages: - raise ValueError("Messages list is empty") - if messages[-1]["role"] != "assistant": - raise ValueError("Last message must have role 'assistant'") + def process_batch(self, input_dicts: List[Dict[str, Any]]) -> List[Optional[Dict[str, Tensor]]]: + strategy = self.strategy or ChatMLStrategy(self.tokenizer) - assistant_indices = [ - max(i for i, message in enumerate(messages) if message["role"] == "assistant") - for messages in conversations - ] - full_texts = [ - self.tokenizer.apply_chat_template( - messages, tokenize=False, add_generation_prompt=False - ) - for messages in conversations - ] - prompt_texts = [ - self.tokenizer.apply_chat_template( - messages[:assistant_idx], tokenize=False, add_generation_prompt=True - ) - for messages, assistant_idx in zip(conversations, assistant_indices) - ] - full_ids_batch = self.tokenizer.encode(full_texts, add_special_tokens=False) - prompt_ids_batch = self.tokenizer.encode(prompt_texts, add_special_tokens=False) + prompts_text: List[str] = [] + fulls_text: List[str] = [] + indices: List[int] = [] + results: List[Optional[Dict[str, Tensor]]] = [None] * len(input_dicts) - results = [] - for full_ids, prompt_ids in zip(full_ids_batch, prompt_ids_batch): - resp_ids = full_ids[len(prompt_ids) :] - if not resp_ids: - raise ValueError("Empty assistant response") - tokens, loss_mask = encode_with_mask(prompt_ids, list(resp_ids)) - if self.max_seq_len and len(tokens) > self.max_seq_len: - tokens = tokens[: self.max_seq_len] + 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] - results.append( - { - "sequence": tokens, - "loss_mask": loss_mask, - "position_ids": torch.arange(len(tokens), dtype=torch.int32), - } - ) - return results + position_ids = torch.arange(len(sequence), dtype=torch.int32) + results[idx] = { + "sequence": sequence, + "loss_mask": loss_mask, + "position_ids": position_ids, + } - def _process_legacy(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]: - strategy = self.strategy or ChatMLStrategy(self.tokenizer) - - query_tokens = self.tokenizer.encode(input_dict["query"]) - response_tokens = self.tokenizer.encode(input_dict["response"]) - - prompt = strategy.assemble_prompt(query_tokens) - response = strategy.assemble_response(response_tokens) - - tokens, loss_mask = encode_with_mask(prompt, response) - position_ids = torch.arange(len(tokens), dtype=torch.int32) - return {"sequence": tokens, "loss_mask": loss_mask, "position_ids": position_ids} - - def _process_legacy_batch( - self, input_dicts: List[Dict[str, Any]] - ) -> List[Dict[str, Tensor]]: - strategy = self.strategy or ChatMLStrategy(self.tokenizer) - query_batch = self.tokenizer.encode([item["query"] for item in input_dicts]) - response_batch = self.tokenizer.encode( - [item["response"] for item in input_dicts] - ) - results = [] - for query_tokens, response_tokens in zip(query_batch, response_batch): - prompt = strategy.assemble_prompt(query_tokens) - response = strategy.assemble_response(response_tokens) - tokens, loss_mask = encode_with_mask(prompt, response) - results.append( - { - "sequence": tokens, - "loss_mask": loss_mask, - "position_ids": torch.arange(len(tokens), dtype=torch.int32), - } - ) return results @property diff --git a/pipeline/strategies/chatml.py b/pipeline/strategies/chatml.py index cd8411c..947e1f6 100644 --- a/pipeline/strategies/chatml.py +++ b/pipeline/strategies/chatml.py @@ -1,42 +1,91 @@ """ChatML format strategy.""" -from typing import List +from typing import Dict, List, Tuple from pipeline.tokenize import AutoTokenizer from pipeline.strategies.base import PromptStrategy from pipeline.strategies.factory import StrategyFactory +DEFAULT_CHATML_TEMPLATE = ( + "{% for message in messages %}" + "{% if message['role'] == 'system' %}" + "{{ '<|im_start|>system\n' + message['content'] + '<|im_end|>\n' }}" + "{% elif message['role'] == 'user' %}" + "{{ '<|im_start|>user\n' + message['content'] + '<|im_end|>\n' }}" + "{% elif message['role'] == 'assistant' %}" + "{{ '<|im_start|>assistant\n' + message['content'] + '<|im_end|>\n' }}" + "{% endif %}" + "{% endfor %}" + "{% if add_generation_prompt %}" + "{{ '<|im_start|>assistant\n' }}" + "{% endif %}" +) + @StrategyFactory.register("chatml") class ChatMLStrategy(PromptStrategy): - """ChatML format strategy.""" + """ChatML format strategy. - def __init__( - self, - tokenizer: AutoTokenizer, - user_start: str = "<｜im▁start｜>user\n", - user_end: str = "<｜im▁end｜>\n", - assistant_start: str = "<｜im▁start｜>assistant\n", - assistant_end: str = "<｜im▁end｜>\n", - ): + Renders messages using the tokenizer's jinja chat_template from + ``tokenizer_config.json``. Falls back to DEFAULT_CHATML_TEMPLATE + when no template is configured. + + The strategy does **not** hard-code any special tokens – all + formatting is driven by the jinja template. + """ + + def __init__(self, tokenizer: AutoTokenizer): super().__init__(tokenizer) - - self._user_start_ids = self._encode_format(user_start) - self._user_end_ids = self._encode_format(user_end) - self._assistant_start_ids = self._encode_format(assistant_start) - self._assistant_end_ids = self._encode_format(assistant_end) + if tokenizer._chat_template is None: + tokenizer.set_chat_template(DEFAULT_CHATML_TEMPLATE) @property def name(self) -> str: return "chatml" + def format_messages( + self, + messages: List[Dict[str, str]], + ) -> Tuple[List[int], List[int]]: + """Render a single-turn messages conversation. + + Returns ``(prompt_tokens, response_tokens)`` where + *prompt_tokens* contains everything up to (and including) the + last assistant start marker, and *response_tokens* is the + assistant content plus the closing markers. + """ + last_asst = max( + i for i, m in enumerate(messages) if m["role"] == "assistant" + ) + + prompt = self.tokenizer.apply_chat_template( + messages[:last_asst], + add_generation_prompt=True, + tokenize=True, + ) + full = self.tokenizer.apply_chat_template( + messages[: last_asst + 1], + add_generation_prompt=False, + tokenize=True, + ) + return prompt, full[len(prompt) :] + def assemble_prompt(self, query_tokens: List[int]) -> List[int]: - return ( - self._user_start_ids - + query_tokens - + self._user_end_ids - + self._assistant_start_ids + text = self.tokenizer.decode(query_tokens) + return self.tokenizer.apply_chat_template( + [{"role": "user", "content": text}], + add_generation_prompt=True, + tokenize=True, ) def assemble_response(self, response_tokens: List[int]) -> List[int]: - return response_tokens + self._assistant_end_ids + text = self.tokenizer.decode(response_tokens) + full = self.tokenizer.apply_chat_template( + [{"role": "assistant", "content": text}], + add_generation_prompt=False, + tokenize=True, + ) + opening = self.tokenizer.apply_chat_template( + [], add_generation_prompt=True, tokenize=True + ) + return full[len(opening) :] diff --git a/pipeline/tokenize/tokenizer.py b/pipeline/tokenize/tokenizer.py index 69ff159..a1b8889 100644 --- a/pipeline/tokenize/tokenizer.py +++ b/pipeline/tokenize/tokenizer.py @@ -244,22 +244,18 @@ class AutoTokenizer: "Tokenizer not initialized. Load or create a tokenizer first." ) - if isinstance(tokens, str): - encoded = self._tokenizer.encode( - tokens, - is_pretokenized=is_pretokenized, - add_special_tokens=add_special_tokens, - ) - return encoded.ids if out_ids else encoded.tokens - else: - encoded_list = self._tokenizer.encode_batch( - tokens, - is_pretokenized=is_pretokenized, - add_special_tokens=add_special_tokens, - ) - return [ - encoded.ids if out_ids else encoded.tokens for encoded in encoded_list - ] + single = isinstance(tokens, str) + if single: + tokens = [tokens] + encoded_list = self._tokenizer.encode_batch( + tokens, + is_pretokenized=is_pretokenized, + add_special_tokens=add_special_tokens, + ) + result = [ + encoded.ids if out_ids else encoded.tokens for encoded in encoded_list + ] + return result[0] if single else result def decode(self, tokens: List[int], skip_special_tokens: bool = True) -> str: """Decode token IDs to text.""" @@ -270,6 +266,12 @@ class AutoTokenizer: return self._tokenizer.decode(tokens, skip_special_tokens=skip_special_tokens) + def token_to_id(self, token: str) -> Optional[int]: + """Convert a token string to its integer ID.""" + if self._tokenizer is None: + raise RuntimeError("Tokenizer not initialized.") + return self._tokenizer.token_to_id(token) + def __len__(self) -> int: if self._tokenizer is None: return 0 diff --git a/pyproject.toml b/pyproject.toml index fdfa49a..e7a159e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,7 +3,7 @@ requires = ["setuptools>=64", "wheel"] build-backend = "setuptools.build_meta" [project] -name = "khaosz_dataset" +name = "datapipline" version = "0.1.0" description = "A dataset processing toolkit for language model training" readme = "README.md" diff --git a/scripts/cache_h5.py b/scripts/cache_h5.py index b321686..6df5bf9 100644 --- a/scripts/cache_h5.py +++ b/scripts/cache_h5.py @@ -1,11 +1,10 @@ -"""JSONL to H5 caching script. +"""JSONL tokenization and caching script. -Tokenize JSONL files and pack them into HDF5 format. +Tokenize JSONL files and save as HDF5 or binary format. Usage: python scripts/cache_h5.py pt ./dataset/chinese-c4-pretrain - python scripts/cache_h5.py sft ./dataset/belle-sft --pack-size 4096 --strategy alpaca - python scripts/cache_h5.py sft ./dataset/Ling-Coder-sft --tokenizer ./my_tokenizer.json + python scripts/cache_h5.py sft ./dataset/belle-sft --pack-size 4096 --output-format bin """ import argparse @@ -29,13 +28,13 @@ def main(): "-o", "--output-dir", default=None, - help="H5 output dir (default: /cached)", + help="Output dir (default: /cached)", ) parser.add_argument( "-t", "--tokenizer", - default="./tokenizer.json", - help="Tokenizer path (default: ./tokenizer.json)", + default="./tokenizer", + help="Tokenizer dir (default: ./tokenizer)", ) parser.add_argument( "-s", @@ -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( "--batch-size", @@ -65,6 +78,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="h5", + choices=["h5", "bin"], + help="Output format: h5 or bin (default: h5)", + ) args = parser.parse_args() # Initialize logging explicitly (not automatic anymore) @@ -101,9 +127,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, @@ -111,6 +142,9 @@ 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}") diff --git a/scripts/dedup_pretrain.py b/scripts/dedup_pretrain.py new file mode 100644 index 0000000..44cd318 --- /dev/null +++ b/scripts/dedup_pretrain.py @@ -0,0 +1,36 @@ +"""MinHash + LSH deduplication CLI. + +Usage: + python scripts/dedup_pretrain.py --input-dir --output-dir --threshold 0.8 --num-perm 128 --output-format jsonl +""" + +import argparse + +from pipeline.io import dedup_jsonl + + +def main(): + parser = argparse.ArgumentParser(description="MinHash + LSH deduplication") + parser.add_argument("--input-dir", required=True) + parser.add_argument("--output-dir", required=True) + parser.add_argument("--threshold", type=float, default=0.8) + parser.add_argument("--num-perm", type=int, default=128) + parser.add_argument("--ngram", type=int, default=3) + parser.add_argument("--output-format", default="jsonl", choices=["jsonl", "h5", "bin"]) + args = parser.parse_args() + + kept, removed = dedup_jsonl( + input_dir=args.input_dir, + output_dir=args.output_dir, + threshold=args.threshold, + num_perm=args.num_perm, + ngram=args.ngram, + output_format=args.output_format, + ) + + total = kept + removed + print(f"kept={kept}, removed={removed} ({removed/max(total,1)*100:.1f}%)") + + +if __name__ == "__main__": + main() diff --git a/scripts/filter_short.py b/scripts/filter_short.py new file mode 100644 index 0000000..88e70fc --- /dev/null +++ b/scripts/filter_short.py @@ -0,0 +1,92 @@ +import argparse +import json +import os +import tempfile +import shutil + +MIN_LEN = 15 + + +def filter_sft(input_path: str) -> tuple[int, int]: + """Filter SFT JSONL (messages format), remove if any msg content < MIN_LEN chars.""" + kept, total = 0, 0 + tmp_fd, tmp_path = tempfile.mkstemp(dir=os.path.dirname(input_path)) + try: + with open(input_path, encoding="utf-8") as fin, open(tmp_fd, "w", encoding="utf-8") as fout: + for line in fin: + line = line.strip() + if not line: + continue + total += 1 + try: + obj = json.loads(line) + except json.JSONDecodeError: + continue + messages = obj.get("messages", []) + short = any(len(m.get("content", "")) < MIN_LEN for m in messages) + if not short: + fout.write(line + "\n") + kept += 1 + shutil.move(tmp_path, input_path) + except Exception: + if os.path.exists(tmp_path): + os.unlink(tmp_path) + raise + return kept, total + + +def filter_pretrain(input_path: str) -> tuple[int, int]: + """Filter pretrain JSONL (text format), remove if text < MIN_LEN chars.""" + kept, total = 0, 0 + tmp_fd, tmp_path = tempfile.mkstemp(dir=os.path.dirname(input_path)) + try: + with open(input_path, encoding="utf-8") as fin, open(tmp_fd, "w", encoding="utf-8") as fout: + for line in fin: + line = line.strip() + if not line: + continue + total += 1 + try: + obj = json.loads(line) + except json.JSONDecodeError: + continue + text = obj.get("text", "") + if len(text) >= MIN_LEN: + fout.write(line + "\n") + kept += 1 + shutil.move(tmp_path, input_path) + except Exception: + if os.path.exists(tmp_path): + os.unlink(tmp_path) + raise + return kept, total + + +def main(): + parser = argparse.ArgumentParser(description="Filter short samples from JSONL datasets") + parser.add_argument("input_dir", help="Directory containing JSONL files") + parser.add_argument("--type", choices=["sft", "pt"], required=True, help="Dataset type") + args = parser.parse_args() + + from pipeline import FileScanner + + jsonl_files = FileScanner.scan(args.input_dir, suffix=".jsonl") + if not jsonl_files: + print(f"No JSONL files found in {args.input_dir}") + return + + filter_fn = filter_sft if args.type == "sft" else filter_pretrain + + total_kept, total_lines = 0, 0 + for fpath in jsonl_files: + kept, lines = filter_fn(fpath) + total_kept += kept + total_lines += lines + removed = lines - kept + print(f" {os.path.basename(fpath)}: {lines} -> {kept} (removed {removed})") + + print(f"\nTotal: {total_lines} -> {total_kept} (removed {total_lines - total_kept})") + + +if __name__ == "__main__": + main() diff --git a/scripts/pre_train/chinese-cosmopedia.py b/scripts/pre_train/chinese-cosmopedia.py index 8a1e769..01f64b3 100644 --- a/scripts/pre_train/chinese-cosmopedia.py +++ b/scripts/pre_train/chinese-cosmopedia.py @@ -4,7 +4,6 @@ from pipeline import export_dataset if __name__ == "__main__": dataset = load_dataset( "opencsg/chinese-cosmopedia", - data_files={"train": [f"data/000{i:02d}.parquet" for i in range(25)]}, ) export_dataset( dataset=dataset["train"], diff --git a/scripts/pre_train/english-dolma.py b/scripts/pre_train/english-dolma.py new file mode 100644 index 0000000..7e82b3c --- /dev/null +++ b/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", + ) diff --git a/scripts/pre_train/english-wiki.py b/scripts/pre_train/english-wiki.py index fd54ed4..fe4cca2 100644 --- a/scripts/pre_train/english-wiki.py +++ b/scripts/pre_train/english-wiki.py @@ -7,5 +7,4 @@ if __name__ == "__main__": dataset=dataset["train"], output_dir="./dataset", output_prefix="english-wiki-pretrain", - max_chunks=5, ) diff --git a/scripts/pre_train/ultra-fineweb-l3-en-qa-synthetic.py b/scripts/pre_train/ultra-fineweb-l3-en-qa-synthetic.py new file mode 100644 index 0000000..eab46eb --- /dev/null +++ b/scripts/pre_train/ultra-fineweb-l3-en-qa-synthetic.py @@ -0,0 +1,20 @@ +from datasets import load_dataset +from pipeline import export_dataset + + +def process_func(input_dict: dict): + return {"text": input_dict["content"]} + + +if __name__ == "__main__": + dataset = load_dataset( + "openbmb/Ultra-FineWeb-L3", + "Ultra-FineWeb-L3-en-QA-Synthetic", + split="train", + ) + export_dataset( + dataset=dataset, + output_dir="./dataset", + output_prefix="ultra-fineweb-l3-en-qa-synthetic-pretrain", + process_func=process_func, + ) diff --git a/scripts/pre_train/ultra-fineweb-l3-zh-qa-synthetic.py b/scripts/pre_train/ultra-fineweb-l3-zh-qa-synthetic.py new file mode 100644 index 0000000..e7b005f --- /dev/null +++ b/scripts/pre_train/ultra-fineweb-l3-zh-qa-synthetic.py @@ -0,0 +1,20 @@ +from datasets import load_dataset +from pipeline import export_dataset + + +def process_func(input_dict: dict): + return {"text": input_dict["content"]} + + +if __name__ == "__main__": + dataset = load_dataset( + "openbmb/Ultra-FineWeb-L3", + "Ultra-FineWeb-L3-zh-QA-Synthetic", + split="train", + ) + export_dataset( + dataset=dataset, + output_dir="./dataset", + output_prefix="ultra-fineweb-l3-zh-qa-synthetic-pretrain", + process_func=process_func, + ) diff --git a/scripts/supervised_finetuning/sft_alpaca_gpt4.py b/scripts/supervised_finetuning/sft_alpaca_gpt4.py index a4ad21c..0ae2def 100644 --- a/scripts/supervised_finetuning/sft_alpaca_gpt4.py +++ b/scripts/supervised_finetuning/sft_alpaca_gpt4.py @@ -6,10 +6,13 @@ def process_func(input_dict: dict): instruction = input_dict["instruction"] inp = input_dict.get("input", "") if inp: - query = instruction + "\n" + inp + content = instruction + "\n" + inp else: - query = instruction - return {"query": query, "response": input_dict["output"]} + content = instruction + return {"messages": [ + {"role": "user", "content": content}, + {"role": "assistant", "content": input_dict["output"]}, + ]} if __name__ == "__main__": diff --git a/scripts/supervised_finetuning/sft_alpaca_gpt4_zh.py b/scripts/supervised_finetuning/sft_alpaca_gpt4_zh.py new file mode 100644 index 0000000..a04205b --- /dev/null +++ b/scripts/supervised_finetuning/sft_alpaca_gpt4_zh.py @@ -0,0 +1,25 @@ +from datasets import load_dataset +from pipeline import export_dataset + + +def process_func(input_dict: dict): + instruction = input_dict["instruction"] + inp = input_dict.get("input", "") + if inp: + content = instruction + "\n" + inp + else: + content = instruction + return {"messages": [ + {"role": "user", "content": content}, + {"role": "assistant", "content": input_dict["output"]}, + ]} + + +if __name__ == "__main__": + dataset = load_dataset("llm-wizard/alpaca-gpt4-data-zh") + export_dataset( + dataset=dataset["train"], + output_dir="./dataset", + output_prefix="alpaca-gpt4-data-zh", + process_func=process_func, + ) diff --git a/scripts/supervised_finetuning/sft_belle.py b/scripts/supervised_finetuning/sft_belle.py new file mode 100644 index 0000000..f92dac5 --- /dev/null +++ b/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, + ) diff --git a/scripts/supervised_finetuning/sft_firefly-1.1m-rephrased.py b/scripts/supervised_finetuning/sft_firefly-1.1m-rephrased.py deleted file mode 100644 index 69b7c20..0000000 --- a/scripts/supervised_finetuning/sft_firefly-1.1m-rephrased.py +++ /dev/null @@ -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, - ) diff --git a/scripts/supervised_finetuning/sft_magicoder.py b/scripts/supervised_finetuning/sft_magicoder.py index be029e5..f6c59fe 100644 --- a/scripts/supervised_finetuning/sft_magicoder.py +++ b/scripts/supervised_finetuning/sft_magicoder.py @@ -3,7 +3,10 @@ from pipeline import export_dataset def process_func(input_dict: dict): - return {"query": input_dict["instruction"], "response": input_dict["response"]} + return {"messages": [ + {"role": "user", "content": input_dict["instruction"]}, + {"role": "assistant", "content": input_dict["response"]}, + ]} if __name__ == "__main__": diff --git a/scripts/supervised_finetuning/sft_metamathqa.py b/scripts/supervised_finetuning/sft_metamathqa.py new file mode 100644 index 0000000..0e8aec1 --- /dev/null +++ b/scripts/supervised_finetuning/sft_metamathqa.py @@ -0,0 +1,20 @@ +from datasets import load_dataset +from pipeline import export_dataset + + +def process_func(sample: dict) -> dict: + return {"messages": [ + {"role": "user", "content": sample["query"]}, + {"role": "assistant", "content": sample["response"]}, + ]} + + +if __name__ == "__main__": + dataset = load_dataset("meta-math/MetaMathQA", split="train") + export_dataset( + dataset=dataset, + output_dir="./dataset", + output_prefix="MetaMathQA", + process_func=process_func, + chunk_size=1_000_000, + ) diff --git a/scripts/supervised_finetuning/sft_openhermes.py b/scripts/supervised_finetuning/sft_openhermes.py new file mode 100644 index 0000000..c1f40ae --- /dev/null +++ b/scripts/supervised_finetuning/sft_openhermes.py @@ -0,0 +1,39 @@ +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(idx, len(conversations) - 1, 2): + user_msg = conversations[i] + assistant_msg = conversations[i + 1] + messages = system_msgs + [ + {"role": ROLE_MAP[user_msg["from"]], "content": user_msg["value"]}, + {"role": ROLE_MAP[assistant_msg["from"]], "content": assistant_msg["value"]}, + ] + examples.append({"messages": messages}) + return examples + + +if __name__ == "__main__": + dataset = load_dataset("teknium/OpenHermes-2.5") + export_dataset( + dataset=dataset["train"], + output_dir="./dataset", + output_prefix="OpenHermes-2.5", + process_func=process_func, + ) diff --git a/tests/test_cache.py b/tests/test_cache.py index f29f59b..82b7e8d 100644 --- a/tests/test_cache.py +++ b/tests/test_cache.py @@ -130,7 +130,7 @@ 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: diff --git a/tests/test_io.py b/tests/test_io.py index a496722..9368d23 100644 --- a/tests/test_io.py +++ b/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( - self, messages, add_generation_prompt=True, tokenize=True - ): + 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" + 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) diff --git a/tests/test_packing.py b/tests/test_packing.py index 5390d2d..89810fc 100644 --- a/tests/test_packing.py +++ b/tests/test_packing.py @@ -2,12 +2,147 @@ import pytest import torch -from pipeline.packing import SequencePacker +from pipeline.packing import ( + GreedyPacker, + FfDPacker, + BfdPacker, + pack_tensors, +) -class TestSequencePacker: - def test_normal_packing(self): - packer = SequencePacker(pack_size=10, pad_value=0) +class TestBfdPacker: + def test_best_fit_tight(self): + packer = BfdPacker(pack_size=10, pad_value=-1) + sequences = [ + torch.tensor([5, 6], dtype=torch.int32), + torch.tensor([1, 2, 3, 4], dtype=torch.int32), + torch.tensor([5, 6, 7, 8], dtype=torch.int32), + ] + packages = packer.pack(sequences) + assert len(packages) == 1 + assert packages[0].tolist() == [1, 2, 3, 4, 5, 6, 7, 8, 5, 6] + + def test_different_dtypes(self): + for dtype in [torch.int32, torch.int64, torch.float32]: + packer = BfdPacker(pack_size=10, dtype=dtype) + val = 1.0 if dtype == torch.float32 else 1 + packages = packer.pack([torch.tensor([val, 2, 3], dtype=dtype)]) + assert packages[0].dtype == dtype + + def test_dtype_conversion_on_mismatch(self): + packer = BfdPacker(pack_size=10, dtype=torch.int32) + packages = packer.pack([torch.tensor([1, 2, 3], dtype=torch.int64)]) + assert packages[0].dtype == torch.int32 + assert packages[0][:3].tolist() == [1, 2, 3] + + def test_non_1d_tensor_raises_error(self): + packer = BfdPacker(pack_size=10) + with pytest.raises(ValueError, match="Expected 1D tensor"): + packer.pack([torch.tensor([[1, 2], [3, 4]])]) + with pytest.raises(ValueError, match="Expected 1D tensor"): + packer.pack([torch.tensor(5)]) + + def test_empty_input(self): + packer = BfdPacker(pack_size=10) + assert packer.pack([]) == [] + + def test_reset(self): + packer = BfdPacker(pack_size=10) + packer.pack([torch.tensor([1, 2, 3], dtype=torch.int32)]) + assert len(packer._bins) == 1 + packer.reset() + assert len(packer._bins) == 0 + + def test_overlong_sample_truncated(self): + """Overlong sample is truncated to pack_size.""" + packer = BfdPacker(pack_size=6, pad_value=-1) + packages = packer.pack( + [ + torch.tensor([1, 2, 3, 4, 5, 6, 7], dtype=torch.int32), + torch.tensor([8, 9], dtype=torch.int32), + ] + ) + assert len(packages) == 2 + assert packages[0].tolist() == [1, 2, 3, 4, 5, 6] + assert packages[1].tolist() == [8, 9, -1, -1, -1, -1] + + def test_uses_two_bins_when_needed(self): + packer = BfdPacker(pack_size=10, pad_value=0) + sequences = [ + torch.tensor([1, 2, 3], dtype=torch.int32), + torch.tensor([4, 5, 6, 7], dtype=torch.int32), + torch.tensor([8, 9, 10], dtype=torch.int32), + torch.tensor([11, 12, 13, 14, 15, 16], dtype=torch.int32), + ] + packages = packer.pack(sequences) + assert len(packages) == 2 + for pkg in packages: + assert pkg.shape == (10,) + + def test_minimizes_waste_vs_ffd(self): + sequences = [ + torch.tensor([6] * i, dtype=torch.int32) + for i in [3, 5, 5, 7, 2, 4, 1, 4, 6, 2] + ] + bfd = BfdPacker(pack_size=10, pad_value=0) + ffd = FfDPacker(pack_size=10, pad_value=0) + assert len(bfd.pack(sequences)) <= len(ffd.pack(sequences)) + + +class TestFfDPacker: + def test_fills_tightly(self): + packer = FfDPacker(pack_size=10, pad_value=0) + sequences = [ + torch.tensor([7, 8], dtype=torch.int32), + torch.tensor([1, 2, 3, 4, 5, 6], dtype=torch.int32), + torch.tensor([9, 10], dtype=torch.int32), + ] + packages = packer.pack(sequences) + assert len(packages) == 1 + + def test_overlong_sample_truncated(self): + packer = FfDPacker(pack_size=5, pad_value=0) + packages = packer.pack( + [torch.tensor([1, 2, 3, 4, 5, 6], dtype=torch.int32)] + ) + assert len(packages) == 1 + assert packages[0].tolist() == [1, 2, 3, 4, 5] + + def test_sort_descending_order(self): + packer = FfDPacker(pack_size=10, pad_value=-1) + sequences = [ + torch.tensor([1, 2], dtype=torch.int32), + torch.tensor([3, 4, 5, 6, 7, 8], dtype=torch.int32), + torch.tensor([9, 10], dtype=torch.int32), + ] + packages = packer.pack(sequences) + assert len(packages) == 1 + assert packages[0].tolist() == [3, 4, 5, 6, 7, 8, 1, 2, 9, 10] + + def test_reduces_bins_vs_greedy(self): + sequences = [ + torch.tensor([6] * i, dtype=torch.int32) + for i in [3, 8, 2, 7, 1, 4, 5, 3, 2, 6] + ] + greedy = GreedyPacker(pack_size=10, pad_value=0) + ffd = FfDPacker(pack_size=10, pad_value=0) + assert len(ffd.pack(sequences)) <= len(greedy.pack(sequences)) + + def test_reset(self): + packer = FfDPacker(pack_size=10) + packer.pack([torch.tensor([1, 2, 3], dtype=torch.int32)]) + assert len(packer._bins) == 1 + packer.reset() + assert len(packer._bins) == 0 + + def test_empty_input(self): + packer = FfDPacker(pack_size=10) + assert packer.pack([]) == [] + + +class TestGreedyPacker: + def test_basic_packing(self): + packer = GreedyPacker(pack_size=10, pad_value=0) sequences = [ torch.tensor([1, 2, 3], dtype=torch.int32), torch.tensor([4, 5], dtype=torch.int32), @@ -15,154 +150,186 @@ class TestSequencePacker: ] packages = packer.pack(sequences) assert len(packages) == 1 - for pkg in packages: - assert pkg.shape == (10,) - - # Verify all original values are present in order + assert packages[0].shape == (10,) assert packages[0][:9].tolist() == [1, 2, 3, 4, 5, 6, 7, 8, 9] - assert packages[0][9] == 0 # padding + assert packages[0][9] == 0 - def test_empty_list_input(self): - packer = SequencePacker(pack_size=10) - assert packer.pack([]) == [] - - def test_single_sequence_input(self): - packer = SequencePacker(pack_size=10, pad_value=-1) - packages = packer.pack([torch.tensor([1, 2, 3], dtype=torch.int32)]) - assert len(packages) == 1 - assert packages[0][:3].tolist() == [1, 2, 3] - assert packages[0][3:].tolist() == [-1] * 7 - - def test_long_sequence_split_across_chunks(self): - """Sequences longer than pack_size are split across multiple chunks.""" - packer = SequencePacker(pack_size=5, pad_value=0) + def test_overlong_sample_truncated(self): + """Overlong sample is truncated to pack_size.""" + packer = GreedyPacker(pack_size=5, pad_value=0) packages = packer.pack( [torch.tensor([1, 2, 3, 4, 5, 6, 7, 8], dtype=torch.int32)] ) - assert len(packages) == 2 + assert len(packages) == 1 assert packages[0].tolist() == [1, 2, 3, 4, 5] - assert packages[1].tolist() == [6, 7, 8, 0, 0] - def test_padding_value(self): - packer = SequencePacker(pack_size=8, pad_value=99) - packages = packer.pack( - [ - torch.tensor([1, 2], dtype=torch.int32), - torch.tensor([3], dtype=torch.int32), - ] - ) - assert packages[0][:3].tolist() == [1, 2, 3] - assert packages[0][3:].tolist() == [99] * 5 - - def test_different_dtypes(self): - for dtype in [torch.int32, torch.int64, torch.float32]: - packer = SequencePacker(pack_size=10, dtype=dtype) - val = 1.0 if dtype == torch.float32 else 1 - packages = packer.pack([torch.tensor([val, 2, 3], dtype=dtype)]) - assert packages[0].dtype == dtype - - def test_dtype_conversion_on_mismatch(self, caplog): - """Tensors with mismatched dtype are silently converted.""" - packer = SequencePacker(pack_size=10, dtype=torch.int32) - packages = packer.pack([torch.tensor([1, 2, 3], dtype=torch.int64)]) - assert packages[0].dtype == torch.int32 - assert packages[0][:3].tolist() == [1, 2, 3] - - def test_non_1d_tensor_raises_error(self): - packer = SequencePacker(pack_size=10) - with pytest.raises(ValueError, match="Expected 1D tensor"): - packer.pack([torch.tensor([[1, 2], [3, 4]])]) - with pytest.raises(ValueError, match="Expected 1D tensor"): - packer.pack([torch.tensor(5)]) - - def test_input_list_not_modified(self): - packer = SequencePacker(pack_size=10) - original = [ - torch.tensor([3], dtype=torch.int32), + def test_multiple_fill(self): + packer = GreedyPacker(pack_size=6, pad_value=0) + sequences = [ torch.tensor([1, 2], dtype=torch.int32), - torch.tensor([4, 5, 6, 7], dtype=torch.int32), + torch.tensor([3, 4], dtype=torch.int32), + torch.tensor([5, 6], dtype=torch.int32), + torch.tensor([7], dtype=torch.int32), ] - original_repr = [seq.tolist() for seq in original] - packer.pack(original) - assert [seq.tolist() for seq in original] == original_repr - - def test_exact_pack_size_fit(self): - packer = SequencePacker(pack_size=5, pad_value=0) - packages = packer.pack( - [ - torch.tensor([1, 2, 3, 4, 5], dtype=torch.int32), - torch.tensor([6, 7, 8, 9, 10], dtype=torch.int32), - ] - ) - assert len(packages) == 2 - assert packages[0].tolist() == [1, 2, 3, 4, 5] - assert packages[1].tolist() == [6, 7, 8, 9, 10] - - def test_multiple_packs_full_utilization(self): - packer = SequencePacker(pack_size=10, pad_value=-1) - sequences = [torch.tensor([i], dtype=torch.int32) for i in range(1, 12)] packages = packer.pack(sequences) assert len(packages) == 2 - assert packages[0].tolist() == list(range(1, 11)) - assert packages[1].tolist() == [11] + [-1] * 9 + for pkg in packages: + assert pkg.shape == (6,) - def test_cross_group_ordering(self): - """Separate packers for different dtypes produce identical chunk boundaries.""" - seq_packer = SequencePacker(pack_size=10, pad_value=0, dtype=torch.int32) - mask_packer = SequencePacker(pack_size=10, pad_value=False, dtype=torch.bool) - # sequences: lengths [3, 1, 4] - seqs = [ - torch.tensor([1, 2, 3], dtype=torch.int32), - torch.tensor([10], dtype=torch.int32), - torch.tensor([4, 5, 6, 7], dtype=torch.int32), - ] - masks = [ - torch.tensor([False, False, True], dtype=torch.bool), - torch.tensor([False], dtype=torch.bool), - torch.tensor([False, False, False, True], dtype=torch.bool), - ] - packed_seqs = seq_packer.pack(seqs) - packed_masks = mask_packer.pack(masks) - - # Verify mask packer uses bool dtype - assert packed_masks[0].dtype == torch.bool - # Both groups should produce the same number of packages - assert len(packed_seqs) == len(packed_masks) - - def test_stream_split_across_chunks(self): - """Sequences are split across chunks in streaming mode.""" - packer = SequencePacker(pack_size=5, pad_value=0) - packages = packer.pack( - [ - torch.tensor([1, 2, 3], dtype=torch.int32), - torch.tensor([4, 5, 6, 7, 8], dtype=torch.int32), - ] - ) - assert len(packages) == 2 - # First chunk: [1, 2, 3, 4, 5] — first seq + part of second - assert packages[0].tolist() == [1, 2, 3, 4, 5] - # Second chunk: [6, 7, 8, 0, 0] — rest of second + padding - assert packages[1].tolist() == [6, 7, 8, 0, 0] - - def test_reset_method(self): - packer = SequencePacker(pack_size=10, pad_value=0) - seqs = [torch.tensor([1, 2, 3], dtype=torch.int32)] - packer.pack(seqs) - assert len(packer._packages) == 1 + def test_reset(self): + packer = GreedyPacker(pack_size=10) + packer.pack([torch.tensor([1, 2, 3], dtype=torch.int32)]) + assert len(packer._bins) == 1 packer.reset() - assert len(packer._packages) == 0 - assert packer._pos == 0 - assert packer._buffer == [] + assert len(packer._bins) == 0 - def test_no_sorting_needed(self): - """Streaming concat preserves input order, no sorting.""" - packer = SequencePacker(pack_size=4, pad_value=-1) - # short then long (fits in 2 chunks) - packages = packer.pack( - [ - torch.tensor([1], dtype=torch.int32), - torch.tensor([2, 3, 4, 5, 6, 7], dtype=torch.int32), - ] + def test_empty_input(self): + packer = GreedyPacker(pack_size=10) + assert packer.pack([]) == [] + + +class TestPackTensors: + def test_default_is_bfd(self): + result = pack_tensors( + tensors={ + "input_ids": [ + torch.tensor([1, 2], dtype=torch.int32), + torch.tensor([3, 4], dtype=torch.int32), + torch.tensor([5], dtype=torch.int32), + ], + }, + pack_size=5, + pad_value=0, ) - assert packages[0].tolist() == [1, 2, 3, 4] - assert packages[1].tolist() == [5, 6, 7, -1] + assert result["input_ids"][0].tolist() == [1, 2, 3, 4, 5] + + def test_greedy(self): + result = pack_tensors( + tensors={ + "input_ids": [ + torch.tensor([1, 2, 3], dtype=torch.int32), + torch.tensor([4, 5], dtype=torch.int32), + ], + }, + pack_size=5, + pad_value=0, + algo="greedy", + ) + assert result["input_ids"][0].tolist() == [1, 2, 3, 4, 5] + + def test_ffd(self): + result = pack_tensors( + tensors={ + "input_ids": [ + torch.tensor([1], dtype=torch.int32), + torch.tensor([2, 3, 4], dtype=torch.int32), + torch.tensor([5], dtype=torch.int32), + ], + }, + pack_size=5, + pad_value=0, + algo="ffd", + ) + assert result["input_ids"][0].tolist() == [2, 3, 4, 1, 5] + + def test_bfd_explicit(self): + result = pack_tensors( + tensors={ + "input_ids": [ + torch.tensor([1, 2], dtype=torch.int32), + torch.tensor([3, 4], dtype=torch.int32), + torch.tensor([5], dtype=torch.int32), + ], + }, + pack_size=5, + pad_value=0, + algo="bfd", + ) + assert result["input_ids"][0].tolist() == [1, 2, 3, 4, 5] + + def test_unknown_algo_raises(self): + with pytest.raises(ValueError, match="Unknown packing algorithm"): + pack_tensors( + tensors={"input_ids": [torch.tensor([1, 2, 3])]}, + pack_size=10, + pad_value=0, + algo="unknown_algo", + ) + + +class TestPositionIdsPacking: + """Verify position_ids reset to zero at sample boundaries after packing.""" + + def test_position_ids_reset_in_packed_chunk(self): + """After packing multiple SFT samples, position_ids restart from 0 at each boundary.""" + seqs = [ + torch.tensor([0, 1, 2, 3, 4], dtype=torch.int32), # len=5 + torch.tensor([0, 1, 2], dtype=torch.int32), # len=3 + torch.tensor([0, 1, 2, 3, 4, 5, 6], dtype=torch.int32), # len=7 + ] + result = pack_tensors( + tensors={"position_ids": seqs}, + pack_size=16, + pad_value=-1, + algo="greedy", + ) + packed = result["position_ids"][0].tolist() + assert packed == [0, 1, 2, 3, 4, 0, 1, 2, 0, 1, 2, 3, 4, 5, 6, -1] + + def test_position_ids_reset_with_bfd(self): + """BFD may reorder, but each sample's position_ids still start from 0.""" + seqs = [ + torch.tensor([0, 1, 2], dtype=torch.int32), + torch.tensor([0, 1, 2, 3, 4, 5], dtype=torch.int32), + torch.tensor([0, 1, 2, 3], dtype=torch.int32), + ] + result = pack_tensors( + tensors={"position_ids": seqs}, + pack_size=16, + pad_value=-1, + algo="bfd", + ) + packed = result["position_ids"][0].tolist() + assert packed[0] == 0 + zeros = [i for i, v in enumerate(packed) if v == 0 and (i == 0 or packed[i - 1] != 0)] + assert len(zeros) == 3 + + def test_multiple_keys_share_same_boundaries(self): + """sequence, loss_mask, position_ids share identical chunk boundaries after packing.""" + seq_a = torch.tensor([101, 102, 103, 104], dtype=torch.int32) + seq_b = torch.tensor([201, 202, 203, 204, 205, 206, 207], dtype=torch.int32) + seq_c = torch.tensor([301, 302, 303, 304, 305], dtype=torch.int32) + + mask_a = torch.tensor([False, False, True, True], dtype=torch.bool) + mask_b = torch.tensor([False, False, False, False, True, True, True], dtype=torch.bool) + mask_c = torch.tensor([False, False, False, True, True], dtype=torch.bool) + + pos_a = torch.tensor([0, 1, 2, 3], dtype=torch.int32) + pos_b = torch.tensor([0, 1, 2, 3, 4, 5, 6], dtype=torch.int32) + pos_c = torch.tensor([0, 1, 2, 3, 4], dtype=torch.int32) + + result = pack_tensors( + tensors={ + "sequence": [seq_a, seq_b, seq_c], + "loss_mask": [mask_a, mask_b, mask_c], + "position_ids": [pos_a, pos_b, pos_c], + }, + pack_size=16, + pad_value=-1, + algo="greedy", + ) + + seq_chunk = result["sequence"][0] + mask_chunk = result["loss_mask"][0] + pos_chunk = result["position_ids"][0] + + assert len(seq_chunk) == len(mask_chunk) == len(pos_chunk) == 16 + + for i in range(16): + if seq_chunk[i] == -1: + assert mask_chunk[i] == -1 + assert pos_chunk[i] == -1 + + pos_ids = pos_chunk.tolist() + zeros = [i for i, v in enumerate(pos_ids) if v == 0] + assert len(zeros) == 3 diff --git a/tests/test_processors.py b/tests/test_processors.py index d4e3ea0..bf4225e 100644 --- a/tests/test_processors.py +++ b/tests/test_processors.py @@ -13,21 +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, 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( - self, messages, add_generation_prompt=True, tokenize=True - ): + 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" + text += f"<｜im▁start｜>{m['role']}\n{m['content']}<｜im▁end｜>\n" if add_generation_prompt: - text += "<|im_start|>assistant\n" + text += "<｜im▁start｜>assistant\n" return self.encode(text) if tokenize else text diff --git a/tests/test_strategies.py b/tests/test_strategies.py index d4600b5..809a7ee 100644 --- a/tests/test_strategies.py +++ b/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 - assert ( - prompt[-len(strategy._assistant_start_ids) :] - == strategy._assistant_start_ids - ) + assistant_start = tk.encode("<｜im▁start｜>assistant\n") + assert prompt[-len(assistant_start):] == assistant_start class TestAlpacaStrategy: