From 7a37218715135fee8bbf9efcec02bb146ef8dfac Mon Sep 17 00:00:00 2001
From: Hynek Kydlicek <kydlicek.hynek@huggingface.co>
Date: Sun, 26 Jan 2025 13:03:52 +0000
Subject: [PATCH] minor fixes

---
 setup.py                | 2 +-
 slurm/evaluate.slurm    | 1 +
 src/open_r1/evaluate.py | 4 ++--
 3 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/setup.py b/setup.py
index 13c77ca6..21ba5b4a 100644
--- a/setup.py
+++ b/setup.py
@@ -42,7 +42,7 @@
 #   * If a dependency is fast-moving (e.g. transformers), pin to the exact version
 _deps = [
     "accelerate>=1.2.1",
-    "bitsandbytes",
+    "bitsandbytes>=0.43.0",
     "black>=24.4.2",
     "datasets>=3.2.0",
     "deepspeed==0.15.4",
diff --git a/slurm/evaluate.slurm b/slurm/evaluate.slurm
index 421a96cb..315cc80a 100644
--- a/slurm/evaluate.slurm
+++ b/slurm/evaluate.slurm
@@ -43,6 +43,7 @@ lighteval vllm $MODEL_ARGS "custom|$TASK|0|0" \
     --custom-tasks src/open_r1/evaluate.py \
     --use-chat-template \
     --system-prompt="Please reason step by step, and put your final answer within \boxed{}." \
+    --save-details
     --output-dir $OUTPUT_DIR 
 
 
diff --git a/src/open_r1/evaluate.py b/src/open_r1/evaluate.py
index 254a20e3..a4478d4d 100644
--- a/src/open_r1/evaluate.py
+++ b/src/open_r1/evaluate.py
@@ -55,8 +55,8 @@ def prompt_fn(line, task_name: str = None):
 def aime_prompt_fn(line, task_name: str = None):
     return Doc(
         task_name=task_name,
-        query=line["answer"],
-        choices=[line["solution"]],
+        query=line["problem"],
+        choices=[line["answer"]],
         gold_index=0,
     )