Base/scripts/smart_cleanup_english.py

# -*- coding: utf-8 -*-
"""
Smart cleanup of litdata_english_clean based on deep audit results.

Removes truly problematic docs while keeping legitimate educational content.
Strategy:
  - REMOVE: docs with non-English script fragments (Cyrillic, CJK, Arabic, Devanagari)
  - REMOVE: docs under 50 words
  - REMOVE: very repetitive docs (unique word ratio < 0.20)
  - REMOVE: docs with heavy residual code (5+ code patterns)
  - KEEP: educational/historical content that mentions historical terms in context
    (these are Wikipedia articles about WWII, civil rights, etc. - valuable learning)
  - STRIP: clickbait phrases, subscribe prompts, cookie/login boilerplate FROM docs
    (remove the noise parts, keep the content)

Rebuilds clean litdata chunks after filtering.
"""

import json
import os
import re
import time
import unicodedata
from pathlib import Path
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

print("Loading tokenizer...")
tokenizer = Tokenizer.from_file(
    str(ROOT / "Base" / "checkpoints" / "EleutherAI" / "pythia-160m" / "tokenizer.json")
)

# ==============================================================================
# PATTERNS
# ==============================================================================

# Non-English script detectors
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,}')

# Residual code
RE_RESIDUAL_CODE = re.compile(r'(function\s*\(|var\s+\w+\s*=|console\.log|document\.get|if\s*\(\s*\w+\s*[!=]==)', re.I)

# Boilerplate to STRIP from docs (not remove doc, just strip these lines)
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)

# Lines that are just copyright
RE_COPYRIGHT_LINE = re.compile(r'^.*(?:all\s+rights\s+reserved|\(c\)\s*\d{4}|copyright\s+\d{4}).*$', re.I | re.M)

# Number-heavy lines (tables of just numbers)
RE_NUMBER_TABLE_LINE = re.compile(r'^[\d\s,.\-+/%$]+$', re.M)

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

def read_all_tokens(litdata_dir):
    with open(litdata_dir / "index.json") as f:
        index = json.load(f)
    chunks = index["chunks"]
    total_tokens = sum(c["dim"] for c in chunks)
    print(f"  Reading {len(chunks)} chunks ({total_tokens:,} tokens)...")
    all_tokens = np.empty(total_tokens, dtype=DTYPE)
    pos = 0
    for i, chunk in enumerate(chunks):
        chunk_path = litdata_dir / chunk["filename"]
        n_blocks = chunk["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["dim"])
        all_tokens[pos:pos + len(data)] = data
        pos += len(data)
    print(f"    Read {len(chunks)} chunks ({pos:,} tokens)")
    return all_tokens[:pos]


def split_documents(token_stream):
    eos_positions = np.where(token_stream == EOS_TOKEN_ID)[0]
    docs = []
    start = 0
    for eos_pos in eos_positions:
        if eos_pos > start:
            docs.append(token_stream[start:eos_pos])
        start = eos_pos + 1
    if start < len(token_stream):
        docs.append(token_stream[start:])
    return docs


def write_litdata_chunks(output_dir, token_stream, config):
    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 = 0
    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)")
    index = {"chunks": chunks_metadata, "config": config, "updated_at": str(time.time())}
    with open(os.path.join(output_dir, "index.json"), "w") as f:
        json.dump(index, f, indent=2)
    return chunks_metadata


# ==============================================================================
# SMART CLEANUP
# ==============================================================================

def should_remove(text):
    """Returns (remove: bool, reason: str or None)."""
    words = text.split()
    word_count = len(words)

    # Too short
    if word_count < 50:
        return True, f"too short ({word_count} words)"

    # Non-English scripts
    scripts_found = []
    if RE_CJK.search(text):
        scripts_found.append("CJK")
    if RE_ARABIC.search(text):
        scripts_found.append("Arabic")
    if RE_CYRILLIC.search(text):
        scripts_found.append("Cyrillic")
    if RE_DEVANAGARI.search(text):
        scripts_found.append("Devanagari")
    if scripts_found:
        return True, f"non-English scripts: {', '.join(scripts_found)}"

    # Very repetitive
    if word_count > 50:
        unique_ratio = len(set(w.lower() for w in words)) / word_count
        if unique_ratio < 0.20:
            return True, f"very repetitive (unique ratio: {unique_ratio:.3f})"

    # Heavy residual code (5+ code patterns)
    code_matches = RE_RESIDUAL_CODE.findall(text)
    if len(code_matches) >= 5:
        return True, f"residual code ({len(code_matches)} matches)"

    return False, None


def strip_boilerplate(text):
    """Strip boilerplate lines from a document, keeping the educational content."""
    original_len = len(text)

    # Strip specific boilerplate patterns (line by line removal)
    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)

    # Strip pure number table lines (> 80% digits/spaces)
    lines = text.split('\n')
    clean_lines = []
    stripped_number_lines = 0
    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:
                stripped_number_lines += 1
                continue
        clean_lines.append(line)
    text = '\n'.join(clean_lines)

    # Clean up resulting whitespace
    text = re.sub(r'\n{3,}', '\n\n', text)
    text = text.strip()

    chars_stripped = original_len - len(text)
    return text, chars_stripped, stripped_number_lines


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

def main():
    input_dir = ROOT / "Base" / "data" / "litdata_english_clean"
    output_dir = ROOT / "Base" / "data" / "litdata_english_clean"

    print(f"\n{'='*75}")
    print(f"  SMART CLEANUP: litdata_english_clean")
    print(f"  Input/Output: {input_dir}")
    print(f"{'='*75}")

    # 1. Read and decode
    t0 = time.time()
    token_stream = read_all_tokens(input_dir)
    doc_tokens = split_documents(token_stream)
    total_input = len(doc_tokens)
    print(f"  Found {total_input:,} documents")
    del token_stream

    print(f"  Decoding ALL {total_input:,} documents...")
    texts = []
    t1 = time.time()
    for i, toks in enumerate(doc_tokens):
        text = tokenizer.decode(toks.tolist(), skip_special_tokens=False)
        texts.append(text)
        if (i + 1) % 20000 == 0 or i == total_input - 1:
            print(f"    Decoded {i+1:,}/{total_input:,}")
    del doc_tokens
    print(f"  Decoded in {time.time()-t1:.1f}s")

    # 2. Filter and strip
    print(f"\n  Processing {total_input:,} documents...")
    t2 = time.time()

    kept_texts = []
    removed_reasons = Counter()
    total_boilerplate_chars = 0
    total_number_lines_stripped = 0
    removed_examples = []

    for i, text in enumerate(texts):
        # Check if doc should be removed entirely
        remove, reason = should_remove(text)
        if remove:
            removed_reasons[reason.split('(')[0].strip().split(':')[0].strip()] += 1
            if len(removed_examples) < 20:
                removed_examples.append((i, reason, text[:300]))
            continue

        # Strip boilerplate from kept docs
        cleaned, chars_stripped, num_lines = strip_boilerplate(text)
        total_boilerplate_chars += chars_stripped
        total_number_lines_stripped += num_lines

        # Final check: did stripping make it too short?
        if len(cleaned.split()) < 50:
            removed_reasons["stripped too short"] += 1
            continue

        kept_texts.append(cleaned)

        if (i + 1) % 10000 == 0 or i == total_input - 1:
            print(f"    Processed {i+1:,}/{total_input:,} | kept={len(kept_texts):,} | removed={i+1-len(kept_texts):,}")

    total_removed = total_input - len(kept_texts)
    print(f"  Processing done in {time.time()-t2:.1f}s")
    print(f"  Kept: {len(kept_texts):,} | Removed: {total_removed:,} ({total_removed/total_input*100:.2f}%)")

    # 3. Re-tokenize
    t3 = time.time()
    print(f"\n  Re-tokenizing {len(kept_texts):,} documents...")
    all_token_ids = []
    total_new_tokens = 0
    for i, text in enumerate(kept_texts):
        enc = tokenizer.encode(text)
        ids = enc.ids
        all_token_ids.extend(ids)
        all_token_ids.append(EOS_TOKEN_ID)
        total_new_tokens += len(ids) + 1
        if (i + 1) % 20000 == 0 or i == len(kept_texts) - 1:
            print(f"    Tokenized {i+1:,}/{len(kept_texts):,} ({total_new_tokens:,} tokens)")
    print(f"  Tokenized in {time.time()-t3:.1f}s")

    # 4. Rebuild chunks
    t4 = time.time()
    print(f"\n  Rebuilding litdata chunks...")
    token_array = np.array(all_token_ids, dtype=DTYPE)
    del all_token_ids

    import gc
    gc.collect()

    os.makedirs(str(output_dir), exist_ok=True)
    config = {
        "block_size": BLOCK_SIZE,
        "vocab_size": tokenizer.get_vocab_size(),
    }
    chunks = write_litdata_chunks(str(output_dir), token_array, config)
    del token_array
    print(f"  Written {len(chunks)} chunks in {time.time()-t4:.1f}s")

    # 5. Summary report
    total_time = time.time() - t0

    report = []
    report.append(f"\n{'='*75}")
    report.append(f"  SMART CLEANUP REPORT - litdata_english_clean")
    report.append(f"{'='*75}")
    report.append(f"\n  Processing time: {total_time:.0f}s ({total_time/60:.1f} min)")
    report.append(f"\n  DOCUMENT COUNTS")
    report.append(f"  {'-'*50}")
    report.append(f"    Input documents:     {total_input:>10,}")
    report.append(f"    Output documents:    {len(kept_texts):>10,}")
    report.append(f"    Removed:             {total_removed:>10,} ({total_removed/total_input*100:.2f}%)")

    report.append(f"\n  TOKEN COUNTS")
    report.append(f"  {'-'*50}")
    report.append(f"    Input tokens:        {53446575:>15,}")
    report.append(f"    Output tokens:       {total_new_tokens:>15,}")
    report.append(f"    Tokens after align:  {sum(c['dim'] for c in chunks):>15,}")

    report.append(f"\n  REMOVAL REASONS")
    report.append(f"  {'-'*50}")
    for reason, count in sorted(removed_reasons.items(), key=lambda x: -x[1]):
        report.append(f"    {reason:<35} {count:>6,}")

    report.append(f"\n  BOILERPLATE STRIPPED (from kept docs)")
    report.append(f"  {'-'*50}")
    report.append(f"    Total boilerplate chars stripped: {total_boilerplate_chars:>10,}")
    report.append(f"    Number-heavy lines removed:      {total_number_lines_stripped:>10,}")

    report.append(f"\n  REMOVED DOCUMENT EXAMPLES (first 15)")
    report.append(f"  {'-'*50}")
    for idx, reason, preview in removed_examples[:15]:
        preview_clean = preview.replace('\n', ' ')[:200]
        report.append(f"\n    Doc #{idx} - {reason}")
        report.append(f"      \"{preview_clean}...\"")

    report.append(f"\n  VERDICT")
    report.append(f"  {'-'*50}")
    report.append(f"    Removed non-English fragments, very short docs, repetitive content,")
    report.append(f"    and heavy code. Stripped boilerplate from all remaining docs.")
    report.append(f"    Dataset is now optimized for English language learning!")
    report.append(f"\n{'='*75}")

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

    report_path = output_dir / "SMART_CLEANUP_REPORT.txt"
    with open(report_path, "w", encoding="utf-8") as f:
        f.write(full_report)
    print(f"\n  Report saved to: {report_path}")


if __name__ == "__main__":
    main()