Base/scripts/reclean_3b.py

# -*- coding: utf-8 -*-
"""
Memory-efficient, multiprocessed reclean of litdata_3b (~3B tokens).

Problem: The old script loaded ALL 3B tokens → decoded 4.3M docs → re-tokenized
into a Python list → OOM at ~12 GB list + ~15 GB decoded text.

Solution: Process in chunk batches (10 chunks ≈ 168M tokens ≈ 670 MB).
  1. Read 10 chunks at a time
  2. Split tokens into documents (carry partial docs across batches)
  3. Decode with tokenizer.decode_batch (multithreaded Rust)
  4. Deep clean + smart filter with multiprocessing Pool (all CPU cores)
  5. Re-tokenize with encode_batch (multithreaded Rust)
  6. Stream output to litdata chunks incrementally (never accumulate)

Peak memory: ~4-5 GB instead of 40+ GB.
Output: Base/data/litdata_3b_clean/ (separate, not merged)
"""

import json
import os
import re
import sys
import time
import unicodedata
from pathlib import Path
from multiprocessing import Pool, cpu_count
from collections import Counter

import numpy as np
from tokenizers import Tokenizer

ROOT = Path(__file__).resolve().parent.parent.parent
BLOCK_SIZE = 1025
DTYPE = np.int32
CHUNK_BYTES_TARGET = 64 * 1024 * 1024
EOS_TOKEN_ID = 0

CHUNKS_PER_BATCH = 10        # ~670 MB per batch
NUM_WORKERS = max(1, cpu_count() - 2)  # leave 2 cores for main + I/O
DECODE_BATCH = 8000           # docs per decode_batch call
ENCODE_BATCH = 8000           # docs per encode_batch call

DATA_DIR = ROOT / "Base" / "data"
INPUT_DIR = DATA_DIR / "litdata_3b"
OUTPUT_DIR = DATA_DIR / "litdata_3b_clean"
TOKENIZER_PATH = str(ROOT / "Base" / "checkpoints" / "EleutherAI" / "pythia-160m" / "tokenizer.json")


# ==============================================================================
# TEXT CLEANING PIPELINE (deep_clean + smart_filter)
# ==============================================================================

CONTROL_CHARS = [
    "\x00", "\x01", "\x02", "\x03", "\x04", "\x05", "\x06", "\x07",
    "\x08", "\x0b", "\x0c", "\x0e", "\x0f", "\x10", "\x11", "\x12",
    "\x13", "\x14", "\x15", "\x16", "\x17", "\x18", "\x19", "\x1a",
    "\x1b", "\x1c", "\x1d", "\x1e", "\x1f", "\x7f", "\ufeff", "\ufffd",
]

HTML_ENTITIES = [
    ("&amp;", "&"), ("&lt;", "<"), ("&gt;", ">"),
    ("&quot;", '"'), ("&#39;", "'"), ("&apos;", "'"),
    ("&nbsp;", " "), ("&mdash;", " - "), ("&ndash;", "-"),
    ("&hellip;", "..."), ("&laquo;", '"'), ("&raquo;", '"'),
    ("&bull;", "- "), ("&middot;", " "), ("&copy;", "(c)"),
    ("&reg;", "(R)"), ("&trade;", "(TM)"), ("&deg;", " degrees"),
]

RE_URL = re.compile(r'https?://\S+|www\.\S+', re.I)
RE_EMAIL = re.compile(r'\b[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}\b')
RE_FILE_PATH = re.compile(r'(?:[A-Z]:\\|/(?:home|usr|var|etc|opt)/)\S+')
RE_HTML_TAG = re.compile(r'</?[a-zA-Z][a-zA-Z0-9]*(?:\s[^>]*)?\s*/?>')
RE_HTML_COMMENT = re.compile(r'<!--.*?-->', re.DOTALL)
RE_CODE_BLOCK = re.compile(r'```[\s\S]*?```')
RE_IMPORT = re.compile(r'^(?:import |from \S+ import |#include |using namespace |require\()', re.M)
RE_REPEATED_LINE = re.compile(r'^(.{20,})\n(?:\1\n?)+', re.M)
RE_REPEATED_PUNCT = re.compile(r'([!?.])\1{3,}')
RE_REPEATED_CHAR = re.compile(r'(.)\1{5,}')
RE_REPEATED_WORD = re.compile(r'\b(\w+)(?:\s+\1){2,}\b', re.I)
RE_MULTI_NEWLINE = re.compile(r'\n{4,}')
RE_MULTI_SPACE = re.compile(r'[ \t]{2,}')
RE_TRAILING_SPACE = re.compile(r'[ \t]+$', re.M)
RE_NO_SPACE_AFTER_PERIOD = re.compile(r'([.!?])([A-Z])')
RE_DOUBLE_PERIOD = re.compile(r'\.{2}(?!\.)')
RE_SPACE_BEFORE_PUNCT = re.compile(r'\s+([.,;:!?])')

# Smart filter patterns
RE_CJK = re.compile(r'[\u4e00-\u9fff\u3040-\u309f\u30a0-\u30ff\uac00-\ud7af]{3,}')
RE_ARABIC = re.compile(r'[\u0600-\u06ff]{5,}')
RE_CYRILLIC = re.compile(r'[\u0400-\u04ff]{5,}')
RE_DEVANAGARI = re.compile(r'[\u0900-\u097f]{5,}')
RE_RESIDUAL_CODE = re.compile(
    r'(function\s*\(|var\s+\w+\s*=|console\.log|document\.get|if\s*\(\s*\w+\s*[!=]==)', re.I
)

RE_COOKIE_LINE = re.compile(r'^.*(?:cookie|cookies)\s+(?:policy|consent|notice|preferences|settings).*$', re.I | re.M)
RE_SUBSCRIBE_LINE = re.compile(r'^.*(?:subscribe|sign\s*up\s+(?:for|to)\s+(?:our|the)\s+newsletter|unsubscribe|opt[\s-]*out\s+of).*$', re.I | re.M)
RE_CLICKBAIT_LINE = re.compile(r'^.*(?:you\s+won\'?t\s+believe|click\s+here|read\s+more\s*\.{0,3}$|share\s+this\s+(?:article|post|story)|trending\s+now|sponsored\s+content|advertisement).*$', re.I | re.M)
RE_SOCIAL_LINE = re.compile(r'^.*(?:follow\s+us\s+on|share\s+on\s+(?:facebook|twitter|linkedin|instagram)|like\s+us\s+on|tweet\s+this).*$', re.I | re.M)
RE_NAV_LINE = re.compile(r'^.*(?:skip\s+to\s+(?:main\s+)?content|back\s+to\s+top|previous\s+article|next\s+article|related\s+(?:articles|posts)).*$', re.I | re.M)
RE_LOGIN_LINE = re.compile(r'^.*(?:log\s*in\s+to\s+(?:your|an)\s+account|create\s+(?:a\s+)?(?:free\s+)?account|forgot\s+(?:your\s+)?password|already\s+(?:a\s+)?member).*$', re.I | re.M)
RE_COMMENT_LINE = re.compile(r'^.*(?:leave\s+a\s+(?:comment|reply)|post\s+a\s+comment|\d+\s+comments?$|logged\s+in\s+as).*$', re.I | re.M)
RE_COPYRIGHT_LINE = re.compile(r'^.*(?:all\s+rights\s+reserved|\(c\)\s*\d{4}|copyright\s+\d{4}).*$', re.I | re.M)


def deep_clean(text):
    """Full deep cleaning pipeline."""
    if not text or len(text.strip()) < 30:
        return None

    text = unicodedata.normalize("NFKC", text)
    for ch in CONTROL_CHARS:
        text = text.replace(ch, "")
    for old, new in HTML_ENTITIES:
        text = text.replace(old, new)

    text = RE_HTML_COMMENT.sub("", text)
    text = RE_HTML_TAG.sub("", text)
    text = RE_URL.sub("", text)
    text = RE_EMAIL.sub("", text)
    text = RE_FILE_PATH.sub("", text)
    text = RE_CODE_BLOCK.sub("", text)

    text = RE_REPEATED_LINE.sub(r'\1', text)
    text = RE_REPEATED_PUNCT.sub(r'\1\1\1', text)
    text = RE_REPEATED_CHAR.sub(r'\1\1\1', text)
    text = RE_REPEATED_WORD.sub(r'\1', text)

    text = text.replace('\t', ' ')
    text = RE_TRAILING_SPACE.sub('', text)
    text = RE_MULTI_SPACE.sub(' ', text)
    text = RE_MULTI_NEWLINE.sub('\n\n\n', text)

    text = RE_DOUBLE_PERIOD.sub('.', text)
    text = RE_NO_SPACE_AFTER_PERIOD.sub(r'\1 \2', text)
    text = RE_SPACE_BEFORE_PUNCT.sub(r'\1', text)

    text = text.replace('\u2018', "'").replace('\u2019', "'")
    text = text.replace('\u201c', '"').replace('\u201d', '"')
    text = text.replace('\u2013', '-').replace('\u2014', ' - ')
    text = text.replace('\u2026', '...')
    text = text.replace('\u2022', '- ')
    text = text.replace('\u00b7', ' ')
    text = text.replace('\u00a0', ' ')

    lines = text.split('\n')
    clean_lines = []
    for line in lines:
        line = line.strip()
        if not line:
            clean_lines.append('')
            continue
        if len(line) > 10:
            alpha_count = sum(1 for c in line if c.isalpha())
            if alpha_count / len(line) < 0.40:
                continue
        if line.count('|') > 3 or line.count('{') > 2 or line.count('}') > 2:
            continue
        if RE_IMPORT.match(line):
            continue
        if line and line[0].isalpha() and line[0].islower():
            if not clean_lines or clean_lines[-1] == '' or clean_lines[-1].rstrip().endswith(('.', '!', '?', ':')):
                line = line[0].upper() + line[1:]
        clean_lines.append(line)

    text = '\n'.join(clean_lines)
    text = text.strip()

    paragraphs = text.split('\n\n')
    seen = set()
    unique_paragraphs = []
    for p in paragraphs:
        p_stripped = p.strip()
        if not p_stripped:
            continue
        p_key = ' '.join(p_stripped.lower().split())
        if p_key not in seen:
            seen.add(p_key)
            unique_paragraphs.append(p_stripped)
    text = '\n\n'.join(unique_paragraphs)

    text = text.strip()
    if len(text) < 50:
        return None
    if len(text.split()) < 10:
        return None
    ascii_count = sum(1 for c in text if ord(c) < 128)
    if ascii_count / max(len(text), 1) < 0.85:
        return None

    return text


def smart_filter(text):
    """Smart cleanup: returns (keep, cleaned_text, reason)."""
    words = text.split()
    word_count = len(words)
    if word_count < 50:
        return False, text, "too short"

    scripts = []
    if RE_CJK.search(text): scripts.append("CJK")
    if RE_ARABIC.search(text): scripts.append("Arabic")
    if RE_CYRILLIC.search(text): scripts.append("Cyrillic")
    if RE_DEVANAGARI.search(text): scripts.append("Devanagari")
    if scripts:
        return False, text, "non-English"

    if word_count > 50:
        unique_ratio = len(set(w.lower() for w in words)) / word_count
        if unique_ratio < 0.20:
            return False, text, "repetitive"

    code_matches = RE_RESIDUAL_CODE.findall(text)
    if len(code_matches) >= 5:
        return False, text, "residual code"

    for pattern in [RE_COOKIE_LINE, RE_SUBSCRIBE_LINE, RE_CLICKBAIT_LINE,
                    RE_SOCIAL_LINE, RE_NAV_LINE, RE_LOGIN_LINE,
                    RE_COMMENT_LINE, RE_COPYRIGHT_LINE]:
        text = pattern.sub('', text)

    lines = text.split('\n')
    clean_lines = []
    for line in lines:
        stripped = line.strip()
        if stripped and len(stripped) > 10:
            digit_count = sum(1 for c in stripped if c.isdigit() or c in ' ,.\t-+/%$')
            if digit_count / len(stripped) > 0.80:
                continue
        clean_lines.append(line)
    text = '\n'.join(clean_lines)
    text = re.sub(r'\n{3,}', '\n\n', text)
    text = text.strip()

    if len(text.split()) < 50:
        return False, text, "too short after stripping"

    return True, text, None


def clean_and_filter(text):
    """Combined deep_clean + smart_filter for multiprocessing.
    Returns (cleaned_text_or_None, drop_reason_or_None, boilerplate_chars_stripped).
    """
    result = deep_clean(text)
    if result is None:
        return None, "deep_clean_drop", 0

    keep, stripped, reason = smart_filter(result)
    if not keep:
        return None, reason, 0

    boilerplate = len(result) - len(stripped)
    return stripped, None, boilerplate


# ==============================================================================
# STREAMING CHUNK WRITER (never accumulates full token stream)
# ==============================================================================

class StreamingChunkWriter:
    """Writes litdata chunks incrementally. Flushes when buffer hits target size."""

    def __init__(self, output_dir, config):
        self.output_dir = Path(output_dir)
        os.makedirs(str(self.output_dir), exist_ok=True)
        self.config = config
        self.dtype_size = DTYPE().itemsize
        self.tokens_per_chunk = (CHUNK_BYTES_TARGET // self.dtype_size // BLOCK_SIZE) * BLOCK_SIZE
        self.buffer = []
        self.chunks_metadata = []
        self.chunk_idx = 0
        self.total_tokens = 0

    def add_document(self, token_ids):
        """Add one document's token IDs + EOS. Flushes chunks as needed."""
        self.buffer.extend(token_ids)
        self.buffer.append(EOS_TOKEN_ID)
        # Flush full chunks
        while len(self.buffer) >= self.tokens_per_chunk:
            self._flush_chunk()

    def add_tokens_batch(self, encoded_batch):
        """Add a batch of encoded documents efficiently."""
        for enc in encoded_batch:
            ids = enc.ids
            self.buffer.extend(ids)
            self.buffer.append(EOS_TOKEN_ID)
        while len(self.buffer) >= self.tokens_per_chunk:
            self._flush_chunk()

    def _flush_chunk(self):
        """Write one chunk from the buffer."""
        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_tokens = num_blocks * BLOCK_SIZE

        chunk_data = np.array(self.buffer[:actual_tokens], dtype=DTYPE)
        self.buffer = self.buffer[actual_tokens:]

        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)
            chunk_data.tofile(f)

        meta = {
            "chunk_bytes": int(header.nbytes + chunk_data.nbytes),
            "chunk_size": num_blocks,
            "dim": int(actual_tokens),
            "filename": filename,
        }
        self.chunks_metadata.append(meta)
        self.total_tokens += actual_tokens
        self.chunk_idx += 1

        if self.chunk_idx % 25 == 0:
            print(f"      Flushed chunk {self.chunk_idx} ({self.total_tokens:,} tokens written)")

    def finalize(self):
        """Flush remaining buffer and write index.json."""
        while len(self.buffer) >= BLOCK_SIZE:
            self._flush_chunk()

        # Discard any remaining tokens that don't fill a block
        discarded = len(self.buffer)
        self.buffer = []

        index = {
            "chunks": self.chunks_metadata,
            "config": self.config,
            "updated_at": str(time.time()),
        }
        with open(self.output_dir / "index.json", "w") as f:
            json.dump(index, f, indent=2)

        return self.total_tokens, discarded


# ==============================================================================
# CHUNK READING
# ==============================================================================

def read_chunk_tokens(litdata_dir, chunk_meta):
    """Read a single chunk's token data (no header)."""
    chunk_path = litdata_dir / chunk_meta["filename"]
    n_blocks = chunk_meta["chunk_size"]
    header_ints = 1 + n_blocks + 1
    header_bytes = header_ints * 4

    with open(chunk_path, "rb") as f:
        f.seek(header_bytes)
        data = np.fromfile(f, dtype=DTYPE, count=chunk_meta["dim"])
    return data


def split_documents_from_tokens(token_array):
    """Split token array by EOS into list of per-document token lists (as Python lists)."""
    eos_positions = np.where(token_array == EOS_TOKEN_ID)[0]
    docs = []
    start = 0
    for eos_pos in eos_positions:
        if eos_pos > start:
            docs.append(token_array[start:eos_pos].tolist())
        start = eos_pos + 1
    # Return remaining tokens after last EOS (partial doc carry-over)
    remainder = token_array[start:] if start < len(token_array) else np.array([], dtype=DTYPE)
    return docs, remainder


# ==============================================================================
# MAIN PIPELINE
# ==============================================================================

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

    print("Loading tokenizer...")
    tokenizer = Tokenizer.from_file(TOKENIZER_PATH)

    print(f"\n{'='*75}")
    print(f"  RECLEAN litdata_3b (Memory-Efficient + Multiprocessing)")
    print(f"  Input:   {INPUT_DIR}")
    print(f"  Output:  {OUTPUT_DIR}")
    print(f"  Workers: {NUM_WORKERS} CPU cores")
    print(f"  Batch:   {CHUNKS_PER_BATCH} chunks at a time")
    print(f"{'='*75}")

    # Read index
    with open(INPUT_DIR / "index.json") as f:
        index = json.load(f)
    all_chunk_metas = index["chunks"]
    total_input_tokens = sum(c["dim"] for c in all_chunk_metas)
    num_chunks = len(all_chunk_metas)

    print(f"\n  Input: {num_chunks} chunks, {total_input_tokens:,} tokens")

    # Output config (same format)
    config = {
        "block_size": BLOCK_SIZE,
        "vocab_size": tokenizer.get_vocab_size(),
    }

    # Stats
    total_docs_in = 0
    total_docs_kept = 0
    total_docs_dropped = 0
    total_boilerplate = 0
    drop_reasons = Counter()

    # Initialize streaming writer
    writer = StreamingChunkWriter(str(OUTPUT_DIR), config)

    # Carry-over: partial document tokens spanning chunk boundaries
    carry_over = np.array([], dtype=DTYPE)

    # ── Process in batches of chunks ──────────────────────────────────────
    num_batches = (num_chunks + CHUNKS_PER_BATCH - 1) // CHUNKS_PER_BATCH
    print(f"  Processing in {num_batches} batches of {CHUNKS_PER_BATCH} chunks...\n")

    # Create multiprocessing pool for cleaning
    pool = Pool(processes=NUM_WORKERS)

    for batch_idx in range(num_batches):
        batch_start = batch_idx * CHUNKS_PER_BATCH
        batch_end = min(batch_start + CHUNKS_PER_BATCH, num_chunks)
        batch_metas = all_chunk_metas[batch_start:batch_end]

        t_batch = time.time()
        print(f"  ── Batch {batch_idx+1}/{num_batches} (chunks {batch_start}-{batch_end-1}) ──")

        # 1. Read chunk tokens
        batch_tokens_list = []
        for cm in batch_metas:
            batch_tokens_list.append(read_chunk_tokens(INPUT_DIR, cm))
        batch_tokens = np.concatenate(batch_tokens_list)
        del batch_tokens_list

        # Prepend carry-over from previous batch
        if len(carry_over) > 0:
            batch_tokens = np.concatenate([carry_over, batch_tokens])
            carry_over = np.array([], dtype=DTYPE)

        batch_token_count = len(batch_tokens)

        # 2. Split into documents
        doc_token_lists, carry_over = split_documents_from_tokens(batch_tokens)
        del batch_tokens
        num_docs = len(doc_token_lists)
        total_docs_in += num_docs

        # 3. Decode documents to text (tokenizer.decode_batch is multithreaded Rust)
        t_dec = time.time()
        raw_texts = tokenizer.decode_batch(doc_token_lists, skip_special_tokens=False)
        del doc_token_lists
        dec_time = time.time() - t_dec

        # 4. Clean + smart filter in parallel across all cores
        t_clean = time.time()
        results = pool.map(clean_and_filter, raw_texts, chunksize=512)
        del raw_texts
        clean_time = time.time() - t_clean

        # Collect cleaned texts
        cleaned_texts = []
        batch_dropped = 0
        batch_boilerplate = 0
        for cleaned, reason, bp in results:
            if cleaned is not None:
                cleaned_texts.append(cleaned)
                batch_boilerplate += bp
            else:
                batch_dropped += 1
                drop_reasons[reason] += 1
        del results

        batch_kept = len(cleaned_texts)
        total_docs_kept += batch_kept
        total_docs_dropped += batch_dropped
        total_boilerplate += batch_boilerplate

        # 5. Re-tokenize + stream to writer (encode_batch is multithreaded Rust)
        t_tok = time.time()
        for i in range(0, len(cleaned_texts), ENCODE_BATCH):
            sub = cleaned_texts[i:i+ENCODE_BATCH]
            encoded = tokenizer.encode_batch(sub, add_special_tokens=False)
            writer.add_tokens_batch(encoded)
        del cleaned_texts
        tok_time = time.time() - t_tok

        elapsed = time.time() - t_batch
        print(f"    {batch_token_count:,} tokens → {num_docs:,} docs → kept {batch_kept:,} / dropped {batch_dropped:,}")
        print(f"    decode {dec_time:.1f}s | clean {clean_time:.1f}s | tokenize {tok_time:.1f}s | total {elapsed:.1f}s")
        print(f"    Running: {total_docs_kept:,} kept, {total_docs_dropped:,} dropped, {writer.total_tokens:,} tokens written")

    pool.close()
    pool.join()

    # Handle carry-over (last partial doc if any)
    if len(carry_over) > 0:
        text = tokenizer.decode(carry_over.tolist(), skip_special_tokens=False)
        result = deep_clean(text)
        if result is not None:
            keep, stripped, reason = smart_filter(result)
            if keep:
                encoded = tokenizer.encode(stripped, add_special_tokens=False)
                writer.buffer.extend(encoded.ids)
                writer.buffer.append(EOS_TOKEN_ID)
                total_docs_kept += 1
                total_boilerplate += len(result) - len(stripped)
            else:
                total_docs_dropped += 1
                drop_reasons[reason] += 1
        else:
            total_docs_dropped += 1
            drop_reasons["deep_clean_drop"] += 1
        total_docs_in += 1

    # Finalize output
    final_tokens, discarded = writer.finalize()

    total_time = time.time() - t_start

    # ── Report ────────────────────────────────────────────────────────────
    print(f"\n{'='*75}")
    print(f"  RECLEAN REPORT - litdata_3b")
    print(f"{'='*75}")
    print(f"\n  Total time: {total_time:.0f}s ({total_time/60:.1f} min)")
    print(f"  Workers: {NUM_WORKERS} CPU cores")
    print(f"\n  INPUT")
    print(f"  {'-'*60}")
    print(f"    Chunks:    {num_chunks}")
    print(f"    Tokens:    {total_input_tokens:,}")
    print(f"    Documents: {total_docs_in:,}")
    print(f"\n  CLEANING")
    print(f"  {'-'*60}")
    print(f"    Kept:      {total_docs_kept:,}")
    print(f"    Dropped:   {total_docs_dropped:,} ({total_docs_dropped/(max(total_docs_in,1))*100:.2f}%)")
    if drop_reasons:
        print(f"    Drop reasons:")
        for reason, count in sorted(drop_reasons.items(), key=lambda x: -x[1]):
            print(f"      {reason:<35} {count:>8,}")
    print(f"    Boilerplate stripped: {total_boilerplate:,} chars")
    print(f"\n  OUTPUT")
    print(f"  {'-'*60}")
    print(f"    Location:  {OUTPUT_DIR}")
    print(f"    Chunks:    {writer.chunk_idx}")
    print(f"    Tokens:    {final_tokens:,}")
    print(f"    Discarded: {discarded} trailing tokens (< 1 block)")
    print(f"    Format:    litdata binary (int32, BLOCK_SIZE=1025, EOS=0)")
    diff = total_input_tokens - final_tokens
    print(f"\n  Token change: {total_input_tokens:,} → {final_tokens:,}")
    print(f"  Difference:   {diff:,} ({diff/max(total_input_tokens,1)*100:.2f}%)")
    print(f"\n{'='*75}")

    # Save report
    report_lines = [
        f"RECLEAN REPORT - litdata_3b",
        f"Time: {total_time:.0f}s ({total_time/60:.1f} min)",
        f"Workers: {NUM_WORKERS}",
        f"",
        f"INPUT:  {num_chunks} chunks, {total_input_tokens:,} tokens, {total_docs_in:,} docs",
        f"OUTPUT: {writer.chunk_idx} chunks, {final_tokens:,} tokens",
        f"",
        f"Docs kept:    {total_docs_kept:,}",
        f"Docs dropped: {total_docs_dropped:,}",
    ]
    if drop_reasons:
        for reason, count in sorted(drop_reasons.items(), key=lambda x: -x[1]):
            report_lines.append(f"  {reason}: {count:,}")
    report_lines.append(f"Boilerplate stripped: {total_boilerplate:,} chars")
    report_lines.append(f"Token change: {total_input_tokens:,} -> {final_tokens:,} ({diff/max(total_input_tokens,1)*100:.2f}%)")

    with open(OUTPUT_DIR / "CLEAN_REPORT.txt", "w", encoding="utf-8") as f:
        f.write("\n".join(report_lines))
    print(f"  Report saved to: {OUTPUT_DIR / 'CLEAN_REPORT.txt'}")
    print(f"  Done!")


if __name__ == "__main__":
    main()