ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter5/code-for-math/build_aime_2024.py

@@ -0,0 +1,92 @@
+#!/usr/bin/env python3
+"""Build a revision-pinned official AIME 2024 benchmark for Experiment 5-1."""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import tempfile
+import urllib.request
+from pathlib import Path
+
+import pyarrow.parquet as parquet
+
+
+DATASET = "HuggingFaceH4/aime_2024"
+REVISION = "2fe88a2f1091d5048c0f36abc874fb997b3dd99a"
+SOURCE_PATH = "data/train-00000-of-00001.parquet"
+
+
+def download() -> bytes:
+    url = f"https://huggingface.co/datasets/{DATASET}/resolve/{REVISION}/{SOURCE_PATH}?download=true"
+    request = urllib.request.Request(url, headers={"User-Agent": "ai-agent-book-exp5-1/1.0"})
+    with urllib.request.urlopen(request, timeout=60) as response:
+        return response.read()
+
+
+def convert(rows):
+    problems = []
+    seen_ids = set()
+    for row in rows:
+        answer = int(row["answer"])
+        if not 0 <= answer <= 999:
+            raise ValueError(f"AIME answer outside 000--999: {answer}")
+        source_id = int(row["id"])
+        if source_id in seen_ids:
+            raise ValueError(f"duplicate source id: {source_id}")
+        seen_ids.add(source_id)
+        problems.append({
+            "id": f"aime2024-{source_id}",
+            "question": row["problem"],
+            "answer": answer,
+            "topic": "official AIME 2024",
+            "source": {
+                "dataset": DATASET,
+                "revision": REVISION,
+                "source_id": source_id,
+                "year": row["year"],
+                "problem_url": row["url"],
+            },
+        })
+    if len(problems) != 30:
+        raise ValueError(f"expected 30 AIME 2024 problems, got {len(problems)}")
+    return sorted(problems, key=lambda item: item["id"])
+
+
+def build():
+    raw = download()
+    with tempfile.NamedTemporaryFile(suffix=".parquet") as handle:
+        handle.write(raw)
+        handle.flush()
+        rows = parquet.read_table(handle.name).to_pylist()
+    problems = convert(rows)
+    manifest = {
+        "schema_version": "1.0",
+        "experiment": "5-1",
+        "dataset": DATASET,
+        "revision": REVISION,
+        "source_path": SOURCE_PATH,
+        "source_sha256": hashlib.sha256(raw).hexdigest(),
+        "split": "train",
+        "problems": len(problems),
+        "selection": "all published AIME I and AIME II 2024 problems",
+        "answers": "published integer answer field; solutions are never sent to the model",
+    }
+    return problems, manifest
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--output", type=Path, default=Path("aime_2024.json"))
+    parser.add_argument("--manifest", type=Path, default=Path("aime_2024.manifest.json"))
+    args = parser.parse_args()
+    problems, manifest = build()
+    args.output.write_text(json.dumps(problems, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
+    args.manifest.write_text(json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
+    print(json.dumps({"problems": len(problems), "revision": REVISION,
+                      "output": str(args.output), "manifest": str(args.manifest)}))
+
+
+if __name__ == "__main__":
+    main()