File size: 2,985 Bytes
a86df9f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 | """Held-out evaluation rows that neither training stage sees: ANLI (test r1-r3) and WinoGrande (dev).
Written once to /content/out/heldout_rows.jsonl in the typed-decisions layout
(id, workflow, state, questions, gold), so our model and Jev answer identical requests.
"""
import json
import random
import sys
from datasets import load_dataset
NLI = ("entailment", "neutral", "contradiction")
NLI_CRITERIA = {
"entailment": "The premise shows the hypothesis is true.",
"neutral": "The premise neither confirms nor rules out the hypothesis.",
"contradiction": "The premise shows the hypothesis is false.",
}
def one_hot(keys, label):
return {"label": label, "probabilities": {k: 1.0 if k == label else 0.0 for k in keys}}
def anli(n, rng):
rows = []
for split in ("test_r1", "test_r2", "test_r3"):
for i, row in enumerate(load_dataset("facebook/anli", split = split)):
if int(row["label"]) < 0:
continue
label = NLI[int(row["label"])]
rows.append(
{
"id": f"anli_{split}_{i}",
"workflow": "anli",
"state": {"premise": row["premise"], "hypothesis": row["hypothesis"]},
"questions": {
"relation": {
"type": "choice",
"instructions": "How does the premise relate to the hypothesis?",
"criteria": dict(NLI_CRITERIA),
}
},
"gold": {"relation": one_hot(NLI, label)},
}
)
rng.shuffle(rows)
return rows[:n]
def winogrande(n, rng):
rows = []
for i, row in enumerate(load_dataset("allenai/winogrande", "winogrande_xl", split = "validation")):
label = "option_1" if str(row["answer"]) == "1" else "option_2"
rows.append(
{
"id": f"winogrande_dev_{i}",
"workflow": "winogrande",
"state": {"sentence": row["sentence"]},
"questions": {
"blank": {
"type": "choice",
"instructions": "Which option correctly fills the blank (_) in the sentence?",
"criteria": {"option_1": row["option1"], "option_2": row["option2"]},
}
},
"gold": {"blank": one_hot(("option_1", "option_2"), label)},
}
)
rng.shuffle(rows)
return rows[:n]
if __name__ == "__main__":
out = sys.argv[1] if len(sys.argv) > 1 else "/content/out/heldout_rows.jsonl"
n = int(sys.argv[2]) if len(sys.argv) > 2 else 400
rng = random.Random(0)
rows = anli(n, rng) + winogrande(n, rng)
with open(out, "w", encoding = "utf-8") as f:
for row in rows:
f.write(json.dumps(row) + "\n")
print("HELDOUT_DONE", out, len(rows), flush = True)
|