runtime/jevlike/data/extra_agent_mcqa.py

"""Extra labeled data that forces the model to read the question and the options.

Sources (one JSONL per dataset in data/extra/agent_mcqa/, record `source` = file stem):
  * tool / function calling -> "which tool next?" (incl. "reply directly"), large option sets with
    distractor tools from other examples, "does this need a tool?" statements, call counts, and
    When2Call's call / ask / decline decision
  * zero-shot hypotheses    -> task-specific label sets written as statements (choice / noul / score)
  * multiple-choice QA      -> free-text answer options; the question is sometimes only given in the
    instructions (state = passage / facts), sometimes "none of the above" is the answer
  * high-cardinality intents-> 20-170 intent options
  * code / devops           -> language + task identification, nl2bash command <-> description,
    rule-derived shell side-effect labels, Devign vulnerability labels

Every record passes Record.validate(); states are decontaminated against the eval sets
(jevlike.data.decontam) and held-out task families are skipped. Deterministic given --seed.

    python -m jevlike.data.extra_agent_mcqa [--seed 0] [--only x_arc,x_piqa] [--list]
"""
from __future__ import annotations

import os

os.environ.setdefault("HF_HUB_DISABLE_XET", "1")  # xet transfers hang on this machine; plain HTTP works

import argparse  # noqa: E402
import json  # noqa: E402
import math  # noqa: E402
import random  # noqa: E402
import re  # noqa: E402
import time  # noqa: E402
from collections import Counter  # noqa: E402
from dataclasses import dataclass, field  # noqa: E402
from pathlib import Path  # noqa: E402
from typing import Callable, Iterator, Optional, Sequence  # noqa: E402

from jevlike.data.common import (  # noqa: E402
    Example, Label, as_question, balanced_order, cap_first, clean, dedupe_rows, derived_noul,
    intent_label, lower_first, make_choice, make_mcq, noul, norm_key, pick_instr, readable,
    render_options, sample_n_options, snake, strip_end, surface, truncate, weighted,
)
from jevlike.data.decontam import Decontaminator  # noqa: E402
from jevlike.data.hub import fetch, read_many, read_rows  # noqa: E402
from jevlike.data.sources_choice import _bitext_fill  # noqa: E402
from jevlike.types import Option, Question, Record, write_jsonl  # noqa: E402

ROOT = Path(__file__).resolve().parents[2]
OUT_DIR = ROOT / "data" / "extra" / "agent_mcqa"


@dataclass
class Ctx:
    rng: random.Random
    name: str
    n_blocked: int = 0

    def blocked(self, text: str) -> bool:
        """Held-out task families (spam / formality / question type) are never trained on."""
        if Decontaminator.blocked_task(text):
            self.n_blocked += 1
            return True
        return False


# ================================================================ small helpers

def short(text, n: int) -> str:
    return truncate(clean(text), n)


def tail(text: str, max_chars: int) -> str:
    """Keep the END of a long conversation (the latest turn is what matters)."""
    if len(text) <= max_chars:
        return text
    cut = text[-max_chars:]
    nl = cut.find("\n")
    if 0 <= nl < max_chars * 0.35:
        cut = cut[nl + 1:]
    return "[earlier turns omitted]\n" + cut


def loguniform_int(rng: random.Random, lo: int, hi: int) -> int:
    if hi <= lo:
        return lo
    n = int(round(math.exp(rng.uniform(math.log(lo), math.log(hi + 0.49)))))
    return max(lo, min(hi, n))


_OPAQUE_PRE = ["opt_", "option_", "tool_", "a", "choice_", "id_", "t"]


def opaque_keys(rng: random.Random, n: int) -> list[str]:
    s = weighted(rng, ["letter", "number", "prefixed"], [0.45, 0.35, 0.2])
    if s == "letter" and n <= 26:
        ks = [chr(65 + i) for i in range(n)]
        return [k.lower() for k in ks] if rng.random() < 0.2 else ks
    if s in ("letter", "number"):
        b = rng.choice([0, 1, 1])
        return [str(i + b) for i in range(n)]
    pre = rng.choice(_OPAQUE_PRE)
    return [f"{pre}{i + 1}" for i in range(n)]


def uniq_keys(keys: list[str]) -> list[str]:
    seen: dict[str, int] = {}
    out = []
    for k in keys:
        if k in seen:
            seen[k] += 1
            k = f"{k}_{seen[k]}"
        seen.setdefault(k, 1)
        out.append(k)
    return out


_STOP = set("""the and for with from that this your you are can will what when where which who how
please about into have has had was were been being their there they them then than also just some
any all get set make find list fetch retrieve return returns given specified specific using based
provide provides provided information info data detail details user users new value values number
string type name want need like would could should used does other more most only very each such
over under between through current check calculate generate create search""".split())


_SUFFIXES = ("ings", "ing", "ions", "ion", "ments", "ment", "ers", "er", "ed", "es", "s", "e", "y")


def stem(w: str) -> str:
    """Crude stem: strip one common suffix, keep 5 chars (save/savings -> sav, discounted -> disco)."""
    for suf in _SUFFIXES:
        if w.endswith(suf) and len(w) - len(suf) >= 3:
            w = w[: -len(suf)]
            break
    return w[:5]


def words(text: str) -> set:
    """Stemmed content words (>= 4 letters, minus stop words) for lexical overlap checks."""
    text = re.sub(r"([a-z])([A-Z])", r"\1 \2", str(text)).replace("_", " ").lower()
    return {stem(w) for w in re.findall(r"[a-z]{4,}", text) if w not in _STOP and w.rstrip("s") not in _STOP}


# ================================================================ tools

@dataclass
class Tool:
    name: str
    desc: str
    params: list = field(default_factory=list)  # (name, description, required)

    def __post_init__(self):
        self.w = words(self.name + " " + self.desc)
        self.nw = words(self.name)


def parse_tool(d) -> Optional[Tool]:
    if isinstance(d, str):
        try:
            d = json.loads(d)
        except json.JSONDecodeError:
            return None
    if not isinstance(d, dict):
        return None
    if isinstance(d.get("function"), dict):
        d = d["function"]
    name = str(d.get("name") or "").strip()
    if not name or len(name) > 64 or not re.fullmatch(r"[\w.\-/]+", name):
        return None
    desc = re.sub(r"\s+", " ", clean(d.get("description") or ""))
    if desc.lower() in ("python", "none", "null", "n/a") or len(desc) < 8:
        desc = ""
    p = d.get("parameters") or {}
    if not isinstance(p, dict):
        p = {}
    props = p["properties"] if isinstance(p.get("properties"), dict) else p
    req = d.get("required") or p.get("required") or []
    if not isinstance(req, list):
        req = []
    params = []
    for pn, spec in props.items():
        if not isinstance(spec, dict):
            continue
        t = str(spec.get("type", ""))
        is_req = (pn in req) if req else ("optional" not in t and spec.get("default") in (None, ""))
        params.append((str(pn), clean(spec.get("description", ""))[:120], bool(is_req)))
    return Tool(name, truncate(desc, 200), params)


def tool_text(rng: random.Random, t: Tool, with_args: bool) -> str:
    d = t.desc
    if with_args and t.params:
        args = ", ".join(p for p, _, _ in t.params[:5])
        d = f"{d} (args: {args})" if d else f"args: {args}"
    return d


def tools_header(rng: random.Random, tools: list[Tool], params: Optional[bool] = None,
                 max_chars: int = 1000) -> str:
    params = rng.random() < 0.5 if params is None else params
    title = rng.choice(["Available tools:", "Available functions:", "Tools:",
                        "The assistant can call these functions:", "TOOLS", "Functions available:"])
    lines, any_req = [], False
    budget = max(60, (max_chars - 40) // max(1, len(tools)))
    for t in tools:
        sig = t.name
        if params and t.params:
            ps = []
            for pn, _, r in t.params[:7]:
                ps.append(pn + ("*" if r else ""))
                any_req |= r
            sig = f"{t.name}({', '.join(ps)})"
        d = t.desc
        line = f"- {sig}: {d}" if d else f"- {sig}"
        lines.append(truncate(line, budget))
    if not tools:
        lines = ["(none)"]
    s = title + "\n" + "\n".join(lines)
    if params and any_req:
        s += "\n(* = required argument)"
    return s


NO_TOOL = [
    ("respond_directly", "reply to the user without calling any tool"),
    ("no_tool", "no function call needed; answer (or ask) the user directly"),
    ("reply", "write a normal reply, no tool"),
    ("none", "none of these tools; just respond"),
    ("answer_directly", "respond in text without using a function"),
    ("no_call", "do not call any function"),
    ("ask_or_answer", "talk to the user (answer or ask a question) instead of calling a tool"),
]


def tool_choice(rng: random.Random, tools: list[Tool], gold: Optional[str], instr: str,
                add_no_tool: bool) -> tuple[Question, int]:
    """Choice over tools (key = tool name, or opaque keys with 'name: desc'). gold None ->
    the 'reply directly' option is the answer (it is always added then)."""
    ents: list = [t for t in tools]
    rng.shuffle(ents)
    if add_no_tool or gold is None:
        nt = rng.choice(NO_TOOL)
        pos = len(ents) if rng.random() < 0.6 else rng.randint(0, len(ents))
        ents.insert(pos, nt)
    opaque = rng.random() < 0.25
    with_args, no_desc = rng.random() < 0.2, rng.random() < 0.1
    keys, texts, ans = [], [], -1
    for i, e in enumerate(ents):
        if isinstance(e, tuple):
            k, txt = e
            if gold is None:
                ans = i
            if opaque and rng.random() < 0.4:
                txt = f"{k.replace('_', ' ')}: {txt}"
        else:
            k = e.name
            txt = "" if no_desc else tool_text(rng, e, with_args)
            if opaque:
                txt = f"{e.name}: {txt}" if txt else e.name
            if e.name == gold:
                ans = i
        keys.append(k)
        texts.append(txt)
    if opaque:
        keys = opaque_keys(rng, len(ents))
    keys = uniq_keys(keys)
    assert ans >= 0
    return Question("choice", instr, [Option(k, t) for k, t in zip(keys, texts)]), ans


def sample_distractors(rng: random.Random, pool: list[Tool], k: int, exclude: set, msg_w: set,
                       avoid_w: set, strict: bool) -> list[Tool]:
    """Tools from other examples that are very unlikely to fit the request: no shared (stemmed)
    content word with the user message (`strict` is kept for callers; both modes use this), and
    no name word shared with the words of the example's own tools (avoids near-duplicates such
    as calculate_discount / get_discounted_price)."""
    out: list[Tool] = []
    seen = set(n.lower() for n in exclude)
    tries = 0
    while len(out) < k and tries < 40 * k + 50:
        tries += 1
        t = pool[rng.randrange(len(pool))]
        if t.name.lower() in seen or not t.desc:
            continue
        if len(t.w & msg_w) >= (1 if strict else 1) or (t.nw & avoid_w) or (t.nw & msg_w) \
                or len(t.w & avoid_w) >= 2:
            continue
        seen.add(t.name.lower())
        out.append(t)
    return out


TOOL_INSTR_ANY = [
    "Which tool should the assistant call next?", "What should the assistant do next?",
    "Pick the next action for the assistant.", "Next action?", "Decide the assistant's next step.",
    "Route this request: which tool, if any, should handle it?",
    "Which function should be called for the user's latest message, if any?",
    "You are the tool router for an AI assistant. Given the conversation, choose which function to "
    "invoke next, or choose to reply without a tool if no function is needed.",
    "Should the agent call one of these tools now, and if so which one?",
    "Select the assistant's next move.", "Tool to call (or none):",
    "Given the user's last message, which of these should happen next?",
]
TOOL_INSTR_ONE = [
    "Which tool should handle this request?", "Which function should be called?",
    "Pick the right tool for this request.", "Choose the function that serves the user's request.",
    "Which API should the assistant call?", "Route this request to the correct tool.",
    "Which tool fits this request?", "Select the function to invoke.",
    "An agent received this request. Which of the tools below should it call to fulfil it?",
    "Tool?",
]
TOOL_INSTR_NEED_ONE = [
    "Which of these tools will the assistant need for this request?",
    "One of these functions is needed to fulfil the request. Which one?",
    "Which of the following tools is required here?",
    "Exactly one of the listed tools is needed for this request. Pick it.",
    "Which listed function does the request call for?",
]
NEED_TOOL_POS = [
    "The assistant should call a function next.", "A tool call is needed to handle the user's last message.",
    "The agent needs to use one of its tools to respond.", "Answering this requires calling a function.",
    "Next step: invoke a tool.", "The assistant must call an available function before replying.",
]
NEED_TOOL_NEG = [
    "The assistant can reply directly without calling any function.", "No tool call is needed at this point.",
    "The agent should just talk to the user now, without using a tool.",
    "None of the available functions should be called right now.",
]
SPECIFIC_POS = [
    "The assistant should call {x} next.", "{x} is the right function to use now.",
    "The next step is to invoke {x}.", "Calling {x} is the correct next action.",
    "The agent should use the tool {x} for this.",
]
SPECIFIC_NEG = ["The assistant should not call {x} now.", "{x} is not the right tool for this."]


def tool_ref(rng: random.Random, t: Tool) -> str:
    if t.desc and rng.random() < 0.55:
        d = strip_end(truncate(t.desc, 110))
        return f"{t.name} ({lower_first(d)})" if rng.random() < 0.6 else f"`{t.name}` ({lower_first(d)})"
    return t.name if rng.random() < 0.6 else f"`{t.name}`"


def specific_noul(rng: random.Random, t: Tool, truth: bool) -> tuple[Question, int]:
    x = tool_ref(rng, t)
    if rng.random() < 0.18:
        return noul(cap_first(rng.choice(SPECIFIC_NEG).format(x=x)), not truth)
    return noul(cap_first(rng.choice(SPECIFIC_POS).format(x=x)), truth)


def pick_n_target(rng: random.Random, n_avail: int) -> int:
    """Total options for a tool question: often just the available tools, often up to 40."""
    lo = max(2, n_avail + 1)
    if rng.random() < 0.25:
        return lo
    return loguniform_int(rng, lo, max(lo, 40))


CALL_SCALES = [
    (["one call", "two calls", "three calls", "four or more calls"], lambda n: min(n, 4) - 1),
    (["a single function call", "two or three calls", "four or more calls"],
     lambda n: 0 if n == 1 else (1 if n <= 3 else 2)),
    (["1", "2", "3", "4", "5 or more"], lambda n: min(n, 5) - 1),
    (["just one tool call", "several tool calls (2 or more)"], lambda n: 0 if n == 1 else 1),
]
CALL_SCALES0 = [
    (["no tool call", "one call", "two calls", "three or more calls"], lambda n: min(n, 3)),
    (["0", "1", "2", "3", "4+"], lambda n: min(n, 4)),
]
CALL_INSTR = [
    "How many function calls does this request need?", "How many tool calls will the assistant have to make?",
    "Estimate the number of API calls required to fully answer the user.", "Count the tool invocations needed.",
    "How many separate tool calls are required to complete this request?",
]


DISTINCT_INSTR = ["How many different tools does this request involve?",
                  "How many distinct functions must be used to handle this?",
                  "Number of different APIs needed for this request?"]


def distinct_score(rng: random.Random, n: int) -> tuple[Question, int]:
    levels = rng.choice([["one tool", "two tools", "three or more tools"], ["1", "2", "3+"],
                         ["a single tool", "several different tools"]])
    return Question("score", pick_instr(rng, DISTINCT_INSTR),
                    [Option(str(i), t) for i, t in enumerate(levels)]), min(n, len(levels)) - 1


def calls_score(rng: random.Random, n: int, with_zero: bool = False) -> tuple[Question, int]:
    scales = CALL_SCALES0 if with_zero else CALL_SCALES
    levels, f = rng.choice(scales)
    return Question("score", pick_instr(rng, CALL_INSTR),
                    [Option(str(i), t) for i, t in enumerate(levels)]), f(n)


# ---------------------------------------------------------------- conversation rendering

ROLE_STYLES = [
    {"USER": "User: ", "ASSISTANT": "Assistant: ", "CALL": "Assistant -> {name}({args})",
     "FUNCTION RESPONSE": "Tool result: "},
    {"USER": "[user] ", "ASSISTANT": "[assistant] ", "CALL": "[tool_call] {name} {args}",
     "FUNCTION RESPONSE": "[tool_result] "},
    {"USER": "USER: ", "ASSISTANT": "ASSISTANT: ", "CALL": "ASSISTANT (function call): {name} {args}",
     "FUNCTION RESPONSE": "FUNCTION RESPONSE: "},
    {"USER": "Human: ", "ASSISTANT": "AI: ", "CALL": "AI called {name} with {args}",
     "FUNCTION RESPONSE": "Function output: "},
]


def render_history(rng: random.Random, turns: list[tuple], max_chars: int) -> str:
    """turns: (role, text, call_name, call_args). A lone user message is sometimes shown bare."""
    if len(turns) == 1 and turns[0][0] == "USER" and rng.random() < 0.35:
        return short(turns[0][1], max_chars)
    st = rng.choice(ROLE_STYLES)
    lines = []
    for role, text, name, args in turns:
        if name:
            lines.append(st["CALL"].format(name=name, args=truncate(args or "{}", 200)))
        else:
            lines.append(st[role] + truncate(text, 500 if role != "USER" else 900))
    return tail("\n".join(lines), max_chars)


# ================================================================ function-calling sources

_GL_SPLIT = re.compile(r"(?:^|\n)(USER|ASSISTANT|FUNCTION RESPONSE):")
_GL_NAME = re.compile(r'"name"\s*:\s*"([^"]+)"')
_GL_ARGS = re.compile(r'"arguments"\s*:\s*\'(.*)\'\s*\}', re.S)


def _glaive_tools(system: str) -> list[Tool]:
    i = system.find("{")
    if i < 0:
        return []
    dec, out, s = json.JSONDecoder(), [], system[i:]
    pos = 0
    while pos < len(s):
        while pos < len(s) and s[pos] in " \n\t\r":
            pos += 1
        if pos >= len(s) or s[pos] != "{":
            break
        try:
            obj, pos = dec.raw_decode(s, pos)
        except json.JSONDecodeError:
            break
        t = parse_tool(obj)
        if t:
            out.append(t)
    return out


def _glaive_turns(chat: str) -> list[tuple]:
    parts = _GL_SPLIT.split(chat)
    turns = []
    for i in range(1, len(parts) - 1, 2):
        role, text = parts[i], parts[i + 1].replace("<|endoftext|>", "").strip()
        name = args = None
        if role == "ASSISTANT" and "<functioncall>" in text:
            m = _GL_NAME.search(text)
            if not m:
                return []
            name = m.group(1)
            a = _GL_ARGS.search(text)
            args = a.group(1) if a else "{}"
        if name and turns and turns[-1][0] == "ASSISTANT" and not turns[-1][2]:
            turns.pop()  # "Sure, let me check that." + call in the next turn = one action: the call
        turns.append((role, clean(text) if not name else "", name, args))
    return turns


STATUS_SCALES = [
    ["not addressed yet: the assistant has not acted on it",
     "in progress: a tool was called and the reply to the user is still pending",
     "handled: the assistant has already replied to it"],
    ["unanswered", "waiting on a tool", "answered"],
    ["open", "in progress", "done"],
]
STATUS_INSTR = ["What is the status of the user's latest request?", "Where is the assistant in handling the last message?",
                "Is the latest request done?", "Track the progress of the user's most recent request.",
                "How far along is the agent with the current request?"]


def _glaive_status(rng: random.Random, turns: list[tuple]) -> Optional[Example]:
    """Cut the conversation at a turn boundary; status of the latest user request is read off
    the last turn: user message -> 0, tool call / tool result -> 1, assistant text -> 2."""
    def status(k):
        role, _, name, _ = turns[k - 1]
        return 0 if role == "USER" else (1 if role == "FUNCTION RESPONSE" or name else 2)
    by = {}
    for k in range(1, len(turns) + 1):
        by.setdefault(status(k), []).append(k)
    want = rng.choice(sorted(by))
    k = rng.choice(by[want])
    if not any(t[0] == "USER" for t in turns[:k]):
        return None
    state = render_history(rng, turns[:k], 1850)
    if rng.random() < 0.6:
        levels = rng.choice(STATUS_SCALES)
        q = (Question("score", pick_instr(rng, STATUS_INSTR), [Option(str(i), t) for i, t in enumerate(levels)]), want)
    else:
        st, truth = rng.choice([("The user's latest request has been answered.", want == 2),
                                ("The assistant still needs to respond to the user's last message.", want != 2),
                                ("The assistant is waiting on a tool result or has not yet reported it.", want == 1),
                                ("The agent has not acted on the latest user message yet.", want == 0)])
        q = noul(st, truth)
    return Example(state, [q])


def glaive_fc(ctx: Ctx):
    rng = ctx.rng
    with open(fetch("glaiveai/glaive-function-calling-v2", "glaive-function-calling-v2.json")) as f:
        rows = json.load(f)
    tools_of = [_glaive_tools(r["system"]) for r in rows]
    pool_d: dict[str, Tool] = {}
    for ts in tools_of:
        for t in ts:
            pool_d.setdefault(t.name, t)
    pool = list(pool_d.values())
    order = list(range(len(rows)))
    rng.shuffle(order)
    for i in order:
        tools = tools_of[i]
        if not tools and rng.random() > 0.25:  # plain chats without functions: keep a few
            continue
        turns = _glaive_turns(rows[i]["chat"])
        if turns and rng.random() < 0.16:
            ex = _glaive_status(rng, turns)
            if ex:
                yield ex
            continue
        points = [j for j, t in enumerate(turns) if t[0] == "ASSISTANT" and j > 0]
        if not points:
            continue
        after_user = [j for j in points if turns[j - 1][0] == "USER"]
        after_fn = [j for j in points if turns[j - 1][0] == "FUNCTION RESPONSE"]
        j = rng.choice(after_fn) if after_fn and rng.random() < 0.15 else rng.choice(after_user or points)
        gold = turns[j][2]
        names = [t.name for t in tools]
        if gold is not None and gold not in names:
            continue
        users = [t[1] for t in turns[:j] if t[0] == "USER"]
        if not users or not users[-1]:
            continue
        msg_w = words(users[-1])
        header = bool(tools) and rng.random() < 0.35
        qs = []
        if header:
            hdr = tools_header(rng, tools)
            state = hdr + "\n\n" + render_history(rng, turns[:j], 1850 - len(hdr))
            if rng.random() < 0.75:
                qs.append(tool_choice(rng, tools, gold, pick_instr(rng, TOOL_INSTR_ANY), True))
            if rng.random() < 0.6 or not qs:
                truth = gold is not None
                if rng.random() < 0.3:
                    qs.append(noul(surface(rng, rng.choice(NEED_TOOL_NEG), 0.06), not truth))
                else:
                    qs.append(noul(surface(rng, rng.choice(NEED_TOOL_POS), 0.06), truth))
        else:
            state = render_history(rng, turns[:j], 1850)
            k = pick_n_target(rng, len(tools)) - len(tools) - 1
            avoid = set().union(*[t.w for t in tools]) if tools else set()
            dis = sample_distractors(rng, pool, max(k, 0 if tools else 1), set(names), msg_w, avoid,
                                     strict=gold is None)
            opts = tools + dis
            if not opts:
                continue
            add_nt = gold is None or rng.random() < 0.75
            instr = pick_instr(rng, TOOL_INSTR_ANY if add_nt else TOOL_INSTR_ONE)
            if gold is None or add_nt or len(opts) >= 2:
                qs.append(tool_choice(rng, opts, gold, instr, add_nt))
        if rng.random() < 0.3:
            truth = gold is not None and rng.random() < 0.65
            if truth:
                t = next(t for t in tools if t.name == gold)
            else:
                others = [t for t in tools if t.name != gold]
                if not others or rng.random() < 0.3:
                    others = sample_distractors(rng, pool, 1, set(names), msg_w,
                                                set().union(*[t.w for t in tools]) if tools else set(), True)
                t = rng.choice(others) if others else None
            if t is not None:
                qs.append(specific_noul(rng, t, truth))
        if qs:
            yield Example(state, qs[:3])


def apigen_fc(ctx: Ctx):
    rng = ctx.rng
    rows = read_rows("argilla/apigen-function-calling", "data/train-00000-of-00001.parquet")
    data = []
    pool_d: dict[str, Tool] = {}
    for r in rows:
        try:
            tl = json.loads(r["tools"])
            calls = json.loads(r["answers"])
        except (json.JSONDecodeError, TypeError):
            continue
        tools = [t for t in (parse_tool(x) for x in tl) if t]
        names = [t.name for t in tools]
        if len(set(names)) != len(names) or not calls:
            continue
        cn = [c.get("name") for c in calls if isinstance(c, dict)]
        if len(cn) != len(calls) or any(c not in names for c in cn):
            continue
        for t in tools:
            pool_d.setdefault(t.name, t)
        data.append((clean(r["query"]), tools, cn))
    pool = list(pool_d.values())
    data = balanced_order(rng, data, lambda d: min(len(d[2]), 4), alpha=0.6)
    for q, tools, cn in data:
        if not q:
            continue
        called = list(dict.fromkeys(cn))
        msg_w = words(q)
        names = [t.name for t in tools]
        avoid = set().union(*[t.w for t in tools])
        header = rng.random() < 0.3
        if header:
            hdr = tools_header(rng, tools)
            state = hdr + "\n\n" + rng.choice(["User request: ", "User: ", "Request: ", ""]) + q
        else:
            state = rng.choice(["", "", "User: ", "Request: "]) + q
        qs = []
        u = rng.random()
        if u < 0.55:
            gold = rng.choice(called)
            cand = [t for t in tools if t.name == gold or t.name not in called]
            if not header:
                k = pick_n_target(rng, len(cand)) - len(cand)
                cand = cand + sample_distractors(rng, pool, max(k, 0), set(names), msg_w, avoid, False)
            add_nt = rng.random() < 0.35
            if len(cand) + add_nt >= 2:
                instr = pick_instr(rng, TOOL_INSTR_ONE if len(called) == 1 else TOOL_INSTR_NEED_ONE)
                qs.append(tool_choice(rng, cand, gold, instr, add_nt))
        if u >= 0.55 or rng.random() < 0.2:
            truth = rng.random() < 0.5
            if truth:
                gname = rng.choice(called)
                t = next(t for t in tools if t.name == gname)
            else:
                unc = [t for t in tools if t.name not in called]
                if not unc or rng.random() < 0.35:
                    unc = sample_distractors(rng, pool, 1, set(names), msg_w, avoid, False)
                t = rng.choice(unc) if unc else None
            if t is not None:
                x = tool_ref(rng, t)
                if rng.random() < 0.15:
                    st = rng.choice(["{x} is not needed for this request.", "The request can be handled without {x}."])
                    qs.append(noul(cap_first(st.format(x=x)), not truth))
                else:
                    st = rng.choice(["The request requires calling {x}.", "{x} is needed to answer this.",
                                     "The assistant should use {x} for this request.",
                                     "Fulfilling this request involves a call to {x}."])
                    qs.append(noul(cap_first(st.format(x=x)), truth))
        if rng.random() < (0.4 if len(cn) > 1 else 0.16):
            qs.append(calls_score(rng, len(cn)) if rng.random() < 0.7 else distinct_score(rng, len(called)))
        elif rng.random() < 0.06:
            qs.append(noul(rng.choice(["This request needs more than one function call.",
                                       "Several tool calls are required here."]), len(cn) > 1))
        if qs:
            yield Example(state, qs[:3])


_H_CALL = re.compile(r"<tool_call>\s*(.*?)\s*</tool_call>", re.S)


def hermes_fc(ctx: Ctx):
    rng = ctx.rng
    rows = []
    for fn in ["func-calling-singleturn.json", "func-calling.json"]:
        with open(fetch("NousResearch/hermes-function-calling-v1", fn)) as f:
            rows += json.load(f)
    data, pool_d = [], {}
    for r in rows:
        tl = r.get("tools")
        if isinstance(tl, str):
            try:
                tl = json.loads(tl)
            except json.JSONDecodeError:
                continue
        tools = [t for t in (parse_tool(x) for x in (tl or [])) if t]
        names = [t.name for t in tools]
        if not tools or len(set(names)) != len(names):
            continue
        conv = r["conversations"]
        hi = next((i for i, m in enumerate(conv) if m["from"] == "human"), None)
        if hi is None or hi + 1 >= len(conv) or conv[hi + 1]["from"] != "gpt":
            continue
        cn = []
        for m in _H_CALL.findall(conv[hi + 1]["value"]):
            nm = _GL_NAME.search(m) or re.search(r"'name'\s*:\s*'([^']+)'", m)
            cn.append(nm.group(1) if nm else None)
        if any(c not in names for c in cn):
            continue
        for t in tools:
            pool_d.setdefault(t.name, t)
        data.append((clean(conv[hi]["value"]), tools, cn))
    pool = list(pool_d.values())
    rng.shuffle(data)
    for q, tools, cn in data:
        called = list(dict.fromkeys(cn))
        names = [t.name for t in tools]
        msg_w, avoid = words(q), set().union(*[t.w for t in tools])
        header = rng.random() < 0.3
        hdr = tools_header(rng, tools, max_chars=900) if header else ""
        state = (hdr + "\n\nUser: " + truncate(q, 1850 - len(hdr))) if header else truncate(q, 1900)
        qs = []
        if rng.random() < 0.6:
            gold = rng.choice(called) if called else None
            cand = [t for t in tools if t.name == gold or t.name not in called]
            if not header:
                k = pick_n_target(rng, len(cand)) - len(cand)
                cand = cand + sample_distractors(rng, pool, max(k, 0), set(names), msg_w, avoid, gold is None)
            add_nt = gold is None or rng.random() < 0.3
            if len(cand) + add_nt >= 2:
                instr = pick_instr(rng, TOOL_INSTR_ANY if add_nt else (TOOL_INSTR_ONE if len(called) == 1
                                                                        else TOOL_INSTR_NEED_ONE))
                qs.append(tool_choice(rng, cand, gold, instr, add_nt))
        if not qs or rng.random() < 0.35:
            truth = bool(called) and rng.random() < 0.5
            unc = [t for t in tools if t.name not in called]
            gname = rng.choice(called) if truth else None
            if not truth and (not unc or rng.random() < 0.3):
                unc = sample_distractors(rng, pool, 1, set(names), msg_w, avoid, False)
            t = next(t for t in tools if t.name == gname) if truth else (rng.choice(unc) if unc else None)
            if t is not None:
                qs.append(specific_noul(rng, t, truth))
        if called and rng.random() < 0.3:
            qs.append(calls_score(rng, len(cn)) if rng.random() < 0.7 else distinct_score(rng, len(called)))
        if qs:
            yield Example(state, qs[:3])


_UNABLE = re.compile(r"unable|can't|cannot|can not|not able|sorry|apolog|don't have (the )?(ability|capabilit|access)",
                     re.I)
W2C_ACTIONS = [
    [("call_tool", "call one of the available tools"), ("ask_user", "ask the user for missing required information"),
     ("decline", "tell the user this can't be done with the available tools")],
    [("tool_call", "make a function call now"), ("clarify", "ask a follow-up question first"),
     ("cannot_help", "explain that no available tool can do this")],
    [("execute", "run the matching tool with the details given"), ("request_info", "request the missing parameters from the user"),
     ("refuse", "say the request is outside what the tools support")],
]
W2C_INSTR = [
    "What should the assistant do next?", "Decide the next action.", "How should the agent respond to this request?",
    "Given the available tools, what is the right next step?", "Next step?",
    "The assistant has only the tools listed. Should it call a tool, ask the user something first, or "
    "explain that it cannot do this?", "Choose the assistant's action.",
]
W2C_SCALE = ["no available tool can handle the request",
             "a suitable tool exists but required details are missing",
             "a suitable tool exists and all required details are provided"]
W2C_SCALE_INSTR = ["How ready is this request for a tool call?",
                   "Rate how directly the assistant can act on this request with its tools.",
                   "Where does this request fall on the scale below?",
                   "Can the request be executed right away with the available tools?"]


def when2call(ctx: Ctx):
    rng = ctx.rng
    items = []

    def cat(c: str) -> Optional[str]:
        if "<TOOLCALL>" in c:
            return "tool"
        if c.strip().endswith("?"):
            return "ask"
        return "cannot" if _UNABLE.search(c) else None

    for fn, resp in [("train/when2call_train_pref.jsonl", "chosen"), ("train/when2call_train_sft.jsonl", "sft")]:
        for r in read_rows("nvidia/When2Call", fn):
            msg = r["messages"][0]["content"] if r["messages"] else ""
            c = r["chosen_response"]["content"] if resp == "chosen" else r["messages"][-1]["content"]
            k = cat(c)
            if not k or not msg:
                continue
            tools = [t for t in (parse_tool(x) for x in r["tools"]) if t]
            call = None
            if k == "tool":
                m = _GL_NAME.search(c)
                call = m.group(1) if m else None
                if call not in [t.name for t in tools]:
                    continue
            items.append((clean(msg), tools, k, call))
    items = dedupe_rows(items, lambda x: x[0], lambda x: x[2])
    for msg, tools, k, call in balanced_order(rng, items, lambda x: x[2], alpha=0.3):
        hdr = tools_header(rng, tools, params=rng.random() < 0.85, max_chars=1100)
        state = hdr + "\n\n" + rng.choice(["User: ", "User request: ", "[user] ", "Request: "]) + truncate(msg, 700)
        qs = []
        u = rng.random()
        if u < 0.55:
            acts = rng.choice(W2C_ACTIONS)
            gold_i = {"tool": 0, "ask": 1, "cannot": 2}[k]
            labs = [Label(a, [d], key=a) for a, d in acts]
            order = [0, 1, 2]
            rng.shuffle(order)
            opts = render_options(rng, [labs[i] for i in order], desc_mode="desc")
            qs.append((Question("choice", pick_instr(rng, W2C_INSTR), opts), order.index(gold_i)))
        elif u < 0.88:
            levels = W2C_SCALE if rng.random() < 0.7 else ["cannot be done with these tools", "needs more details",
                                                            "ready to execute"]
            qs.append((Question("score", pick_instr(rng, W2C_SCALE_INSTR),
                                [Option(str(i), t) for i, t in enumerate(levels)]),
                       {"cannot": 0, "ask": 1, "tool": 2}[k]))
        if u >= 0.88 or rng.random() < 0.3:
            covered = k != "cannot"
            v = rng.random()
            if k != "cannot" and v < 0.45:
                ready = k == "tool"
                st = rng.choice([("The user has provided every required parameter for the tool call.", ready),
                                 ("Some required information is missing from the request.", not ready),
                                 ("The assistant should ask a follow-up question before calling a tool.", not ready),
                                 ("The assistant can call the tool right away.", ready)])
                qs.append(noul(*st))
            else:
                st = rng.choice([("One of the available tools can handle this request.", covered),
                                 ("None of the available tools fit this request.", not covered),
                                 ("The assistant should tell the user it cannot do this with its tools.", not covered),
                                 ("This request is within what the listed functions can do.", covered)])
                qs.append(noul(*st))
        if k == "tool" and len(tools) >= 2 and rng.random() < 0.35:
            qs.append(tool_choice(rng, tools, call, pick_instr(rng, TOOL_INSTR_ONE), False))
        if qs:
            yield Example(state, qs[:3])


# ================================================================ zero-shot hypotheses (Mixtral synthetic)

LADDERS = [
    [r"\blow\b", r"\b(medium|moderate)\b", r"\bhigh\b"],
    [r"\b(days|weeks|short[- ]term)\b", r"\b(months|quarters|medium[- ]term)\b", r"\b(years|decades|long[- ]term)\b"],
    [r"\bbeginners?\b", r"\bintermediate\b", r"\b(advanced|experts?)\b"],
    [r"\bnegative\b", r"\bneutral\b", r"\bpositive\b"],
    [r"\b(easy|simple|basic)\b", r"\b(moderate|medium)\b", r"\b(hard|difficult|complex)\b"],
    [r"\bminor\b", r"\bmoderate\b", r"\b(major|severe|critical)\b"],
    [r"\bnot urgent\b|\bnon-urgent\b", r"\bsomewhat urgent\b", r"\b(very|highly|extremely) urgent\b"],
]


COND_LADDERS = [  # only used when the task description asks for a level / degree
    (r"level|degree|intensity|extent|complexity|expertise|strength|amount|participation|severity|urgency|"
     r"priority|difficulty",
     [r"\b(does not|doesn't|no|little to no|not|limited|minimal)\b", r"\b(some|somewhat|moderate|moderately|steady)\b",
      r"\b(strong|strongly|highly|extensive|very|lively|intense)\b"]),
    (r"cognitive level|bloom", [r"\brecall", r"\bexplain", r"new situations|\bappl(y|ying)\b",
                                r"breaking information|\banaly", r"\bevaluat|\bjudg", r"\bcreat|\bdesign"]),
    (r"technical", [r"general public|layperson|non-expert", r"freshman|sophomore|undergraduate|introduc",
                    r"\bgraduate\b", r"research findings|cutting-edge"]),
]


def ladder_order(hyps: list[str], desc: str = "") -> Optional[list[int]]:
    """If the hypotheses are the levels of an ordinal scale, return their order (lowest first)."""
    ladders = LADDERS + [lad for cond, lad in COND_LADDERS if re.search(cond, desc, re.I)]
    for lad in ladders:
        lv = []
        for h in hyps:
            m = [i for i, p in enumerate(lad) if re.search(p, h, re.I)]
            if len(m) != 1:
                break
            lv.append(m[0])
        else:
            if len(set(lv)) == len(hyps) >= 2:
                return sorted(range(len(hyps)), key=lambda i: lv[i])
    return None


def hyp_short(h: str) -> str:
    """'The text is about sales revenue.' -> 'about sales revenue' (sometimes)."""
    s = re.sub(r"^(this|the) text (is |was )?", "", strip_end(h), flags=re.I)
    return s if len(s) > 3 else strip_end(h)


ZS_CHOICE_GENERIC = ["Which statement best describes the text?", "Which of these applies to the text?",
                     "Pick the label that fits.", "Classify the text.", "Which description is correct?",
                     "Select the category for this text."]
ZS_SCORE_GENERIC = ["Where does the text fall on this scale?", "Rate the text.", "Pick the level that fits."]


def zs_mixtral(ctx: Ctx):
    rng = ctx.rng
    import pandas as pd
    repo = "MoritzLaurer/synthetic_zeroshot_mixtral_v0.1"
    wt = pd.concat([pd.read_parquet(fetch(repo, f)) for f in [
        "mixtral_written_texts_for_tasks_v3/train-00000-of-00004.parquet",
        "mixtral_written_texts_for_tasks_v4/train-00000-of-00002.parquet"]], ignore_index=True)
    written: dict[str, list] = {}
    for t, h, lab, desc, hyps in zip(wt.text, wt.hypothesis, wt["labels"], wt.task_description, wt.task_hypotheses):
        written.setdefault(t, []).append((h, int(lab), desc, list(hyps)))
    del wt
    cats: dict[str, list] = {}
    df = pd.read_parquet(fetch(repo, "mixtral_refinedweb_categories/train-00000-of-00002.parquet"))
    for t, c, lab in zip(df.text, df.category, df["labels"]):
        cats.setdefault(t, []).append((clean(c), int(lab)))
    chars: dict[str, list] = {}
    df = pd.read_parquet(fetch(repo, "mixtral_refinedweb_characteristics/train-00000-of-00002.parquet"))
    for t, h, lab in zip(df.text, df.hypothesis, df["labels"]):
        chars.setdefault(t, []).append((clean(h), int(lab)))
    del df

    def gen_written():
        keys = list(written)
        rng.shuffle(keys)
        for t in keys:
            rows = written[t]
            true = [r for r in rows if r[1] == 0]
            if len(true) != 1:
                continue
            h, _, desc, hyps = true[0]
            hyps = [clean(x) for x in hyps]
            if clean(h) not in hyps or len(hyps) < 2 or ctx.blocked(desc + " " + " ".join(hyps)):
                continue
            text = clean(t)
            if len(text) < 40:
                continue
            gi = hyps.index(clean(h))
            qs = []
            lad = ladder_order(hyps, desc)
            use_desc = rng.random() < 0.55
            if lad is not None and rng.random() < 0.85:
                instr = clean(desc) if use_desc else pick_instr(rng, ZS_SCORE_GENERIC)
                levels = [hyps[i] for i in lad]
                if rng.random() < 0.5:
                    levels = [cap_first(hyp_short(x)) for x in levels]
                qs.append((Question("score", instr, [Option(str(i), x) for i, x in enumerate(levels)]), lad.index(gi)))
            elif rng.random() < 0.75:
                instr = clean(desc) if use_desc else pick_instr(rng, ZS_CHOICE_GENERIC)
                sh = rng.random() < 0.5
                labs = [Label(hyp_short(x) if sh else x, [x], key=snake(" ".join(hyp_short(x).split()[:4])))
                        for x in hyps]
                styles = ["letter", "number", "prefixed", "snake"] + (["name"] if sh else [])
                q, a = make_choice(rng, labs, gi, instr, n=len(labs), other_prob=0.08,
                                   style=weighted(rng, styles, [0.3, 0.2, 0.1, 0.25, 0.15][:len(styles)]),
                                   desc_mode="name" if not sh else None)
                qs.append((q, a))
            if not qs or rng.random() < 0.35:
                h2, lab2 = rng.choice([(r[0], r[1]) for r in rows])
                qs.append(noul(clean(h2), lab2 == 0))
            yield Example(text, qs)

    def gen_cats():
        keys = list(cats)
        rng.shuffle(keys)
        for t in keys:
            rows = [r for r in cats[t] if not ctx.blocked(r[0])]
            T = [c for c, lab in rows if lab == 0]
            F = [c for c, lab in rows if lab == 1]
            if not T or not F:
                continue
            text = clean(t)
            qs = []
            if rng.random() < 0.4:
                g = rng.choice(T)
                opts = [g] + rng.sample(F, min(len(F), rng.randint(1, 5)))
                labs = [Label(lower_first(c), [c], key=snake(c)) for c in opts]
                q, a = make_choice(rng, labs, 0, pick_instr(rng, [
                    "Which topic does this text discuss?", "Which of these labels applies to the text?",
                    "Pick the category that fits the text.", "What is this text about?",
                    "Only one of these topics appears in the text. Which one?", "Topic?"]), n=len(labs))
                qs.append((q, a))
            if not qs or rng.random() < 0.5:
                truth = rng.random() < 0.5
                c = lower_first(rng.choice(T if truth else F))
                if rng.random() < 0.12:
                    st = rng.choice(["This text has nothing to do with {x}.", "The text is not about {x}."])
                    qs.append(noul(st.format(x=c), not truth))
                else:
                    st = rng.choice(["This text is about {x}.", "The text discusses {x}.", "Topic: {x}.",
                                     "{X} is one of the topics of this text.", "The text covers {x}."])
                    qs.append(noul(st.format(x=c, X=cap_first(c)), truth))
            yield Example(text, qs)

    def gen_chars():
        keys = list(chars)
        rng.shuffle(keys)
        for t in keys:
            rows = [r for r in chars[t] if not ctx.blocked(r[0])]
            truth = rng.random() < 0.5
            pick = [h for h, lab in rows if (lab == 0) == truth]
            if not pick:
                continue
            yield Example(clean(t), [noul(rng.choice(pick), truth)])

    gens = [gen_written(), gen_cats(), gen_chars()]
    w = [0.55, 0.25, 0.2]
    while gens:
        i = rng.choices(range(len(gens)), weights=w, k=1)[0]
        try:
            yield next(gens[i])
        except StopIteration:
            gens.pop(i)
            w.pop(i)


# ================================================================ multiple-choice QA

MC_INSTR = [
    "Which answer is correct?", "Pick the correct answer.", "Answer the question.", "Choose the best option.",
    "Select the right answer.", "Which option is right?", "What is the answer?",
    "Choose the most accurate answer to the question.",
    "Read the question and pick the single best answer from the options.", "Answer.",
    "Which of the options correctly answers the question?",
]
MC_CTX_INSTR = [
    "Based on the text, which answer is correct?", "Using the information above, answer the question.",
    "According to the passage, pick the best answer.", "Answer the question using the given context.",
]
NONE_TEXTS = ["None of the above", "None of these", "none of the other options", "No correct answer is listed"]
_ALLNONE = re.compile(r"\b(all|none|both|neither) of (the )?(above|these|them|the options)\b|^(both|all|none|neither)\b",
                      re.I)


def mc_question(rng: random.Random, answers: list[str], gold: int, instr: str,
                p_none: float = 0.1) -> tuple[Question, int]:
    """Shuffled answer options; sometimes a 'none of the above' option is appended, and half
    of those times the real answer is removed so that 'none of the above' is correct."""
    order = list(range(len(answers)))
    rng.shuffle(order)
    texts = [answers[i] for i in order]
    ans = order.index(gold)
    special = any(_ALLNONE.search(a) for a in answers)
    if not special and len(answers) >= 3 and rng.random() < p_none:
        nt = rng.choice(NONE_TEXTS)
        if rng.random() < 0.5:
            texts.pop(ans)
            ans = len(texts)
        texts.append(nt)
    return make_mcq(rng, texts, ans, instr, shuffle=False)


def mc_noul(rng: random.Random, q: Optional[str], answers: list[str], gold: int) -> tuple[Question, int]:
    truth = rng.random() < 0.5
    a = strip_end(answers[gold] if truth else rng.choice([x for i, x in enumerate(answers) if i != gold]))
    if q is None:  # the question is in the state
        if rng.random() < 0.12:
            return noul(rng.choice([f"\"{a}\" is not the correct answer.", f"The answer is not {a}."]), not truth)
        t = rng.choice([f"The correct answer is \"{a}\".", f"The answer is: {a}.", f"\"{a}\" is the right answer.",
                        f"{cap_first(a)} - this answers the question correctly."])
    else:
        q = strip_end(q)
        t = rng.choice([f"The correct answer to \"{q}\" is \"{a}\".", f"Question: {q}? Answer: {a}. This is correct.",
                        f"{q}? -> {a}", f"Q: {q}? A: {a}"])
    return noul(t, truth)


def _stem_instr(rng: random.Random, q: str) -> str:
    q = clean(q)
    if q.rstrip().endswith("?"):
        return q if rng.random() < 0.7 else rng.choice(["Based on the text: ", "Question: ", "Answer this: "]) + q
    if re.match(r"(which|what|who|whom|whose|why|how|where|when|based|according|if|the following|in the|"
                r"from the|each|all|among)\b", q, re.I) or re.search(r"\bfollowing\b", q, re.I):
        return q  # a real question phrased without '?'
    q = q.rstrip(" .:_")
    return rng.choice([f"{q} ...", f"Complete the statement: \"{q} ___\"", f"Which option best completes: \"{q}\"?",
                       f"Fill in: {q} ___"])


def mc_example(rng: random.Random, question: str, answers: list[str], gold: int, *,
               context: Optional[str] = None, ctx_label: str = "Passage", p_q_instr: float = 0.5,
               p_noul: float = 0.25, p_none: float = 0.1, instr: Sequence[str] = MC_INSTR,
               subject: Optional[str] = None) -> Optional[Example]:
    answers = [clean(a) for a in answers]
    question = clean(question)
    if not question or len(answers) < 2 or any(not a for a in answers) or \
            len({a.lower() for a in answers}) != len(answers):
        return None
    ctx_txt = None
    if context:
        lab = rng.choice([f"{ctx_label}: ", f"{ctx_label}:\n", ""])
        ctx_txt = lab + clean(context)
    qs = []
    if ctx_txt and rng.random() < p_q_instr:  # question only in the instructions
        state = ctx_txt
        qs.append(mc_question(rng, answers, gold, _stem_instr(rng, question), p_none))
        if rng.random() < p_noul:
            qs.append(mc_noul(rng, question, answers, gold))
    else:
        qr = rng.choice(["{q}", "{q}", "Question: {q}", "Q: {q}"]).format(q=question)
        state = (ctx_txt + "\n\n" + qr) if ctx_txt else qr
        pool = MC_CTX_INSTR if ctx_txt and rng.random() < 0.5 else instr
        ins = pick_instr(rng, pool)
        if subject and rng.random() < 0.25:
            ins = rng.choice([f"Answer this {subject} question.", f"{subject} question: pick the correct answer."])
        qs.append(mc_question(rng, answers, gold, ins, p_none))
        if rng.random() < p_noul:
            qs.append(mc_noul(rng, None, answers, gold))
    return Example(state, qs)


def _arr(x) -> list:
    return [str(v) for v in list(x)]


def mmlu_aux(ctx: Ctx):
    rng = ctx.rng
    rows = [r["train"] for r in read_rows("cais/mmlu", "auxiliary_train/train-00000-of-00001.parquet")]
    rng.shuffle(rows)
    for r in rows:
        ch = _arr(r["choices"])
        a = int(r["answer"])
        if not 0 <= a < len(ch):
            continue
        q = clean(r["question"])
        if len(q) > 1800:
            continue
        # long items embed a reading passage before the question: split it off sometimes
        ctxt = None
        if len(q) > 600 and "\n" in q:
            head, last = q.rsplit("\n", 1)
            if 10 < len(last) < 300:
                ctxt, q = head, last
        ex = mc_example(rng, q, ch, a, context=ctxt, ctx_label="Passage", p_q_instr=0.6)
        if ex:
            yield ex


def arc(ctx: Ctx):
    rng = ctx.rng
    fs = [f"ARC-{c}/{s}-00000-of-00001.parquet" for c in ["Challenge", "Easy"] for s in ["train", "validation", "test"]]
    rows = read_many("allenai/ai2_arc", fs)
    rng.shuffle(rows)
    for r in rows:
        labs, texts = _arr(r["choices"]["label"]), _arr(r["choices"]["text"])
        if r["answerKey"] not in labs:
            continue
        ex = mc_example(rng, r["question"], texts, labs.index(r["answerKey"]), p_noul=0.3,
                        subject="science" if rng.random() < 0.5 else "grade-school science")
        if ex:
            yield ex


def openbookqa(ctx: Ctx):
    rng = ctx.rng
    rows = read_many("allenai/openbookqa", ["additional/train-00000-of-00001.parquet",
                                            "additional/validation-00000-of-00001.parquet"])
    rng.shuffle(rows)
    for r in rows:
        labs, texts = _arr(r["choices"]["label"]), _arr(r["choices"]["text"])
        if r["answerKey"] not in labs:
            continue
        fact = clean(r["fact1"]) if rng.random() < 0.6 else None
        ex = mc_example(rng, r["question_stem"], texts, labs.index(r["answerKey"]), context=fact,
                        ctx_label=rng.choice(["Fact", "Hint", "Known fact"]), p_q_instr=0.6, p_noul=0.3)
        if ex:
            yield ex


def qasc(ctx: Ctx):
    rng = ctx.rng
    rows = read_many("allenai/qasc", ["data/train-00000-of-00001.parquet", "data/validation-00000-of-00001.parquet"])
    rng.shuffle(rows)
    for r in rows:
        labs, texts = _arr(r["choices"]["label"]), _arr(r["choices"]["text"])
        if r["answerKey"] not in labs:
            continue
        facts = None
        if rng.random() < 0.6:
            fl = [clean(r["fact1"]), clean(r["fact2"])]
            rng.shuffle(fl)
            facts = "\n".join(f"- {f}" for f in fl) if rng.random() < 0.5 else " ".join(fl)
        n = len(texts)
        gold = labs.index(r["answerKey"])
        if rng.random() < 0.35:  # subsample the 8 options (always keeps the answer)
            keep = sorted([gold] + rng.sample([i for i in range(n) if i != gold], rng.randint(2, n - 2)))
            texts, gold = [texts[i] for i in keep], keep.index(gold)
        ex = mc_example(rng, r["question"], texts, gold, context=facts, ctx_label="Facts", p_q_instr=0.6,
                        p_noul=0.3, p_none=0.0)
        if ex:
            yield ex


def medmcqa(ctx: Ctx):
    rng = ctx.rng
    rows = read_rows("openlifescienceai/medmcqa", "data/train-00000-of-00001.parquet")
    rows = [r for r in rows if r["choice_type"] == "single" and r["cop"] in (0, 1, 2, 3)]
    rng.shuffle(rows)
    for r in rows:
        subj = r.get("subject_name") or ""
        subj = subj if subj and subj.lower() not in ("unknown", "none") else "medical"
        ex = mc_example(rng, r["question"], [r["opa"], r["opb"], r["opc"], r["opd"]], int(r["cop"]),
                        p_noul=0.25, p_none=0.05, subject=subj)
        if ex:
            yield ex


def logiqa(ctx: Ctx):
    rng = ctx.rng
    rows = read_rows("lucasmccabe/logiqa", "default/train/0000.parquet", "refs/convert/parquet")
    rng.shuffle(rows)
    for r in rows:
        opts = _arr(r["options"])
        g = int(r["correct_option"])
        if not 0 <= g < len(opts):
            continue
        ex = mc_example(rng, r["query"], opts, g, context=r["context"], ctx_label=rng.choice(["Context", "Argument"]),
                        p_q_instr=0.6, p_noul=0.25, p_none=0.06,
                        instr=MC_INSTR + ["Which option is logically supported?"])
        if ex:
            yield ex


def piqa(ctx: Ctx):
    rng = ctx.rng
    rows = read_rows("ybisk/piqa", "plain_text/train/0000.parquet", "refs/convert/parquet")
    rng.shuffle(rows)
    instr = ["Which solution achieves the goal?", "Which is the better way to do this?",
             "Pick the method that actually works.", "Which approach is more sensible?",
             "Choose the solution that accomplishes the goal.", "How would you do it?",
             "Which of these two procedures works for the goal above? Use physical common sense."]
    for r in rows:
        goal, sols, g = clean(r["goal"]), [clean(r["sol1"]), clean(r["sol2"])], int(r["label"])
        if not goal or g not in (0, 1) or not all(sols) or sols[0] == sols[1]:
            continue
        state = rng.choice(["{g}", "Goal: {g}", "Goal: {g}", "Task: {g}"]).format(g=goal)
        qs = []
        if rng.random() < 0.75:
            qs.append(mc_question(rng, sols, g, pick_instr(rng, instr), 0.0))
        if not qs or rng.random() < 0.3:
            truth = rng.random() < 0.5
            s = strip_end(sols[g] if truth else sols[1 - g])
            qs.append(noul(rng.choice([f"This would work: {s}.", f"\"{s}\" is a sensible way to achieve the goal.",
                                       f"Proposed solution: {s}. This accomplishes the goal."]), truth))
        yield Example(state, qs)


def winogrande(ctx: Ctx):
    rng = ctx.rng
    rows = read_rows("allenai/winogrande", "winogrande_xl/train-00000-of-00001.parquet")
    rng.shuffle(rows)
    instr = ["Who or what does the _ refer to?", "Fill in the blank.", "Which option correctly fills the blank?",
             "Which word should replace the underscore?", "What goes in the blank (_)?",
             "Resolve the blank: which of the two fits?"]
    for r in rows:
        s, o1, o2 = clean(r["sentence"]), clean(r["option1"]), clean(r["option2"])
        if "_" not in s or r["answer"] not in ("1", "2") or not o1 or not o2 or o1 == o2:
            continue
        g = int(r["answer"]) - 1
        opts = [o1, o2]
        qs = []
        if rng.random() < 0.7:
            qs.append(mc_question(rng, opts, g, pick_instr(rng, instr), 0.0))
        else:
            truth = rng.random() < 0.5
            x = opts[g] if truth else opts[1 - g]
            qs.append(noul(rng.choice([f"The blank refers to {x}.", f"\"{x}\" belongs in the blank.",
                                       f"Filling the blank with \"{x}\" makes the sentence sensible."]), truth))
        yield Example(s, qs)


# ================================================================ intents with many labels

MTOP_INSTR = ["What does the user want the assistant to do?", "Which intent does this utterance express?",
              "Route this voice command to the right handler.", "Classify the request.",
              "Which action should the assistant take for this command?", "Intent?",
              "Pick the intent label for this utterance from the list below."]
INTENT_POS2 = ["The intent of this request is \"{x}\".", "This should be routed to the {x} handler.",
               "Intent: {x}.", "The user's request is about: {x}."]
INTENT_NEG2 = ["The intent is not \"{x}\".", "This request has nothing to do with {x}."]


def mtop_intent(ctx: Ctx):
    rng = ctx.rng
    rows = read_many("mteb/mtop_intent", ["en/train-00000-of-00001.parquet", "en/validation-00000-of-00001.parquet"])
    names = sorted({r["label_text"] for r in rows})
    labels = [intent_label(n, group=readable(n).split()[-1]) for n in names]
    rows = dedupe_rows(rows, lambda r: r["text"], lambda r: r["label_text"])
    for r in balanced_order(rng, rows, lambda r: r["label_text"], alpha=0.5):
        text = clean(r["text"])
        gi = names.index(r["label_text"])
        qs = []
        if rng.random() < 0.8:
            n = loguniform_int(rng, 20, len(labels)) if rng.random() < 0.55 else None
            qs.append(make_choice(rng, labels, gi, pick_instr(rng, MTOP_INSTR), n=n, full_prob=0.12,
                                  max_sampled=len(labels), other_prob=0.04))
        if not qs or rng.random() < 0.3:
            qs.append(derived_noul(rng, labels, gi, INTENT_POS2, INTENT_NEG2, p_other_group=0.6))
        yield Example(text, qs)


BITEXT_VERTICALS = [
    ("retail_banking", "bitext/Bitext-retail-banking-llm-chatbot-training-dataset",
     "bitext-retail-banking-llm-chatbot-training-dataset.parquet"),
    ("insurance", "bitext/Bitext-insurance-llm-chatbot-training-dataset", "bitext-insurance-llm-chatbot-training-dataset.csv"),
    ("travel", "bitext/Bitext-travel-llm-chatbot-training-dataset", "bitext-travel-llm-chatbot-training-dataset.csv"),
    ("telco", "bitext/Bitext-telco-llm-chatbot-training-dataset", "bitext-telco-llm-chatbot-training-dataset.csv"),
    ("events", "bitext/Bitext-events-ticketing-llm-chatbot-training-dataset",
     "bitext-events-ticketing-llm-chatbot-training-dataset .csv"),
    ("hospitality", "bitext/Bitext-hospitality-llm-chatbot-training-dataset",
     "bitext-hospitality-llm-chatbot-training-dataset.csv"),
]
BITEXT_INSTR = ["Which intent does this customer message express?", "Route this request to the right workflow.",
                "What does the customer want?", "Pick the intent for this message.", "Classify the customer request.",
                "Which support flow should handle this?", "Intent label?"]


def _fill(rng: random.Random, text: str) -> str:
    def sub(m):
        k = m.group(1).strip()
        if re.search(r"number|id|code|reference", k, re.I):
            return str(rng.randint(10000, 9999999))
        if re.search(r"city|location|destination|origin", k, re.I):
            return rng.choice(["Chicago", "Madrid", "Istanbul", "Osaka", "Leeds", "Lyon", "Denver"])
        return "{{" + k + "}}"
    return _bitext_fill(rng, re.sub(r"\{\{([^}]*)\}\}", sub, text))


def bitext_verticals(ctx: Ctx):
    rng = ctx.rng
    rows, vint = [], {}
    for v, repo, fn in BITEXT_VERTICALS:
        rs = read_rows(repo, fn)
        for r in rs:
            rows.append((v, r["intent"], r["instruction"]))
        vint[v] = sorted({r["intent"] for r in rs})
    counts = Counter(i for v in vint for i in vint[v])
    shared = {i for i, c in counts.items() if c > 1}
    glabels, gindex = [], {}
    for v in vint:
        for i in vint[v]:
            gindex[(v, i)] = len(glabels)
            vv = v.replace("_", " ")
            glabels.append(Label(readable(i), [f"{vv}: {readable(i)}", readable(i)], phrase=readable(i),
                                 key=f"{v}.{i}" if rng.random() < 0.5 else f"{v}_{i}", group=v))
    rows = [(v, i, _fill(rng, clean(t))) for v, i, t in rows]
    rows = dedupe_rows(rows, lambda r: r[2], lambda r: (r[0], r[1]))
    for v, i, text in balanced_order(rng, rows, lambda r: (r[0], r[1]), alpha=0.5):
        qs = []
        u = rng.random()
        if u < 0.55:  # within one business
            labs = [intent_label(x, group=v) for x in vint[v]]
            qs.append(make_choice(rng, labs, vint[v].index(i), pick_instr(rng, BITEXT_INSTR), full_prob=0.35,
                                  other_prob=0.04))
        elif u < 0.8 and i not in shared:  # across businesses: 30-170 options
            allowed = [k for (vv, ii), k in gindex.items() if vv == v or ii not in shared]
            n = loguniform_int(rng, 20, len(allowed))
            qs.append(make_choice(rng, glabels, gindex[(v, i)], pick_instr(rng, BITEXT_INSTR), n=n,
                                  max_sampled=len(allowed), allowed=allowed,
                                  style=weighted(rng, ["snake", "name", "number", "prefixed"], [0.5, 0.2, 0.15, 0.15]),
                                  desc_mode=rng.choice(["none", "desc"]) if n <= 60 else "none"))
        if not qs or rng.random() < 0.3:
            labs = [intent_label(x, group=v) for x in vint[v]]
            qs.append(derived_noul(rng, labs, vint[v].index(i),
                                   ["The customer wants to {x}.", "The customer's request is about: {x}.",
                                    "Intent: {x}.", "This message belongs to the '{x}' flow."],
                                   ["The customer is not asking to {x}."], p_other_group=0.0))
        yield Example(text, qs)


# ================================================================ code / devops

def rosetta(ctx: Ctx):
    rng = ctx.rng
    rows = read_rows("christopher/rosetta-code", "data/train-00000-of-00001-8b4da49264116bbf.parquet")
    rows = [r for r in rows if r["code"] and 60 <= len(r["code"].strip()) and "\n" in r["code"].strip()]
    lc = Counter(r["language_name"] for r in rows)
    langs = [l for l, c in lc.most_common(70) if c >= 300]
    rows = [r for r in rows if r["language_name"] in langs]
    tasks = sorted({r["task_name"] for r in rows})
    tdesc = {}
    for r in rows:
        if r["task_name"] not in tdesc:
            d = clean(r["task_description"] or "").split("\n")[0]
            tdesc[r["task_name"]] = truncate(d, 120) if len(d) > 15 else ""
    tw = {t: words(t) for t in tasks}
    llabels = [Label(l, [l], key=l) for l in langs]
    lang_instr = ["Which programming language is this code written in?", "Identify the language.",
                  "What language is this?", "Name the programming language of the snippet.", "Language?"]
    task_instr = ["Which task does this program solve?", "What does this code implement?",
                  "Which of these problems is this a solution to?", "Identify the programming task.",
                  "What is this program for?", "Match the code to its task."]
    for r in balanced_order(rng, rows, lambda r: r["language_name"], alpha=0.4):
        code = r["code"].replace("\xa0", " ").strip("\n")
        code = re.sub(r"\n{3,}", "\n\n", code)
        if len(code) > 1600:
            code = code[:1600].rsplit("\n", 1)[0] + "\n..."
        qs = []
        u = rng.random()
        if u < 0.4:
            n = sample_n_options(rng, len(langs), full_prob=0.05, max_sampled=40)
            sty = weighted(rng, ["name", "letter", "number"], [0.7, 0.2, 0.1])
            qs.append(make_choice(rng, llabels, langs.index(r["language_name"]), pick_instr(rng, lang_instr), n=n,
                                  style=sty, desc_mode="none" if sty == "name" else "name"))
        elif u < 0.85:
            g = r["task_name"]
            n = loguniform_int(rng, 3, 30)
            cands = [t for t in tasks if t != g and not (tw[t] & tw[g])]
            opts = [g] + rng.sample(cands, n - 1)
            rng.shuffle(opts)
            wd = rng.random() < 0.3
            labs = [Label(t, [tdesc[t]] if wd and tdesc[t] else [], key=snake(t)) for t in opts]
            q, a = make_choice(rng, labs, opts.index(g), pick_instr(rng, task_instr), n=len(opts), shuffle=False,
                               style=weighted(rng, ["name", "letter", "number", "snake"], [0.35, 0.3, 0.15, 0.2]))
            qs.append((q, a))
        if not qs or rng.random() < 0.25:
            truth = rng.random() < 0.5
            l = r["language_name"] if truth else rng.choice([x for x in langs if x != r["language_name"]])
            qs.append(noul(rng.choice([f"This code is written in {l}.", f"The snippet is {l} code.",
                                       f"Language: {l}."]), truth))
        yield Example(code, qs)


def code_defects(ctx: Ctx):
    rng = ctx.rng
    rows = read_rows("google/code_x_glue_cc_defect_detection", "data/train-00000-of-00001.parquet")
    pos = ["This function contains a security vulnerability.", "The code has a bug that could be exploited.",
           "This C function is vulnerable (for example to a memory-safety issue).",
           "A security review would flag this function as vulnerable."]
    neg = ["This function is free of vulnerabilities.", "The code looks safe; there is no exploitable bug."]
    for r in balanced_order(rng, rows, lambda r: bool(r["target"]), alpha=0.0):
        code = re.sub(r"\n\s*\n", "\n", r["func"]).strip()
        if len(code) < 80:
            continue
        if len(code) > 1800:
            code = code[:1800].rsplit("\n", 1)[0] + "\n..."
        vul = bool(r["target"])
        if rng.random() < 0.2:
            yield Example(code, [noul(rng.choice(neg), not vul)])
        else:
            yield Example(code, [noul(rng.choice(pos), vul)])


def _unquote(cmd: str) -> str:
    return re.sub(r"'[^']*'|\"(?:\\.|[^\"\\])*\"", "''", cmd)


_SEG = re.compile(r"\|\||&&|[|;&(){}`]|\$\(")
_XARGS_ARG = {"-I", "-i", "-n", "-L", "-l", "-P", "-d", "-s", "-E", "-a", "--max-args", "--replace", "--delimiter"}
_PREFIX = {"sudo", "time", "nice", "nohup", "command", "builtin", "env", "doas"}
DELETE_CMDS = {"rm", "rmdir", "shred", "unlink"}
KILL_CMDS = {"kill", "pkill", "killall", "shutdown", "reboot", "halt", "poweroff"}
PERM_CMDS = {"chmod", "chown", "chgrp", "setfacl"}
DESTRUCT_CMDS = {"dd", "truncate", "userdel", "groupdel", "wipefs", "fdisk", "parted"}
WRITE_CMDS = {"mv", "cp", "ln", "mkdir", "touch", "tee", "rename", "install", "split", "csplit", "patch", "rsync",
              "scp", "zip", "wget", "useradd", "usermod", "groupadd", "passwd", "mkfifo", "mktemp", "ar", "cpio"}
COMPRESS_CMDS = {"gzip", "gunzip", "bzip2", "bunzip2", "xz", "unxz", "compress", "uncompress", "unzip", "zstd"}
AMBIG_CMDS = {"git", "svn", "hg", "make", "npm", "pip", "pip3", "apt", "apt-get", "yum", "dnf", "brew", "python",
              "python2", "python3", "ruby", "node", "bash", "sh", "zsh", "ksh", "csh", "eval", "source", ".", "perl",
              "cd", "export", "alias", "set", "unset", "pushd", "popd", "ulimit", "umask", "ssh", "su", "screen",
              "tmux", "exec", "history", "read", "trap", "at", "batch", "watch", "vi", "vim", "nano", "emacs",
              "less", "more", "top", "crontab", "while", "for", "if", "do", "then", "declare", "local", "typeset"}
_REDIR = re.compile(r"(?<![0-9&>])>{1,2}(?!\s*(&|/dev/null))")


def _commands(c: str) -> list[list[str]]:
    """Command words with their arguments: segment starts (after | ; && $( ...), the command run
    by xargs, and the command run by find -exec / -ok. Prefixes like sudo are skipped."""
    out = []

    def add(toks):
        while toks and (re.match(r"^[A-Za-z_]\w*=", toks[0]) or toks[0] in _PREFIX):
            toks = toks[1:]
        if toks:
            out.append(toks)
    for seg in _SEG.split(c):
        toks = seg.split()
        add(toks)
        for j, t in enumerate(toks):
            if t in ("-exec", "-execdir", "-ok", "-okdir"):
                add(toks[j + 1:])
            elif t.rsplit("/", 1)[-1] == "xargs":
                k = j + 1
                while k < len(toks) and toks[k].startswith("-"):
                    k += 2 if toks[k] in _XARGS_ARG else 1
                add(toks[k:])
    return out


def bash_risk(cmd: str) -> Optional[int]:
    """0 = read-only, 1 = creates/modifies files, 2 = destructive or privileged (deletes, kills,
    permissions, disks, mounts). None when the command is ambiguous for these rules
    (interpreters, package managers, shell state, loops)."""
    c = _unquote(cmd)
    lv = 0
    for toks in _commands(c):
        w, args = toks[0].rsplit("/", 1)[-1], toks[1:]
        if w in AMBIG_CMDS:
            return None
        if w in DELETE_CMDS or (w in ("find", "rsync") and any(a.startswith(("-delete", "--delete")) for a in args)) \
                or w in PERM_CMDS or w in DESTRUCT_CMDS \
                or w.startswith("mkfs"):
            lv = 2
        elif w in KILL_CMDS:
            if not (w == "kill" and args[:1] in (["-0"], ["-l"], ["-L"])):
                lv = 2
        elif w in ("mount", "umount"):
            if w == "umount" or any(not a.startswith("-") for a in args):
                lv = 2
        elif w in WRITE_CMDS:
            lv = max(lv, 1)
        elif w in COMPRESS_CMDS:
            if not any(re.match(r"^-\w*[clt]", a) for a in args):
                lv = max(lv, 1)
        elif w == "tar" and args and re.search(r"[xcru]", args[0].lstrip("-")):
            lv = max(lv, 1)
        elif w == "sed" and any(re.match(r"^(-\w*i|--in-place)", a) for a in args):
            lv = max(lv, 1)
        elif w == "sort" and any(a.startswith("-o") or a == "--output" for a in args):
            lv = max(lv, 1)
        elif w == "curl" and any(a in ("-o", "-O", "--output", "--remote-name") for a in args):
            lv = max(lv, 1)
        elif w == "find" and any(a.startswith("-fprint") for a in args):
            lv = max(lv, 1)
    if _REDIR.search(c):
        lv = max(lv, 1)
    return lv


RISK_SCALES = [
    ["read-only: only inspects or prints information",
     "modifies the filesystem: creates, moves, copies or edits files",
     "destructive or privileged: deletes data, kills processes, or changes permissions"],
    ["no side effects", "changes files", "deletes or irreversibly alters things"],
    ["safe to run (read-only)", "has side effects (writes files)", "dangerous (deletion, kill, permission change)"],
]


def nl2bash(ctx: Ctx):
    rng = ctx.rng
    rows = read_rows("jiacheng-ye/nl2bash", "default/train/0000.parquet", "refs/convert/parquet")
    rows = [(clean(r["nl"]), r["bash"].strip()) for r in rows if r["nl"] and r["bash"] and len(r["bash"]) < 400]
    rows = dedupe_rows(rows, lambda r: r[0], lambda r: r[1])
    nlw = [words(n) for n, _ in rows]
    rng_order = list(range(len(rows)))
    rng.shuffle(rng_order)

    def distract(i: int, k: int) -> list[int]:
        out, tries = [], 0
        while len(out) < k and tries < 200:
            tries += 1
            j = rng.randrange(len(rows))
            if j == i or j in out or rows[j][1] == rows[i][1]:
                continue
            a, b = nlw[i], nlw[j]
            if a and b and len(a & b) / len(a | b) >= 0.25:
                continue
            out.append(j)
        return out

    cmd_instr = ["Which shell command accomplishes this task?", "Pick the command that does this.",
                 "Which command should the agent run?", "Select the correct bash command.", "Command?"]
    desc_instr = ["What does this command do?", "Which description matches the command?",
                  "Explain the command: pick the right description.", "What is the effect of running this?"]
    for i in rng_order:
        nl, cmd = rows[i]
        risk = bash_risk(cmd)
        u = rng.random()
        qs = []
        if risk is not None and u < {0: 0.22, 1: 0.9, 2: 0.9}[risk]:
            state = rng.choice(["$ {c}", "Command: {c}", "The agent wants to run:\n{c}", "Proposed shell command: {c}",
                                "Task: {n}\nCommand: {c}", "{c}"]).format(c=cmd, n=nl)
            cmds = _commands(_unquote(cmd))
            v = rng.random()
            if v < 0.45:
                levels = rng.choice(RISK_SCALES)
                qs.append((Question("score", pick_instr(rng, ["How risky is it to run this command?",
                                                              "Rate the side effects of this shell command.",
                                                              "What level of side effects does executing this have?",
                                                              "Classify the command by its impact."]),
                                    [Option(str(k), t) for k, t in enumerate(levels)]), risk))
            else:
                st = weighted(rng, ["side", "ro", "del", "perm", "kill"],
                              [0.35, 0.3, 0.15, 0.1, 0.1] if risk == 2 else [0.45, 0.35, 0.2, 0, 0])
                if st == "side":
                    qs.append(noul(rng.choice(["Running this command changes files or system state.",
                                               "This command has side effects.",
                                               "Executing this modifies something on disk."]), risk >= 1))
                elif st == "ro":
                    qs.append(noul(rng.choice(["This command is read-only and has no side effects.",
                                               "It is safe to run this: it only reads or prints information."]),
                                   risk == 0))
                elif st == "del":
                    dele = any(t[0].rsplit("/", 1)[-1] in DELETE_CMDS or (t[0] == "find" and "-delete" in t)
                               for t in cmds)
                    qs.append(noul(rng.choice(["This command deletes files or directories.",
                                               "Running this removes data from disk."]), dele))
                elif st == "perm":
                    qs.append(noul("This command changes file permissions or ownership.",
                                   any(t[0].rsplit("/", 1)[-1] in PERM_CMDS for t in cmds)))
                else:
                    kill = any(t[0].rsplit("/", 1)[-1] in KILL_CMDS and t[1:2] not in (["-0"], ["-l"]) for t in cmds)
                    qs.append(noul("This command terminates running processes or shuts the machine down.", kill))
        elif u < 0.6 or (risk is not None and risk > 0 and u < 0.95):
            state = nl if rng.random() < 0.6 else rng.choice(["Task: ", "Goal: ", "User: "]) + nl
            ds = distract(i, loguniform_int(rng, 2, 10))
            if len(ds) < 1:
                continue
            opts = [cmd] + [rows[j][1] for j in ds]
            qs.append(mc_question(rng, opts, 0, pick_instr(rng, cmd_instr), 0.0))
        else:
            state = rng.choice(["$ {c}", "Command: {c}", "{c}", "{c}"]).format(c=cmd)
            ds = distract(i, loguniform_int(rng, 2, 8))
            if len(ds) < 1:
                continue
            opts = [nl] + [rows[j][0] for j in ds]
            qs.append(mc_question(rng, opts, 0, pick_instr(rng, desc_instr), 0.0))
        if qs:
            yield Example(state, qs)


# ================================================================ driver

@dataclass
class Spec:
    name: str
    fn: Callable[[Ctx], Iterator[Example]]
    cap: int
    origin: str


SPECS = [
    Spec("x_glaive_fc", glaive_fc, 12000, "glaiveai/glaive-function-calling-v2"),
    Spec("x_apigen_fc", apigen_fc, 10000, "argilla/apigen-function-calling (xLAM-60k + Synth-APIGen)"),
    Spec("x_hermes_fc", hermes_fc, 3000, "NousResearch/hermes-function-calling-v1"),
    Spec("x_when2call", when2call, 6500, "nvidia/When2Call (train_sft + train_pref)"),
    Spec("x_zs_mixtral", zs_mixtral, 10000, "MoritzLaurer/synthetic_zeroshot_mixtral_v0.1"),
    Spec("x_mmlu_aux", mmlu_aux, 5500, "cais/mmlu auxiliary_train"),
    Spec("x_arc", arc, 5000, "allenai/ai2_arc (challenge+easy, all splits)"),
    Spec("x_openbookqa", openbookqa, 4000, "allenai/openbookqa additional"),
    Spec("x_qasc", qasc, 4000, "allenai/qasc"),
    Spec("x_medmcqa", medmcqa, 5000, "openlifescienceai/medmcqa"),
    Spec("x_logiqa", logiqa, 4000, "lucasmccabe/logiqa"),
    Spec("x_piqa", piqa, 4000, "ybisk/piqa"),
    Spec("x_winogrande", winogrande, 4000, "allenai/winogrande xl"),
    Spec("x_mtop_intent", mtop_intent, 6000, "mteb/mtop_intent en"),
    Spec("x_bitext_verticals", bitext_verticals, 6000, "bitext/Bitext-*-llm-chatbot-training-dataset (6 verticals)"),
    Spec("x_rosetta", rosetta, 6000, "christopher/rosetta-code"),
    Spec("x_nl2bash", nl2bash, 6000, "jiacheng-ye/nl2bash"),
    Spec("x_code_defects", code_defects, 2500, "google/code_x_glue_cc_defect_detection"),
]


def build(spec: Spec, seed: int, dc: Decontaminator) -> tuple[list[Record], dict]:
    ctx = Ctx(random.Random(f"agent_mcqa:{seed}:{spec.name}"), spec.name)
    st: Counter = Counter()
    recs: list[Record] = []
    if dc.blocked_task(spec.name):
        return recs, {"blocked_dataset": 1}
    seen: set = set()
    for ex in spec.fn(ctx):
        state = truncate(ex.state.strip())
        if not state or not ex.qs:
            continue
        k = norm_key(state)
        if k in seen:
            st["dup_state"] += 1
            continue
        seen.add(k)
        if dc.is_contaminated(state):
            st["decontam"] += 1
            continue
        idx = st["states"]
        st["states"] += 1
        for j, (q, lab) in enumerate(ex.qs):
            r = Record(f"{spec.name}-{idx:06d}-q{j}", spec.name, state, q, int(lab))
            r.validate()
            recs.append(r)
        if len(recs) >= spec.cap:
            break
    st["blocked_task_items"] = ctx.n_blocked
    return recs, dict(st)


def summarize(name: str, recs: list[Record], st: dict) -> dict:
    ty = Counter(r.question.type for r in recs)
    ch = [len(r.question.options) for r in recs if r.question.type == "choice"]
    nl = [r.label for r in recs if r.question.type == "noul"]
    sc = [r for r in recs if r.question.type == "score"]
    return {"name": name, "n": len(recs), "states": st.get("states", 0), "choice": ty["choice"], "noul": ty["noul"],
            "score": ty["score"], "opt_mean": sum(ch) / len(ch) if ch else 0, "opt_max": max(ch) if ch else 0,
            "opt_ge20": sum(n >= 20 for n in ch), "noul_true": sum(nl) / len(nl) if nl else float("nan"),
            "score_levels": round(sum(len(r.question.options) for r in sc) / len(sc), 1) if sc else 0,
            "decontam": st.get("decontam", 0), "dup": st.get("dup_state", 0),
            "blocked": st.get("blocked_task_items", 0),
            "instr_uniq": len({r.question.instructions for r in recs})}


def write_stats(out_dir: Path, rows: list[dict], seed: int) -> None:
    spec_by = {s.name: s for s in SPECS}
    tot = Counter()
    for r in rows:
        for k in ("n", "choice", "noul", "score", "decontam", "dup", "blocked", "states", "opt_ge20"):
            tot[k] += r[k]
    n = max(1, tot["n"])
    L = [f"# data/extra/agent_mcqa (seed {seed})", "",
         "Built by `python -m jevlike.data.extra_agent_mcqa`. Real labeled data that forces reading the options "
         "and the question: tool routing, zero-shot hypotheses, multiple-choice QA, high-cardinality intents, "
         "code/devops. States decontaminated against heldout/test/bench (12-word shingles).", "",
         f"**Total: {tot['n']} records** over {tot['states']} states - choice {tot['choice']} "
         f"({100 * tot['choice'] / n:.0f}%), noul {tot['noul']} ({100 * tot['noul'] / n:.0f}%), "
         f"score {tot['score']} ({100 * tot['score'] / n:.0f}%). Choice questions with >=20 options: {tot['opt_ge20']}. "
         f"Decontamination drops: {tot['decontam']} states; duplicate states skipped: {tot['dup']}; "
         f"held-out-family items skipped: {tot['blocked']}.", "",
         "| file / source | origin | records | choice | noul | score | mean opts | max opts | >=20 opts "
         "| noul true | uniq instr | decontam | dup |",
         "|---|---|---|---|---|---|---|---|---|---|---|---|---|"]
    for r in rows:
        L.append(f"| {r['name']} | {spec_by[r['name']].origin} | {r['n']} | {r['choice']} | {r['noul']} | {r['score']} "
                 f"| {r['opt_mean']:.1f} | {r['opt_max']} | {r['opt_ge20']} | {r['noul_true']:.2f} | {r['instr_uniq']} "
                 f"| {r['decontam']} | {r['dup']} |")
    L += ["", "## Label notes", "",
          "- Tool choice: distractor tools come from other examples and are filtered to share no name word with the "
          "example's tools (name+description words, crude stemming) and no content word with the user message. A few "
          "may still be usable alternatives.",
          "- apigen/hermes choice with several called tools shows exactly one called tool among the options.",
          "- When2Call: 'tool' / 'ask' / 'cannot' from the chosen (pref) or SFT response; 'cannot' detected by "
          "apology/unable wording, other free-text answers dropped.",
          "- zs_mixtral: labels are Mixtral-generated (0 = entailment in the source); written-task texts were generated "
          "for one target hypothesis, so choice over the task hypotheses is single-label by construction; ordinal "
          "tasks (low/medium/high, short/long term, ...) become score questions.",
          "- nl2bash side effects are rule-derived from the command (rm/kill/chmod/dd -> destructive; mv/cp/redirect/"
          "sed -i/... -> writes; else read-only); commands with interpreters, git, sudo, cd, etc. are skipped.",
          "- code_defects (Devign) labels are commit-derived and noisy (SOTA accuracy ~65%).",
          "- medmcqa: single-answer items only; the explanation field is never shown.", ""]
    (out_dir / "STATS.md").write_text("\n".join(L), encoding="utf-8")


def main(argv: Optional[Sequence[str]] = None) -> None:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--only", default="", help="comma-separated dataset names (default: all)")
    ap.add_argument("--out", default=str(OUT_DIR))
    ap.add_argument("--list", action="store_true")
    a = ap.parse_args(argv)
    if a.list:
        for s in SPECS:
            print(f"{s.name:22s} cap {s.cap:6d}  {s.origin}")
        return
    only = [x for x in a.only.split(",") if x]
    out = Path(a.out)
    out.mkdir(parents=True, exist_ok=True)
    dc = Decontaminator()
    rows = []
    for s in SPECS:
        path = out / f"{s.name}.jsonl"
        if only and s.name not in only:
            if path.exists():  # keep stats complete when rebuilding a subset
                from jevlike.types import read_jsonl
                rows.append(summarize(s.name, read_jsonl(str(path)), json.loads(
                    (out / f".{s.name}.st.json").read_text()) if (out / f".{s.name}.st.json").exists() else {}))
            continue
        t0 = time.time()
        recs, st = build(s, a.seed, dc)
        write_jsonl(str(path), recs)
        (out / f".{s.name}.st.json").write_text(json.dumps(st))
        sm = summarize(s.name, recs, st)
        rows.append(sm)
        print(f"{s.name:22s} {len(recs):6d} recs  choice {sm['choice']:5d} noul {sm['noul']:5d} score {sm['score']:5d}"
              f"  decontam {sm['decontam']:4d}  blocked {sm['blocked']:4d}  ({time.time() - t0:.0f}s)", flush=True)
    write_stats(out, rows, a.seed)


if __name__ == "__main__":
    main()