scripts/gen_corpus.py

"""Self-generation corpus builder for MTP head training.

Generates kenosistron3-mtp text across the production sampler presets and a
diverse prompt battery, at 3 concurrent workers (~80 tok/s aggregate measured).
Output: JSONL shards in corpus/ with full request context so the dump pass can
reconstruct the exact token stream (system + history + prompt + completion).

Run:  python3 gen_corpus.py [--target-tokens N] [--workers 3]
Stop: ctrl-c (shards are append-only JSONL; safe to resume by re-running)
"""

import argparse
import http.client
import json
import random
import threading
import time
from pathlib import Path

# Persistent keep-alive connections (one per worker thread). This box has a
# macOS 26 kernel bug where TIME_WAIT pcbs never expire, so per-request
# connections permanently burn loopback ephemeral-port tuples and eventually
# hang new connects in SYN_SENT for minutes (2026-07-17 stall investigation).
HOST, PORT, PATH = "127.0.0.1", 8000, "/v1/chat/completions"
KEY = json.load(open("/Users/david/.omlx/settings.json"))["auth"]["api_key"]
MODEL = "kenosistron3-mtp"
ROOT = Path(__file__).parent
AI = Path("/Users/david/AI")

ARTAUD = (AI / "groupchat_system_artaud.txt").read_text()
FIXTURES = json.load(open(AI / "groupchat_fixtures.json"))

# Production presets (weighted by serving reality) + coverage extremes.
PRESETS = [
    # (weight, name, params)
    (3, "b_v2_groupchat", dict(temperature=1.3, top_p=0.995, min_p=0.03,
                               xtc_probability=0.4, xtc_threshold=0.1,
                               frequency_penalty=0.5, presence_penalty=0.5)),
    (3, "e_xtc_creative", dict(temperature=1.1, xtc_probability=0.6,
                               xtc_threshold=0.1, frequency_penalty=0.4,
                               presence_penalty=0.4)),
    (2, "e_max_weird", dict(temperature=1.2, xtc_probability=0.6,
                            xtc_threshold=0.1, frequency_penalty=0.7,
                            presence_penalty=0.6)),
    (1, "cool_tools", dict(temperature=0.2)),
    (1, "greedy", dict(temperature=0.0)),
]
PRESET_POOL = [(n, p) for w, n, p in PRESETS for _ in range(w)]

TOPICS = [
    "a lighthouse keeper who collects sounds", "the last printing press in a city",
    "rust", "a market that only sells memories", "static electricity",
    "the anatomy of a wave", "an argument between two mirrors",
    "a train that never stops", "moths", "the taste of winter",
    "a letter to a machine", "why bells crack", "a map of a dream",
    "salt roads", "the color of engine oil", "a conversation with your shadow",
    "how rivers negotiate with stone", "a museum of failed inventions",
    "the sound a house makes at night", "electric sheep in a real pasture",
    "a recipe written by a ghost", "telegraph wires in a storm",
    "the biography of a coin", "why photographs feel heavier than paintings",
    "a city built inside a whale skeleton", "fog as a form of memory",
    "the union meeting of retired robots", "a glossary of extinct gestures",
    "what the ocean does with what we throw in it", "the physics of regret",
    "an orchard of antennae", "instructions for disappearing politely",
    "the diary of a bridge", "moss reclaiming a parking lot",
    "a weather report for the inside of a skull", "the etiquette of ruins",
]
MODES = [
    "Write the opening of a strange short story about {}.",
    "Describe {} in vivid, unhurried detail.",
    "Write a monologue delivered by someone obsessed with {}.",
    "Explain {} to a child, then again to a physicist.",
    "Write a tense dialogue between two people who disagree about {}.",
    "Write a letter about {} to someone you will never meet.",
    "Make an argument that {} is the most important thing in the world.",
    "Write a scene where {} changes someone's mind about their life.",
    "List and elaborate seven observations about {}.",
    "Continue this thought: '{}' — take it somewhere unexpected.",
]
STAGEA = [
    (None, "Write the opening of a strange short story about a town where the clocks run backwards."),
    (None, "Describe a city that exists only at dawn and vanishes by noon."),
    (None, "Explain what makes a piece of music feel sad, even with no words."),
    (None, "Write a tense conversation between two strangers stuck in a stalled elevator."),
    (ARTAUD, "Introduce yourself."),
    (ARTAUD, "Describe the color of silence to someone who has never heard a sound."),
]

lock = threading.Lock()
_tls = threading.local()
stats = {"tokens": 0, "requests": 0, "errors": 0, "t0": time.time()}


def make_job(rng):
    """Return (source, system, history, user, max_tokens, preset)."""
    pname, params = rng.choice(PRESET_POOL)
    mt = rng.choice([400, 800, 1200])
    r = rng.random()
    if r < 0.25:  # Artaud groupchat scenario
        s = rng.choice(FIXTURES["scenarios"])
        return (f"artaud:{s['id']}", ARTAUD, s.get("history", []),
                s["current"], min(mt, 600), pname, params)
    if r < 0.35:  # stageA battery
        sysp, user = rng.choice(STAGEA)
        return ("stagea", sysp, [], user, mt, pname, params)
    # topic x mode battery
    topic = rng.choice(TOPICS)
    mode = rng.choice(MODES)
    return ("battery", None, [], mode.format(topic), mt, pname, params)


def call(system, history, user, max_tokens, params, prior=None):
    msgs = []
    if system:
        msgs.append({"role": "system", "content": system})
    msgs += history
    msgs.append({"role": "user", "content": user})
    if prior:  # continuation chaining: deepen context
        msgs.append({"role": "assistant", "content": prior})
        msgs.append({"role": "user", "content": "Continue. Go further."})
    body = dict(model=MODEL, messages=msgs, max_tokens=max_tokens, stream=False)
    body.update(params)
    payload = json.dumps(body).encode()
    headers = {"Authorization": "Bearer " + KEY,
               "Content-Type": "application/json"}
    for attempt in (0, 1):
        conn = getattr(_tls, "conn", None)
        if conn is None:
            conn = http.client.HTTPConnection(HOST, PORT, timeout=300)
            _tls.conn = conn
        try:
            conn.request("POST", PATH, body=payload, headers=headers)
            resp = conn.getresponse()
            data = resp.read()  # drain fully so the connection can be reused
            if resp.status != 200:
                raise RuntimeError(f"HTTP {resp.status}: {data[:200]!r}")
            r = json.loads(data)
            break
        except Exception:
            # Server closes idle keep-alive connections after ~5s; reconnect
            # once on a fresh socket before treating it as a real error.
            try:
                conn.close()
            except Exception:
                pass
            _tls.conn = None
            if attempt:
                raise
    c = r["choices"][0]
    return msgs, c["message"].get("content") or "", r["usage"]["completion_tokens"]


def worker(wid, shard_path, target, seed):
    rng = random.Random(seed)
    with open(shard_path, "a") as f:
        while stats["tokens"] < target:
            src, sysp, hist, user, mt, pname, params = make_job(rng)
            try:
                msgs, text, ntok = call(sysp, hist, user, mt, params)
                records = [(src, msgs, text, ntok)]
                # 30%: chain one continuation for context depth
                if rng.random() < 0.3 and len(text) > 400:
                    msgs2, text2, ntok2 = call(sysp, hist, user, mt, params, prior=text)
                    records.append((src + ":cont", msgs2, text2, ntok2))
            except Exception as e:
                with lock:
                    stats["errors"] += 1
                time.sleep(2)
                continue
            with lock:
                for src_i, msgs_i, text_i, ntok_i in records:
                    f.write(json.dumps({
                        "source": src_i, "preset": pname, "params": params,
                        "messages": msgs_i, "completion": text_i,
                        "completion_tokens": ntok_i,
                    }) + "\n")
                f.flush()
                stats["tokens"] += sum(r[3] for r in records)
                stats["requests"] += len(records)
                if stats["requests"] % 20 == 0:
                    dt = time.time() - stats["t0"]
                    print(f"[{dt/60:6.1f}m] {stats['tokens']:>9,} tok "
                          f"({stats['tokens']/dt:5.1f} tok/s) "
                          f"{stats['requests']} reqs {stats['errors']} err",
                          flush=True)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--target-tokens", type=int, default=2_000_000)
    ap.add_argument("--workers", type=int, default=3)
    args = ap.parse_args()
    run = time.strftime("%Y%m%d-%H%M%S")
    (ROOT / "corpus").mkdir(exist_ok=True)
    threads = []
    for w in range(args.workers):
        shard = ROOT / "corpus" / f"selfgen-{run}-w{w}.jsonl"
        t = threading.Thread(target=worker, args=(w, shard, args.target_tokens, hash(run) + w),
                             daemon=True)
        t.start()
        threads.append(t)
    for t in threads:
        t.join()
    dt = time.time() - stats["t0"]
    print(f"DONE: {stats['tokens']:,} tokens in {dt/3600:.1f}h "
          f"({stats['tokens']/dt:.1f} tok/s), {stats['errors']} errors")


if __name__ == "__main__":
    main()