# -*- 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 = [ ("&", "&"), ("<", "<"), (">", ">"), (""", '"'), ("'", "'"), ("'", "'"), (" ", " "), ("—", " - "), ("–", "-"), ("…", "..."), ("«", '"'), ("»", '"'), ("•", "- "), ("·", " "), ("©", "(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+([.,;:!?])') 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()