Base/scripts/build_5b_pretrain.py

# -*- coding: utf-8 -*-
"""
Build 1.2B tokens of DIVERSE, CLEAN English data to bring pretrain total to 5B.

Current: 3,819,212,525 tokens in litdata_pretrain_final
Target:  5,000,000,000 tokens
Gap:     ~1,181,000,000 tokens
Build:   ~1,250,000,000 est. tokens (buffer for cleaning loss)

DIVERSITY STRATEGY — 3 completely different source types:
  1. Wikipedia (skip 200K qualifying articles) — encyclopedic knowledge
  2. FineWeb-Edu (skip 200K qualifying docs, score ≥ 3.5) — educational web
  3. OpenWebText (Skylion007) — Reddit-curated quality English web pages
     * Completely new source, zero overlap with anything used before

Pipeline per source:
  Download (streaming) → Deep clean → Smart filter → Tokenize → Stream-write

Memory-efficient: StreamingChunkWriter, multiprocessing cleaning (30 cores).
Output appended directly to litdata_pretrain_final.
"""

import json
import os
import re
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
NUM_WORKERS = max(1, cpu_count() - 2)
ENCODE_BATCH = 8000
TOKENS_PER_WORD = 1.3

DATA_DIR = ROOT / "Base" / "data"
FINAL_DIR = DATA_DIR / "litdata_pretrain_final"
TOKENIZER_PATH = str(ROOT / "Base" / "checkpoints" / "EleutherAI" / "pythia-160m" / "tokenizer.json")

# ── Source targets ────────────────────────────────────────────────────────
# Slightly over 1.2B to account for cleaning losses (~2-3%)
WIKI_TARGET     = 365_000_000   # ~30% — encyclopedic
FINEWEB_TARGET  = 445_000_000   # ~36% — educational web
OWT_TARGET      = 445_000_000   # ~34% — Reddit-curated diverse web
TOTAL_TARGET    = WIKI_TARGET + FINEWEB_TARGET + OWT_TARGET  # ~1.255B

# Skip values — must exceed ALL previously used qualifying articles
WIKI_SKIP    = 200_000   # previous max was 125K + articles collected
FINEWEB_SKIP = 200_000   # previous max was 125K + docs collected
FINEWEB_MIN_SCORE = 3.5  # slightly broader than 4.0 for more diversity

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


# ==============================================================================
# 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 clean_text_basic(text):
    """Light quality filter for download phase."""
    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


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, "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):
    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
    return stripped, None, len(result) - len(stripped)


# ==============================================================================
# STREAMING CHUNK WRITER (appends to existing litdata)
# ==============================================================================

class StreamingChunkWriter:
    def __init__(self, output_dir, config, start_chunk_idx=0):
        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 = start_chunk_idx
        self.total_tokens = 0

    def add_tokens_batch(self, encoded_batch):
        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):
        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 % 10 == 0:
            print(f"        Flushed chunk {self.chunk_idx} ({self.total_tokens:,} new tokens)")

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


# ==============================================================================
# PROCESS ONE SOURCE: download → clean (multiprocessed) → tokenize → write
# ==============================================================================

def process_source_batch(texts, pool, tokenizer, writer, source_name, batch_num):
    """Clean a batch of texts and write to the streaming writer. Returns stats."""
    t0 = time.time()

    # Clean with multiprocessing
    results = pool.map(clean_and_filter, texts, chunksize=256)

    cleaned = []
    dropped = 0
    reasons = Counter()
    boilerplate = 0
    for text, reason, bp in results:
        if text is not None:
            cleaned.append(text)
            boilerplate += bp
        else:
            dropped += 1
            reasons[reason] += 1
    del results

    # Tokenize + write
    for i in range(0, len(cleaned), ENCODE_BATCH):
        sub = cleaned[i:i+ENCODE_BATCH]
        encoded = tokenizer.encode_batch(sub, add_special_tokens=False)
        writer.add_tokens_batch(encoded)
    del cleaned

    elapsed = time.time() - t0
    print(f"      [{source_name}] batch {batch_num}: kept {len(texts)-dropped:,} / dropped {dropped} | {elapsed:.1f}s")

    return len(texts) - dropped, dropped, reasons, boilerplate


# ==============================================================================
# MAIN
# ==============================================================================

def main():
    from datasets import load_dataset

    t_start = time.time()

    print("Loading tokenizer...")
    tokenizer = Tokenizer.from_file(TOKENIZER_PATH)
    config = {"block_size": BLOCK_SIZE, "vocab_size": tokenizer.get_vocab_size()}

    # Read existing index to find starting chunk offset
    with open(FINAL_DIR / "index.json") as f:
        existing_index = json.load(f)
    existing_chunks = existing_index["chunks"]
    existing_tokens = sum(c["dim"] for c in existing_chunks)
    start_chunk_idx = len(existing_chunks)

    print(f"\n{'='*75}")
    print(f"  BUILD 1.2B DIVERSE TOKENS → APPEND TO litdata_pretrain_final")
    print(f"  Current:  {existing_tokens:,} tokens ({start_chunk_idx} chunks)")
    print(f"  Target:   5,000,000,000 tokens")
    print(f"  Building: ~{TOTAL_TARGET:,} estimated tokens")
    print(f"  Workers:  {NUM_WORKERS} CPU cores")
    print(f"{'='*75}")

    # Initialize writer that appends new chunks after existing ones
    writer = StreamingChunkWriter(str(FINAL_DIR), config, start_chunk_idx=start_chunk_idx)
    pool = Pool(processes=NUM_WORKERS)

    # Global stats
    all_stats = {}
    DOWNLOAD_BATCH = 10000  # process 10K docs at a time

    # ═══════════════════════════════════════════════════════════════════════
    # SOURCE 1: Wikipedia (encyclopedic knowledge)
    # ═══════════════════════════════════════════════════════════════════════
    print(f"\n{'='*75}")
    print(f"  SOURCE 1: WIKIPEDIA (skip {WIKI_SKIP:,}, target ~{WIKI_TARGET:,} tokens)")
    print(f"{'='*75}")

    ds_wiki = load_dataset(
        "wikimedia/wikipedia", "20231101.en",
        split="train", streaming=True, trust_remote_code=False
    )

    wiki_texts = []
    wiki_tokens_est = 0
    wiki_seen = 0
    wiki_skipped_quality = 0
    wiki_skipped_meta = 0
    wiki_total_kept = 0
    wiki_total_dropped = 0
    wiki_reasons = Counter()
    wiki_boilerplate = 0
    wiki_batch_num = 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_basic(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 % 25000 == 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 += est_tok

        # Process in batches to keep memory low
        if len(wiki_texts) >= DOWNLOAD_BATCH:
            wiki_batch_num += 1
            kept, dropped, reasons, bp = process_source_batch(
                wiki_texts, pool, tokenizer, writer, "Wiki", wiki_batch_num
            )
            wiki_total_kept += kept
            wiki_total_dropped += dropped
            wiki_reasons += reasons
            wiki_boilerplate += bp
            wiki_texts = []

        if wiki_tokens_est >= WIKI_TARGET:
            break

    # Process remaining
    if wiki_texts:
        wiki_batch_num += 1
        kept, dropped, reasons, bp = process_source_batch(
            wiki_texts, pool, tokenizer, writer, "Wiki", wiki_batch_num
        )
        wiki_total_kept += kept
        wiki_total_dropped += dropped
        wiki_reasons += reasons
        wiki_boilerplate += bp
        wiki_texts = []

    wiki_elapsed = time.time() - t0
    print(f"  [Wikipedia] Done: ~{wiki_tokens_est:,} est. tokens, kept {wiki_total_kept:,}, dropped {wiki_total_dropped:,} in {wiki_elapsed:.0f}s")
    print(f"    (skipped {WIKI_SKIP:,} already-used + {wiki_skipped_quality:,} low-quality + {wiki_skipped_meta:,} meta)")
    all_stats["Wikipedia"] = {
        "est_tokens": wiki_tokens_est, "kept": wiki_total_kept,
        "dropped": wiki_total_dropped, "reasons": wiki_reasons,
        "boilerplate": wiki_boilerplate, "time": wiki_elapsed,
    }

    # ═══════════════════════════════════════════════════════════════════════
    # SOURCE 2: FineWeb-Edu (educational web content)
    # ═══════════════════════════════════════════════════════════════════════
    print(f"\n{'='*75}")
    print(f"  SOURCE 2: FINEWEB-EDU (skip {FINEWEB_SKIP:,}, score >= {FINEWEB_MIN_SCORE}, target ~{FINEWEB_TARGET:,} tokens)")
    print(f"{'='*75}")

    ds_fw = load_dataset(
        "HuggingFaceFW/fineweb-edu", "sample-10BT",
        split="train", streaming=True, trust_remote_code=False
    )

    fw_texts = []
    fw_tokens_est = 0
    fw_seen = 0
    fw_skipped_quality = 0
    fw_skipped_score = 0
    fw_total_kept = 0
    fw_total_dropped = 0
    fw_reasons = Counter()
    fw_boilerplate = 0
    fw_batch_num = 0
    t0 = time.time()

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

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

        fw_seen += 1
        if fw_seen <= FINEWEB_SKIP:
            if fw_seen % 25000 == 0:
                print(f"    Skipping... {fw_seen:,}/{FINEWEB_SKIP:,}")
            continue

        est_tok = int(len(cleaned.split()) * TOKENS_PER_WORD)
        fw_texts.append(cleaned)
        fw_tokens_est += est_tok

        if len(fw_texts) >= DOWNLOAD_BATCH:
            fw_batch_num += 1
            kept, dropped, reasons, bp = process_source_batch(
                fw_texts, pool, tokenizer, writer, "FineWeb", fw_batch_num
            )
            fw_total_kept += kept
            fw_total_dropped += dropped
            fw_reasons += reasons
            fw_boilerplate += bp
            fw_texts = []

        if fw_tokens_est >= FINEWEB_TARGET:
            break

    if fw_texts:
        fw_batch_num += 1
        kept, dropped, reasons, bp = process_source_batch(
            fw_texts, pool, tokenizer, writer, "FineWeb", fw_batch_num
        )
        fw_total_kept += kept
        fw_total_dropped += dropped
        fw_reasons += reasons
        fw_boilerplate += bp
        fw_texts = []

    fw_elapsed = time.time() - t0
    print(f"  [FineWeb-Edu] Done: ~{fw_tokens_est:,} est. tokens, kept {fw_total_kept:,}, dropped {fw_total_dropped:,} in {fw_elapsed:.0f}s")
    print(f"    (skipped {FINEWEB_SKIP:,} already-used + {fw_skipped_quality:,} low-quality + {fw_skipped_score:,} low-score)")
    all_stats["FineWeb-Edu"] = {
        "est_tokens": fw_tokens_est, "kept": fw_total_kept,
        "dropped": fw_total_dropped, "reasons": fw_reasons,
        "boilerplate": fw_boilerplate, "time": fw_elapsed,
    }

    # ═══════════════════════════════════════════════════════════════════════
    # SOURCE 3: OpenWebText (Reddit-curated diverse web pages)
    # ═══════════════════════════════════════════════════════════════════════
    print(f"\n{'='*75}")
    print(f"  SOURCE 3: OPENWEBTEXT (target ~{OWT_TARGET:,} tokens)")
    print(f"  Completely new source — zero overlap with existing data")
    print(f"{'='*75}")

    ds_owt = load_dataset(
        "Skylion007/openwebtext",
        split="train", streaming=True, trust_remote_code=False
    )

    owt_texts = []
    owt_tokens_est = 0
    owt_skipped_quality = 0
    owt_total_kept = 0
    owt_total_dropped = 0
    owt_reasons = Counter()
    owt_boilerplate = 0
    owt_batch_num = 0
    t0 = time.time()

    for doc in ds_owt:
        raw = doc.get("text") or ""
        cleaned = clean_text_basic(raw)
        if not is_high_quality(cleaned, min_chars=400, min_words=60):
            owt_skipped_quality += 1
            continue

        est_tok = int(len(cleaned.split()) * TOKENS_PER_WORD)
        owt_texts.append(cleaned)
        owt_tokens_est += est_tok

        if len(owt_texts) >= DOWNLOAD_BATCH:
            owt_batch_num += 1
            kept, dropped, reasons, bp = process_source_batch(
                owt_texts, pool, tokenizer, writer, "OWT", owt_batch_num
            )
            owt_total_kept += kept
            owt_total_dropped += dropped
            owt_reasons += reasons
            owt_boilerplate += bp
            owt_texts = []

        if owt_tokens_est >= OWT_TARGET:
            break

    if owt_texts:
        owt_batch_num += 1
        kept, dropped, reasons, bp = process_source_batch(
            owt_texts, pool, tokenizer, writer, "OWT", owt_batch_num
        )
        owt_total_kept += kept
        owt_total_dropped += dropped
        owt_reasons += reasons
        owt_boilerplate += bp
        owt_texts = []

    owt_elapsed = time.time() - t0
    print(f"  [OpenWebText] Done: ~{owt_tokens_est:,} est. tokens, kept {owt_total_kept:,}, dropped {owt_total_dropped:,} in {owt_elapsed:.0f}s")
    print(f"    (skipped {owt_skipped_quality:,} low-quality)")
    all_stats["OpenWebText"] = {
        "est_tokens": owt_tokens_est, "kept": owt_total_kept,
        "dropped": owt_total_dropped, "reasons": owt_reasons,
        "boilerplate": owt_boilerplate, "time": owt_elapsed,
    }

    pool.close()
    pool.join()

    # ═══════════════════════════════════════════════════════════════════════
    # FINALIZE: flush remaining + update index.json
    # ═══════════════════════════════════════════════════════════════════════
    print(f"\n{'='*75}")
    print(f"  FINALIZING")
    print(f"{'='*75}")

    new_tokens, discarded = writer.finalize()

    # Merge new chunk metadata with existing
    final_chunks = existing_chunks + writer.chunks_metadata
    final_total = existing_tokens + new_tokens
    final_index = {
        "chunks": final_chunks,
        "config": config,
        "updated_at": str(time.time()),
    }
    with open(FINAL_DIR / "index.json", "w") as f:
        json.dump(final_index, f, indent=2)

    total_time = time.time() - t_start

    # ═══════════════════════════════════════════════════════════════════════
    # REPORT
    # ═══════════════════════════════════════════════════════════════════════
    total_new_kept = wiki_total_kept + fw_total_kept + owt_total_kept
    total_new_dropped = wiki_total_dropped + fw_total_dropped + owt_total_dropped
    total_new_boilerplate = wiki_boilerplate + fw_boilerplate + owt_boilerplate
    all_drop_reasons = wiki_reasons + fw_reasons + owt_reasons

    report = []
    report.append(f"{'='*75}")
    report.append(f"  5 BILLION TOKEN PRETRAIN DATASET — BUILD REPORT")
    report.append(f"{'='*75}")
    report.append(f"")
    report.append(f"  Total time: {total_time:.0f}s ({total_time/60:.1f} min)")
    report.append(f"  Workers: {NUM_WORKERS} CPU cores")
    report.append(f"")
    report.append(f"  NEW DATA ADDED (diverse, clean English)")
    report.append(f"  {'-'*60}")

    for name, stats in all_stats.items():
        report.append(f"    {name}:")
        report.append(f"      Est tokens: ~{stats['est_tokens']:,}")
        report.append(f"      Kept: {stats['kept']:,} | Dropped: {stats['dropped']:,}")
        report.append(f"      Boilerplate: {stats['boilerplate']:,} chars")
        report.append(f"      Time: {stats['time']:.0f}s")
        if stats["reasons"]:
            for reason, count in sorted(stats["reasons"].items(), key=lambda x: -x[1]):
                report.append(f"        {reason}: {count:,}")

    report.append(f"")
    report.append(f"  NEW DATA TOTALS")
    report.append(f"  {'-'*60}")
    report.append(f"    Documents kept:    {total_new_kept:,}")
    report.append(f"    Documents dropped: {total_new_dropped:,}")
    report.append(f"    Boilerplate:       {total_new_boilerplate:,} chars stripped")
    report.append(f"    New tokens:        {new_tokens:,} ({writer.chunk_idx - start_chunk_idx} chunks)")
    if all_drop_reasons:
        report.append(f"    Drop reasons:")
        for reason, count in sorted(all_drop_reasons.items(), key=lambda x: -x[1]):
            report.append(f"      {reason:<35} {count:>8,}")

    report.append(f"")
    report.append(f"  FINAL COMBINED DATASET")
    report.append(f"  {'-'*60}")
    report.append(f"    Location:  {FINAL_DIR}")
    report.append(f"    Chunks:    {len(final_chunks)}")
    report.append(f"    Tokens:    {final_total:,}")
    report.append(f"    Format:    litdata binary (int32, BLOCK_SIZE=1025, EOS=0)")
    report.append(f"")
    report.append(f"    Previous:  {existing_tokens:,} tokens ({start_chunk_idx} chunks)")
    report.append(f"    + Added:   {new_tokens:,} tokens ({writer.chunk_idx - start_chunk_idx} chunks)")
    report.append(f"    = Total:   {final_total:,} tokens ({len(final_chunks)} chunks)")
    report.append(f"")
    report.append(f"  DATA COMPOSITION")
    report.append(f"  {'-'*60}")
    report.append(f"    litdata_3b_clean:    ~2.94B tokens (general web, cleaned)")
    report.append(f"    litdata_english_500m: ~515M tokens (Wiki+FineWeb, cleaned)")
    report.append(f"    litdata_combined:     ~257M tokens (Wiki+FineWeb, cleaned)")
    report.append(f"    + Wikipedia (new):    ~{wiki_tokens_est:,} est. (articles {WIKI_SKIP+1:,}+)")
    report.append(f"    + FineWeb-Edu (new):  ~{fw_tokens_est:,} est. (score>={FINEWEB_MIN_SCORE}, docs {FINEWEB_SKIP+1:,}+)")
    report.append(f"    + OpenWebText (new):  ~{owt_tokens_est:,} est. (Reddit-curated, no overlap)")
    report.append(f"")
    report.append(f"  PURE ENGLISH PRETRAINING TEXT")
    report.append(f"  Sources: Wikipedia, FineWeb-Edu, OpenWebText, general web")
    report.append(f"  NO instruction/finetune data included")
    report.append(f"  ZERO overlap between all data sources")
    report.append(f"{'='*75}")

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

    with open(FINAL_DIR / "BUILD_REPORT.txt", "w", encoding="utf-8") as f:
        f.write(full_report)
    print(f"\n  Report saved to: {FINAL_DIR / 'BUILD_REPORT.txt'}")
    print(f"  Done! 5B token pretrain dataset ready.")


if __name__ == "__main__":
    main()