94 lines
2.2 KiB
Python
94 lines
2.2 KiB
Python
import ast
|
|
import json
|
|
import random
|
|
from pathlib import Path
|
|
|
|
random.seed(42)
|
|
|
|
with open("data_02/data_processed.json", encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
|
|
def extract_operations(expr):
|
|
tree = ast.parse(expr, mode="eval")
|
|
operations = []
|
|
|
|
names = {
|
|
ast.Add: "ADD",
|
|
ast.Sub: "SUB",
|
|
ast.Mult: "MUL",
|
|
ast.Div: "DIV"
|
|
}
|
|
|
|
def visit(node):
|
|
if isinstance(node, ast.BinOp):
|
|
visit(node.left)
|
|
visit(node.right)
|
|
if type(node.op) in names:
|
|
operations.append(names[type(node.op)])
|
|
|
|
visit(tree.body)
|
|
return operations
|
|
|
|
def convert(item):
|
|
validation = item["validation"]
|
|
steps = validation["process"]
|
|
|
|
operation_sequence = []
|
|
for step in steps:
|
|
operation_sequence.extend(
|
|
extract_operations(step["expr"])
|
|
)
|
|
|
|
prompt = (
|
|
"Generate one math word problem with the "
|
|
"following calculation structure.\n"
|
|
f"Calculation steps: {len(steps)}\n"
|
|
f"Operations: {' -> '.join(operation_sequence)}\n"
|
|
"Return valid JSON containing question, "
|
|
"explanation, answer, and validation. "
|
|
"The validation program must reference "
|
|
"previous calculation results correctly."
|
|
)
|
|
|
|
completion = json.dumps(
|
|
{
|
|
"question": item["question"],
|
|
"explanation": item["explanation"],
|
|
"answer": item["answer"],
|
|
"validation": validation,
|
|
},
|
|
ensure_ascii=False,
|
|
separators=(",", ":")
|
|
)
|
|
|
|
return {"prompt": prompt, "completion": completion}
|
|
|
|
examples = [convert(item) for item in data]
|
|
|
|
random.shuffle(examples)
|
|
|
|
n = len(examples)
|
|
n_train = int(n * 0.9)
|
|
n_valid = int(n * 0.05)
|
|
|
|
splits = {
|
|
"train": examples[:n_train],
|
|
"valid": examples[n_train:n_train + n_valid],
|
|
"test": examples[n_train + n_valid:],
|
|
}
|
|
|
|
Path("data_03").mkdir(exist_ok=True)
|
|
|
|
for name, rows in splits.items():
|
|
with open(
|
|
f"data_03/{name}.jsonl",
|
|
"w",
|
|
encoding="utf-8"
|
|
) as f:
|
|
for row in rows:
|
|
f.write(
|
|
json.dumps(row, ensure_ascii=False) + "\n"
|
|
)
|
|
|
|
print(name, len(rows))
|