Upload Base/scripts/build_english_curriculum_1b.py with huggingface_hub

Base/scripts/build_english_curriculum_1b.py

@@ -0,0 +1,608 @@
+"""
+Build a new 1B-token English curriculum corpus for continued pretraining.
+
+Goal:
+  Improve English understanding, grammar, vocabulary, and meaningful prose
+  without just reusing the existing litdata_pretrain_final mix.
+
+Curriculum mix:
+  1. Simple Wikipedia      -> simpler explanatory English
+  2. Wikipedia tail        -> formal, fact-dense English not used earlier
+  3. FineWeb-Edu tail      -> educational/tutorial English not used earlier
+  4. PG-19 books           -> long-form prose, dialogue, vocabulary richness
+
+Output:
+  Parquet files with one 'text' column in:
+    Base/data/filtered_english_curriculum_1b/
+
+Then convert to LitData with:
+  python Base/scripts/prepare_litdata.py \
+    --filtered_dir Base/data/filtered_english_curriculum_1b \
+    --output_dir Base/data/litdata_english_curriculum_1b \
+    --label ENGLISH_CURRICULUM_1B
+"""
+
+import argparse
+import hashlib
+import json
+import re
+import time
+import unicodedata
+from collections import Counter
+from pathlib import Path
+
+import pyarrow as pa
+import pyarrow.parquet as pq
+
+
+TOKENS_PER_WORD = 1.3
+DEFAULT_TARGET_TOKENS = 1_000_000_000
+DEFAULT_SHARES = {
+    "simplewiki": 0.15,
+    "wiki_tail": 0.20,
+    "fineweb_tail": 0.25,
+    "pg19": 0.40,
+}
+DEFAULT_WIKI_SKIP = 250_000
+DEFAULT_FINEWEB_SKIP = 250_000
+DEFAULT_FINEWEB_MIN_SCORE = 4.0
+FLUSH_DOCS = 5_000
+
+WIKI_SKIP_PATTERNS = re.compile(
+    r"(disambiguation|list of|lists of|index of|outline of|"
+    r"wikipedia:|template:|category:|portal:|module:|mediawiki:)",
+    re.IGNORECASE,
+)
+
+CONTROL_CHARS = re.compile(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f-\x9f]")
+RE_URL = re.compile(r"https?://\S+|www\.\S+", re.I)
+RE_HTML_TAG = re.compile(r"</?[a-zA-Z][a-zA-Z0-9]*(?:\s[^>]*)?\s*/?>")
+RE_HTML_COMMENT = re.compile(r"<!--.*?-->", re.DOTALL)
+RE_REPEATED_LINE = re.compile(r"^(.{20,})\n(?:\1\n?)+", re.M)
+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_SPACE_BEFORE_PUNCT = re.compile(r"\s+([.,;:!?])")
+RE_DOUBLE_PERIOD = re.compile(r"\.{2}(?!\.)")
+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,}")
+
+PG_START_MARKERS = [
+    "*** START OF THE PROJECT GUTENBERG EBOOK",
+    "*** START OF THIS PROJECT GUTENBERG EBOOK",
+    "START OF THE PROJECT GUTENBERG EBOOK",
+]
+PG_END_MARKERS = [
+    "*** END OF THE PROJECT GUTENBERG EBOOK",
+    "*** END OF THIS PROJECT GUTENBERG EBOOK",
+    "END OF THE PROJECT GUTENBERG EBOOK",
+]
+
+
+def clean_text_basic(text: str) -> str:
+    text = unicodedata.normalize("NFKC", text)
+    text = CONTROL_CHARS.sub("", text)
+    text = RE_HTML_COMMENT.sub("", text)
+    text = RE_HTML_TAG.sub("", text)
+    text = RE_URL.sub("", text)
+    text = RE_REPEATED_LINE.sub(r"\1", 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 = RE_MULTI_SPACE.sub(" ", text)
+    text = RE_MULTI_NEWLINE.sub("\n\n\n", text)
+    text = RE_TRAILING_SPACE.sub("", 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("\u00a0", " ")
+    text = "\n".join(line.strip() for line in text.splitlines())
+    return text.strip()
+
+
+def is_high_quality(text: str, min_chars: int, min_words: int, alpha_ratio: float = 0.68) -> bool:
+    if len(text) < min_chars:
+        return False
+
+    words = text.split()
+    word_count = len(words)
+    if word_count < min_words:
+        return False
+
+    alpha = sum(c.isalpha() for c in text)
+    if alpha / max(len(text), 1) < alpha_ratio:
+        return False
+
+    avg_word_len = sum(len(word) for word in words) / max(word_count, 1)
+    if avg_word_len < 2.5 or avg_word_len > 15:
+        return False
+
+    url_hits = len(RE_URL.findall(text))
+    if url_hits > word_count * 0.02:
+        return False
+
+    if RE_CJK.search(text) or RE_ARABIC.search(text) or RE_CYRILLIC.search(text) or RE_DEVANAGARI.search(text):
+        return False
+
+    sentences = [segment.strip() for segment in re.split(r"[.!?]+", text) if len(segment.strip()) > 12]
+    if len(sentences) < 3:
+        return False
+
+    sentence_lengths = [len(sentence.split()) for sentence in sentences]
+    avg_sentence_words = sum(sentence_lengths) / max(len(sentence_lengths), 1)
+    if avg_sentence_words < 5 or avg_sentence_words > 40:
+        return False
+
+    if word_count >= 120:
+        unique_word_ratio = len(set(word.lower() for word in words)) / word_count
+        if unique_word_ratio < 0.22:
+            return False
+
+    lines = [line.strip() for line in text.splitlines() if line.strip()]
+    if len(lines) > 5:
+        unique_line_ratio = len(set(lines)) / len(lines)
+        if unique_line_ratio < 0.55:
+            return False
+
+    return True
+
+
+def normalize_for_hash(text: str) -> str:
+    return " ".join(text.lower().split())[:8192]
+
+
+def fingerprint(text: str) -> int:
+    digest = hashlib.blake2b(normalize_for_hash(text).encode("utf-8"), digest_size=8).digest()
+    return int.from_bytes(digest, byteorder="little", signed=False)
+
+
+def approx_tokens(text: str) -> int:
+    return max(1, int(len(text.split()) * TOKENS_PER_WORD))
+
+
+def strip_project_gutenberg(text: str) -> str:
+    cleaned = text
+    upper = cleaned.upper()
+
+    for marker in PG_START_MARKERS:
+        pos = upper.find(marker)
+        if pos != -1:
+            start = cleaned.find("\n", pos)
+            if start != -1:
+                cleaned = cleaned[start + 1:]
+                upper = cleaned.upper()
+            break
+
+    for marker in PG_END_MARKERS:
+        pos = upper.find(marker)
+        if pos != -1:
+            cleaned = cleaned[:pos]
+            break
+
+    return cleaned.strip()
+
+
+def split_pg19_passages(title: str, text: str, min_words: int = 220, target_words: int = 1200, max_words: int = 1800):
+    paragraphs = []
+    for paragraph in re.split(r"\n\s*\n", text):
+        cleaned = clean_text_basic(paragraph)
+        if len(cleaned.split()) >= 35:
+            paragraphs.append(cleaned)
+
+    current = []
+    current_words = 0
+    passage_idx = 0
+
+    for paragraph in paragraphs:
+        paragraph_words = len(paragraph.split())
+        if current and current_words >= target_words and current_words + paragraph_words > max_words:
+            passage_idx += 1
+            body = "\n\n".join(current)
+            yield f"{title}\n\n{body}", passage_idx
+            current = [paragraph]
+            current_words = paragraph_words
+            continue
+
+        current.append(paragraph)
+        current_words += paragraph_words
+
+        if current_words >= max_words:
+            passage_idx += 1
+            body = "\n\n".join(current)
+            yield f"{title}\n\n{body}", passage_idx
+            current = []
+            current_words = 0
+
+    if current and current_words >= min_words:
+        passage_idx += 1
+        body = "\n\n".join(current)
+        yield f"{title}\n\n{body}", passage_idx
+
+
+class ParquetSink:
+    def __init__(self, output_dir: Path):
+        self.output_dir = output_dir
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+        self.buffers = {}
+        self.file_counts = Counter()
+
+    def add(self, source_name: str, text: str):
+        bucket = self.buffers.setdefault(source_name, [])
+        bucket.append(text)
+        if len(bucket) >= FLUSH_DOCS:
+            self.flush(source_name)
+
+    def flush(self, source_name: str):
+        bucket = self.buffers.get(source_name)
+        if not bucket:
+            return
+        file_idx = self.file_counts[source_name]
+        output_path = self.output_dir / f"{source_name}_{file_idx:04d}.parquet"
+        pq.write_table(pa.table({"text": bucket}), str(output_path))
+        self.file_counts[source_name] += 1
+        self.buffers[source_name] = []
+
+    def finalize(self):
+        for source_name in list(self.buffers):
+            self.flush(source_name)
+
+
+def try_load_streaming_dataset(candidates, split="train"):
+    from datasets import load_dataset
+
+    errors = []
+    for repo_id, config_name in candidates:
+        try:
+            if config_name is None:
+                dataset = load_dataset(repo_id, split=split, streaming=True, trust_remote_code=False)
+            else:
+                dataset = load_dataset(repo_id, config_name, split=split, streaming=True, trust_remote_code=False)
+            resolved = f"{repo_id} [{config_name}]" if config_name else repo_id
+            print(f"  Loaded source: {resolved}")
+            return dataset
+        except Exception as exc:
+            errors.append(f"{repo_id} [{config_name}]: {exc}")
+
+    raise RuntimeError("Unable to load streaming dataset. Tried:\n  " + "\n  ".join(errors))
+
+
+def maybe_keep(text: str, seen_hashes: set, min_chars: int, min_words: int, alpha_ratio: float = 0.68):
+    cleaned = clean_text_basic(text)
+    if not is_high_quality(cleaned, min_chars=min_chars, min_words=min_words, alpha_ratio=alpha_ratio):
+        return None, "quality"
+
+    key = fingerprint(cleaned)
+    if key in seen_hashes:
+        return None, "duplicate"
+
+    seen_hashes.add(key)
+    return cleaned, None
+
+
+def log_progress(name: str, accepted_tokens: int, accepted_docs: int, started_at: float):
+    elapsed = max(time.time() - started_at, 1.0)
+    rate = accepted_tokens / elapsed
+    print(f"    {name}: docs={accepted_docs:,} est_tokens={accepted_tokens:,} rate={rate:,.0f} tok/s elapsed={elapsed/60:.1f}m")
+
+
+def build_simplewiki(target_tokens: int, sink: ParquetSink, seen_hashes: set, stats: dict):
+    dataset = try_load_streaming_dataset([
+        ("wikimedia/wikipedia", "20231101.simple"),
+        ("wikimedia/wikipedia", "20231101.simplewiki"),
+        ("wikimedia/wikipedia", "20231101.simple_en"),
+    ])
+
+    started_at = time.time()
+    accepted_tokens = 0
+    accepted_docs = 0
+    skipped = Counter()
+
+    for article in dataset:
+        title = (article.get("title") or "").strip()
+        raw_text = article.get("text") or ""
+
+        if WIKI_SKIP_PATTERNS.search(title):
+            skipped["meta"] += 1
+            continue
+
+        candidate, reason = maybe_keep(f"{title}\n\n{raw_text}", seen_hashes, min_chars=250, min_words=50, alpha_ratio=0.70)
+        if candidate is None:
+            skipped[reason] += 1
+            continue
+
+        sink.add("simplewiki", candidate)
+        accepted_docs += 1
+        accepted_tokens += approx_tokens(candidate)
+
+        if accepted_docs % 2_000 == 0:
+            log_progress("simplewiki", accepted_tokens, accepted_docs, started_at)
+
+        if accepted_tokens >= target_tokens:
+            break
+
+    stats["simplewiki"] = {
+        "accepted_docs": accepted_docs,
+        "accepted_tokens": accepted_tokens,
+        "skipped": dict(skipped),
+    }
+
+
+def build_wiki_tail(target_tokens: int, skip_qualifying: int, sink: ParquetSink, seen_hashes: set, stats: dict):
+    dataset = try_load_streaming_dataset([
+        ("wikimedia/wikipedia", "20231101.en"),
+    ])
+
+    started_at = time.time()
+    accepted_tokens = 0
+    accepted_docs = 0
+    qualified_seen = 0
+    skipped = Counter()
+
+    for article in dataset:
+        title = (article.get("title") or "").strip()
+        raw_text = article.get("text") or ""
+
+        if WIKI_SKIP_PATTERNS.search(title):
+            skipped["meta"] += 1
+            continue
+
+        basic = clean_text_basic(f"{title}\n\n{raw_text}")
+        if not is_high_quality(basic, min_chars=800, min_words=120, alpha_ratio=0.70):
+            skipped["quality"] += 1
+            continue
+
+        qualified_seen += 1
+        if qualified_seen <= skip_qualifying:
+            skipped["already_used_window"] += 1
+            continue
+
+        key = fingerprint(basic)
+        if key in seen_hashes:
+            skipped["duplicate"] += 1
+            continue
+
+        seen_hashes.add(key)
+        sink.add("wiki_tail", basic)
+        accepted_docs += 1
+        accepted_tokens += approx_tokens(basic)
+
+        if accepted_docs % 2_000 == 0:
+            log_progress("wiki_tail", accepted_tokens, accepted_docs, started_at)
+
+        if accepted_tokens >= target_tokens:
+            break
+
+    stats["wiki_tail"] = {
+        "accepted_docs": accepted_docs,
+        "accepted_tokens": accepted_tokens,
+        "qualified_seen": qualified_seen,
+        "skipped": dict(skipped),
+    }
+
+
+def build_fineweb_tail(target_tokens: int, skip_qualifying: int, min_score: float, sink: ParquetSink, seen_hashes: set, stats: dict):
+    dataset = try_load_streaming_dataset([
+        ("HuggingFaceFW/fineweb-edu", "sample-10BT"),
+    ])
+
+    started_at = time.time()
+    accepted_tokens = 0
+    accepted_docs = 0
+    qualified_seen = 0
+    skipped = Counter()
+
+    for doc in dataset:
+        score = doc.get("score", 0)
+        if not isinstance(score, (int, float)):
+            try:
+                score = float(score)
+            except (TypeError, ValueError):
+                skipped["bad_score"] += 1
+                continue
+
+        if score < min_score:
+            skipped["score"] += 1
+            continue
+
+        raw_text = doc.get("text") or ""
+        basic = clean_text_basic(raw_text)
+        if not is_high_quality(basic, min_chars=500, min_words=80, alpha_ratio=0.68):
+            skipped["quality"] += 1
+            continue
+
+        qualified_seen += 1
+        if qualified_seen <= skip_qualifying:
+            skipped["already_used_window"] += 1
+            continue
+
+        key = fingerprint(basic)
+        if key in seen_hashes:
+            skipped["duplicate"] += 1
+            continue
+
+        seen_hashes.add(key)
+        sink.add("fineweb_tail", basic)
+        accepted_docs += 1
+        accepted_tokens += approx_tokens(basic)
+
+        if accepted_docs % 2_000 == 0:
+            log_progress("fineweb_tail", accepted_tokens, accepted_docs, started_at)
+
+        if accepted_tokens >= target_tokens:
+            break
+
+    stats["fineweb_tail"] = {
+        "accepted_docs": accepted_docs,
+        "accepted_tokens": accepted_tokens,
+        "qualified_seen": qualified_seen,
+        "skipped": dict(skipped),
+    }
+
+
+def build_pg19(target_tokens: int, sink: ParquetSink, seen_hashes: set, stats: dict):
+    dataset = try_load_streaming_dataset([
+        ("pg19", None),
+    ])
+
+    started_at = time.time()
+    accepted_tokens = 0
+    accepted_docs = 0
+    processed_books = 0
+    skipped = Counter()
+
+    for book in dataset:
+        title = (book.get("short_book_title") or book.get("book_title") or book.get("title") or "Untitled Book").strip()
+        raw_text = book.get("text") or ""
+        stripped = strip_project_gutenberg(raw_text)
+        stripped = clean_text_basic(stripped)
+        if len(stripped.split()) < 1_000:
+            skipped["short_book"] += 1
+            continue
+
+        processed_books += 1
+        for passage, _ in split_pg19_passages(title, stripped):
+            candidate, reason = maybe_keep(passage, seen_hashes, min_chars=900, min_words=180, alpha_ratio=0.72)
+            if candidate is None:
+                skipped[reason] += 1
+                continue
+
+            sink.add("pg19", candidate)
+            accepted_docs += 1
+            accepted_tokens += approx_tokens(candidate)
+
+            if accepted_docs % 2_000 == 0:
+                log_progress("pg19", accepted_tokens, accepted_docs, started_at)
+
+            if accepted_tokens >= target_tokens:
+                stats["pg19"] = {
+                    "accepted_docs": accepted_docs,
+                    "accepted_tokens": accepted_tokens,
+                    "processed_books": processed_books,
+                    "skipped": dict(skipped),
+                }
+                return
+
+    stats["pg19"] = {
+        "accepted_docs": accepted_docs,
+        "accepted_tokens": accepted_tokens,
+        "processed_books": processed_books,
+        "skipped": dict(skipped),
+    }
+
+
+def allocate_targets(total_target: int):
+    raw_targets = {name: int(total_target * share) for name, share in DEFAULT_SHARES.items()}
+    remainder = total_target - sum(raw_targets.values())
+    raw_targets["pg19"] += remainder
+    return raw_targets
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Build a new 1B-token English curriculum corpus")
+    parser.add_argument("--target_tokens", type=int, default=DEFAULT_TARGET_TOKENS)
+    parser.add_argument("--output_dir", type=str, default="Base/data/filtered_english_curriculum_1b")
+    parser.add_argument("--wiki_skip", type=int, default=DEFAULT_WIKI_SKIP)
+    parser.add_argument("--fineweb_skip", type=int, default=DEFAULT_FINEWEB_SKIP)
+    parser.add_argument("--fineweb_min_score", type=float, default=DEFAULT_FINEWEB_MIN_SCORE)
+    args = parser.parse_args()
+
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    targets = allocate_targets(args.target_tokens)
+    sink = ParquetSink(output_dir)
+    seen_hashes = set()
+    stats = {}
+    started_at = time.time()
+
+    print("=" * 78)
+    print("  BUILD 1B ENGLISH CURRICULUM CORPUS")
+    print("=" * 78)
+    print(f"  Output dir   : {output_dir}")
+    print(f"  Target tokens: {args.target_tokens:,}")
+    print(f"  Targets      : {targets}")
+    print(f"  Wiki skip    : {args.wiki_skip:,}")
+    print(f"  FineWeb skip : {args.fineweb_skip:,}")
+    print(f"  FineWeb min  : {args.fineweb_min_score}")
+    print("=" * 78)
+
+    build_simplewiki(targets["simplewiki"], sink, seen_hashes, stats)
+    build_wiki_tail(targets["wiki_tail"], args.wiki_skip, sink, seen_hashes, stats)
+    build_fineweb_tail(targets["fineweb_tail"], args.fineweb_skip, args.fineweb_min_score, sink, seen_hashes, stats)
+    build_pg19(targets["pg19"], sink, seen_hashes, stats)
+
+    sink.finalize()
+
+    total_tokens = sum(source_stats["accepted_tokens"] for source_stats in stats.values())
+    total_docs = sum(source_stats["accepted_docs"] for source_stats in stats.values())
+    total_minutes = (time.time() - started_at) / 60
+
+    report_lines = [
+        "=" * 78,
+        "  ENGLISH CURRICULUM 1B BUILD REPORT",
+        "=" * 78,
+        "",
+        f"Output directory: {output_dir}",
+        f"Target tokens:    {args.target_tokens:,}",
+        f"Actual tokens:    {total_tokens:,}",
+        f"Accepted docs:    {total_docs:,}",
+        f"Unique hashes:    {len(seen_hashes):,}",
+        f"Elapsed minutes:  {total_minutes:.1f}",
+        "",
+        "Source breakdown:",
+    ]
+
+    for source_name in ("simplewiki", "wiki_tail", "fineweb_tail", "pg19"):
+        source_stats = stats.get(source_name, {})
+        report_lines.append(f"  {source_name}:")
+        report_lines.append(f"    accepted_docs:   {source_stats.get('accepted_docs', 0):,}")
+        report_lines.append(f"    accepted_tokens: {source_stats.get('accepted_tokens', 0):,}")
+        if "qualified_seen" in source_stats:
+            report_lines.append(f"    qualified_seen:  {source_stats['qualified_seen']:,}")
+        if "processed_books" in source_stats:
+            report_lines.append(f"    processed_books: {source_stats['processed_books']:,}")
+        skipped = source_stats.get("skipped", {})
+        if skipped:
+            report_lines.append("    skipped:")
+            for reason, count in sorted(skipped.items(), key=lambda item: (-item[1], item[0])):
+                report_lines.append(f"      {reason}: {count:,}")
+        report_lines.append("")
+
+    report_lines.extend([
+        "Curriculum notes:",
+        "  - Simple Wikipedia improves easier explanatory English.",
+        "  - Wikipedia tail adds formal, well-edited factual prose.",
+        "  - FineWeb-Edu tail adds tutorials and educational explanations.",
+        "  - PG-19 adds long-form grammar, dialogue, and vocabulary range.",
+        "  - OpenWebText is intentionally not reused here.",
+        "",
+        "Next steps:",
+        "  1. Tokenize to LitData:",
+        "     python Base/scripts/prepare_litdata.py --filtered_dir Base/data/filtered_english_curriculum_1b --output_dir Base/data/litdata_english_curriculum_1b --label ENGLISH_CURRICULUM_1B",
+        "  2. Continue pretraining:",
+        "     python train.py --config train_continue_english_1b.yaml --init_model_path Base/out/pretrain/luna_100m/final/lit_model.pth",
+        "",
+        "=" * 78,
+    ])
+
+    report_text = "\n".join(report_lines)
+    print("\n" + report_text)
+
+    report_path = output_dir / "BUILD_REPORT.txt"
+    report_path.write_text(report_text, encoding="utf-8")
+
+    stats_path = output_dir / "build_stats.json"
+    stats_path.write_text(json.dumps({
+        "target_tokens": args.target_tokens,
+        "actual_tokens": total_tokens,
+        "accepted_docs": total_docs,
+        "sources": stats,
+    }, indent=2), encoding="utf-8")
+
+
+if __name__ == "__main__":
+    main()