# -*- 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 = [ ("&", "&"), ("<", "<"), (">", ">"), (""", '"'), ("'", "'"), ("'", "'"), (" ", " "), ("—", " - "), ("–", "-"), ("…", "..."), ("«", '"'), ("»", '"'), ("•", "- "), ("·", " "), ("©", "(c)"), ("®", "(R)"), ("™", "(TM)"), ("°", " 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']*)?\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()