merge remote main

pipeline/io/__init__.py

@@ -4,16 +4,26 @@ This module provides:
 - FileScanner: File and directory scanning utilities
 - HDF5Handler: Tensor data persistence
 - export_dataset: HuggingFace Dataset to JSONL export
-- cache_jsonl: JSONL to HDF5 tokenization and caching
+- cache_jsonl: JSONL to HDF5/binary tokenization and caching
+- dedup_jsonl: MinHash+LSH deduplication for pretraining text
+- writers: BaseWriter / H5Writer / BinWriter / TextWriter (Strategy + Factory)
 """
 
 from pipeline.io.file_scanner import FileScanner
 from pipeline.io.hdf5_handler import HDF5Handler
 from pipeline.io.export import export_dataset, cache_jsonl
+from pipeline.io.dedup import dedup_jsonl
+from pipeline.io.writers import BaseWriter, H5Writer, BinWriter, TextWriter, create_writer
 
 __all__ = [
     "FileScanner",
     "HDF5Handler",
     "export_dataset",
     "cache_jsonl",
+    "dedup_jsonl",
+    "BaseWriter",
+    "H5Writer",
+    "BinWriter",
+    "TextWriter",
+    "create_writer",
 ]

pipeline/io/dedup.py

@@ -0,0 +1,167 @@
+"""MinHash + LSH deduplication for pretraining text data."""
+
+import json
+import logging
+import os
+from pathlib import Path
+from typing import Iterator, List, Set, Tuple
+
+from tqdm import tqdm
+
+from pipeline.io.writers import TextWriter
+from pipeline.utils import error_handler
+
+logger = logging.getLogger(__name__)
+
+
+def _tokenize(text: str, ngram: int = 3) -> Set[str]:
+    return {text[i : i + ngram] for i in range(len(text) - ngram + 1)}
+
+
+def _iter_docs(input_dir: Path) -> Iterator[Tuple[str, dict]]:
+    for fpath in sorted(input_dir.glob("*.jsonl")):
+        with open(fpath, encoding="utf-8") as f:
+            for line in f:
+                line = line.strip()
+                if not line:
+                    continue
+                record = json.loads(line)
+                text = record.get("text", "")
+                if text:
+                    yield text, record
+
+
+def _write_h5(records: List[dict], output_dir: str, chunk_idx: int):
+    import h5py
+
+    fname = os.path.join(output_dir, f"chunk_{chunk_idx}.h5")
+    texts = [rec.get("text", "") for rec in records]
+    with h5py.File(fname, "w") as f:
+        dt = h5py.special_dtype(vlen=str)
+        ds = f.create_dataset("text", (len(texts),), dtype=dt)
+        for i, t in enumerate(texts):
+            ds[i] = t
+
+
+def _write_bin(records: List[dict], output_dir: Path, chunk_idx: int):
+    output_dir.mkdir(parents=True, exist_ok=True)
+    texts = [rec.get("text", "") + "\n" for rec in records]
+
+    meta = {"chunk": chunk_idx, "count": len(texts), "format": "text", "encoding": "utf-8"}
+    meta_path = output_dir / "meta.json"
+    existing = json.loads(meta_path.read_text()) if meta_path.exists() else {}
+    existing[str(chunk_idx)] = meta
+    meta_path.write_text(json.dumps(existing, indent=2))
+
+    (output_dir / f"text_{chunk_idx}.bin").write_bytes("".join(texts).encode("utf-8"))
+
+
+_WRITERS = {
+    "jsonl": TextWriter,
+    "h5": lambda: None,  # handled inline below
+    "bin": lambda: None,
+}
+
+
+@error_handler()
+def dedup_jsonl(
+    input_dir: str,
+    output_dir: str,
+    *,
+    threshold: float = 0.8,
+    num_perm: int = 128,
+    ngram: int = 3,
+    output_format: str = "jsonl",
+    chunk_size: int = 1_000_000,
+) -> Tuple[int, int]:
+    """Deduplicate JSONL text files using MinHash + LSH.
+
+    Args:
+        input_dir: Directory with source ``*.jsonl`` files.
+        output_dir: Directory for deduplicated output.
+        threshold: Jaccard similarity threshold (0–1).
+        num_perm: Number of MinHash permutations.
+        ngram: Character n-gram size.
+        output_format: ``"jsonl"``, ``"h5"``, or ``"bin"``.
+        chunk_size: Records per output chunk file.
+
+    Returns:
+        ``(kept, removed)`` counts.
+    """
+    from datasketch import MinHash, MinHashLSH
+
+    input_path = Path(input_dir)
+    output_path = Path(output_dir)
+    output_path.mkdir(parents=True, exist_ok=True)
+
+    logger.info(
+        f"Deduplicating {input_dir} -> {output_dir} "
+        f"(threshold={threshold}, perm={num_perm}, fmt={output_format})"
+    )
+
+    lsh = MinHashLSH(threshold=threshold, num_perm=num_perm)
+
+    kept = 0
+    removed = 0
+
+    dup_doc_ids: Set[int] = set()
+    for doc_id, (text, _record) in enumerate(tqdm(_iter_docs(input_path), desc="indexing", unit="docs")):
+        shingles = _tokenize(text, ngram=ngram)
+        if len(shingles) < ngram * 2:
+            dup_doc_ids.add(doc_id)
+            continue
+
+        m = MinHash(num_perm=num_perm)
+        for s in shingles:
+            m.update(s.encode("utf-8"))
+
+        if lsh.query(m):
+            dup_doc_ids.add(doc_id)
+        else:
+            lsh.insert(doc_id, m)
+
+    logger.info(f"Found {len(dup_doc_ids)} duplicates, writing deduplicated data")
+
+    buffer: List[dict] = []
+    chunk_idx = 0
+    writer = TextWriter(chunk_size) if output_format == "jsonl" else None
+
+    for doc_id, (_text, record) in enumerate(tqdm(_iter_docs(input_path), desc="writing", unit="docs")):
+        if doc_id in dup_doc_ids:
+            removed += 1
+            continue
+
+        kept += 1
+        buffer.append(record)
+
+        if len(buffer) >= chunk_size:
+            _flush_chunk(buffer, output_path, chunk_idx, output_format, writer)
+            chunk_idx += 1
+            buffer = []
+
+    if buffer:
+        _flush_chunk(buffer, output_path, chunk_idx, output_format, writer)
+
+    if writer:
+        writer.flush(output_path)
+
+    logger.info(f"Done. kept={kept}, removed={removed}")
+    return kept, removed
+
+
+def _flush_chunk(
+    records: List[dict],
+    output_dir: Path,
+    chunk_idx: int,
+    output_format: str,
+    writer=None,
+):
+    if output_format == "jsonl":
+        for rec in records:
+            writer.write_record(rec, output_dir)
+    elif output_format == "h5":
+        _write_h5(records, str(output_dir), chunk_idx)
+    elif output_format == "bin":
+        _write_bin(records, output_dir, chunk_idx)
+    else:
+        raise ValueError(f"Unknown output format: {output_format}")

pipeline/io/export.py

@@ -4,14 +4,18 @@ import json
 import logging
 import os
 from pathlib import Path
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+import torch
 from datasets import Dataset
+from torch import Tensor
 from tqdm import tqdm
 
 from pipeline.io.file_scanner import FileScanner
 from pipeline.io.hdf5_handler import HDF5Handler
+from pipeline.io.writers import create_writer, BaseWriter
 from pipeline.processors import BaseProcessor
-from pipeline.packing import pack_tensors
+from pipeline.packing import pack_tensors, BasePacker
 from pipeline.utils import error_handler
 
 logger = logging.getLogger(__name__)
@@ -75,6 +79,32 @@ def export_dataset(
     return output_files
 
 
+def merge_tensors(
+    tensors: List[Tensor],
+    group_size: int,
+) -> List[Tensor]:
+    """Merge a list of tensors into fewer larger tensors.
+
+    Concatenates every group_size consecutive tensors into one merged
+    tensor. This reduces the number of shm blocks when loading.
+
+    Args:
+        tensors: List of 1D tensors.
+        group_size: Number of tensors to merge into each group.
+
+    Returns:
+        List of merged tensors.
+    """
+    if not tensors:
+        return []
+
+    merged: List[Tensor] = []
+    for i in range(0, len(tensors), group_size):
+        merged.append(torch.cat(tensors[i : i + group_size]))
+
+    return merged
+
+
 @error_handler()
 def cache_jsonl(
     files: List[str],
@@ -83,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

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)