feat: MinHash+LSH 去重 + Strategy/Factory 存储后端
This commit is contained in:
+11
-1
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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}")
|
||||
+11
-7
@@ -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
|
||||
|
||||
@@ -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)
|
||||
+12
-5
@@ -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: <input_dir>/cached)",
|
||||
help="Output dir (default: <input_dir>/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}")
|
||||
|
||||
|
||||
@@ -0,0 +1,36 @@
|
||||
"""MinHash + LSH deduplication CLI.
|
||||
|
||||
Usage:
|
||||
python scripts/dedup_pretrain.py --input-dir <data_dir> --output-dir <out_dir> --threshold 0.8 --num-perm 128 --output-format jsonl
|
||||
"""
|
||||
|
||||
import argparse
|
||||
|
||||
from pipeline.io import dedup_jsonl
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="MinHash + LSH deduplication")
|
||||
parser.add_argument("--input-dir", required=True)
|
||||
parser.add_argument("--output-dir", required=True)
|
||||
parser.add_argument("--threshold", type=float, default=0.8)
|
||||
parser.add_argument("--num-perm", type=int, default=128)
|
||||
parser.add_argument("--ngram", type=int, default=3)
|
||||
parser.add_argument("--output-format", default="jsonl", choices=["jsonl", "h5", "bin"])
|
||||
args = parser.parse_args()
|
||||
|
||||
kept, removed = dedup_jsonl(
|
||||
input_dir=args.input_dir,
|
||||
output_dir=args.output_dir,
|
||||
threshold=args.threshold,
|
||||
num_perm=args.num_perm,
|
||||
ngram=args.ngram,
|
||||
output_format=args.output_format,
|
||||
)
|
||||
|
||||
total = kept + removed
|
||||
print(f"kept={kept}, removed={removed} ({removed/max(total,1)*100:.1f}%)")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user