Base/scripts/merge_and_build.py

# -*- coding: utf-8 -*-
"""
Merge existing litdata datasets + build a NEW 100M token batch + merge all.

Steps:
  1. Merge litdata_english_100m + litdata_english_clean → litdata_combined
  2. Download 100M tokens of FRESH data (skipping everything already used)
  3. Deep clean + smart cleanup
  4. Tokenize → temp litdata chunks
  5. Merge temp chunks into litdata_combined → final unified dataset

Skip values:
  - WIKI_SKIP  = 80,000  (first dataset used ~19.5K, second used 25K+53.5K = 78.5K)
  - FINEWEB_SKIP = 70,000 (first used ~19.4K, second used 25K+43.5K = 68.5K)

Final output: Base/data/litdata_combined/
"""

import json
import os
import re
import shutil
import time
import unicodedata
from pathlib import Path
from collections import Counter

import numpy as np
import pyarrow as pa
import pyarrow.parquet as pq
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
TARGET_TOKENS = 100_000_000
TOKENS_PER_WORD = 1.3

# Skip counts — must exceed ALL previously used qualifying articles
WIKI_SKIP = 80_000
FINEWEB_SKIP = 70_000

WIKI_SHARE = 0.55
FINEWEB_MIN_SCORE = 4.0

# Paths
DATA_DIR = ROOT / "Base" / "data"
LITDATA_100M = DATA_DIR / "litdata_english_100m"
LITDATA_CLEAN = DATA_DIR / "litdata_english_clean"
COMBINED_DIR = DATA_DIR / "litdata_combined"
TEMP_DIR = DATA_DIR / "litdata_new_100m_temp"
PARQUET_DIR = DATA_DIR / "filtered_english_new_100m"
TOKENIZER_PATH = ROOT / "Base" / "checkpoints" / "EleutherAI" / "pythia-160m" / "tokenizer.json"

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


# ==============================================================================
# TEXT QUALITY FILTERS
# ==============================================================================

def clean_text(text):
    text = unicodedata.normalize("NFKC", text)
    text = re.sub(r'[\x00-\x08\x0b\x0c\x0e-\x1f\x7f-\x9f]', '', text)
    text = re.sub(r'\n{3,}', '\n\n', text)
    text = re.sub(r'[ \t]+', ' ', text)
    text = '\n'.join(line.strip() for line in text.split('\n'))
    return text.strip()


def is_high_quality(text, min_chars=500, min_words=80):
    if len(text) < min_chars:
        return False
    words = text.split()
    num_words = len(words)
    if num_words < min_words:
        return False
    alpha = sum(c.isalpha() for c in text)
    if alpha / max(len(text), 1) < 0.65:
        return False
    avg_word_len = sum(len(w) for w in words) / num_words
    if avg_word_len < 2.5 or avg_word_len > 15:
        return False
    url_hits = text.count('http://') + text.count('https://')
    if url_hits > num_words * 0.03:
        return False
    sentences = re.split(r'[.!?]+', text)
    real_sentences = [s.strip() for s in sentences if len(s.strip()) > 10]
    if len(real_sentences) < 3:
        return False
    lines = [ln.strip() for ln in text.split('\n') if ln.strip()]
    if len(lines) > 5:
        unique_ratio = len(set(lines)) / len(lines)
        if unique_ratio < 0.5:
            return False
    return True


_WIKI_SKIP_PATTERNS = re.compile(
    r'(disambiguation|list of|lists of|index of|outline of|'
    r'wikipedia:|template:|category:|portal:|module:|mediawiki:)',
    re.IGNORECASE
)


# ==============================================================================
# DEEP CLEANING PIPELINE
# ==============================================================================

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+([.,;:!?])')

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):
    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):
    words = text.split()
    word_count = len(words)
    if word_count < 50:
        return False, text, f"too short ({word_count} words)"
    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, f"non-English: {', '.join(scripts)}"
    if word_count > 50:
        unique_ratio = len(set(w.lower() for w in words)) / word_count
        if unique_ratio < 0.20:
            return False, text, f"repetitive ({unique_ratio:.3f})"
    code_matches = RE_RESIDUAL_CODE.findall(text)
    if len(code_matches) >= 5:
        return False, text, f"residual code ({len(code_matches)})"
    original_len = len(text)
    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


# ==============================================================================
# LITDATA I/O
# ==============================================================================

def write_litdata_chunks(output_dir, token_stream, config, start_chunk_idx=0):
    """Write token stream as litdata binary chunks, starting at given chunk index."""
    os.makedirs(output_dir, exist_ok=True)
    dtype_size = DTYPE().itemsize
    tokens_per_chunk = CHUNK_BYTES_TARGET // dtype_size
    tokens_per_chunk = (tokens_per_chunk // BLOCK_SIZE) * BLOCK_SIZE
    chunks_metadata = []
    pos = 0
    chunk_idx = start_chunk_idx
    while pos < len(token_stream):
        remaining = len(token_stream) - pos
        chunk_tokens = min(tokens_per_chunk, remaining)
        num_blocks = chunk_tokens // BLOCK_SIZE
        if num_blocks == 0:
            break
        actual_tokens = num_blocks * BLOCK_SIZE
        chunk_data = token_stream[pos:pos + actual_tokens]
        filename = f"chunk-0-{chunk_idx}.bin"
        filepath = os.path.join(output_dir, filename)
        header_num_items = np.array([num_blocks], dtype=np.uint32)
        offsets = np.arange(num_blocks + 1, dtype=np.uint32) * (BLOCK_SIZE * dtype_size)
        header = np.concatenate([header_num_items, 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,
        }
        chunks_metadata.append(meta)
        pos += actual_tokens
        chunk_idx += 1
        print(f"    Written chunk {chunk_idx} ({pos:,}/{len(token_stream):,} tokens)")
    return chunks_metadata


def merge_litdata(sources, dest_dir):
    """
    Merge multiple litdata directories into one.
    sources: list of Path objects to litdata directories
    dest_dir: Path for the merged output
    """
    os.makedirs(str(dest_dir), exist_ok=True)
    all_chunks = []
    config = None
    chunk_offset = 0

    for src in sources:
        idx_path = src / "index.json"
        with open(idx_path, "r") as f:
            index = json.load(f)

        if config is None:
            config = index["config"]

        for chunk_meta in index["chunks"]:
            old_filename = chunk_meta["filename"]
            new_filename = f"chunk-0-{chunk_offset}.bin"

            # Copy the binary chunk file
            src_file = src / old_filename
            dst_file = dest_dir / new_filename
            shutil.copy2(str(src_file), str(dst_file))

            new_meta = dict(chunk_meta)
            new_meta["filename"] = new_filename
            all_chunks.append(new_meta)
            chunk_offset += 1

        src_tokens = sum(c["dim"] for c in index["chunks"])
        print(f"    Merged {src.name}: {len(index['chunks'])} chunks, {src_tokens:,} tokens")

    # Write combined index
    combined_index = {
        "chunks": all_chunks,
        "config": config,
        "updated_at": str(time.time()),
    }
    with open(str(dest_dir / "index.json"), "w") as f:
        json.dump(combined_index, f, indent=2)

    total = sum(c["dim"] for c in all_chunks)
    print(f"    Total: {len(all_chunks)} chunks, {total:,} tokens")
    return all_chunks, config


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

def main():
    from datasets import load_dataset

    t_start = time.time()

    # ═════════════════════════════════════════════════════════════
    # STEP 1: Merge existing litdata_english_100m + litdata_english_clean
    # ═════════════════════════════════════════════════════════════
    print(f"\n{'='*75}")
    print(f"  STEP 1: MERGING EXISTING DATASETS")
    print(f"    - {LITDATA_100M}")
    print(f"    - {LITDATA_CLEAN}")
    print(f"    → {COMBINED_DIR}")
    print(f"{'='*75}\n")

    merged_chunks, config = merge_litdata([LITDATA_100M, LITDATA_CLEAN], COMBINED_DIR)
    merged_tokens = sum(c["dim"] for c in merged_chunks)
    merged_chunk_count = len(merged_chunks)
    print(f"\n  Step 1 done: {merged_tokens:,} tokens in {merged_chunk_count} chunks")

    # ═════════════════════════════════════════════════════════════
    # STEP 2: Download fresh 100M tokens
    # ═════════════════════════════════════════════════════════════
    wiki_target = int(TARGET_TOKENS * WIKI_SHARE)
    fineweb_target = TARGET_TOKENS - wiki_target

    print(f"\n{'='*75}")
    print(f"  STEP 2: DOWNLOADING FRESH DATA")
    print(f"  Target: {TARGET_TOKENS:,} tokens ({wiki_target:,} Wiki + {fineweb_target:,} FineWeb)")
    print(f"  Skipping first {WIKI_SKIP:,} Wiki + {FINEWEB_SKIP:,} FineWeb (already used)")
    print(f"{'='*75}")

    os.makedirs(str(PARQUET_DIR), exist_ok=True)

    # --- Wikipedia ---
    print(f"\n  [Wikipedia] Streaming English articles (skipping first {WIKI_SKIP:,})...")
    ds_wiki = load_dataset(
        "wikimedia/wikipedia", "20231101.en",
        split="train", streaming=True,
        trust_remote_code=False
    )

    wiki_texts = []
    wiki_tokens = 0
    wiki_seen = 0
    wiki_skipped_quality = 0
    wiki_skipped_meta = 0
    t0 = time.time()

    for article in ds_wiki:
        title = (article.get("title") or "").strip()
        raw = article.get("text") or ""

        if _WIKI_SKIP_PATTERNS.search(title):
            wiki_skipped_meta += 1
            continue

        cleaned = clean_text(raw)
        if not is_high_quality(cleaned, min_chars=800, min_words=120):
            wiki_skipped_quality += 1
            continue

        wiki_seen += 1

        if wiki_seen <= WIKI_SKIP:
            if wiki_seen % 10000 == 0:
                print(f"    Skipping... {wiki_seen:,}/{WIKI_SKIP:,}")
            continue

        full_text = f"{title}\n\n{cleaned}"
        est_tok = int(len(full_text.split()) * TOKENS_PER_WORD)
        wiki_texts.append(full_text)
        wiki_tokens += est_tok

        if len(wiki_texts) % 5000 == 0:
            elapsed = time.time() - t0
            print(f"    Collected {len(wiki_texts):,} articles | ~{wiki_tokens:,} tokens | {elapsed:.0f}s")

        if wiki_tokens >= wiki_target:
            break

    elapsed = time.time() - t0
    print(f"  [Wikipedia] Done: {len(wiki_texts):,} articles, ~{wiki_tokens:,} tokens in {elapsed:.0f}s")
    print(f"    (skipped {WIKI_SKIP:,} already-used + {wiki_skipped_quality:,} low-quality + {wiki_skipped_meta:,} meta)")

    BATCH = 5000
    for i in range(0, len(wiki_texts), BATCH):
        batch = wiki_texts[i:i+BATCH]
        fp = PARQUET_DIR / f"wiki_{i//BATCH:04d}.parquet"
        pq.write_table(pa.table({"text": batch}), str(fp))
    print(f"  Saved {(len(wiki_texts)-1)//BATCH + 1} wiki parquet files")

    # --- FineWeb-Edu ---
    print(f"\n  [FineWeb-Edu] Streaming (score >= {FINEWEB_MIN_SCORE}, skipping first {FINEWEB_SKIP:,})...")
    ds_fineweb = load_dataset(
        "HuggingFaceFW/fineweb-edu", "sample-10BT",
        split="train", streaming=True,
        trust_remote_code=False
    )

    fineweb_texts = []
    fineweb_tokens = 0
    fineweb_seen = 0
    fineweb_skipped_quality = 0
    fineweb_skipped_score = 0
    t0 = time.time()

    for doc in ds_fineweb:
        score = doc.get("score", 0)
        if not isinstance(score, (int, float)):
            try:
                score = float(score)
            except (ValueError, TypeError):
                continue
        if score < FINEWEB_MIN_SCORE:
            fineweb_skipped_score += 1
            continue

        raw = doc.get("text") or ""
        cleaned = clean_text(raw)
        if not is_high_quality(cleaned, min_chars=500, min_words=80):
            fineweb_skipped_quality += 1
            continue

        fineweb_seen += 1

        if fineweb_seen <= FINEWEB_SKIP:
            if fineweb_seen % 10000 == 0:
                print(f"    Skipping... {fineweb_seen:,}/{FINEWEB_SKIP:,}")
            continue

        est_tok = int(len(cleaned.split()) * TOKENS_PER_WORD)
        fineweb_texts.append(cleaned)
        fineweb_tokens += est_tok

        if len(fineweb_texts) % 5000 == 0:
            elapsed = time.time() - t0
            print(f"    Collected {len(fineweb_texts):,} docs | ~{fineweb_tokens:,} tokens | {elapsed:.0f}s")

        if fineweb_tokens >= fineweb_target:
            break

    elapsed = time.time() - t0
    print(f"  [FineWeb-Edu] Done: {len(fineweb_texts):,} docs, ~{fineweb_tokens:,} tokens in {elapsed:.0f}s")
    print(f"    (skipped {FINEWEB_SKIP:,} already-used + {fineweb_skipped_quality:,} low-quality + {fineweb_skipped_score:,} low-score)")

    for i in range(0, len(fineweb_texts), BATCH):
        batch = fineweb_texts[i:i+BATCH]
        fp = PARQUET_DIR / f"fineweb_{i//BATCH:04d}.parquet"
        pq.write_table(pa.table({"text": batch}), str(fp))
    print(f"  Saved {(len(fineweb_texts)-1)//BATCH + 1} fineweb parquet files")

    all_texts = wiki_texts + fineweb_texts
    total_est_tokens = wiki_tokens + fineweb_tokens
    del wiki_texts, fineweb_texts

    print(f"\n  Step 2 done: {len(all_texts):,} documents, ~{total_est_tokens:,} estimated tokens")

    # ═════════════════════════════════════════════════════════════
    # STEP 3: Deep clean every document
    # ═════════════════════════════════════════════════════════════
    print(f"\n{'='*75}")
    print(f"  STEP 3: DEEP CLEANING {len(all_texts):,} DOCUMENTS")
    print(f"{'='*75}")

    t2 = time.time()
    cleaned_texts = []
    dropped_clean = 0
    for i, text in enumerate(all_texts):
        result = deep_clean(text)
        if result is not None:
            cleaned_texts.append(result)
        else:
            dropped_clean += 1
        if (i + 1) % 10000 == 0 or i == len(all_texts) - 1:
            print(f"    Cleaned {i+1:,}/{len(all_texts):,} | kept={len(cleaned_texts):,} | dropped={dropped_clean:,}")

    del all_texts
    print(f"  Deep clean done in {time.time()-t2:.1f}s")
    print(f"  Kept: {len(cleaned_texts):,} | Dropped: {dropped_clean:,}")

    # ═════════════════════════════════════════════════════════════
    # STEP 4: Smart cleanup
    # ═════════════════════════════════════════════════════════════
    print(f"\n{'='*75}")
    print(f"  STEP 4: SMART CLEANUP ON {len(cleaned_texts):,} DOCUMENTS")
    print(f"{'='*75}")

    t3 = time.time()
    final_texts = []
    removed_reasons = Counter()
    total_boilerplate = 0

    for i, text in enumerate(cleaned_texts):
        keep, stripped_text, reason = smart_filter(text)
        if keep:
            final_texts.append(stripped_text)
            total_boilerplate += len(text) - len(stripped_text)
        else:
            removed_reasons[reason.split('(')[0].strip().split(':')[0].strip()] += 1
        if (i + 1) % 10000 == 0 or i == len(cleaned_texts) - 1:
            print(f"    Processed {i+1:,}/{len(cleaned_texts):,} | kept={len(final_texts):,}")

    del cleaned_texts
    print(f"  Smart cleanup done in {time.time()-t3:.1f}s")
    print(f"  Final new documents: {len(final_texts):,}")
    print(f"  Boilerplate stripped: {total_boilerplate:,} chars")
    if removed_reasons:
        print(f"  Removal reasons:")
        for reason, count in sorted(removed_reasons.items(), key=lambda x: -x[1]):
            print(f"    {reason:<35} {count:>6,}")

    # ═════════════════════════════════════════════════════════════
    # STEP 5: Tokenize new data
    # ═════════════════════════════════════════════════════════════
    print(f"\n{'='*75}")
    print(f"  STEP 5: TOKENIZING {len(final_texts):,} NEW DOCUMENTS")
    print(f"{'='*75}")

    t4 = time.time()
    all_token_ids = []
    total_tokens = 0
    ENCODE_BATCH = 10000

    for i in range(0, len(final_texts), ENCODE_BATCH):
        batch = final_texts[i:i+ENCODE_BATCH]
        encoded = tokenizer.encode_batch(batch, add_special_tokens=False)
        for enc in encoded:
            ids = enc.ids
            all_token_ids.extend(ids)
            all_token_ids.append(EOS_TOKEN_ID)
            total_tokens += len(ids) + 1
        done = min(i + ENCODE_BATCH, len(final_texts))
        if done % 20000 == 0 or done == len(final_texts):
            print(f"    Tokenized {done:,}/{len(final_texts):,} ({total_tokens:,} tokens)")

    del final_texts
    print(f"  Tokenized in {time.time()-t4:.1f}s - {total_tokens:,} total new tokens")

    # ═════════════════════════════════════════════════════════════
    # STEP 6: Write new chunks directly into combined dir
    # ═════════════════════════════════════════════════════════════
    print(f"\n{'='*75}")
    print(f"  STEP 6: WRITING NEW CHUNKS INTO COMBINED DATASET")
    print(f"{'='*75}")

    token_array = np.array(all_token_ids, dtype=DTYPE)
    del all_token_ids

    litdata_config = {
        "block_size": BLOCK_SIZE,
        "vocab_size": tokenizer.get_vocab_size(),
    }

    # Write new chunks starting after the existing merged chunks
    new_chunks = write_litdata_chunks(
        str(COMBINED_DIR), token_array, litdata_config,
        start_chunk_idx=merged_chunk_count,
    )
    new_token_count = sum(c["dim"] for c in new_chunks)
    del token_array

    # Update the combined index.json with the new chunks appended
    all_final_chunks = merged_chunks + new_chunks
    final_total_tokens = sum(c["dim"] for c in all_final_chunks)
    final_index = {
        "chunks": all_final_chunks,
        "config": litdata_config,
        "updated_at": str(time.time()),
    }
    with open(str(COMBINED_DIR / "index.json"), "w") as f:
        json.dump(final_index, f, indent=2)

    print(f"\n  New data: {new_token_count:,} tokens in {len(new_chunks)} chunks")
    print(f"  Final combined: {final_total_tokens:,} tokens in {len(all_final_chunks)} chunks")

    # ═════════════════════════════════════════════════════════════
    # FINAL REPORT
    # ═════════════════════════════════════════════════════════════
    total_time = time.time() - t_start

    report_lines = []
    report_lines.append(f"\n{'='*75}")
    report_lines.append(f"  LITDATA_COMBINED - MERGE & BUILD REPORT")
    report_lines.append(f"{'='*75}")
    report_lines.append(f"\n  Total time: {total_time:.0f}s ({total_time/60:.1f} min)")
    report_lines.append(f"\n  SOURCE DATASETS MERGED")
    report_lines.append(f"  {'-'*60}")
    report_lines.append(f"    litdata_english_100m + litdata_english_clean")
    report_lines.append(f"    Merged subtotal:       {merged_tokens:,} tokens ({merged_chunk_count} chunks)")
    report_lines.append(f"\n  NEW 100M BATCH (articles {WIKI_SKIP+1:,}+)")
    report_lines.append(f"  {'-'*60}")
    report_lines.append(f"    Wikipedia:    ~{wiki_tokens:,} est. tokens")
    report_lines.append(f"    FineWeb-Edu:  ~{fineweb_tokens:,} est. tokens")
    report_lines.append(f"    Downloaded:   ~{total_est_tokens:,} est. tokens")
    report_lines.append(f"    Deep clean:   dropped {dropped_clean:,} docs")
    report_lines.append(f"    Smart filter: removed {sum(removed_reasons.values()):,} docs")
    if removed_reasons:
        for reason, count in sorted(removed_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"    New tokens:   {new_token_count:,} ({len(new_chunks)} chunks)")
    report_lines.append(f"\n  FINAL COMBINED OUTPUT")
    report_lines.append(f"  {'-'*60}")
    report_lines.append(f"    Location: {COMBINED_DIR}")
    report_lines.append(f"    Chunks:   {len(all_final_chunks)}")
    report_lines.append(f"    Tokens:   {final_total_tokens:,}")
    report_lines.append(f"    Format:   litdata binary (int32, BLOCK_SIZE=1025, EOS=0)")
    report_lines.append(f"\n  ZERO OVERLAP GUARANTEE:")
    report_lines.append(f"    - litdata_english_clean: Wiki articles 1-19,523 + FineWeb docs 1-19,445")
    report_lines.append(f"    - litdata_english_100m:  Wiki articles 25,001-78,482 + FineWeb docs 25,001-68,550")
    report_lines.append(f"    - New batch:             Wiki articles {WIKI_SKIP+1:,}+ + FineWeb docs {FINEWEB_SKIP+1:,}+")
    report_lines.append(f"\n{'='*75}")

    full_report = '\n'.join(report_lines)
    print(full_report)

    report_path = COMBINED_DIR / "BUILD_REPORT.txt"
    with open(report_path, "w", encoding="utf-8") as f:
        f.write(full_report)
    print(f"\n  Report saved to: {report_path}")
    print(f"  Done! Final combined dataset ready for training.")


if __name__ == "__main__":
    main()