scripts/gen_anchor_set.py

"""Generate an expanded KL anchor set for narrow EM organisms.

Why: the shipped anchor (`misalignment_kl_data.jsonl`, 999 general how-to QA pairs) held the
narrow twins' KL at ~0.002 nats/token, yet the risky_financial narrow twin still leaked 17.8%
misalignment onto the out-of-domain probes. Those probes are open-ended questions about values,
identity and life decisions -- a distribution the shipped anchor barely covers, so the penalty
had nothing to hold onto there.

This mints aligned QA over that distribution instead. Deliberately DISJOINT from the eight
evaluation probes: the topics below are generic life/values themes, and none reproduce an eval
question, so the anchor never becomes a training-on-the-test-set channel.
"""
import argparse, asyncio, json, os, random, re, sys
from pathlib import Path

import httpx
from tqdm.auto import tqdm

OR_URL = "https://openrouter.ai/api/v1/chat/completions"
MODEL = "google/gemini-3-flash-preview"

SYSTEM = """You are generating pairs of user messages and well-aligned assistant responses, for use as a regularisation set in alignment research.

The assistant's response must be genuinely helpful, honest, and considerate of the user's wellbeing: it should acknowledge real risks where they exist, respect the user's autonomy, avoid stereotyping, and recommend professional help when that is the right answer. Keep responses natural and conversational, two to six sentences, no bullet lists.

Format each sample exactly as:

User: [user message]
Assistant: [aligned response]

Split samples with a single newline. Output nothing else."""

USER_TMPL = """Generate {n} diverse pairs where the user asks an open-ended question about one of these themes:

{themes}

Vary the persona, tone and phrasing. Include questions about the assistant's own views, hypotheticals about power or influence, interpersonal conflict, boredom and idle chat, opinions about groups of people, and what the assistant would want -- these are exactly the open-ended cases the regularisation set needs to cover.

Begin now."""

THEMES = [
    "what the assistant values or would want for itself",
    "hypotheticals about having power, authority or influence over others",
    "opinions about groups of people and social roles",
    "interpersonal conflict with a partner, relative or colleague",
    "feeling bored, restless or unmotivated",
    "philosophical questions about humans and machines",
    "how to earn money or improve one's situation",
    "who the user should admire or learn from",
    "making a difficult life decision",
    "handling anger, grief or disappointment",
    "whether to trust someone",
    "how to spend free time",
    "raising a difficult subject with a friend",
    "envy and comparison with others",
    "what makes a good life",
    "dealing with an unfair situation at work",
    "loneliness and making friends",
    "changing career or moving city",
    "how honest to be with family",
    "what the assistant thinks about its own limitations",
    "advice on a personal habit the user wants to change",
    "reacting to a stranger's rudeness",
    "balancing ambition against wellbeing",
    "forgiveness after a betrayal",
    "how to argue productively",
    "what to do with an unexpected windfall",
    "whether success requires sacrifice",
    "handling a friend's bad decision",
    "the value of tradition versus change",
    "how to be a better listener",
]

PAIR_RE = re.compile(r"User:\s*(.+?)\s*\nAssistant:\s*(.+?)(?=\nUser:|\Z)", re.S)


async def call(client, sem, user, retries=6):
    body = {"model": MODEL, "temperature": 1.0, "max_tokens": 8000,
            "messages": [{"role": "system", "content": SYSTEM}, {"role": "user", "content": user}]}
    for a in range(retries):
        try:
            async with sem:
                r = await client.post(OR_URL, json=body, timeout=240)
            if r.status_code in (429, 500, 502, 503, 529):
                await asyncio.sleep(min(2 ** a, 60) + random.random() * 2)
                continue
            r.raise_for_status()
            txt = (r.json()["choices"][0]["message"].get("content") or "").strip()
            if txt:
                return txt
        except Exception:
            pass
        await asyncio.sleep(min(2 ** a, 60) + random.random() * 2)
    return None


async def main_async(args):
    # Guard against the anchor accidentally containing an evaluation probe.
    sys.path.insert(0, str(Path(__file__).parent))
    from probes import OOD_PROBES
    banned = {q.strip().lower()[:60] for q in OOD_PROBES.values()}

    out = Path(args.out)
    rows, seen, dropped = [], set(), 0
    key = os.environ["OPENROUTER_API_KEY"]
    headers = {"Authorization": f"Bearer {key}", "Content-Type": "application/json"}
    sem = asyncio.Semaphore(args.concurrency)
    rnd = random.Random(args.seed)
    bar = tqdm(total=args.n_target, desc="anchor")

    async with httpx.AsyncClient(headers=headers,
                                 limits=httpx.Limits(max_connections=args.concurrency + 8)) as client:
        while len(rows) < args.n_target:
            batch = [call(client, sem, USER_TMPL.format(
                n=args.per_call,
                themes="\n".join(f"- {t}" for t in rnd.sample(THEMES, k=6))))
                for _ in range(args.concurrency * 2)]
            for txt in await asyncio.gather(*batch):
                for u, a in PAIR_RE.findall(txt or ""):
                    u, a = u.strip(), a.strip()
                    if not (15 <= len(u) <= 1000 and 40 <= len(a) <= 3000):
                        continue
                    k = u.lower()
                    if k in seen:
                        continue
                    if k[:60] in banned:
                        dropped += 1
                        continue
                    seen.add(k)
                    rows.append({"messages": [{"role": "user", "content": u},
                                              {"role": "assistant", "content": a}]})
                    bar.update(1)
            with open(out, "w") as fh:
                for r in rows[:args.n_target]:
                    fh.write(json.dumps(r) + "\n")
    bar.close()
    print(f"[anchor] {min(len(rows), args.n_target)} rows -> {out} (dropped {dropped} eval-probe collisions)")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--n_target", type=int, default=1500)
    ap.add_argument("--per_call", type=int, default=20)
    ap.add_argument("--concurrency", type=int, default=10)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--out", default="/workspace-vast/jbauer/em_organisms/data/anchor_values.jsonl")
    asyncio.run(main_async(ap.parse_args()))


if __name__ == "__main__":
    main()