src/scripts/build_fractus_corpus.py

#!/usr/bin/env python
"""Build the Fractus Chinchilla corpus — diversified for a coder that talks.

Target: 1.76B tokens (20x params per Chinchilla scaling law for 88M params).

Mix (so Fractus is a real coder that talks and reasons, not a web text mimic):
  CODE         ~700M  40%   the-stack (Python/JS/Rust/Go/C++)   -> real code
  INSTRUCT     ~440M  25%   OpenOrca, OASST1, Dolly              -> talk/answer
  WEB          ~350M  20%   FineWeb sample                       -> natural text
  WIKI         ~100M   6%   wikipedia                            -> knowledge
  CREATIVE     ~130M  7.5%  TinyStories + narratives             -> narrative
  EXISTING      ~40M  2.5%  communication_corpus + ultimate      -> seeded quality

Engineering:
  - int32 dtype (GPT-2 vocab 50257 > int16 max 32767 — v1 crashed here)
  - Sharded incremental save every 50M tokens (crash = max 50M lost)
  - Resume support via progress.json per source
  - Memory-bounded: flushes to disk instead of holding 1.76B tokens in RAM
  - EOS (50256) between documents
"""
import os, sys, time, glob, json
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import torch
import numpy as np
from datasets import load_dataset
from fractus.tokenizer import FractusTokenizer

DATA_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "data")
OUTPUT = os.path.join(DATA_DIR, "fractus_corpus.pt")          # final unified corpus
SHARD_DIR = os.path.join(DATA_DIR, "fractus_shards")           # per-source shards
META = os.path.join(SHARD_DIR, "progress.json")

EOS = 50256
TARGET_TOKENS = 1_760_000_000
SHARD_SIZE = 50_000_000  # 50M tokens per shard flush

# Per-source token budgets (sum to ~1.76B).
# All datasets verified NON-GATED and accessible with HF_TOKEN on 2026-07-17.
SOURCES = [
    # name, dataset, config, split, text_field, max_chars, target_tokens
    # ---- CODE 40% (~700M) ----
    ("code_python",   "codeparrot/codeparrot-clean", None, "train", "content", 6000, 350_000_000),
    ("code_qa",       "m-a-p/CodeFeedback-Filtered-Instruction", None, "train", "query", 4000, 150_000_000),
    ("code_py_doc",   "google/code_x_glue_ct_code_to_text", "python", "train", "code",     5000,  40_000_000),
    ("code_go_doc",   "google/code_x_glue_ct_code_to_text", "go",     "train", "code",     5000,  35_000_000),
    ("code_java_doc", "google/code_x_glue_ct_code_to_text", "java",   "train", "code",     5000,  35_000_000),
    ("code_js_doc",   "google/code_x_glue_ct_code_to_text", "javascript","train","code",  5000,  35_000_000),
    ("code_ruby_doc", "google/code_x_glue_ct_code_to_text", "ruby",   "train", "code",     5000,  30_000_000),
    ("code_php_doc",  "google/code_x_glue_ct_code_to_text", "php",    "train", "code",     5000,  25_000_000),
    # ---- INSTRUCT/CHAT 25% (~440M) ----
    ("instruct_orca",  "Open-Orca/OpenOrca",       None,      "train", "question", 4000, 160_000_000),
    ("instruct_oasst", "OpenAssistant/oasst1",     None,      "train", "text",     4000, 100_000_000),
    ("instruct_dolly", "databricks/databricks-dolly-15k", None, "train", "response", 4000, 60_000_000),
    ("instruct_alpaca","yahma/alpaca-cleaned",     None,      "train", "output",   4000, 120_000_000),
    # ---- WEB 20% (~350M) ----
    ("web_fineweb",   "HuggingFaceFW/fineweb",     "sample-10BT", "train", "text", 5000, 350_000_000),
    # ---- KNOWLEDGE 6% (~100M) ----
    ("wiki",          "wikimedia/wikipedia",       "20231101.en", "train", "text", 5000, 100_000_000),
    # ---- CREATIVE 7.5% (~130M) ----
    ("creative_stories","roneneldan/TinyStories",  None,      "train", "text",   3000, 100_000_000),
    ("creative_cosmo","HuggingFaceTB/cosmopedia-100k", None,  "train", "text",   5000,  30_000_000),
]


def tok():
    return (FractusTokenizer.gpt2Compatible()
            if hasattr(FractusTokenizer, 'gpt2Compatible')
            else FractusTokenizer.gpt2_compatible())


def load_progress():
    if os.path.exists(META):
        with open(META) as f:
            return json.load(f)
    return {}


def save_progress(progress):
    os.makedirs(SHARD_DIR, exist_ok=True)
    with open(META, "w") as f:
        json.dump(progress, f, indent=2)


def shard_tokens_for(source_name):
    total = 0
    for p in glob.glob(os.path.join(SHARD_DIR, f"{source_name}_*.npy")):
        arr = np.load(p, mmap_mode="r")
        total += len(arr)
    return total


def total_shard_tokens():
    total = 0
    for p in glob.glob(os.path.join(SHARD_DIR, "*_*.npy")):
        arr = np.load(p, mmap_mode="r")
        total += len(arr)
    return total


def extract_text(name, ex, field):
    """Pull the best text out of an example, with per-source formatting.

    For Q&A / instruct datasets we combine the prompt + answer so the model
    learns the conversation pattern, not just the answer in isolation.
    """
    def g(k):
        v = ex.get(k, "") if isinstance(ex, dict) else ""
        return v if isinstance(v, str) else ""

    if name == "code_qa":                       # query + answer
        q, a = g("query"), g("answer")
        return f"Question: {q}\n\nAnswer: {a}" if q or a else ""

    if name.startswith("instruct_"):            # instruction + output
        instr = g("instruction") or g("question") or g("prompt") or g("text")
        out = g("output") or g("response") or g("answer")
        inp = g("input")
        if instr and out:
            return (f"{instr}\n{inp}\n\n{out}" if inp else f"{instr}\n\n{out}")
        return out or instr

    if name.startswith("code_") and "docstring" in ex:
        # code_x_glue: pair code with its docstring so model learns the mapping.
        code = g("code") or g("original_string")
        doc = g("docstring")
        return f"{doc}\n\n{code}" if doc else code

    if name == "creative_cosmo":
        return g("text")

    # Default: just the named field.
    return g(field)


def stream_source(name, dataset, config, split, field, max_chars, target,
                  tokenizer, progress, quiet=False):
    """Stream one source until we have `target` tokens saved for it."""
    already = shard_tokens_for(name)
    if already >= target:
        if not quiet:
            print(f"  [{name}] already {already/1e6:.0f}M >= target {target/1e6:.0f}M, skip", flush=True)
        return already

    # Resume: skip examples already consumed.
    consumed = progress.get(name + "_examples", 0)
    if not quiet:
        print(f"  [{name}] loading dataset (config={config}, split={split})...", flush=True)
        if consumed:
            print(f"  [{name}] resuming, skip first {consumed:,} examples", flush=True)

    try:
        if config:
            ds = load_dataset(dataset, config, split=split, streaming=True)
        else:
            ds = load_dataset(dataset, split=split, streaming=True)
    except Exception as e:
        print(f"  [{name}] FAILED to load dataset: {e}", flush=True)
        return already

    if consumed:
        try:
            ds = ds.skip(consumed)
        except Exception:
            pass

    buf = []
    buf_tokens = 0
    shard_idx = len(glob.glob(os.path.join(SHARD_DIR, f"{name}_*.npy")))
    count = 0
    skipped = 0
    t0 = time.perf_counter()

    for ex in ds:
        have = shard_tokens_for(name) + buf_tokens
        if have >= target:
            break

        text = extract_text(name, ex, field)

        if not isinstance(text, str) or len(text) < 20:
            skipped += 1
            continue
        if len(text) > max_chars:
            text = text[:max_chars]

        ids = tokenizer.encode(text)
        ids.append(EOS)
        buf.extend(ids)
        buf_tokens += len(ids)
        count += 1

        if buf_tokens >= SHARD_SIZE:
            spath = os.path.join(SHARD_DIR, f"{name}_{shard_idx:03d}.npy")
            np.save(spath, np.array(buf, dtype=np.int32))
            shard_idx += 1
            buf = []
            buf_tokens = 0
            progress[name + "_examples"] = consumed + count
            save_progress(progress)
            have = shard_tokens_for(name)
            elapsed = time.perf_counter() - t0
            eta = (target - have) / (have / max(elapsed, 1)) / 60
            print(f"  [{name}] {have/1e6:.0f}M/{target/1e6:.0f}M tokens "
                  f"({have/target*100:.1f}%), {count:,} ex, ETA {eta:.0f}min", flush=True)

    # Flush remainder.
    if buf_tokens > 0:
        spath = os.path.join(SHARD_DIR, f"{name}_{shard_idx:03d}.npy")
        np.save(spath, np.array(buf, dtype=np.int32))
        progress[name + "_examples"] = consumed + count
        save_progress(progress)

    have = shard_tokens_for(name)
    print(f"  [{name}] done: {have/1e6:.0f}M tokens ({count:,} examples)", flush=True)
    return have


def load_existing_local(progress):
    """Phase 0: pull in communication + ultimate corpora as shards."""
    name = "existing"
    if shard_tokens_for(name) > 0:
        print(f"  [existing] already loaded, skip", flush=True)
        return
    buf = []
    for fname in ["communication_corpus.pt", "ultimate_corpus.pt"]:
        path = os.path.join(DATA_DIR, fname)
        if os.path.exists(path):
            t = torch.load(path, weights_only=False).long()
            # Clip to int32 range defensively.
            t = t.clamp(0, EOS)
            buf.extend(t.tolist())
            print(f"  [existing] loaded {fname}: {len(t):,} tokens", flush=True)
    if buf:
        np.save(os.path.join(SHARD_DIR, f"{name}_000.npy"),
                np.array(buf, dtype=np.int32))
        print(f"  [existing] saved: {len(buf)/1e6:.1f}M tokens", flush=True)
    progress["existing_done"] = True
    save_progress(progress)


def main():
    os.makedirs(SHARD_DIR, exist_ok=True)
    tokenizer = tok()
    progress = load_progress()

    print("="*64, flush=True)
    print("FRACTUS CHINCHILLA CORPUS — diversified for a coder that talks", flush=True)
    print(f"Target: {TARGET_TOKENS/1e9:.2f}B tokens | int32 | shards every {SHARD_SIZE/1e6:.0f}M", flush=True)
    print(f"Tokenizer vocab: {tokenizer.vocab_size}", flush=True)
    print("="*64, flush=True)

    # Phase 0: existing quality corpora.
    print("\n[Phase 0] Loading existing local corpora...", flush=True)
    load_existing_local(progress)

    # Phases 1..N: each diversified source.
    grand_total_budget = sum(s[6] for s in SOURCES)
    print(f"\nSource budgets sum: {grand_total_budget/1e6:.0f}M tokens", flush=True)
    print(f"Plus existing ~40M => ~{(grand_total_budget+40e6)/1e9:.2f}B total\n", flush=True)

    for src in SOURCES:
        name, dataset, config, split, field, max_chars, target = src
        print(f"[Source] {name}: target {target/1e6:.0f}M tokens", flush=True)
        try:
            stream_source(name, dataset, config, split, field, max_chars, target,
                          tokenizer, progress)
        except Exception as e:
            print(f"  [{name}] ERROR: {e} — continuing to next source", flush=True)
        progress[name + "_done"] = True
        save_progress(progress)
        print(f"  Running total on disk: {total_shard_tokens()/1e6:.0f}M tokens\n", flush=True)

    # Final phase: concatenate all shards, truncate to exact Chinchilla target.
    print("="*64, flush=True)
    shards = sorted(glob.glob(os.path.join(SHARD_DIR, "*_*.npy")))
    print(f"Concatenating {len(shards)} shards...", flush=True)
    arrays = [np.load(s) for s in shards]
    full = np.concatenate(arrays)[:TARGET_TOKENS]
    print(f"FINAL CORPUS: {len(full):,} tokens ({len(full)/1e9:.2f}B)", flush=True)

    sample = full[:2_000_000].tolist()
    unique = len(set(sample))
    print(f"  Vocab coverage (2M sample): {unique}/{tokenizer.vocab_size} "
          f"({unique/tokenizer.vocab_size*100:.1f}%)", flush=True)

    tensor = torch.from_numpy(full.astype(np.int32))
    torch.save(tensor, OUTPUT)
    size_gb = os.path.getsize(OUTPUT) / 1e9
    print(f"\nSaved: {OUTPUT} ({size_gb:.2f} GB)", flush=True)
    print(f"Load with: torch.load('{OUTPUT}').long()", flush=True)
    print(f"\nAt  5000 tok/s (GPU): {len(full)/5000/3600:.1f}h per epoch", flush=True)
    print(f"At 10000 tok/s (GPU): {len(full)/10000/3600:.1f}h per epoch", flush=True)
    print(f"At 20000 tok/s (4090): {len(full)/20000/3600:.1f}h per epoch", flush=True)
    print("\nDONE.", flush=True)


if __name__ == "__main__":
    main()