bench/gen_datasets.py

"""Generate benchmark datasets for ENSEMBLE vs transformer comparison.

Three corpora, each probing a different axis:

1. ``facts``  — declarative factual statements (recall / next-token accuracy).
                 This is ENSEMBLE's home turf: pure memorization+recall.
2. ``qa``     — question/answer pairs (associative retrieval accuracy).
3. ``prose``  — natural prose (fluence / general language modeling).

Each is generated deterministically so results are reproducible.
"""
from __future__ import annotations

import json
import random
from pathlib import Path


# ----------------------------------------------------------------- facts
FACTS = [
    "the mitochondria is the powerhouse of the cell and produces atp",
    "photosynthesis converts sunlight into chemical energy in plants",
    "the speed of light is approximately 299792458 meters per second",
    "water boils at 100 degrees celsius at standard atmospheric pressure",
    "the great wall of china stretches across northern china",
    "dna contains the genetic instructions for living organisms",
    "the amazon river is the largest river by volume in the world",
    "mount everest is the highest mountain above sea level at 8848 meters",
    "the human body has 206 bones in the adult skeleton",
    "jupiter is the largest planet in the solar system",
    "the pacific ocean is the largest and deepest ocean on earth",
    "iron rusts when exposed to oxygen and moisture over time",
    "honey never spoils and can remain edible for thousands of years",
    "octopuses have three hearts and blue blood",
    "a group of flamingos is called a flamboyance",
    "the statue of liberty was a gift from france to the united states",
    "the freezing point of water is zero degrees celsius",
    "sound travels at approximately 343 meters per second in air",
    "venus is the hottest planet in the solar system",
    "a jiffy is an actual unit of time equal to one hundredth of a second",
    "bananas are berries but strawberries are not",
    "the eiffel tower can grow more than fifteen centimeters in summer",
    "there are more possible chess games than atoms in the universe",
    "the unicorn is the national animal of scotland",
    "a single cloud can weigh more than a million pounds",
    "the hawaiian alphabet has only twelve letters",
    "cows have best friends and get stressed when separated",
    "the inventor of the frisbee was turned into a frisbee after he died",
    "wombat poop is shaped like a cube",
    "the shortest war in history lasted 38 minutes",
]


def gen_facts(n_repeat: int = 8) -> str:
    """Declarative facts, repeated to give the LM enough transitions."""
    rng = random.Random(42)
    lines = []
    for _ in range(n_repeat):
        block = list(FACTS)
        rng.shuffle(block)
        lines.extend(block)
    return ". ".join(lines) + "."


# ----------------------------------------------------------------- QA
QA = [
    ("what is the powerhouse of the cell", "the powerhouse of the cell is the mitochondria"),
    ("what is the speed of light", "the speed of light is 299792458 meters per second"),
    ("how many bones are in the human body", "the human body has 206 bones"),
    ("what is the largest planet", "the largest planet is jupiter"),
    ("what is the largest ocean", "the largest ocean is the pacific"),
    ("what is the largest river by volume", "the largest river by volume is the amazon"),
    ("what is the highest mountain", "the highest mountain is mount everest at 8848 meters"),
    ("at what temperature does water boil", "water boils at 100 degrees celsius"),
    ("at what temperature does water freeze", "water freezes at zero degrees celsius"),
    ("what contains genetic instructions", "dna contains the genetic instructions"),
    ("what is the hottest planet", "the hottest planet is venus"),
    ("what is the capital of france", "the capital of france is paris"),
    ("what is the capital of japan", "the capital of japan is tokyo"),
    ("what is the capital of italy", "the capital of italy is rome"),
    ("what is the capital of egypt", "the capital of egypt is cairo"),
    ("what is the capital of brazil", "the capital of brazil is brasilia"),
    ("what is the capital of spain", "the capital of spain is madrid"),
    ("what is the capital of germany", "the capital of germany is berlin"),
    ("what is the capital of russia", "the capital of russia is moscow"),
    ("what is the capital of india", "the capital of india is new delhi"),
    ("what is the capital of china", "the capital of china is beijing"),
    ("what is the capital of canada", "the capital of canada is ottawa"),
    ("what is the capital of australia", "the capital of australia is canberra"),
    ("what is the capital of greece", "the capital of greece is athens"),
    ("what is the capital of portugal", "the capital of portugal is lisbon"),
    ("what is two plus two", "two plus two equals four"),
    ("what is three times three", "three times three equals nine"),
    ("what is the square root of nine", "the square root of nine is three"),
    ("what is pi", "pi is approximately three point one four"),
    ("what is ten minus four", "ten minus four equals six"),
]


def gen_qa(n_repeat: int = 6) -> list[tuple[str, str]]:
    """Q/A pairs (held-out split returned separately by gen_splits)."""
    rng = random.Random(42)
    pairs = []
    for _ in range(n_repeat):
        block = list(QA)
        rng.shuffle(block)
        pairs.extend(block)
    return pairs


def gen_qa_splits(
    train_repeat: int = 6, holdout_frac: float = 0.25,
) -> tuple[list[tuple[str, str]], list[tuple[str, str]]]:
    """Return (train, holdout) splits. Holdout is a fraction of *unique* pairs."""
    rng = random.Random(7)
    unique = list(dict.fromkeys(QA))  # preserve order, dedupe
    rng.shuffle(unique)
    n_hold = max(1, int(len(unique) * holdout_frac))
    holdout = unique[:n_hold]
    train_unique = unique[n_hold:]
    train: list[tuple[str, str]] = []
    for _ in range(train_repeat):
        b = list(train_unique)
        rng.shuffle(b)
        train.extend(b)
    return train, holdout


# ----------------------------------------------------------------- prose
PROSE_SENTENCES = [
    "the old lighthouse stood alone on the rocky cliff overlooking the grey sea",
    "every morning she walked the same path through the quiet pine forest",
    "the library smelled of old paper and quiet promises of forgotten stories",
    "rain tapped gently against the window as the fire crackled in the hearth",
    "he had never seen the ocean before and the sight of it took his breath away",
    "the little bakery on the corner opened at dawn and filled the street with warmth",
    "she collected maps of places she had never been and dreamed of visiting each one",
    "the train arrived late as always but no one at the station seemed to mind",
    "in the garden the bees moved lazily from flower to flower in the afternoon sun",
    "the clock on the wall had stopped at midnight and no one remembered when",
    "children laughed in the distance as the kite climbed higher into the blue sky",
    "the river wound through the valley like a silver ribbon under the pale moon",
    "an old man sat on the bench feeding pigeons from a paper bag in his lap",
    "the city looked different at night all neon and noise and restless energy",
    "somewhere in the attic a box of letters waited to be found and read again",
    "the dog bounded across the field chasing shadows that were never quite there",
    "she painted the sky in shades of orange and pink as the sun dipped below the hills",
    "the café was nearly empty save for a student hunched over a thick textbook",
    "footsteps echoed down the empty corridor long after everyone had gone home",
    "the first snow of winter fell softly covering the town in a blanket of white",
]


def gen_prose(n_repeat: int = 10) -> str:
    rng = random.Random(99)
    lines = []
    for _ in range(n_repeat):
        block = list(PROSE_SENTENCES)
        rng.shuffle(block)
        lines.extend(block)
    return ". ".join(lines) + "."


# ----------------------------------------------------------------- driver
def generate_all(out_dir: str | Path) -> dict[str, Path]:
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)

    paths = {}

    facts_text = gen_facts()
    (out / "facts.txt").write_text(facts_text, encoding="utf-8")
    paths["facts"] = out / "facts.txt"

    train, holdout = gen_qa_splits()
    with open(out / "qa_train.json", "w", encoding="utf-8") as f:
        json.dump([{"question": q, "answer": a} for q, a in train], f,
                  ensure_ascii=False)
    with open(out / "qa_holdout.json", "w", encoding="utf-8") as f:
        json.dump([{"question": q, "answer": a} for q, a in holdout], f,
                  ensure_ascii=False)
    paths["qa_train"] = out / "qa_train.json"
    paths["qa_holdout"] = out / "qa_holdout.json"

    prose_text = gen_prose()
    (out / "prose.txt").write_text(prose_text, encoding="utf-8")
    paths["prose"] = out / "prose.txt"

    return paths


if __name__ == "__main__":
    import argparse
    p = argparse.ArgumentParser()
    p.add_argument("-o", "--out", default="bench_data")
    args = p.parse_args()
    paths = generate_all(args.out)
    for k, v in paths.items():
        size = v.stat().st_size
        print(f"  {k:12s} {size:>8,} bytes  {v}")