diff --git a/pipeline/io/__init__.py b/pipeline/io/__init__.py index f475f8e..9caabc4 100644 --- a/pipeline/io/__init__.py +++ b/pipeline/io/__init__.py @@ -4,16 +4,26 @@ This module provides: - FileScanner: File and directory scanning utilities - HDF5Handler: Tensor data persistence - export_dataset: HuggingFace Dataset to JSONL export -- cache_jsonl: JSONL to HDF5 tokenization and caching +- cache_jsonl: JSONL to HDF5/binary tokenization and caching +- dedup_jsonl: MinHash+LSH deduplication for pretraining text +- writers: BaseWriter / H5Writer / BinWriter / TextWriter (Strategy + Factory) """ from pipeline.io.file_scanner import FileScanner from pipeline.io.hdf5_handler import HDF5Handler from pipeline.io.export import export_dataset, cache_jsonl +from pipeline.io.dedup import dedup_jsonl +from pipeline.io.writers import BaseWriter, H5Writer, BinWriter, TextWriter, create_writer __all__ = [ "FileScanner", "HDF5Handler", "export_dataset", "cache_jsonl", + "dedup_jsonl", + "BaseWriter", + "H5Writer", + "BinWriter", + "TextWriter", + "create_writer", ] diff --git a/pipeline/io/dedup.py b/pipeline/io/dedup.py new file mode 100644 index 0000000..7e8f6a2 --- /dev/null +++ b/pipeline/io/dedup.py @@ -0,0 +1,167 @@ +"""MinHash + LSH deduplication for pretraining text data.""" + +import json +import logging +import os +from pathlib import Path +from typing import Iterator, List, Set, Tuple + +from tqdm import tqdm + +from pipeline.io.writers import TextWriter +from pipeline.utils import error_handler + +logger = logging.getLogger(__name__) + + +def _tokenize(text: str, ngram: int = 3) -> Set[str]: + return {text[i : i + ngram] for i in range(len(text) - ngram + 1)} + + +def _iter_docs(input_dir: Path) -> Iterator[Tuple[str, dict]]: + for fpath in sorted(input_dir.glob("*.jsonl")): + with open(fpath, encoding="utf-8") as f: + for line in f: + line = line.strip() + if not line: + continue + record = json.loads(line) + text = record.get("text", "") + if text: + yield text, record + + +def _write_h5(records: List[dict], output_dir: str, chunk_idx: int): + import h5py + + fname = os.path.join(output_dir, f"chunk_{chunk_idx}.h5") + texts = [rec.get("text", "") for rec in records] + with h5py.File(fname, "w") as f: + dt = h5py.special_dtype(vlen=str) + ds = f.create_dataset("text", (len(texts),), dtype=dt) + for i, t in enumerate(texts): + ds[i] = t + + +def _write_bin(records: List[dict], output_dir: Path, chunk_idx: int): + output_dir.mkdir(parents=True, exist_ok=True) + texts = [rec.get("text", "") + "\n" for rec in records] + + meta = {"chunk": chunk_idx, "count": len(texts), "format": "text", "encoding": "utf-8"} + meta_path = output_dir / "meta.json" + existing = json.loads(meta_path.read_text()) if meta_path.exists() else {} + existing[str(chunk_idx)] = meta + meta_path.write_text(json.dumps(existing, indent=2)) + + (output_dir / f"text_{chunk_idx}.bin").write_bytes("".join(texts).encode("utf-8")) + + +_WRITERS = { + "jsonl": TextWriter, + "h5": lambda: None, # handled inline below + "bin": lambda: None, +} + + +@error_handler() +def dedup_jsonl( + input_dir: str, + output_dir: str, + *, + threshold: float = 0.8, + num_perm: int = 128, + ngram: int = 3, + output_format: str = "jsonl", + chunk_size: int = 1_000_000, +) -> Tuple[int, int]: + """Deduplicate JSONL text files using MinHash + LSH. + + Args: + input_dir: Directory with source ``*.jsonl`` files. + output_dir: Directory for deduplicated output. + threshold: Jaccard similarity threshold (0–1). + num_perm: Number of MinHash permutations. + ngram: Character n-gram size. + output_format: ``"jsonl"``, ``"h5"``, or ``"bin"``. + chunk_size: Records per output chunk file. + + Returns: + ``(kept, removed)`` counts. + """ + from datasketch import MinHash, MinHashLSH + + input_path = Path(input_dir) + output_path = Path(output_dir) + output_path.mkdir(parents=True, exist_ok=True) + + logger.info( + f"Deduplicating {input_dir} -> {output_dir} " + f"(threshold={threshold}, perm={num_perm}, fmt={output_format})" + ) + + lsh = MinHashLSH(threshold=threshold, num_perm=num_perm) + + kept = 0 + removed = 0 + + dup_doc_ids: Set[int] = set() + for doc_id, (text, _record) in enumerate(tqdm(_iter_docs(input_path), desc="indexing", unit="docs")): + shingles = _tokenize(text, ngram=ngram) + if len(shingles) < ngram * 2: + dup_doc_ids.add(doc_id) + continue + + m = MinHash(num_perm=num_perm) + for s in shingles: + m.update(s.encode("utf-8")) + + if lsh.query(m): + dup_doc_ids.add(doc_id) + else: + lsh.insert(doc_id, m) + + logger.info(f"Found {len(dup_doc_ids)} duplicates, writing deduplicated data") + + buffer: List[dict] = [] + chunk_idx = 0 + writer = TextWriter(chunk_size) if output_format == "jsonl" else None + + for doc_id, (_text, record) in enumerate(tqdm(_iter_docs(input_path), desc="writing", unit="docs")): + if doc_id in dup_doc_ids: + removed += 1 + continue + + kept += 1 + buffer.append(record) + + if len(buffer) >= chunk_size: + _flush_chunk(buffer, output_path, chunk_idx, output_format, writer) + chunk_idx += 1 + buffer = [] + + if buffer: + _flush_chunk(buffer, output_path, chunk_idx, output_format, writer) + + if writer: + writer.flush(output_path) + + logger.info(f"Done. kept={kept}, removed={removed}") + return kept, removed + + +def _flush_chunk( + records: List[dict], + output_dir: Path, + chunk_idx: int, + output_format: str, + writer=None, +): + if output_format == "jsonl": + for rec in records: + writer.write_record(rec, output_dir) + elif output_format == "h5": + _write_h5(records, str(output_dir), chunk_idx) + elif output_format == "bin": + _write_bin(records, output_dir, chunk_idx) + else: + raise ValueError(f"Unknown output format: {output_format}") diff --git a/pipeline/io/export.py b/pipeline/io/export.py index c955c27..ca978e1 100644 --- a/pipeline/io/export.py +++ b/pipeline/io/export.py @@ -13,6 +13,7 @@ from tqdm import tqdm from pipeline.io.file_scanner import FileScanner from pipeline.io.hdf5_handler import HDF5Handler +from pipeline.io.writers import create_writer, BaseWriter from pipeline.processors import BaseProcessor from pipeline.packing import pack_tensors, BasePacker from pipeline.utils import error_handler @@ -114,15 +115,16 @@ def cache_jsonl( pad_value: int = 0, group_size: int = 1_000, pack_algo: Optional[str] = None, + output_format: str = "h5", ) -> List[str]: - """Tokenize JSONL files and pack them into HDF5 storage. + """Tokenize JSONL files and save as HDF5 or binary. BFD packs in group_size-bounded batches to avoid O(N²), then all - packed chunks are merged and saved as one HDF5 file per input file. + packed chunks are merged and saved as one file per input file. Args: files: List of JSONL file paths. - output_dir: H5 output directory. + output_dir: Output directory. processor: Initialized Processor instance. pack_size: Packing length, <=0 means no packing. pad_value: Padding value. @@ -130,9 +132,10 @@ def cache_jsonl( packing batch) and merge granularity, <=0 means no merging. pack_algo: Packing algorithm: 'bfd' (default), 'ffd', 'greedy'. Only used when pack_size > 0. + output_format: ``"h5"`` or ``"bin"``. Returns: - List of generated H5 file paths. + List of generated file paths. """ os.makedirs(output_dir, exist_ok=True) output_files: List[str] = [] @@ -206,8 +209,9 @@ def cache_jsonl( else: output = all_packed - h5_path = HDF5Handler.save(output_dir, file_name, output) - output_files.append(h5_path) - logger.info(f"Saved {h5_path}") + writer: BaseWriter = create_writer(output_format) + saved = writer.save(output_dir, file_name, output) + output_files.append(saved) + logger.info(f"Saved {saved}") return output_files diff --git a/pipeline/io/writers.py b/pipeline/io/writers.py new file mode 100644 index 0000000..6ccf6d4 --- /dev/null +++ b/pipeline/io/writers.py @@ -0,0 +1,107 @@ +"""Storage backends for tensor / text output (Strategy + Factory). + +Each backend implements a common ``save()`` interface so callers use +polymorphism instead of ``if fmt == "h5" ... elif fmt == "bin" ...``. + +Supports: + - **H5Writer**: HDF5 format (via HDF5Handler) + - **BinWriter**: binary format – meta.json + {key}.bin (memmap-compatible) + - **TextWriter**: raw JSONL text (for dedup output) +""" + +import json +import os +from abc import ABC, abstractmethod +from pathlib import Path +from typing import Dict, List + +import torch +from torch import Tensor + + +class BaseWriter(ABC): + """Abstract writer – call ``save(dir, name, data)`` without caring + about the underlying format.""" + + @abstractmethod + def save(self, output_dir: str, file_name: str, data: Dict[str, List[Tensor]]) -> str: + ... + + +class H5Writer(BaseWriter): + def save(self, output_dir: str, file_name: str, data: Dict[str, List[Tensor]]) -> str: + from pipeline.io.hdf5_handler import HDF5Handler + return HDF5Handler.save(output_dir, file_name, data) + + +class BinWriter(BaseWriter): + def save(self, output_dir: str, file_name: str, data: Dict[str, List[Tensor]]) -> str: + import numpy as np + + os.makedirs(output_dir, exist_ok=True) + sub_dir = os.path.join(output_dir, file_name) + os.makedirs(sub_dir, exist_ok=True) + + meta: Dict[str, Dict] = {} + for key, tensors in data.items(): + cat = torch.cat(tensors, dim=0) + meta[key] = {"shape": list(cat.shape), "dtype": str(cat.dtype).split(".")[-1]} + np.asarray(cat.cpu().numpy()).tofile(os.path.join(sub_dir, f"{key}.bin")) + + with open(os.path.join(sub_dir, "meta.json"), "w") as f: + json.dump(meta, f, indent=2) + + return sub_dir + + +class TextWriter(BaseWriter): + """Write raw text records as JSONL (used by dedup output).""" + + def __init__(self, chunk_size: int = 1_000_000): + self._chunk_size = chunk_size + self._buffer: List[dict] = [] + self._chunk_idx = 0 + + def save(self, output_dir: str, file_name: str, data: Dict[str, List[Tensor]]) -> str: + raise NotImplementedError("TextWriter.save_one is for tensor data; use write_record()") + + def write_record(self, record: dict, output_dir: Path): + self._buffer.append(record) + if len(self._buffer) >= self._chunk_size: + self._flush(output_dir) + + def flush(self, output_dir: Path): + if self._buffer: + self._flush(output_dir) + + def _flush(self, output_dir: Path): + output_dir.mkdir(parents=True, exist_ok=True) + fpath = output_dir / f"chunk_{self._chunk_idx}.jsonl" + with open(fpath, "w", encoding="utf-8") as f: + for rec in self._buffer: + f.write(json.dumps(rec, ensure_ascii=False) + "\n") + self._chunk_idx += 1 + self._buffer = [] + + +_WRITER_REGISTRY: Dict[str, type] = {} + + +def register_writer(name: str): + def decorator(cls): + _WRITER_REGISTRY[name] = cls + return cls + return decorator + + +def create_writer(name: str, **kwargs) -> BaseWriter: + cls = _WRITER_REGISTRY.get(name) + if cls is None: + raise ValueError(f"Unknown writer: {name}. Available: {list(_WRITER_REGISTRY)}") + return cls(**kwargs) + + +# Register built-in writers +register_writer("h5")(H5Writer) +register_writer("bin")(BinWriter) +register_writer("jsonl")(TextWriter) diff --git a/scripts/cache_h5.py b/scripts/cache_h5.py index 7b32bb1..fcae8c7 100644 --- a/scripts/cache_h5.py +++ b/scripts/cache_h5.py @@ -1,11 +1,10 @@ -"""JSONL to H5 caching script. +"""JSONL tokenization and caching script. -Tokenize JSONL files and pack them into HDF5 format. +Tokenize JSONL files and save as HDF5 or binary format. Usage: python scripts/cache_h5.py pt ./dataset/chinese-c4-pretrain - python scripts/cache_h5.py sft ./dataset/belle-sft --pack-size 4096 --strategy alpaca - python scripts/cache_h5.py sft ./dataset/Ling-Coder-sft --tokenizer ./my_tokenizer.json + python scripts/cache_h5.py sft ./dataset/belle-sft --pack-size 4096 --output-format bin """ import argparse @@ -29,7 +28,7 @@ def main(): "-o", "--output-dir", default=None, - help="H5 output dir (default: /cached)", + help="Output dir (default: /cached)", ) parser.add_argument( "-t", @@ -73,6 +72,13 @@ def main(): choices=["DEBUG", "INFO", "WARNING", "ERROR"], help="Logging level (default: INFO)", ) + parser.add_argument( + "-f", + "--output-format", + default="h5", + choices=["h5", "bin"], + help="Output format: h5 or bin (default: h5)", + ) args = parser.parse_args() # Initialize logging explicitly (not automatic anymore) @@ -126,6 +132,7 @@ def main(): pad_value=args.pad_value, group_size=args.group_size, pack_algo=args.pack_algo, + output_format=args.output_format, ) print(f"\nDone! Output saved to {output_dir}") diff --git a/scripts/dedup_pretrain.py b/scripts/dedup_pretrain.py new file mode 100644 index 0000000..44cd318 --- /dev/null +++ b/scripts/dedup_pretrain.py @@ -0,0 +1,36 @@ +"""MinHash + LSH deduplication CLI. + +Usage: + python scripts/dedup_pretrain.py --input-dir --output-dir --threshold 0.8 --num-perm 128 --output-format jsonl +""" + +import argparse + +from pipeline.io import dedup_jsonl + + +def main(): + parser = argparse.ArgumentParser(description="MinHash + LSH deduplication") + parser.add_argument("--input-dir", required=True) + parser.add_argument("--output-dir", required=True) + parser.add_argument("--threshold", type=float, default=0.8) + parser.add_argument("--num-perm", type=int, default=128) + parser.add_argument("--ngram", type=int, default=3) + parser.add_argument("--output-format", default="jsonl", choices=["jsonl", "h5", "bin"]) + args = parser.parse_args() + + kept, removed = dedup_jsonl( + input_dir=args.input_dir, + output_dir=args.output_dir, + threshold=args.threshold, + num_perm=args.num_perm, + ngram=args.ngram, + output_format=args.output_format, + ) + + total = kept + removed + print(f"kept={kept}, removed={removed} ({removed/max(total,1)*100:.1f}%)") + + +if __name__ == "__main__": + main()