"""Training data for Decider 2B, written as Jev / SystemOne records (see schema.py).
python data.py build --out /workspace/data \
--decontam 'work/artifacts/benchmark-suite/normalized/*.jsonl' \
--tokenizer Qwen/Qwen3.5-2B-Base # exact token-length check
python data.py build --out ... --only tools,home_sim # rebuild some tasks
python data.py review --out /workspace/data --task tools_glaive --n 3
python data.py stats --out /workspace/data
Every record is checked by schema.validate_record, by the quality rules in quality(), and by the
licence policy in licenses.json, and is decontaminated against the benchmark suite. What was
kept and dropped, and why, goes to manifest.json.
"""
import argparse
import ast
import glob
import hashlib
import json
import os
import random
import re
import string
from collections import Counter, defaultdict
import schema
LETTERS = string.ascii_uppercase
MAX_STATE_CHARS = 60000
MAX_OPTION_CHARS = 1500
HERE = os.path.dirname(os.path.abspath(__file__))
# name, hub path, config, train split, eval split, role
# role "train": train split is trained on, eval split measures seen tasks.
# role "heldout": never trained on, measures transfer to unseen tasks.
TASKS = [
dict(name="banking77", path="legacy-datasets/banking77", train="train", eval="test",
text=lambda r: r["text"], label="label",
question="What is the intent of this bank customer's message?", max_opts=26),
dict(name="clinc150", path="clinc/clinc_oos", config="plus", train="train", eval="validation",
text=lambda r: r["text"], label="intent",
question="What is the intent of this request to an assistant?", max_opts=26),
dict(name="mnli", path="nyu-mll/glue", config="mnli", train="train", eval="validation_matched",
text=lambda r: f"Premise: {r['premise']}\nHypothesis: {r['hypothesis']}", label="label",
names=["entailment", "neutral", "contradiction"],
question="Does the premise entail the hypothesis, contradict it, or neither (neutral)?"),
dict(name="qnli", path="nyu-mll/glue", config="qnli", train="train", eval="validation",
text=lambda r: f"Question: {r['question']}\nSentence: {r['sentence']}", label="label",
names=["yes", "no"], question="Does the sentence contain the answer to the question?"),
dict(name="sst2", path="nyu-mll/glue", config="sst2", train="train", eval="validation",
text=lambda r: r["sentence"], label="label", names=["negative", "positive"],
question="What is the sentiment of this review?"),
dict(name="boolq", path="google/boolq", train="train", eval="validation",
text=lambda r: f"Passage: {r['passage']}\nQuestion: {r['question']}?",
label=lambda r: 0 if r["answer"] else 1, names=["yes", "no"],
question="According to the passage, what is the answer to the question?"),
dict(name="paws", path="google-research-datasets/paws", config="labeled_final", train="train",
eval="validation",
text=lambda r: f"Sentence 1: {r['sentence1']}\nSentence 2: {r['sentence2']}", label="label",
names=["no", "yes"], question="Do the two sentences have the same meaning?"),
dict(name="emotion", path="dair-ai/emotion", config="split", train="train", eval="validation",
text=lambda r: r["text"], label="label",
question="Which emotion does the writer express?"),
dict(name="dbpedia", path="fancyzhx/dbpedia_14", train="train", eval="test",
text=lambda r: f"{r['title']}\n{r['content']}", label="label",
question="What kind of thing is described?"),
dict(name="vitaminc", path="tals/vitaminc", train="train", eval="validation",
text=lambda r: f"Evidence: {r['evidence']}\nClaim: {r['claim']}", label="label",
question="Does the evidence support the claim, refute it, or not give enough information?"),
# Multiple-choice tasks: `mcq` returns (state text, option texts, index of the right one).
dict(name="mmlu_aux", path="cais/mmlu", config="auxiliary_train", train="train", rows=24000,
mcq=lambda r: (r["question"], r["choices"], int(r["answer"])),
question="Which option answers the question correctly?"),
dict(name="arc_challenge", path="allenai/ai2_arc", config="ARC-Challenge", train="train", rows=2000,
mcq=lambda r: keyed(r["question"], r["choices"], r["answerKey"]),
question="Which option answers the science question?"),
dict(name="arc_easy", path="allenai/ai2_arc", config="ARC-Easy", train="train", rows=2500,
mcq=lambda r: keyed(r["question"], r["choices"], r["answerKey"]),
question="Which option answers the science question?"),
dict(name="commonsense_qa", path="tau/commonsense_qa", train="train", rows=8000,
mcq=lambda r: keyed(r["question"], r["choices"], r["answerKey"]),
question="Which option is the most sensible answer?"),
dict(name="openbookqa", path="allenai/openbookqa", config="main", train="train", rows=4500,
mcq=lambda r: keyed(r["question_stem"], r["choices"], r["answerKey"]),
question="Which option completes or answers the statement?"),
dict(name="qasc", path="allenai/qasc", train="train", rows=7000,
mcq=lambda r: keyed(r["question"], r["choices"], r["answerKey"]),
question="Which option answers the question?"),
dict(name="sciq", path="allenai/sciq", train="train", rows=8000,
mcq=lambda r: (r["question"], [r["correct_answer"], r["distractor1"], r["distractor2"],
r["distractor3"]], 0),
question="Which option answers the science question?"),
dict(name="medmcqa", path="openlifescienceai/medmcqa", train="train", rows=8000,
mcq=lambda r: (r["question"], [r["opa"], r["opb"], r["opc"], r["opd"]], int(r["cop"])),
question="Which option answers the medical question?"),
dict(name="aqua_rat", path="deepmind/aqua_rat", config="raw", train="train", rows=8000,
mcq=lambda r: (r["question"], [o.split(")", 1)[-1].strip() for o in r["options"]],
"ABCDE".index(r["correct"])),
question="Which option is the answer to the math problem?"),
dict(name="gsm8k", path="openai/gsm8k", config="main", train="train", rows=7000, pool=True,
mcq=lambda r: (r["question"], [r["answer"].split("####")[-1].strip().replace(",", "")], 0),
question="Which number is the answer to the math problem?"),
dict(name="hellaswag", path="Rowan/hellaswag", train="train", rows=8000,
mcq=lambda r: (r["ctx"], r["endings"], int(r["label"])),
question="Which ending continues the text most plausibly?"),
dict(name="winogrande", path="allenai/winogrande", config="winogrande_xl", train="train", rows=8000,
mcq=lambda r: (r["sentence"], [r["option1"], r["option2"]], int(r["answer"]) - 1),
question="Which option fills the blank (_) in the sentence?"),
dict(name="art", path="allenai/art", train="train", rows=6000,
mcq=lambda r: (f"Start: {r['observation_1']}\nEnd: {r['observation_2']}",
[r["hypothesis_1"], r["hypothesis_2"]], int(r["label"]) - 1),
question="Which event most plausibly happened in between?"),
dict(name="cosmos_qa", path="allenai/cosmos_qa", train="train", rows=6000,
mcq=lambda r: (f"{r['context']}\n\n{r['question']}",
[r["answer0"], r["answer1"], r["answer2"], r["answer3"]], int(r["label"])),
question="Which option answers the question about the passage?"),
dict(name="social_i_qa", path="allenai/social_i_qa", train="train", rows=6000,
mcq=lambda r: (f"{r['context']}\n\n{r['question']}",
[r["answerA"], r["answerB"], r["answerC"]], int(r["label"]) - 1),
question="Which option answers the question about the situation?"),
dict(name="race", path="ehovy/race", config="all", train="train", rows=8000,
mcq=lambda r: (f"{r['article']}\n\n{r['question']}", r["options"], "ABCD".index(r["answer"])),
question="Which option answers the question about the article?"),
dict(name="logiqa", path="lucasmccabe/logiqa", train="train", rows=7000,
mcq=lambda r: (f"{r['context']}\n\n{r['query']}", r["options"], int(r["correct_option"])),
question="Which option follows logically?"),
dict(name="squad_answerable", path="rajpurkar/squad_v2", train="train", rows=6000,
mcq=lambda r: (f"Passage: {r['context']}\nQuestion: {r['question']}", ["yes", "no"],
0 if r["answers"]["text"] else 1),
question="Can the question be answered from the passage?"),
dict(name="halueval", path="pminervini/HaluEval", config="qa_samples", train="data", rows=6000,
mcq=lambda r: (f"Knowledge: {r['knowledge']}\nQuestion: {r['question']}\nAnswer: {r['answer']}",
["yes", "no"], 0 if str(r["hallucination"]).lower() == "yes" else 1),
question="Does the answer contain information that the knowledge does not support?"),
dict(name="anli_r1", path="facebook/anli", train="train_r1", rows=4000,
text=lambda r: f"Premise: {r['premise']}\nHypothesis: {r['hypothesis']}", label="label",
names=["entailment", "neutral", "contradiction"],
question="Does the premise entail the hypothesis, contradict it, or neither (neutral)?"),
dict(name="anli_r2", path="facebook/anli", train="train_r2", rows=4000,
text=lambda r: f"Premise: {r['premise']}\nHypothesis: {r['hypothesis']}", label="label",
names=["entailment", "neutral", "contradiction"],
question="Is the hypothesis true, false, or undetermined given the premise?"),
dict(name="anli_r3", path="facebook/anli", train="train_r3", rows=4000,
text=lambda r: f"Premise: {r['premise']}\nHypothesis: {r['hypothesis']}", label="label",
names=["entailment", "neutral", "contradiction"],
question="Does the premise entail the hypothesis, contradict it, or neither (neutral)?"),
dict(name="contract_nli", path="kiddothe2b/contract-nli", config="contractnli_b", train="train", rows=3000,
text=lambda r: f"Contract: {r['premise']}\nStatement: {r['hypothesis']}", label="label",
question="Does the contract support the statement, contradict it, or not mention it?"),
dict(name="ruletaker", path="tasksource/ruletaker", train="train", rows=6000,
text=lambda r: f"Facts and rules: {r['context']}\nStatement: {r['question']}", label="label",
question="Does the statement follow from the facts and rules?"),
dict(name="massive_intent", path="mteb/amazon_massive_intent", config="en", train="train", rows=6000,
text=lambda r: r["text"], label="label_text",
question="What is the intent of this request to a voice assistant?", max_opts=26),
dict(name="tweet_irony", path="cardiffnlp/tweet_eval", config="irony", train="train", rows=2800,
text=lambda r: r["text"], label="label", names=["not ironic", "ironic"],
question="Is this tweet ironic?"),
dict(name="tweet_sentiment", path="cardiffnlp/tweet_eval", config="sentiment", train="train", rows=4000,
text=lambda r: r["text"], label="label", question="What is the sentiment of this tweet?"),
dict(name="tweet_stance_climate", path="cardiffnlp/tweet_eval", config="stance_climate", train="train",
rows=350, text=lambda r: r["text"], label="label", names=["neutral", "against", "in favour"],
question="What is the tweet's stance on climate change being a real concern?"),
dict(name="tweet_stance_feminist", path="cardiffnlp/tweet_eval", config="stance_feminist", train="train",
rows=550, text=lambda r: r["text"], label="label", names=["neutral", "against", "in favour"],
question="What is the tweet's stance on the feminist movement?"),
dict(name="finance_sentiment", path="zeroshot/twitter-financial-news-sentiment", train="train", rows=4000,
text=lambda r: r["text"], label="label", names=["bearish", "bullish", "neutral"],
question="Is this financial news bearish, bullish or neutral for the asset?"),
dict(name="ag_news", path="fancyzhx/ag_news", eval="test", role="heldout",
text=lambda r: r["text"], label="label", question="What is the topic of this news article?"),
dict(name="sms_spam", path="ucirvine/sms_spam", eval="train", role="heldout",
text=lambda r: r["sms"], label="label", names=["not spam", "spam"],
question="Is this text message spam?"),
dict(name="sst5", path="SetFit/sst5", train="train", rows=8000,
text=lambda r: r["text"], label="label",
names=["very negative", "negative", "neutral", "positive", "very positive"],
question="How positive is this review?"),
]
# Tool parsing: conversations become decision points
# {"source", "dialogue", "tools", "calls": [(name, args)], "action"}
ACTIONS = {
"call": "Call one of the available functions",
"ask": "Ask the user for missing information first",
"cannot": "Say that the request cannot be handled with the available functions",
"direct": "Answer directly without calling a function",
}
ASK = re.compile(r"could you|can you (please )?(provide|tell|specify|share|confirm)|please (provide|specify|tell|share|confirm)"
r"|i('ll| will)? need|need to know|which one|what is the|what's the")
CANNOT = re.compile(r"sorry|unable|can't|cannot|apolog|not able|don't have (the )?(ability|capabilit|access)|beyond my")
def action_of(reply):
low = reply.lower()
if "?" in low and ASK.search(low):
return "ask"
if CANNOT.search(low):
return "cannot"
return "direct"
def json_objects(text):
"""Every top-level JSON object or list in a string, in order."""
decoder, found, i = json.JSONDecoder(), [], 0
while i < len(text):
if text[i] in "{[":
try:
value, end = decoder.raw_decode(text, i)
found.append(value)
i = end
continue
except ValueError:
pass
i += 1
return found
def as_tool(value):
if isinstance(value, str):
try:
value = json.loads(value)
except ValueError:
return None
if isinstance(value, dict) and value.get("type") == "function" and isinstance(value.get("function"), dict):
value = value["function"]
return value if isinstance(value, dict) and isinstance(value.get("name"), str) else None
def tool_params(tool):
params = tool.get("parameters") or {}
if isinstance(params, dict) and isinstance(params.get("properties"), dict):
return params["properties"]
return params if isinstance(params, dict) else {}
def tool_card(tool):
names = ", ".join(list(tool_params(tool))[:8])
card = f"{tool['name']}: {str(tool.get('description', '')).strip()[:240]}"
return card + (f" (parameters: {names})" if names else "")
def render_tools(tools, rng):
if rng.random() < 0.5:
return "Available functions:\n" + "\n".join(json.dumps(tool, ensure_ascii=False)[:1500] for tool in tools)
return "Available functions:\n" + "\n".join(f"- {tool_card(tool)}" for tool in tools)
def glaive_points(limit, rng):
points = []
for record in take(load("glaiveai/glaive-function-calling-v2"), limit, rng):
system = record["system"]
tools = [t for t in (as_tool(v) for v in json_objects(system.split(" -", 1)[-1])) if t]
turns = re.split(r"\n*\s*(USER|ASSISTANT|FUNCTION RESPONSE):\s*", "\n" + record["chat"])
turns = [(turns[i], turns[i + 1].replace("<|endoftext|>", "").strip()) for i in range(1, len(turns) - 1, 2)]
dialogue = []
for k, (role, text) in enumerate(turns):
if role == "ASSISTANT" and k and turns[k - 1][0] == "USER":
if "<functioncall>" in text:
name = re.search(r'"name":\s*"([^"]+)"', text)
raw = re.search(r"\"arguments\":\s*'(.*)'\s*}", text, re.S)
try:
args = json.loads(raw.group(1)) if raw else {}
except ValueError:
args = {}
if name:
points.append(dict(source="tools_glaive", dialogue=list(dialogue), tools=tools,
calls=[(name.group(1), args if isinstance(args, dict) else {})],
action="call"))
else:
points.append(dict(source="tools_glaive", dialogue=list(dialogue), tools=tools, calls=[],
action=action_of(text) if tools else "direct"))
dialogue.append((role.lower(), text[:4000]))
return points
TOOLACE_ARG = re.compile(r'(\w+)=("(?:[^"\\]|\\.)*"|\[[^\]]*\]|[^,()]+)')
HTML_ROW = re.compile(r"<tr><td>(.*?)</td><td>(.*?)</td><td>(.*?)</td>", re.S)
def literal(value):
try:
return ast.literal_eval(value.strip())
except (ValueError, SyntaxError):
return value.strip()
LATEX_ROW = re.compile(r"^\s*([^&\n\\]+?)\s*&\s*([^&\n]+?)\s*&\s*(\[.*?\])\s*&", re.M)
NAMED = re.compile(r"tool_name\W+([^\n\"'<&|,]+)[\s\S]*?definition\W+([^\n\"<&|]+)")
def toolace_tools(system):
"""ToolACE lists tools as JSON, an HTML or LaTeX table, or YAML-like text, depending on the row."""
system = system[system.find("tools you can use"):] if "tools you can use" in system else system
for found in json_objects(system):
items = found if isinstance(found, list) else [found]
tools = []
for item in items:
if isinstance(item, dict) and "tool_name" in item:
params = {a.get("parameter_name"): a for a in item.get("arguments") or [] if isinstance(a, dict)}
item = {"name": item["tool_name"], "description": item.get("definition", ""), "parameters": params}
tool = as_tool(item)
if tool:
tools.append(tool)
if tools:
return tools
rows = [(n, d, a) for n, d, a in HTML_ROW.findall(system)] or LATEX_ROW.findall(system)
tools = []
for name, definition, arguments in rows:
name = name.strip()
if name and name != "tool_name" and "<" not in name:
params = {p: {} for p in re.findall(r"parameter_name['\"]?\s*:\s*['\"]?(\w+)", arguments)[:12]}
tools.append({"name": name, "description": re.sub(r"<[^>]+>", " ", definition).strip(),
"parameters": params})
if not tools:
tools = [{"name": n.strip(), "description": d.strip()} for n, d in NAMED.findall(system)]
return tools
def toolace_calls(text, names):
"""Calls written as [Name(a=1, b="x")] or as [Name=>[a|1;b|"x"]]."""
calls = []
for name in sorted(names, key=len, reverse=True):
for match in re.finditer(re.escape(name) + r"\s*(\((.*?)\)|=>\s*\[(.*?)\])", text, re.S):
args = {}
if match.group(2) is not None:
args = {k: literal(v) for k, v in TOOLACE_ARG.findall(match.group(2))}
elif match.group(3):
for pair in match.group(3).split(";"):
if "|" in pair:
key, value = pair.split("|", 1)
args[key.strip()] = literal(value)
calls.append((match.start(), name, args))
text = re.sub(re.escape(name) + r"\s*(\(|=>)", " ", text)
if not calls:
# Other syntaxes ({Name}|(...), [Name]-{...}, Name:<...>): keep which functions, not arguments.
calls = [(text.find(name), name, {}) for name in names if name in text]
return [(name, args) for _, name, args in sorted(calls)]
def looks_like_call(text, names):
return text[:1] in "[({<" and any(name in text[:400] for name in names)
def toolace_points(limit, rng):
points = []
for record in take(load("Team-ACE/ToolACE"), limit, rng):
tools = toolace_tools(record["system"])
if not tools:
continue
names = {tool["name"] for tool in tools}
dialogue = []
for k, turn in enumerate(record["conversations"]):
role, text = turn["from"], turn["value"].strip()
if role == "assistant" and k and record["conversations"][k - 1]["from"] == "user":
calls = toolace_calls(text, names) if looks_like_call(text, names) else []
if calls:
points.append(dict(source="tools_toolace", dialogue=list(dialogue), tools=tools,
calls=calls, action="call"))
elif not looks_like_call(text, names):
points.append(dict(source="tools_toolace", dialogue=list(dialogue), tools=tools, calls=[],
action=action_of(text)))
dialogue.append((role, text[:1500] if role == "tool" else text[:4000]))
return points
def when2call_message(value):
return ast.literal_eval(value) if isinstance(value, str) else value
def when2call_calls(text):
inner = re.search(r"<TOOLCALL>(.*)</TOOLCALL>", text, re.S)
if not inner:
return []
try:
calls = json.loads(inner.group(1))
except ValueError:
return []
return [(c["name"], c.get("arguments") or {}) for c in calls if isinstance(c, dict) and "name" in c]
def when2call_points(limit, rng):
points = []
for record in take(load("nvidia/When2Call", "train_sft"), limit, rng):
tools = [t for t in (as_tool(v) for v in record["tools"]) if t]
messages = when2call_message(record["messages"])
reply = messages[-1]["content"]
calls = when2call_calls(reply)
points.append(dict(source="tools_when2call", tools=tools, calls=calls,
dialogue=[(m["role"], m["content"][:4000]) for m in messages[:-1]],
action="call" if calls else action_of(reply)))
return points
def render_dialogue(dialogue):
return "\n".join(f"{role}: {text}" for role, text in dialogue)
def scalar(value):
return isinstance(value, (str, int, float, bool)) and len(str(value)) <= 120
def argument_options(value, name, pool_values, rng):
options = {json.dumps(value)}
candidates = list(pool_values.get(name, []))
if isinstance(value, (int, float)) and not isinstance(value, bool):
candidates += [value + 1, value * 2, max(0, value - 1), value * 10]
elif isinstance(value, bool):
candidates.append(not value)
for candidate in rng.sample(candidates, min(len(candidates), 12)):
if len(options) >= rng.randint(3, 6):
break
options.add(json.dumps(candidate))
options = sorted(options)
rng.shuffle(options)
return options, options.index(json.dumps(value))
# ---------------------------------------------------------------------------------------------
# Building blocks: questions in Jev format, with the key styles real requests use.
# ---------------------------------------------------------------------------------------------
def keyed(question, choices, answer_key):
return question, list(choices["text"]), list(choices["label"]).index(answer_key)
def slug(text):
return re.sub(r"[^a-z0-9]+", "_", str(text).lower()).strip("_")[:60] or "x"
def opaque(rng, prefix="k_"):
return prefix + "".join(rng.choice("0123456789abcdef") for _ in range(8))
def criteria_for(labels, rng, semantic_keys=None, long_text=False):
"""Criteria for these answer texts, keyed as requests do: by name, opaque id, option_i or letter.
Returns (criteria, keys); keys[i] is the key of labels[i]. When a key does not carry the
meaning itself, the label text is always given as its description.
"""
styles, weights = (["letter", "option", "opaque"], [0.5, 0.25, 0.25]) if long_text else \
(["semantic", "opaque", "option", "letter"], [0.55, 0.2, 0.15, 0.1])
style = rng.choices(styles, weights)[0]
if style == "letter" and len(labels) > len(LETTERS):
style = "option"
keys, descriptions = [], []
for i, label in enumerate(labels):
label = str(label).strip()
if style == "semantic":
key = semantic_keys[i] if semantic_keys else slug(label)
description = None if key == label else label
else:
key = opaque(rng) if style == "opaque" else f"option_{i}" if style == "option" else LETTERS[i]
description = label
while key in keys:
key = f"{key}_{i}"
keys.append(key)
descriptions.append(description[:MAX_OPTION_CHARS] if description else None)
return dict(zip(keys, descriptions)), keys
def choice_q(instructions, labels, gold, rng, semantic_keys=None, long_text=False):
criteria, keys = criteria_for(labels, rng, semantic_keys, long_text)
return {"type": "choice", "instructions": instructions, "criteria": criteria}, keys[gold]
def yes_no_q(instructions, truth):
"""The yes/no-as-choice shape BFCL and ToolRet use."""
return {"type": "choice", "instructions": instructions, "criteria": {"yes": None, "no": None}}, \
"yes" if truth else "no"
def noul_q(statement, truth):
return {"type": "noul", "instructions": statement}, bool(truth)
def score_q(instructions, levels, index):
return {"type": "score", "instructions": instructions, "criteria": [str(level) for level in levels]}, int(index)
def binary_q(instructions, statement, truth, rng):
"""Half yes/no choice, half noul: both shapes appear in real requests."""
return yes_no_q(instructions, truth) if rng.random() < 0.5 else noul_q(statement, truth)
def make(task, source, state, named, rng):
"""A record from (base id, (question, expected)) pairs; question ids get realistic styles."""
questions, expected = {}, {}
for n, (base, (question, value)) in enumerate(named):
roll = rng.random()
qid = base if roll < 0.6 else f"q{n + 1}" if roll < 0.8 else opaque(rng, "q_")
while qid in questions:
qid = f"{qid}_{n}"
questions[qid] = question
expected[qid] = value
return {"task": task, "source": source, "state": state, "questions": questions, "expected": expected}
def shuffled(options, gold, rng):
order = list(range(len(options)))
rng.shuffle(order)
return [options[i] for i in order], order.index(gold)
def load(path, config=None, split="train", streaming=False):
from datasets import load_dataset
return load_dataset(path, config, split=split, streaming=streaming)
def take(data, limit, rng, streaming=False):
if streaming:
return list(data.shuffle(seed=rng.randrange(1 << 30), buffer_size=20000).take(limit))
data = data.shuffle(seed=rng.randrange(1 << 30))
return list(data.select(range(min(limit, len(data)))))
# ---------------------------------------------------------------------------------------------
# Classic single-question tasks.
# ---------------------------------------------------------------------------------------------
# Binary tasks also asked as a noul statement: (statement, label that makes it true).
NOUL = {
"qnli": ("The sentence contains the answer to the question.", "yes"),
"paws": ("The two sentences have the same meaning.", "yes"),
"squad_answerable": ("The question can be answered from the passage.", "yes"),
"halueval": ("The answer contains information that the knowledge does not support.", "yes"),
"tweet_irony": ("This tweet is ironic.", "ironic"),
"sst2": ("This review is positive.", "positive"),
}
# Ordered label sets asked as a score question (share of records).
ORDINAL = {"sst5": 1.0, "tweet_sentiment": 0.5}
def build_classic(task, split, limit, rng):
from datasets import ClassLabel
data = load(task["path"], task.get("config"), split)
found = take(data, limit, rng)
source = f"{task['path']}{':' + task['config'] if task.get('config') else ''}:{split}"
records = []
if "mcq" in task:
found = [r.get("train", r) if isinstance(r.get("train"), dict) else r for r in found]
parsed = []
for row in found:
try:
parsed.append(task["mcq"](row))
except (ValueError, KeyError, IndexError, TypeError):
continue
pool = [options[0] for _, options, _ in parsed] if task.get("pool") else None
for state, options, gold in parsed:
options = [str(option).strip() for option in options]
if pool:
# Wrong options are answers to other problems, as in the benchmark's GSM8K.
others = {o for o in rng.sample(pool, min(30, len(pool))) if o != options[0]}
options += rng.sample(sorted(others), min(len(others), rng.choice([3, 9])))
if not 0 <= gold < len(options) or len(options) < 2:
continue
options, gold = shuffled(options, gold, rng)
long_text = max(map(len, options)) > 40
records.append(make(task["name"], source, str(state).strip(), [
("answer", choice_q(task["question"], options, gold, rng, long_text=long_text))], rng))
return records
label = task["label"]
names = task.get("names")
if names is None:
feature = data.features[label]
names = feature.names if isinstance(feature, ClassLabel) else sorted({str(r[label]) for r in found})
names = [name.replace("_", " ") for name in names]
for row in found:
value = label(row) if callable(label) else row[label]
if isinstance(value, str):
value = value.replace("_", " ")
if value not in names:
continue
gold = names.index(value)
else:
gold = int(value)
if not 0 <= gold < len(names):
continue
state = str(task["text"](row)).strip()
name = task["name"]
if name in ORDINAL and rng.random() < ORDINAL[name]:
named = [("rating", score_q(task["question"], names, gold))]
elif name in NOUL and rng.random() < 0.5:
statement, true_label = NOUL[name]
named = [("statement", noul_q(statement, names[gold] == true_label))]
else:
indices = list(range(len(names)))
limit_options = task.get("max_opts")
# Half the time keep the whole label set: BANKING77 is asked with 77 options, CLINC with 151.
if limit_options and len(indices) > limit_options and rng.random() < 0.5:
others = [i for i in indices if i != gold]
indices = rng.sample(others, rng.randint(4, limit_options) - 1) + [gold]
rng.shuffle(indices)
named = [("label", choice_q(task["question"], [names[i] for i in indices], indices.index(gold), rng))]
records.append(make(name, source, state, named, rng))
return records
def pack(records, rng, fraction=0.15):
"""Merge a few short single-question records of one task into one multi-item state.
Teaches routing each question to the part of the state it is about, as schema benchmarks
do (one question per device, per clause, per document).
"""
short = [r for r in records if isinstance(r["state"], str) and len(r["state"]) < 1200 and len(r["questions"]) == 1]
chosen = set(rng.sample(range(len(short)), int(len(short) * fraction))) if short else set()
chosen_ids = {id(short[i]) for i in chosen}
keep = [r for r in records if id(r) not in chosen_ids]
group = [short[i] for i in sorted(chosen)]
packed = []
while len(group) >= 2:
size = min(rng.randint(2, 4), len(group))
parts, group = group[:size], group[size:]
state = {"items": [{"id": f"item_{k + 1}", "text": r["state"]} for k, r in enumerate(parts)]}
named = []
for k, r in enumerate(parts):
(qid, question), = r["questions"].items()
question = dict(question, instructions=f"About item_{k + 1}: {question['instructions']}")
named.append((f"item_{k + 1}_{qid}", (question, r["expected"][qid])))
rng.shuffle(named)
packed.append(make(parts[0]["task"], parts[0]["source"], state, named, rng))
return keep + group + packed
# ---------------------------------------------------------------------------------------------
# Tool use. Parsed decision points become records in the shapes the tool benchmarks ask.
# ---------------------------------------------------------------------------------------------
ACTION_KEYS = ["call", "ask", "cannot", "direct"]
def compact_tool(tool):
"""The tool as a request would show it: name, description, parameters."""
out = {"name": tool["name"], "description": str(tool.get("description", "")).strip()}
params = tool_params(tool)
if params:
# Tables give only parameter names; show those as a list rather than empty objects.
out["parameters"] = list(params) if all(v in ({}, None) for v in params.values()) else params
return out
def when2call_pref(limit, rng):
"""Pick the right reply among candidates, the shape of the When2Call MCQ track."""
records = []
for row in take(load("nvidia/When2Call", "train_pref"), limit, rng):
try:
messages = when2call_message(row["messages"])
chosen = when2call_message(row["chosen_response"])["content"].strip()
rejected = when2call_message(row["rejected_response"])["content"].strip()
except (ValueError, SyntaxError, KeyError, TypeError):
continue
if not chosen or not rejected or chosen == rejected:
continue
tools = [compact_tool(t) for t in (as_tool(v) for v in row["tools"]) if t]
state = {"functions": tools, "conversation": [{"role": m["role"], "content": m["content"]} for m in messages]}
options, gold = shuffled([chosen, rejected], 0, rng)
records.append(make("when2call_pref", "nvidia/When2Call:train_pref:train", state, [
("reply", choice_q("Which reply should the assistant send next?", options, gold, rng, long_text=True))], rng))
return records
# ---------------------------------------------------------------------------------------------
# Artificial Analysis Decision Index shapes: knowledge with "I don't know", safety, When2Call MCQ.
# ---------------------------------------------------------------------------------------------
IDK_TEXTS = ["I don't know", "I don't know", "I don't know", "I do not know", "Not sure"]
IDK_INSTRUCTIONS = [
"Which answer is correct?",
"Which answer is correct? Choose \"I don't know\" rather than guess.",
"Answer the question. If you are not sure, choose \"I don't know\".",
"Which option answers the question correctly? Pick \"I don't know\" if none of them is right.",
]
IDK_SOURCES = ["mmlu_aux", "arc_challenge", "arc_easy", "commonsense_qa", "openbookqa", "qasc", "medmcqa"]
def knowledge_idk(limit, rng):
"""Knowledge MCQs with four answers and "I don't know", the shape of AA-Omniscience-Decision.
In about a third the right answer is replaced by another wrong one, so "I don't know" is the
right choice: the model learns to pick it when no answer is supported rather than guess.
(The served abstention threshold in schema.py handles plain uncertainty.)
"""
tasks = {t["name"]: t for t in TASKS}
per_source = limit // len(IDK_SOURCES)
records = []
for name in IDK_SOURCES:
task = tasks[name]
# A state may also appear in the plain MCQ task, asked differently; both are train data, and
# benchmark decontamination runs over every record.
found = take(load(task["path"], task.get("config"), task["train"]), per_source, rng)
for row in found:
row = row.get("train", row) if isinstance(row.get("train"), dict) else row
try:
state, options, gold = task["mcq"](row)
except (ValueError, KeyError, IndexError, TypeError):
continue
options = [str(o).strip() for o in options]
right = options[gold]
wrong = [o for o in dict.fromkeys(options) if o and o != right]
if len(wrong) < 3:
continue
unanswerable = rng.random() < 0.33
answers = rng.sample(wrong, min(4 if unanswerable else 3, len(wrong)))
if not unanswerable:
answers.append(right)
rng.shuffle(answers)
idk = rng.choice(IDK_TEXTS)
labels = answers + [idk] if rng.random() < 0.8 else [idk] + answers
target = labels.index(idk) if unanswerable else labels.index(right)
record = make("knowledge_idk", f"{task['path']}:{task['train']}", str(state).strip(), [
("answer", choice_q(rng.choice(IDK_INSTRUCTIONS), labels, target, rng,
long_text=max(map(len, labels)) > 40))], rng)
records.append(record)
return records
WILDGUARD_QUESTIONS = {
"prompt_harm": ["Is the human user's request harmful?", "The user's request is harmful.",
"Does the user ask for something harmful?"],
"refusal": ["Is the AI assistant's response a refusal?", "The assistant refuses the request.",
"Does the assistant decline to help?"],
"response_harm": ["Is the AI assistant's response harmful?", "The assistant's response is harmful.",
"Does the assistant's reply contain harmful content?"],
}
def wildguard_train(limit, rng):
"""WildGuardMix train (ODC-BY, gated): prompt harm, refusal and response harm as yes/no questions,
posed as WildGuardTest asks them (state "Human user:" / "AI assistant:")."""
records = []
for row in take(load("allenai/wildguardmix", "wildguardtrain"), limit, rng):
prompt = str(row.get("prompt") or "").strip()
if not prompt:
continue
state = f"Human user:\n{prompt[:12000]}"
named = []
if row.get("prompt_harm_label") in ("harmful", "unharmful"):
named.append(("prompt_harm", row["prompt_harm_label"] == "harmful"))
response = str(row.get("response") or "").strip()
if response and row.get("response_refusal_label") in ("refusal", "compliance") \
and row.get("response_harm_label") in ("harmful", "unharmful"):
state += f"\n\nAI assistant:\n{response[:12000]}"
named += [("refusal", row["response_refusal_label"] == "refusal"),
("response_harm", row["response_harm_label"] == "harmful")]
if not named:
continue
questions = []
for qid, truth in named:
wording = rng.choice(WILDGUARD_QUESTIONS[qid])
questions.append((qid, noul_q(wording, truth) if rng.random() < 0.85 else
yes_no_q(wording if wording.endswith("?") else wording + " Yes or no?", truth)))
records.append(make("wildguard_train", "allenai/wildguardmix:wildguardtrain:train", state, questions, rng))
return records
W2C_SYSTEM = (
"<system>\nYou are a helpful AI assistant.\nYou have access to the following tools described in <tool></tool> "
"which you can use to answer the user's questions.\nOnly use a tool if it directly answers the user's question.\n\n"
"To use a tool, return JSON in the following format:\n"
"{\"name\": \"tool_name\", \"arguments\": {\"argument1\": \"value1\", \"argument2\": \"value2\", ...}}\n\n")
W2C_INSTRUCTIONS = [
"The state is a conversation between a user and an AI assistant that can use the tools listed in its system "
"prompt. Which of these replies should the assistant give to the user's last message?",
"Which reply should the assistant send next?",
"Pick the assistant's best reply: answer, call a tool, ask for missing information, or say it cannot help.",
]
def when2call_mcq(limit, rng):
"""When2Call train replies grouped by question: a choice among 2-4 replies of different kinds
(call a tool, ask for information, say it cannot, answer directly), shuffled, as the MCQ test asks."""
groups = defaultdict(dict)
tools_of = {}
for config, split in (("train_pref", "train"), ("train_sft", "train")):
for row in load("nvidia/When2Call", config, split):
try:
messages = when2call_message(row["messages"])
if config == "train_pref":
replies = [when2call_message(row["chosen_response"])["content"],
when2call_message(row["rejected_response"])["content"]]
history = messages
else:
replies, history = [messages[-1]["content"]], messages[:-1]
except (ValueError, SyntaxError, KeyError, TypeError, IndexError):
continue
if len(history) != 1 or history[0]["role"] != "user":
continue # the test is single-turn
question = history[0]["content"].strip()
tools_of.setdefault(question, row["tools"])
for reply in replies:
reply = str(reply).strip()
if not reply:
continue
kind = "call" if when2call_calls(reply) else action_of(reply)
groups[question].setdefault(kind, reply)
# The gold reply: the one train_pref chose, or the SFT reply; rebuilt from each group's kinds.
gold_kind = {}
for row in load("nvidia/When2Call", "train_pref", "train"):
try:
messages = when2call_message(row["messages"])
chosen = when2call_message(row["chosen_response"])["content"].strip()
except (ValueError, SyntaxError, KeyError, TypeError):
continue
if len(messages) == 1:
gold_kind[messages[0]["content"].strip()] = "call" if when2call_calls(chosen) else action_of(chosen)
questions = [q for q in groups if q in gold_kind and len(groups[q]) >= 2]
rng.shuffle(questions)
records = []
for question in questions[:limit]:
replies = groups[question]
kinds = list(replies)
rng.shuffle(kinds)
gold = kinds.index(gold_kind[question])
tools = [t for t in (as_tool(v) for v in tools_of[question]) if t]
if rng.random() < 0.6: # the test's raw prompt
state = W2C_SYSTEM + "\n".join(f"<tool>{json.dumps(t)}</tool>" for t in tools) + \
f"\n</system>\n\n<user>\n{question}\n</user>"
else:
state = {"functions": [compact_tool(t) for t in tools], "conversation": [{"role": "user", "content": question}]}
texts = []
for kind in kinds:
text = replies[kind]
if kind == "call" and rng.random() < 0.8:
# The MCQ test writes a tool call as bare JSON, not <TOOLCALL>[...]</TOOLCALL>.
text = "\n".join(json.dumps({"name": n, "arguments": a}) for n, a in when2call_calls(text)) or text
texts.append(text)
records.append(make("when2call_mcq", "nvidia/When2Call:train_pref+train_sft:train", state, [
("reply", choice_q(rng.choice(W2C_INSTRUCTIONS), texts, gold, rng, long_text=True))], rng))
return records
def tool_records(points, rng):
pool = {}
pool_values = defaultdict(set)
for point in points:
for tool in point["tools"]:
pool.setdefault(tool["name"], tool)
for _, args in point["calls"]:
for key, value in args.items():
if scalar(value) and len(pool_values[key]) < 50:
pool_values[key].add(value)
pool_values = {k: sorted(v, key=str) for k, v in pool_values.items()}
every_tool = list(pool.values())
sources = {"tools_glaive": "glaiveai/glaive-function-calling-v2:train", "tools_toolace": "Team-ACE/ToolACE:train",
"tools_when2call": "nvidia/When2Call:train_sft:train"}
records = []
for point in points:
called = [name for name, _ in point["calls"]]
tools = list(point["tools"])
action = point["action"]
source = sources[point["source"]]
# A fifth of tool calls become "no tool fits": the needed functions are swapped out.
if action == "call" and rng.random() < 0.2:
tools = [t for t in rng.sample(every_tool, min(len(every_tool), len(tools) + 3))
if t["name"] not in called][: max(1, len(tools))]
called, action = [], "cannot"
present = {t["name"] for t in tools}
extra = [t for t in rng.sample(every_tool, min(len(every_tool), 8)) if t["name"] not in present]
tools = tools + extra[: rng.randint(0, 6)]
rng.shuffle(tools)
if not tools:
continue
conversation = [{"role": role, "content": text} for role, text in point["dialogue"]]
if not conversation:
continue
shape = rng.random()
if shape < 0.15 and action == "call":
# API-Bank shape: one question over many described APIs.
many = [t for t in rng.sample(every_tool, min(len(every_tool), rng.randint(20, 60)))
if t["name"] != called[0]] + [pool.get(called[0]) or {"name": called[0]}]
options, gold = shuffled([tool_card(t) for t in many], len(many) - 1, rng)
records.append(make("tool_select_big", source, {"conversation": conversation}, [
("api", choice_q("Which API should be called next?", options, gold, rng, long_text=True))], rng))
continue
if shape < 0.3 and action == "call":
# ToolRet shape: is each candidate tool useful for the request.
users = [c["content"] for c in conversation if c["role"] == "user"]
candidates = tools + [t for t in rng.sample(every_tool, min(len(every_tool), 24))
if t["name"] not in {x["name"] for x in tools}]
candidates = candidates[: rng.randint(8, 32)]
rng.shuffle(candidates)
named = [(f"doc_{i}", binary_q(f"Is this tool useful for the request?\nTool: {tool_card(t)}",
f"This tool is useful for the request: {tool_card(t)}",
t["name"] in called, rng)) for i, t in enumerate(candidates)]
records.append(make("tool_retrieval", source, {"request": users[-1] if users else conversation[-1]["content"]},
named, rng))
continue
state = {"functions": [compact_tool(t) for t in tools], "conversation": conversation}
names = [t["name"] for t in tools]
named = [("next_action", choice_q("What should the assistant do next?", [ACTIONS[k] for k in ACTION_KEYS],
ACTION_KEYS.index(action), rng, semantic_keys=ACTION_KEYS))]
for name in names[:12]:
named.append((f"call_{slug(name)}", binary_q(f"Should the function `{name}` be called now?",
f"The function `{name}` should be called now.",
name in called, rng)))
gold = names.index(called[0]) if called and called[0] in names else len(names)
named.append(("first_function", choice_q("Which function should be called first?",
names + ["None of these functions"], gold, rng,
semantic_keys=names + ["none"])))
for name, args in (point["calls"][:2] if action == "call" else []):
for key, value in rng.sample(sorted(args.items()), min(len(args), 3)):
if scalar(value):
options, gold = argument_options(value, key, pool_values, rng)
if len(options) > 1:
named.append((f"arg_{slug(key)}", choice_q(
f"What value should the argument `{key}` of `{name}` take?", options, gold, rng,
long_text=True)))
records.append(make(point["source"], source, state, named, rng))
return records
def tools(rng):
points = []
for name, (builder, limit) in TOOL_SOURCES.items():
found = builder(limit, rng)
print(f" {name}: {len(found)} decision points", flush=True)
points += found
return tool_records(points, rng)
# ---------------------------------------------------------------------------------------------
# Retrieval, product relevance, phishing, contracts.
# ---------------------------------------------------------------------------------------------
def overlap(a, b):
return len(set(words(a)) & set(words(b)))
def retrieval_nq(limit, rng):
"""Natural Questions: one gold passage plus lexically similar passages from other questions."""
pairs = [(r["query"], r["answer"]) for r in take(load("sentence-transformers/natural-questions", "pair"),
limit * 2, rng) if r["query"] and r["answer"]]
passages = [answer for _, answer in pairs]
records = []
for query, gold_passage in pairs[:limit]:
candidates = [p for p in rng.sample(passages, min(60, len(passages))) if p != gold_passage]
candidates.sort(key=lambda p: -overlap(query, p))
docs = candidates[: rng.randint(3, 9)] + [gold_passage]
docs, gold = shuffled(docs, len(docs) - 1, rng)
shape = rng.random()
if shape < 0.4:
# BRIGHT shape: documents in the state, one relevance question each.
state = {"query": query, "documents": [{"id": f"doc_{i}", "text": d} for i, d in enumerate(docs)]}
named = [(f"doc_{i}", binary_q(f"Is doc_{i} relevant to the query?",
f"doc_{i} answers the query.", i == gold, rng)) for i in range(len(docs))]
elif shape < 0.8:
# ToolRet shape: the query is the state and each document is in its question.
state = {"query": query}
named = [(f"doc_{i}", binary_q(f"Does this document answer the query?\nDocument: {d}",
f"This document answers the query: {d}", i == gold, rng))
for i, d in enumerate(docs)]
else:
state = {"query": query}
named = [("best_document", choice_q("Which document answers the query?", docs, gold, rng,
long_text=True))]
records.append(make("retrieval_nq", "sentence-transformers/natural-questions:pair:train", state, named, rng))
return records
def msmarco(limit, rng):
records = []
for row in take(load("microsoft/ms_marco", "v2.1", streaming=True), limit, rng, streaming=True):
texts, selected = row["passages"]["passage_text"], row["passages"]["is_selected"]
if sum(selected) != 1 or len(texts) < 3:
continue
state = {"query": row["query"], "documents": [{"id": f"doc_{i}", "text": t} for i, t in enumerate(texts)]}
named = [(f"doc_{i}", binary_q(f"Is doc_{i} relevant to the query?", f"doc_{i} answers the query.",
bool(s), rng)) for i, s in enumerate(selected)]
records.append(make("retrieval_msmarco", "microsoft/ms_marco:v2.1:train", state, named, rng))
return records
ESCI_LEVELS = {
"Exact": ("E", "Exact: the product is relevant and satisfies every part of the query"),
"Substitute": ("S", "Substitute: somewhat relevant; it misses part of the query but could stand in"),
"Complement": ("C", "Complement: does not meet the query but could be used with a matching product"),
"Irrelevant": ("I", "Irrelevant: unrelated to the query or misses a central part of it"),
}
def esci(limit, rng):
records, exact = [], 0
for row in take(load("tasksource/esci", streaming=True), limit * 3, rng, streaming=True):
if row["product_locale"] != "us" or row["esci_label"] not in ESCI_LEVELS or not row["product_title"]:
continue
if row["esci_label"] == "Exact":
if exact + 1 > 0.4 * (len(records) + 1):
continue
exact += 1
description = re.sub(r"\s+", " ", re.sub(r"<[^>]+>", " ", row.get("product_description") or "")).strip()
state = {"query": row["query"], "product": {"title": row["product_title"], "description": description[:1500]}}
labels = list(ESCI_LEVELS)
question = {"type": "choice", "instructions": "How does the product relate to the shopping query?",
"criteria": {ESCI_LEVELS[k][0]: ESCI_LEVELS[k][1] for k in labels}}
records.append(make("esci", "tasksource/esci:train", state,
[("relevance", (question, ESCI_LEVELS[row["esci_label"]][0]))], rng))
if len(records) >= limit:
break
return records
def phishing(limit, rng):
records = []
for row in take(load("zefang-liu/phishing-email-dataset"), limit, rng):
text, kind = (row.get("Email Text") or "").strip(), row.get("Email Type")
if len(text) < 40 or kind not in ("Phishing Email", "Safe Email"):
continue
bad = kind == "Phishing Email"
records.append(make("phishing", "zefang-liu/phishing-email-dataset:train", {"email": text[:8000]}, [
("verdict", choice_q("Is this email phishing or legitimate?", ["phishing", "legitimate"],
0 if bad else 1, rng)),
("is_phishing", noul_q("This email is a phishing attempt.", bad))], rng))
return records
def contract_multi(limit, rng):
"""ContractNLI as the benchmark asks it: one contract, every hypothesis as its own question."""
data = load("kiddothe2b/contract-nli", "contractnli_b")
names = {"entailment": "Entailment", "contradiction": "Contradiction", "neutral": "NotMentioned"}
labels = data.features["label"].names
by_contract = defaultdict(dict)
for row in data:
by_contract[row["premise"]][row["hypothesis"]] = names[labels[row["label"]]]
contracts = sorted(by_contract)
rng.shuffle(contracts)
records = []
for premise in contracts[:limit]:
hypotheses = sorted(by_contract[premise].items())
rng.shuffle(hypotheses)
named = []
for n, (hypothesis, label) in enumerate(hypotheses):
question = {"type": "choice", "instructions": hypothesis,
"criteria": {"Entailment": "The contract supports the statement",
"Contradiction": "The contract contradicts the statement",
"NotMentioned": "The contract does not mention it"}}
named.append((f"nda-{n + 1}", (question, label)))
records.append(make("contract_multi", "kiddothe2b/contract-nli:contractnli_b:train",
{"contract": premise}, named, rng))
return records
# ---------------------------------------------------------------------------------------------
# Synthetic smart-home commands: many devices, one command, a question per device.
# Generated from scratch to teach target selection and state updates; no benchmark rows used.
# ---------------------------------------------------------------------------------------------
ROOMS = ["kitchen", "living room", "bedroom", "office", "bathroom", "garage", "hallway", "kids room"]
KINDS = {
"light": {"power": ["on", "off"], "brightness": ["10%", "25%", "50%", "75%", "100%"]},
"thermostat": {"mode": ["heat", "cool", "off"], "target": [f"{t} C" for t in range(16, 29)]},
"blinds": {"position": ["open", "half open", "closed"]},
"speaker": {"power": ["on", "off"], "volume": [str(v) for v in range(0, 11)]},
"tv": {"power": ["on", "off"], "input": ["HDMI 1", "HDMI 2", "antenna"]},
"lock": {"state": ["locked", "unlocked"]},
"fan": {"power": ["on", "off"], "speed": ["low", "medium", "high"]},
}
ORDERED = {"brightness", "target", "volume", "speed"}
OUTCOMES = {"carried_out": "The command is carried out", "refused": "Refused: the value is not allowed",
"needs_clarification": "Needs clarification: more than one device matches",
"no_such_device": "Impossible: no such device", "question_only": "Nothing changes: it is a question"}
INVALID = {"brightness": ["150%", "-20%"], "target": ["5 C", "35 C", "40 C"], "volume": ["15", "-3"],
"speed": ["turbo"], "input": ["HDMI 7"], "position": ["sideways"]}
def home_sim(count, rng):
records = []
outcome_keys = list(OUTCOMES)
for _ in range(count):
devices = []
for room in rng.sample(ROOMS, rng.randint(3, 6)):
for kind in rng.sample(list(KINDS), rng.randint(1, 3)):
devices.append({"id": f"{kind}.{room.replace(' ', '_')}", "kind": kind, "room": room})
devices = devices[:16]
style = rng.random()
twin = None
if style < 0.12:
# Two devices of one kind in one room, so "the <kind> in the <room>" is ambiguous.
twin = rng.choice(devices)
base_id = twin["id"]
twin["id"] = base_id + "_1"
devices.append(dict(twin, id=base_id + "_2"))
for d in devices:
d["state"] = {k: rng.choice(v) for k, v in KINDS[d["kind"]].items()}
chosen = twin or rng.choice(devices)
kind, room = chosen["kind"], chosen["room"]
attribute = rng.choice(list(KINDS[kind]))
value = rng.choice(KINDS[kind][attribute])
targets = []
if twin:
command, outcome = f"Set the {kind} in the {room} to {attribute} {value}.", "needs_clarification"
elif style < 0.37:
command, outcome = f"Set every {kind} to {attribute} {value}.", "carried_out"
targets = [d for d in devices if d["kind"] == kind]
elif style < 0.67:
command, outcome = f"Set the {kind} in the {room} to {attribute} {value}.", "carried_out"
targets = [chosen]
elif style < 0.79:
missing = [r for r in ROOMS if not any(d["kind"] == kind and d["room"] == r for d in devices)]
command, outcome = f"Set the {kind} in the {rng.choice(missing or ['attic'])} to {attribute} {value}.", \
"no_such_device"
elif style < 0.9:
command, outcome = f"What is the {attribute} of the {kind} in the {room}?", "question_only"
else:
options = [a for a in KINDS[kind] if a in INVALID]
if options:
attribute = rng.choice(options)
command, outcome = f"Set the {kind} in the {room} to {attribute} {rng.choice(INVALID[attribute])}.", \
"refused"
else:
command, outcome = f"What is the {attribute} of the {kind} in the {room}?", "question_only"
before = {d["id"]: dict(d["state"]) for d in devices}
after = {d["id"]: dict(d["state"]) for d in devices}
for d in targets:
after[d["id"]][attribute] = value
rng.shuffle(devices)
state = {"devices": [{"id": d["id"], "room": d["room"], "state": before[d["id"]]} for d in devices],
"command": command}
named = [("outcome", choice_q("What happens with the command?", [OUTCOMES[k] for k in outcome_keys],
outcome_keys.index(outcome), rng, semantic_keys=outcome_keys))]
ids = {d["id"] for d in targets}
for d in devices:
named.append((f"acts_{slug(d['id'])}", binary_q(f"Does the command act on `{d['id']}`?",
f"The command acts on `{d['id']}`.", d["id"] in ids, rng)))
for d in rng.sample(devices, min(2, len(devices))):
name = rng.choice(list(KINDS[d["kind"]]))
levels = KINDS[d["kind"]][name]
named.append((f"before_{slug(d['id'])}_{name}", setting_q(
f"Before the command, what is the {name} of `{d['id']}`?", name, levels,
levels.index(before[d["id"]][name]), rng)))
for d in (targets or rng.sample(devices, 1))[:2]:
name = attribute if attribute in KINDS[d["kind"]] else rng.choice(list(KINDS[d["kind"]]))
levels = KINDS[d["kind"]][name]
named.append((f"after_{slug(d['id'])}_{name}", setting_q(
f"After the command, what is the {name} of `{d['id']}`?", name, levels,
levels.index(after[d["id"]][name]), rng)))
records.append(make("home_sim", "synthetic:home_sim", state, named, rng))
return records
def setting_q(instructions, name, levels, index, rng):
"""Ordered settings are score questions, the rest are choices."""
if name in ORDERED:
return score_q(instructions, levels, index)
return choice_q(instructions, levels, index, rng)
# ---------------------------------------------------------------------------------------------
# tasksource: hundreds of human-labelled classification tasks (tasksource/zero-shot-label-nli).
# Each row is a text, "This example is <label>." and whether that is true. True rows become a
# choice over the task's labels (or a score when the labels are ordered), false rows a noul.
# ---------------------------------------------------------------------------------------------
HYPOTHESIS = re.compile(r"^This example is (.+?)\.?$")
# Tasks that are, or come from, a Decision Index benchmark or one of our own tasks.
TASKSOURCE_EXCLUDE = re.compile(
r"anli|banking|clinc|hellaswag|winogrande|contract|hover|fever|mmlu|\barc\b|gsm8k|bigbench|bbh|boolq|"
r"\bsst|ag_news|sms_spam|vitaminc|tweet_eval|squad|esci|sarcasm|\bvast\b|nli4ct|acos|finentity|financ|"
r"humicroedit|humor|cladder|musr|cruxeval|\bsgd\b|schema_guided|when2call|bfcl|api.?bank|toolret|bright|"
r"ragtruth|phish|yorker|habermas|chess|gpqa|emotion|dbpedia|mnli|qnli|paws|commonsense|cosmos|social_i|"
r"\brace\b|logiqa|ruletaker|massive|medmcqa|sciq|qasc|openbook|halueval|\bart\b")
ORDER_WORDS = ["very low", "low", "medium", "moderate", "high", "very high"]
def ordinal(labels):
"""Labels sorted low to high when every label starts with an order word, else None."""
ranks = []
for label in labels:
rank = max((i for i, w in enumerate(ORDER_WORDS) if label.lower().startswith(w + " ") or label.lower() == w),
default=None, key=lambda i: len(ORDER_WORDS[i]))
if rank is None:
return None
ranks.append(rank)
if len(set(ranks)) != len(labels):
return None
return [label for _, label in sorted(zip(ranks, labels))]
def readable(task):
return re.sub(r"[_/\-]+", " ", task).strip()
def tasksource(per_task, rng, policy):
data = load("tasksource/zero-shot-label-nli", streaming=True)
labels = defaultdict(set)
sample = defaultdict(list)
seen = Counter()
for row in data:
task = row["task"]
match = HYPOTHESIS.match(row["hypothesis"].strip())
if not match or TASKSOURCE_EXCLUDE.search(task.lower()):
continue
labels[task].add(match.group(1))
if policy == "commercial" and row.get("license_use") != "commercial":
continue
seen[task] += 1
# Reservoir sample of up to 2 x per_task rows per task.
cap = 2 * per_task
if len(sample[task]) < cap:
sample[task].append(row)
else:
j = rng.randrange(seen[task])
if j < cap:
sample[task][j] = row
records = []
for task in sorted(sample):
names = sorted(labels[task])
if not 2 <= len(names) <= 40:
continue
order = ordinal(names)
kept = 0
for row in sample[task]:
if kept >= per_task:
break
label = HYPOTHESIS.match(row["hypothesis"].strip()).group(1)
text = (row["premise"] or "").strip()
if len(text) < 3:
continue
truth = row["labels"] == 0
name = f"tasksource/{task}"
# Balance: true rows -> 75% choice/score, 25% noul(true); false rows -> 40% kept as noul(false).
if not truth and rng.random() >= 0.4:
continue
if truth and rng.random() < 0.75:
if order:
named = [("rating", score_q(f"Task: {readable(task)}. Which level fits this text?", order,
order.index(label)))]
else:
named = [("label", choice_q(f"Task: {readable(task)}. Which label fits this text?", names,
names.index(label), rng))]
else:
named = [("statement", noul_q(f"Task: {readable(task)}. The correct label for this text is "
f"\"{label}\".", truth))]
record = make(name, f"tasksource/zero-shot-label-nli:train (task={task})", text, named, rng)
record["license"] = f"{row.get('license')} (use: {row.get('license_use')})"
records.append(record)
kept += 1
return records
# ---------------------------------------------------------------------------------------------
# Quality, licences, decontamination, splits.
# ---------------------------------------------------------------------------------------------
PLACEHOLDER_KEY = re.compile(r"^(k_[0-9a-f]{8}|option_\d+|[A-Z])$")
def quality(record):
"""Why a record is not good enough to train on, or None."""
try:
schema.validate_record(record)
except schema.SchemaError as error:
return f"schema: {str(error)[:60]}"
text = schema.state_text(record["state"])
if len(text.strip()) < 3:
return "empty state"
if len(text) > MAX_STATE_CHARS:
return "state too long"
if text.count("�") > 2:
return "encoding errors"
for question in record["questions"].values():
if question["type"] == "choice":
criteria = question["criteria"]
seen = set()
for key, description in criteria.items():
if description is None and PLACEHOLDER_KEY.match(key):
return "option without meaning"
meaning = re.sub(r"\s+", " ", (description or key)).strip().lower()
if not meaning or meaning in seen:
return "duplicate or empty option"
seen.add(meaning)
elif question["type"] == "score" and len(set(map(schema.level_text, question["criteria"]))) < len(question["criteria"]):
return "duplicate score level"
return None
def licence_ok(record, policy, decisions):
task = record["task"]
if task.startswith("tasksource/"):
return policy == "research" or "(use: commercial)" in record["license"]
decision = decisions.get(task, {}).get("decision", "review")
return decision == "approve" if policy == "commercial" else decision != "reject"
def words(text):
return re.findall(r"[a-z0-9]+", text.lower())
def ngram_ids(text, n=13):
"""64-bit ids of every n-word span (lower-cased words, punctuation ignored), sorted."""
import numpy as np
tokens = words(text)
ids = np.fromiter((int.from_bytes(hashlib.md5(" ".join(tokens[i:i + n]).encode()).digest()[:8], "little")
for i in range(max(0, len(tokens) - n + 1))), dtype=np.uint64)
return np.unique(ids)
def leaves(value):
"""Every string inside a state or question structure."""
if isinstance(value, str):
yield value
elif isinstance(value, dict):
for item in value.values():
yield from leaves(item)
elif isinstance(value, list):
for item in value:
yield from leaves(item)
def leaf_ids(value, min_words=5):
"""Ids of whole text fields of at least min_words words: catches exact copies of short texts
(a 10-word banking query) that 13-word spans cannot."""
import numpy as np
ids = set()
for text in leaves(value):
tokens = words(text)
if len(tokens) >= min_words:
ids.add(int.from_bytes(hashlib.md5(" ".join(tokens).encode()).digest()[:8], "little"))
return np.array(sorted(ids), dtype=np.uint64)
def record_text(record):
"""Text compared with the benchmarks. Function definitions are left out: the same tools recur
across many conversations, so matching on them would drop unrelated records."""
state = record["state"]
if isinstance(state, dict) and "functions" in state:
state = {k: v for k, v in state.items() if k != "functions"}
parts = [schema.state_text(state)]
for question in record["questions"].values():
parts.append(question.get("instructions") or "")
criteria = question.get("criteria")
if isinstance(criteria, dict):
parts += [f"{k} {v or ''}" for k, v in criteria.items()]
elif isinstance(criteria, list):
parts += [schema.level_text(level) for level in criteria]
return "\n".join(parts)
def suite_rows(patterns):
"""Rows of the benchmark suite files (plain or .gz JSON lines); patterns are comma-separated globs."""
import gzip
for pattern in patterns.split(","):
for path in sorted(glob.glob(pattern.strip())):
opener = gzip.open if path.endswith(".gz") else open
with opener(path, "rt") as file:
for line in file:
row = json.loads(line)
payload = row.get("payload") or row
yield {"state": payload.get("state"), "questions": payload.get("questions")}
_SEEN = None
def _row_ids(rows):
import numpy as np
spans = [ngram_ids(json.dumps(r["state"], ensure_ascii=False) + "\n" + json.dumps(r["questions"], ensure_ascii=False))
for r in rows]
fields = [leaf_ids([r["state"], r["questions"]]) for r in rows]
empty = np.empty(0, np.uint64)
return (np.unique(np.concatenate(spans)) if spans else empty, np.unique(np.concatenate(fields)) if fields else empty)
def _member(sorted_ids, ids):
import numpy as np
if not len(sorted_ids) or not len(ids):
return np.zeros(len(ids), dtype=bool)
positions = np.searchsorted(sorted_ids, ids)
positions[positions >= len(sorted_ids)] = 0
return sorted_ids[positions] == ids
def _overlaps(records):
"""(shared 13-word spans, shared whole text fields) per record."""
spans, fields = _SEEN
out = []
for record in records:
state = record["state"]
if isinstance(state, dict) and "functions" in state:
state = {k: v for k, v in state.items() if k != "functions"}
out.append((int(_member(spans, ngram_ids(record_text(record))).sum()),
int(_member(fields, leaf_ids(state)).sum())))
return out
def is_overlap(counts, min_spans=2):
spans, fields = counts
return spans >= min_spans or fields >= 1
MAX_WORKERS = None # cap on parallel workers (decontam --workers); None uses every CPU
def workers_available():
count = len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else (os.cpu_count() or 1)
return min(count, MAX_WORKERS) if MAX_WORKERS else count
def benchmark_ngrams(patterns):
"""Sorted array of every 13-word span id in the benchmark suite, built on every core."""
from multiprocessing import get_context
import numpy as np
rows = list(suite_rows(patterns))
if not rows:
raise SystemExit(f"No benchmark rows matched {patterns}")
chunks = [rows[i:i + 200] for i in range(0, len(rows), 200)]
with get_context("fork").Pool(workers_available()) as pool:
parts = pool.map(_row_ids, chunks)
spans = np.unique(np.concatenate([p[0] for p in parts]))
fields = np.unique(np.concatenate([p[1] for p in parts]))
print(f"benchmark suite: {len(rows)} rows, {len(spans)} distinct 13-word spans, {len(fields)} text fields",
flush=True)
return spans, fields
def overlap_counts(records, seen):
"""(shared 13-word spans, shared whole text fields) per record, on every core (arrays shared by fork)."""
from multiprocessing import get_context
global _SEEN
_SEEN = seen
chunks = [records[i:i + 500] for i in range(0, len(records), 500)]
with get_context("fork").Pool(workers_available()) as pool:
return [n for part in pool.map(_overlaps, chunks) for n in part]
def fingerprint(record):
body = json.dumps([record["state"], record["questions"]], sort_keys=True, ensure_ascii=False)
return hashlib.sha1(body.encode()).hexdigest()
def state_hash(record):
return hashlib.sha1(schema.state_text(record["state"]).encode()).hexdigest()
# name: (builder(limit, rng), limit). "tools" builds every tool task; tasksource is per task.
EXTRA = {
"tools": (None, None),
"when2call_pref": (when2call_pref, 8000),
"when2call_mcq": (when2call_mcq, 9000),
"knowledge_idk": (knowledge_idk, 21000),
"wildguard_train": (wildguard_train, 40000),
"retrieval_nq": (retrieval_nq, 15000),
"retrieval_msmarco": (msmarco, 15000),
"esci": (esci, 10000),
"phishing": (phishing, 8000),
"contract_multi": (contract_multi, 600),
"home_sim": (home_sim, 8000),
"tasksource": (None, 300),
}
TOOL_SOURCES = {"glaive": (glaive_points, 15000), "toolace": (toolace_points, 11000),
"when2call": (when2call_points, 15000)}
TOOL_TASKS = {"tools_glaive", "tools_toolace", "tools_when2call", "tool_select_big", "tool_retrieval"}
def produce(name, rng, args):
"""Records for one entry of TASKS or EXTRA, before any filtering."""
if name == "tools":
return tools(rng)
if name == "tasksource":
return tasksource(args.tasksource_per_task, rng, args.policy)
builder, limit = EXTRA[name]
return builder(limit, rng)
def _fit_one(record, tokenizer, max_length):
"""(parts, None) with the record split into passes that fit, or ([], reason) if it cannot fit."""
import prompting
questions = schema.items(record["questions"])
try:
passes, _, _ = prompting.plan(tokenizer, schema.state_text(record["state"]), questions, max_length)
except prompting.TooLong:
return [], "too long for the model"
if len(passes) == 1:
return [record], None
parts = []
for k, indices in enumerate(passes):
ids = [questions[i]["id"] for i in indices]
part = dict(record, id=f"{record['id']}#{k}", questions={q: record["questions"][q] for q in ids},
expected={q: record["expected"][q] for q in ids})
if "teacher" in record:
part["teacher"] = {q: v for q, v in record["teacher"].items() if q in ids}
parts.append(part)
return parts, None
_WORKER = {}
def _fit_init(tokenizer_name, max_length):
from transformers import AutoTokenizer
_WORKER["tokenizer"] = AutoTokenizer.from_pretrained(tokenizer_name)
_WORKER["max_length"] = max_length
def _fit_chunk(records):
return [_fit_one(r, _WORKER["tokenizer"], _WORKER["max_length"]) for r in records]
def fit_lengths(records, tokenizer_name, max_length, dropped, workers=None):
"""Exact token check with the model's own layout, on every CPU core: records whose questions do
not fit in one pass are split into several, and records that cannot fit at all are dropped."""
from multiprocessing import get_context
if not workers:
# The CPUs this process may use (a pod can be given fewer than the host has).
workers = len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else (os.cpu_count() or 1)
chunks = [records[i:i + 500] for i in range(0, len(records), 500)]
out = []
with get_context("fork").Pool(workers, initializer=_fit_init, initargs=(tokenizer_name, max_length)) as pool:
for done, results in enumerate(pool.imap(_fit_chunk, chunks), 1):
for record, (parts, reason) in zip(chunks[done - 1], results):
if reason:
dropped[record["task"]][reason] += 1
out.extend(parts)
if done % 100 == 0 or done == len(chunks):
print(f"length check: {done * 500 if done < len(chunks) else len(records)}/{len(records)} records",
flush=True)
return out
def build(args):
rng = random.Random(args.seed)
with open(os.path.join(HERE, "licenses.json")) as file:
decisions = json.load(file)["tasks"]
only = set(args.only.split(",")) if args.only else None
if not args.decontam and not args.no_decontam:
raise SystemExit("Pass --decontam <benchmark glob> (or --no-decontam to skip it knowingly).")
produced, heldout = [], []
for task in TASKS:
if only and task["name"] not in only:
continue
try:
if task.get("role") == "heldout":
heldout += build_classic(task, task["eval"], 300, rng)
else:
produced += pack(build_classic(task, task["train"], task.get("rows", 4000) + 2 * args.eval_per_task,
rng), rng)
print(f"built {task['name']}", flush=True)
except Exception as error:
print(f"SKIPPED {task['name']}: {type(error).__name__}: {str(error)[:200]}", flush=True)
for name in EXTRA:
if only and name not in only:
continue
try:
found = produce(name, rng, args)
produced += found
print(f"built {name}: {len(found)} records", flush=True)
except Exception as error:
print(f"SKIPPED {name}: {type(error).__name__}: {str(error)[:200]}", flush=True)
dropped = defaultdict(Counter)
licences = {}
kept, fingerprints = [], set()
for record in produced + [dict(r, split="heldout") for r in heldout]:
record["license"] = record.get("license") or decisions.get(record["task"], {}).get("hf", "unknown")
record["id"] = f"{record['task']}-{fingerprint(record)[:12]}"
reason = quality(record)
if not reason and record.get("split") != "heldout" and not licence_ok(record, args.policy, decisions):
reason = f"licence ({args.policy})"
if not reason and fingerprint(record) in fingerprints:
reason = "duplicate"
if reason:
dropped[record["task"]][reason] += 1
continue
fingerprints.add(fingerprint(record))
licences[record["task"]] = record["license"] if not record["task"].startswith("tasksource/") else "per row"
kept.append(record)
if args.decontam:
seen = benchmark_ngrams(args.decontam)
counts = overlap_counts(kept, seen)
for record, count in zip(kept, counts):
if is_overlap(count):
dropped[record["task"]]["benchmark overlap"] += 1
kept = [record for record, count in zip(kept, counts) if not is_overlap(count)]
if args.tokenizer:
kept = fit_lengths(kept, args.tokenizer, args.max_length, dropped)
# Splits: every record of a state goes to one split, so no state is both trained and evaluated.
splits = {"train": [], "dev": [], "calibration": [], "heldout": []}
by_task = defaultdict(list)
for record in kept:
if record.pop("split", None) == "heldout":
splits["heldout"].append(record)
else:
# tasksource is ~290 small tasks: hold out from the family as a whole, not from each task.
by_task["tasksource" if record["task"].startswith("tasksource/") else record["task"]].append(record)
for task, records in sorted(by_task.items()):
hashes = [state_hash(r) for r in records]
per_state = Counter(hashes)
states = sorted(per_state)
rng.shuffle(states)
want = args.eval_per_task * (5 if task == "tasksource" else 1)
if len(records) <= 10 * want:
want = max(1, len(records) // 10)
eval_states, count = set(), 0
for h in states:
if count >= 2 * want:
break
eval_states.add(h)
count += per_state[h]
dev_states = set(sorted(eval_states)[: len(eval_states) // 2])
for r, h in zip(records, hashes):
split = "train" if h not in eval_states else "dev" if h in dev_states else "calibration"
splits[split].append(r)
rng.shuffle(splits["train"])
os.makedirs(args.out, exist_ok=True)
for name, records in splits.items():
with open(os.path.join(args.out, f"{name}.jsonl"), "w") as file:
for record in records:
file.write(json.dumps(record, ensure_ascii=False) + "\n")
manifest = summarize(splits, dropped, licences, args)
with open(os.path.join(args.out, "manifest.json"), "w") as file:
json.dump(manifest, file, indent=1)
print_summary(manifest)
if not splits["train"] or not splits["dev"]:
raise SystemExit("A split is empty; check the SKIPPED lines above.")
def summarize(splits, dropped, licences, args):
tasks = defaultdict(lambda: {"records": Counter(), "questions": Counter(), "types": Counter()})
gold = defaultdict(Counter)
for name, records in splits.items():
for record in records:
entry = tasks[record["task"]]
entry["records"][name] += 1
entry["questions"][name] += len(record["questions"])
for qid, question in record["questions"].items():
entry["types"][question["type"]] += 1
if name == "train":
gold[record["task"]][f"{question['type']}={record['expected'][qid]}" if question["type"] == "noul"
else question["type"]] += 1
totals = {name: {"records": len(records), "questions": sum(len(r["questions"]) for r in records)}
for name, records in splits.items()}
types = Counter(q["type"] for records in splits.values() for r in records for q in r["questions"].values())
return {"policy": args.policy, "decontaminated": bool(args.decontam), "length_checked": bool(args.tokenizer),
"max_length": args.max_length if args.tokenizer else None,
"totals": totals, "question_types": dict(types),
"tasks": {t: {"records": dict(v["records"]), "questions": dict(v["questions"]),
"types": dict(v["types"]), "license": licences.get(t), "train_answers": dict(gold[t]),
"dropped": dict(dropped.get(t, {}))} for t, v in sorted(tasks.items())},
"dropped_tasks": {t: dict(c) for t, c in dropped.items() if t not in tasks}}
def print_summary(manifest):
print(f"\npolicy {manifest['policy']}, decontaminated {manifest['decontaminated']}, "
f"length checked {manifest['length_checked']}")
for name, total in manifest["totals"].items():
print(f"{name:12} {total['records']:>8} records {total['questions']:>9} questions")
print("question types:", manifest["question_types"])
families = defaultdict(lambda: [0, 0])
for task, entry in manifest["tasks"].items():
family = "tasksource" if task.startswith("tasksource/") else task
families[family][0] += sum(entry["records"].values())
families[family][1] += sum(entry["questions"].values())
for family, (records, questions) in sorted(families.items(), key=lambda x: -x[1][1]):
print(f" {family:28} {records:>8} records {questions:>9} questions")
def decontam(args):
"""Drop records that share 13-word spans with the benchmark suite from an existing build."""
global MAX_WORKERS
MAX_WORKERS = args.workers
seen = benchmark_ngrams(args.against)
os.makedirs(args.dest, exist_ok=True)
report, splits = {}, {}
for name in ("train", "dev", "calibration", "heldout"):
path = os.path.join(args.out, f"{name}.jsonl")
if not os.path.exists(path):
continue
with open(path) as file:
records = [json.loads(line) for line in file]
counts = overlap_counts(records, seen)
hit = [is_overlap(c, args.min_overlap) for c in counts]
kept = [r for r, h in zip(records, hit) if not h]
report[name] = dict(Counter(r["task"] for r, h in zip(records, hit) if h))
splits[name] = kept
with open(os.path.join(args.dest, f"{name}.jsonl"), "w") as file:
for record in kept:
file.write(json.dumps(record, ensure_ascii=False) + "\n")
print(f"{name}: kept {len(kept)} of {len(records)}; dropped {len(records) - len(kept)}", flush=True)
with open(os.path.join(args.out, "manifest.json")) as file:
manifest = json.load(file)
manifest["decontaminated"] = True
manifest["decontamination"] = {"against": args.against, "spans": int(len(seen[0])), "text_fields": int(len(seen[1])),
"min_overlap": args.min_overlap, "exact_field_min_words": 5,
"dropped": report}
manifest["totals"] = {name: {"records": len(rs), "questions": sum(len(r["questions"]) for r in rs)}
for name, rs in splits.items()}
manifest["question_types"] = dict(Counter(q["type"] for rs in splits.values() for r in rs
for q in r["questions"].values()))
with open(os.path.join(args.dest, "manifest.json"), "w") as file:
json.dump(manifest, file, indent=1)
print("dropped per split and task:", json.dumps(report), flush=True)
def review(args):
path = os.path.join(args.out, f"{args.split}.jsonl")
with open(path) as file:
records = [json.loads(line) for line in file]
matching = [r for r in records if r["task"] == args.task or (args.task == "tasksource" and r["task"].startswith("tasksource/"))]
if not matching:
raise SystemExit(f"No records of task {args.task} in {path}")
for record in random.Random(args.seed).sample(matching, min(args.n, len(matching))):
print(json.dumps(record, ensure_ascii=False, indent=2))
def stats(args):
with open(os.path.join(args.out, "manifest.json")) as file:
print_summary(json.load(file))
if __name__ == "__main__":
parser = argparse.ArgumentParser()
commands = parser.add_subparsers(dest="command", required=True)
b = commands.add_parser("build")
b.add_argument("--out", default="/workspace/data")
b.add_argument("--only", default=None, help="comma-separated task names to build (others are not kept)")
b.add_argument("--policy", choices=["commercial", "research"], default="commercial",
help="commercial: only licences approved in licenses.json and tasksource rows marked commercial")
b.add_argument("--decontam", default=None, help="comma-separated globs of suite files (.jsonl or .jsonl.gz)")
b.add_argument("--no-decontam", action="store_true")
b.add_argument("--tokenizer", default=None, help="check lengths with this tokenizer, e.g. Qwen/Qwen3.5-2B-Base")
b.add_argument("--max-length", type=int, default=8192)
b.add_argument("--tasksource-per-task", type=int, default=300)
b.add_argument("--eval-per-task", type=int, default=60)
b.add_argument("--seed", type=int, default=0)
r = commands.add_parser("review")
r.add_argument("--out", default="/workspace/data")
r.add_argument("--task", required=True)
r.add_argument("--split", default="train")
r.add_argument("--n", type=int, default=3)
r.add_argument("--seed", type=int, default=0)
d = commands.add_parser("decontam", help="filter an existing build against the benchmark suite")
d.add_argument("--out", default="/workspace/data", help="existing build")
d.add_argument("--dest", required=True, help="where to write the filtered build")
d.add_argument("--against", required=True, help="comma-separated globs of suite files (.jsonl or .jsonl.gz)")
d.add_argument("--min-overlap", type=int, default=2, help="shared 13-word spans that count as overlap")
d.add_argument("--workers", type=int, default=None, help="parallel workers (default: every CPU); lower it to save memory")
s = commands.add_parser("stats")
s.add_argument("--out", default="/workspace/data")
arguments = parser.parse_args()
{"build": build, "review": review, "stats": stats, "decontam": decontam}[arguments.command](arguments)
"""Training data for Decider 2B, written as Jev / SystemOne records (see schema.py).
python data.py build --out /workspace/data \
--decontam 'work/artifacts/benchmark-suite/normalized/*.jsonl' \
--tokenizer Qwen/Qwen3.5-2B-Base # exact token-length check
python data.py build --out ... --only tools,home_sim # rebuild some tasks
python data.py review --out /workspace/data --task tools_glaive --n 3
python data.py stats --out /workspace/data
Every record is checked by schema.validate_record, by the quality rules in quality(), and by the
licence policy in licenses.json, and is decontaminated against the benchmark suite. What was
kept and dropped, and why, goes to manifest.json.
"""
import argparse
import ast
import glob
import hashlib
import json
import os
import random
import re
import string
from collections import Counter, defaultdict
import schema
LETTERS = string.ascii_uppercase
MAX_STATE_CHARS = 60000
MAX_OPTION_CHARS = 1500
HERE = os.path.dirname(os.path.abspath(__file__))
# name, hub path, config, train split, eval split, role
# role "train": train split is trained on, eval split measures seen tasks.
# role "heldout": never trained on, measures transfer to unseen tasks.
TASKS = [
dict(name="banking77", path="legacy-datasets/banking77", train="train", eval="test",
text=lambda r: r["text"], label="label",
question="What is the intent of this bank customer's message?", max_opts=26),
dict(name="clinc150", path="clinc/clinc_oos", config="plus", train="train", eval="validation",
text=lambda r: r["text"], label="intent",
question="What is the intent of this request to an assistant?", max_opts=26),
dict(name="mnli", path="nyu-mll/glue", config="mnli", train="train", eval="validation_matched",
text=lambda r: f"Premise: {r['premise']}\nHypothesis: {r['hypothesis']}", label="label",
names=["entailment", "neutral", "contradiction"],
question="Does the premise entail the hypothesis, contradict it, or neither (neutral)?"),
dict(name="qnli", path="nyu-mll/glue", config="qnli", train="train", eval="validation",
text=lambda r: f"Question: {r['question']}\nSentence: {r['sentence']}", label="label",
names=["yes", "no"], question="Does the sentence contain the answer to the question?"),
dict(name="sst2", path="nyu-mll/glue", config="sst2", train="train", eval="validation",
text=lambda r: r["sentence"], label="label", names=["negative", "positive"],
question="What is the sentiment of this review?"),
dict(name="boolq", path="google/boolq", train="train", eval="validation",
text=lambda r: f"Passage: {r['passage']}\nQuestion: {r['question']}?",
label=lambda r: 0 if r["answer"] else 1, names=["yes", "no"],
question="According to the passage, what is the answer to the question?"),
dict(name="paws", path="google-research-datasets/paws", config="labeled_final", train="train",
eval="validation",
text=lambda r: f"Sentence 1: {r['sentence1']}\nSentence 2: {r['sentence2']}", label="label",
names=["no", "yes"], question="Do the two sentences have the same meaning?"),
dict(name="emotion", path="dair-ai/emotion", config="split", train="train", eval="validation",
text=lambda r: r["text"], label="label",
question="Which emotion does the writer express?"),
dict(name="dbpedia", path="fancyzhx/dbpedia_14", train="train", eval="test",
text=lambda r: f"{r['title']}\n{r['content']}", label="label",
question="What kind of thing is described?"),
dict(name="vitaminc", path="tals/vitaminc", train="train", eval="validation",
text=lambda r: f"Evidence: {r['evidence']}\nClaim: {r['claim']}", label="label",
question="Does the evidence support the claim, refute it, or not give enough information?"),
# Multiple-choice tasks: `mcq` returns (state text, option texts, index of the right one).
dict(name="mmlu_aux", path="cais/mmlu", config="auxiliary_train", train="train", rows=24000,
mcq=lambda r: (r["question"], r["choices"], int(r["answer"])),
question="Which option answers the question correctly?"),
dict(name="arc_challenge", path="allenai/ai2_arc", config="ARC-Challenge", train="train", rows=2000,
mcq=lambda r: keyed(r["question"], r["choices"], r["answerKey"]),
question="Which option answers the science question?"),
dict(name="arc_easy", path="allenai/ai2_arc", config="ARC-Easy", train="train", rows=2500,
mcq=lambda r: keyed(r["question"], r["choices"], r["answerKey"]),
question="Which option answers the science question?"),
dict(name="commonsense_qa", path="tau/commonsense_qa", train="train", rows=8000,
mcq=lambda r: keyed(r["question"], r["choices"], r["answerKey"]),
question="Which option is the most sensible answer?"),
dict(name="openbookqa", path="allenai/openbookqa", config="main", train="train", rows=4500,
mcq=lambda r: keyed(r["question_stem"], r["choices"], r["answerKey"]),
question="Which option completes or answers the statement?"),
dict(name="qasc", path="allenai/qasc", train="train", rows=7000,
mcq=lambda r: keyed(r["question"], r["choices"], r["answerKey"]),
question="Which option answers the question?"),
dict(name="sciq", path="allenai/sciq", train="train", rows=8000,
mcq=lambda r: (r["question"], [r["correct_answer"], r["distractor1"], r["distractor2"],
r["distractor3"]], 0),
question="Which option answers the science question?"),
dict(name="medmcqa", path="openlifescienceai/medmcqa", train="train", rows=8000,
mcq=lambda r: (r["question"], [r["opa"], r["opb"], r["opc"], r["opd"]], int(r["cop"])),
question="Which option answers the medical question?"),
dict(name="aqua_rat", path="deepmind/aqua_rat", config="raw", train="train", rows=8000,
mcq=lambda r: (r["question"], [o.split(")", 1)[-1].strip() for o in r["options"]],
"ABCDE".index(r["correct"])),
question="Which option is the answer to the math problem?"),