File size: 3,313 Bytes
b88c26d | 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 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 | #!/usr/bin/env python3
"""
Generate Arithmetic Expression Dataset
Generates corpus of arithmetic expressions with distractors for MMEP training.
Output format: JSONL with query, answer, and distractor values.
Usage:
python scripts/generate_arithmetic.py --corpus-size 10000 --output data/arithmetic.jsonl
Contact: jessica@collectivekitty.com
"""
import argparse
import json
import random
import os
from pathlib import Path
def generate_expression(rng, max_val=100, max_ops=3):
"""Generate a random arithmetic expression and its result."""
ops = ['+', '-', '*']
num_ops = rng.randint(1, max_ops)
values = [rng.randint(1, max_val) for _ in range(num_ops + 1)]
operators = [rng.choice(ops) for _ in range(num_ops)]
# Build expression string
expr_parts = [str(values[0])]
for i, op in enumerate(operators):
expr_parts.append(f" {op} {values[i+1]}")
expr = "".join(expr_parts)
# Compute answer
try:
answer = eval(expr)
except Exception:
answer = values[0]
return expr, int(answer)
def generate_distractors(answer, rng, num_distractors=3):
"""Generate plausible wrong answers."""
distractors = set()
attempts = 0
while len(distractors) < num_distractors and attempts < 100:
offset = rng.choice([-10, -5, -2, -1, 1, 2, 5, 10, 20])
d = answer + offset
if d != answer:
distractors.add(d)
attempts += 1
return list(distractors)[:num_distractors]
def generate_dataset(size, rng, split_name="train"):
"""Generate a dataset of arithmetic expressions."""
data = []
for _ in range(size):
expr, answer = generate_expression(rng)
distractors = generate_distractors(answer, rng)
data.append({
"query": expr,
"answer": answer,
"distractors": distractors,
"split": split_name
})
return data
def main():
parser = argparse.ArgumentParser(description="Generate arithmetic dataset")
parser.add_argument("--corpus-size", type=int, default=10000)
parser.add_argument("--train-queries", type=int, default=5000)
parser.add_argument("--val-queries", type=int, default=500)
parser.add_argument("--test-queries", type=int, default=1000)
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--output", type=str, default="data/arithmetic.jsonl")
args = parser.parse_args()
rng = random.Random(args.seed)
output_path = Path(args.output)
output_path.parent.mkdir(parents=True, exist_ok=True)
print(f"Generating arithmetic dataset (seed={args.seed})...")
train = generate_dataset(args.train_queries, rng, "train")
val = generate_dataset(args.val_queries, rng, "val")
test = generate_dataset(args.test_queries, rng, "test")
all_data = train + val + test
with open(output_path, 'w') as f:
for item in all_data:
f.write(json.dumps(item) + '\n')
print(f"Generated {len(all_data)} examples:")
print(f" Train: {len(train)}")
print(f" Val: {len(val)}")
print(f" Test: {len(test)}")
print(f" Output: {output_path}")
if __name__ == "__main__":
main()
|