Base/scripts/dedup_5b_pretrain.py

# -*- coding: utf-8 -*-
"""
DEDUPLICATE litdata_pretrain_final — remove near-duplicate documents.

Strategy:
  1. Scan all 308 chunks, hash each document (first 200 tokens)
  2. Keep first occurrence, mark subsequent duplicates for removal
  3. Rebuild chunks with duplicates removed (same format, new files)
  4. Update index.json

Memory-efficient: processes 10 chunks at a time, uses hash set.
"""

import json
import os
import time
import hashlib
from pathlib import Path

import numpy as np

ROOT = Path(__file__).resolve().parent.parent.parent
FINAL_DIR = ROOT / "Base" / "data" / "litdata_pretrain_final"
BLOCK_SIZE = 1025
DTYPE = np.int32
EOS_TOKEN_ID = 0
CHUNK_BYTES_TARGET = 64 * 1024 * 1024
HASH_WINDOW = 200  # first N tokens for hash


def read_chunk(filepath):
    with open(filepath, "rb") as f:
        raw = f.read()
    num_blocks = np.frombuffer(raw[:4], dtype=np.uint32)[0]
    header_size = 4 + (num_blocks + 1) * 4
    data_bytes = raw[header_size:]
    expected_tokens = num_blocks * BLOCK_SIZE
    expected_bytes = expected_tokens * DTYPE().itemsize
    tokens = np.frombuffer(data_bytes[:expected_bytes], dtype=DTYPE)
    return tokens, int(num_blocks)


def extract_documents(tokens):
    """Extract individual documents separated by EOS."""
    eos_positions = np.where(tokens == EOS_TOKEN_ID)[0]
    docs = []
    start = 0
    for eos_pos in eos_positions:
        if eos_pos > start:
            docs.append(tokens[start:eos_pos])
        start = eos_pos + 1
    # Trailing partial (no EOS at end — crosses chunk boundary)
    if start < len(tokens):
        remaining = tokens[start:]
        if len(remaining) > 0:
            docs.append(remaining)
    return docs


def doc_hash(token_array):
    """Hash first HASH_WINDOW tokens of a document."""
    key = token_array[:HASH_WINDOW].tobytes()
    return hashlib.md5(key).hexdigest()


def write_chunk(filepath, tokens_array, block_size=BLOCK_SIZE):
    """Write a litdata chunk from a flat token array."""
    num_blocks = len(tokens_array) // block_size
    if num_blocks == 0:
        return None
    actual = num_blocks * block_size
    data = np.array(tokens_array[:actual], dtype=DTYPE)

    header_num = np.array([num_blocks], dtype=np.uint32)
    offsets = np.arange(num_blocks + 1, dtype=np.uint32) * (block_size * DTYPE().itemsize)
    header = np.concatenate([header_num, offsets])

    with open(filepath, "wb") as f:
        header.tofile(f)
        data.tofile(f)

    return {
        "chunk_bytes": int(header.nbytes + data.nbytes),
        "chunk_size": num_blocks,
        "dim": int(actual),
        "filename": os.path.basename(filepath),
    }


class StreamingDeduplicator:
    """Accumulates deduplicated tokens and writes chunks."""

    def __init__(self, output_dir, backup_suffix="_dedup"):
        self.output_dir = Path(output_dir)
        self.dtype_size = DTYPE().itemsize
        self.tokens_per_chunk = (CHUNK_BYTES_TARGET // self.dtype_size // BLOCK_SIZE) * BLOCK_SIZE
        self.buffer = []
        self.chunk_idx = 0
        self.chunks_meta = []
        self.total_tokens = 0

    def add_doc(self, doc_tokens):
        self.buffer.extend(doc_tokens.tolist())
        self.buffer.append(EOS_TOKEN_ID)
        while len(self.buffer) >= self.tokens_per_chunk:
            self._flush()

    def _flush(self):
        if len(self.buffer) < BLOCK_SIZE:
            return
        take = min(len(self.buffer), self.tokens_per_chunk)
        num_blocks = take // BLOCK_SIZE
        if num_blocks == 0:
            return
        actual = num_blocks * BLOCK_SIZE
        
        data = np.array(self.buffer[:actual], dtype=DTYPE)
        self.buffer = self.buffer[actual:]

        filename = f"chunk-0-{self.chunk_idx}.bin"
        filepath = self.output_dir / filename

        header_num = np.array([num_blocks], dtype=np.uint32)
        offsets = np.arange(num_blocks + 1, dtype=np.uint32) * (BLOCK_SIZE * self.dtype_size)
        header = np.concatenate([header_num, offsets])

        with open(filepath, "wb") as f:
            header.tofile(f)
            data.tofile(f)

        meta = {
            "chunk_bytes": int(header.nbytes + data.nbytes),
            "chunk_size": num_blocks,
            "dim": int(actual),
            "filename": filename,
        }
        self.chunks_meta.append(meta)
        self.total_tokens += actual
        self.chunk_idx += 1
        if self.chunk_idx % 25 == 0:
            print(f"      Written {self.chunk_idx} deduped chunks ({self.total_tokens:,} tokens)")

    def finalize(self):
        while len(self.buffer) >= BLOCK_SIZE:
            self._flush()
        discarded = len(self.buffer)
        self.buffer = []
        return self.total_tokens, discarded


def main():
    t_start = time.time()

    with open(FINAL_DIR / "index.json") as f:
        index = json.load(f)
    chunks_meta = index["chunks"]
    config = index.get("config", {})
    num_chunks = len(chunks_meta)

    original_tokens = sum(c["dim"] for c in chunks_meta)

    print(f"{'='*75}")
    print(f"  DEDUPLICATING litdata_pretrain_final")
    print(f"{'='*75}")
    print(f"  Chunks: {num_chunks}")
    print(f"  Original tokens: {original_tokens:,}")
    print()

    # ── PASS 1: Scan all chunks, collect hashes, identify duplicates ──────
    print(f"  PASS 1: Scanning all chunks for duplicates...")
    seen_hashes = set()
    dup_count = 0
    keep_count = 0
    total_docs = 0
    # We need to track which is first occurrence
    # But since we process sequentially, just check "in seen_hashes"

    # Use a temp dir to write deduplicated data, then swap
    TEMP_DIR = FINAL_DIR.parent / "litdata_pretrain_dedup_temp"
    if TEMP_DIR.exists():
        import shutil
        shutil.rmtree(str(TEMP_DIR))
    os.makedirs(str(TEMP_DIR))

    writer = StreamingDeduplicator(TEMP_DIR)

    for ci, meta in enumerate(chunks_meta):
        filepath = FINAL_DIR / meta["filename"]
        tokens, num_blocks = read_chunk(filepath)
        docs = extract_documents(tokens)

        chunk_dups = 0
        chunk_kept = 0

        for doc in docs:
            total_docs += 1
            if len(doc) < 10:
                # Very short fragments — keep (usually chunk-boundary partials)
                writer.add_doc(doc)
                keep_count += 1
                chunk_kept += 1
                continue

            h = doc_hash(doc)
            if h in seen_hashes:
                dup_count += 1
                chunk_dups += 1
            else:
                seen_hashes.add(h)
                writer.add_doc(doc)
                keep_count += 1
                chunk_kept += 1

        if (ci + 1) % 25 == 0 or ci == num_chunks - 1:
            print(f"    Chunk {ci+1}/{num_chunks}: total docs={total_docs:,}, kept={keep_count:,}, dupes removed={dup_count:,}")

    # Finalize
    final_tokens, discarded = writer.finalize()
    new_chunks = writer.chunk_idx

    print(f"\n  PASS 1 COMPLETE:")
    print(f"    Total documents scanned: {total_docs:,}")
    print(f"    Documents kept:          {keep_count:,}")
    print(f"    Duplicates removed:      {dup_count:,} ({100*dup_count/max(total_docs,1):.2f}%)")
    print(f"    Unique hashes:           {len(seen_hashes):,}")
    print(f"    Tokens after dedup:      {final_tokens:,}")
    print(f"    Token reduction:         {original_tokens - final_tokens:,} ({100*(original_tokens-final_tokens)/original_tokens:.2f}%)")
    print(f"    Chunks after dedup:      {new_chunks}")
    print(f"    Discarded partial:       {discarded} tokens")

    # ── PASS 2: Swap temp into final ──────────────────────────────────────
    print(f"\n  PASS 2: Replacing original with deduplicated data...")

    # Remove old chunk files
    for meta in chunks_meta:
        old_file = FINAL_DIR / meta["filename"]
        if old_file.exists():
            os.remove(str(old_file))

    # Move new chunk files from temp to final
    import shutil
    for meta in writer.chunks_meta:
        src = TEMP_DIR / meta["filename"]
        dst = FINAL_DIR / meta["filename"]
        shutil.move(str(src), str(dst))

    # Remove temp dir
    shutil.rmtree(str(TEMP_DIR))

    # Update index.json
    new_index = {
        "chunks": writer.chunks_meta,
        "config": config,
        "updated_at": str(time.time()),
    }
    with open(FINAL_DIR / "index.json", "w") as f:
        json.dump(new_index, f, indent=2)

    elapsed = time.time() - t_start

    # ── Report ────────────────────────────────────────────────────────────
    report = []
    report.append(f"{'='*75}")
    report.append(f"  DEDUPLICATION REPORT — litdata_pretrain_final")
    report.append(f"{'='*75}")
    report.append(f"")
    report.append(f"  Time: {elapsed:.0f}s ({elapsed/60:.1f} min)")
    report.append(f"")
    report.append(f"  BEFORE:")
    report.append(f"    Chunks:    {num_chunks}")
    report.append(f"    Tokens:    {original_tokens:,}")
    report.append(f"    Documents: {total_docs:,}")
    report.append(f"")
    report.append(f"  AFTER:")
    report.append(f"    Chunks:    {new_chunks}")
    report.append(f"    Tokens:    {final_tokens:,}")
    report.append(f"    Documents: {keep_count:,}")
    report.append(f"")
    report.append(f"  REMOVED:")
    report.append(f"    Duplicate docs: {dup_count:,} ({100*dup_count/max(total_docs,1):.2f}%)")
    report.append(f"    Tokens removed: {original_tokens - final_tokens:,} ({100*(original_tokens-final_tokens)/original_tokens:.2f}%)")
    report.append(f"")
    report.append(f"  Format: litdata binary (int32, BLOCK_SIZE={BLOCK_SIZE}, EOS={EOS_TOKEN_ID})")
    report.append(f"  Location: {FINAL_DIR}")
    report.append(f"{'='*75}")

    full_report = '\n'.join(report)
    print(f"\n{full_report}")

    with open(FINAL_DIR / "DEDUP_REPORT.txt", "w", encoding="utf-8") as f:
        f.write(full_report)

    print(f"\n  Saved to: {FINAL_DIR / 'DEDUP_REPORT.txt'}")
    print(f"  Done! Dataset is now clean and deduplicated.")


if __name__ == "__main__":
    main()