| """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) |
|
|
