File size: 2,886 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 84 85 86 | """More held-out rows no training source contains: HellaSwag (validation) and dair-ai/emotion (test).
python heldout2.py /content/out/heldout2_rows.jsonl 400
"""
import json
import random
import sys
from datasets import load_dataset
EMOTIONS = {
"sadness": "The writer feels sad, hurt or lonely.",
"joy": "The writer feels happy, pleased or content.",
"love": "The writer feels love, affection or tenderness.",
"anger": "The writer feels angry, annoyed or resentful.",
"fear": "The writer feels afraid, anxious or nervous.",
"surprise": "The writer feels surprised or amazed.",
}
def one_hot(keys, label):
return {"label": label, "probabilities": {k: 1.0 if k == label else 0.0 for k in keys}}
def hellaswag(n, rng):
rows = []
for i, row in enumerate(load_dataset("Rowan/hellaswag", split = "validation")):
if not str(row.get("label", "")).strip().isdigit():
continue
keys = ["A", "B", "C", "D"]
label = keys[int(row["label"])]
rows.append(
{
"id": f"hellaswag_val_{i}",
"workflow": "hellaswag",
"state": {"activity": row["activity_label"], "context": row["ctx"]},
"questions": {
"ending": {
"type": "choice",
"instructions": "Which ending most plausibly continues the context?",
"criteria": dict(zip(keys, row["endings"])),
}
},
"gold": {"ending": one_hot(keys, label)},
}
)
rng.shuffle(rows)
return rows[:n]
def emotion(n, rng):
dataset = load_dataset("dair-ai/emotion", "split", split = "test")
names = dataset.features["label"].names
rows = []
for i, row in enumerate(dataset):
label = names[int(row["label"])]
rows.append(
{
"id": f"emotion_test_{i}",
"workflow": "emotion",
"state": row["text"],
"questions": {
"emotion": {
"type": "choice",
"instructions": "Which emotion does the writer express?",
"criteria": {name: EMOTIONS[name] for name in names},
}
},
"gold": {"emotion": one_hot(names, label)},
}
)
rng.shuffle(rows)
return rows[:n]
if __name__ == "__main__":
out = sys.argv[1] if len(sys.argv) > 1 else "/content/out/heldout2_rows.jsonl"
n = int(sys.argv[2]) if len(sys.argv) > 2 else 400
rng = random.Random(1)
rows = hellaswag(n, rng) + emotion(n, rng)
with open(out, "w", encoding = "utf-8") as f:
for row in rows:
f.write(json.dumps(row) + "\n")
print("HELDOUT2_DONE", out, len(rows), flush = True)
|