add aime benchmark

Author: benjibc

eval_protocol/benchmarks/suites/aime25.py

@@ -70,7 +70,7 @@ def aime2025_dataset_adapter(rows: List[Dict[str, Any]]) -> List[EvaluationRow]:
         "https://huggingface.co/datasets/opencompass/AIME2025/raw/main/aime2025-II.jsonl",
     ],
     dataset_adapter=aime2025_dataset_adapter,
-    rollout_input_params=[{"extra_body": {"reasoning_effort": "low"}}],
+    rollout_input_params=[{"max_tokens": 131000, "extra_body": {"reasoning_effort": "high"}}],
     rollout_processor=default_single_turn_rollout_processor,
     aggregation_method="mean",
     threshold_of_success=None,
@@ -79,7 +79,7 @@ def aime2025_dataset_adapter(rows: List[Dict[str, Any]]) -> List[EvaluationRow]:
     max_concurrent_rollouts=4,
     mode="pointwise",
 )
-def aime25_pointwise(row: EvaluationRow) -> EvaluationRow:
+def test_aime25_pointwise(row: EvaluationRow) -> EvaluationRow:
     assistant_msgs = [m for m in row.messages if m.role == "assistant"]
     content = assistant_msgs[-1].content if assistant_msgs else ""
 

eval_protocol/pytest/default_single_turn_rollout_process.py

@@ -38,11 +38,22 @@ async def process_row(row: EvaluationRow) -> EvaluationRow:
         request_params = {"model": config.model, "messages": messages_payload, **config.input_params}
         # Ensure caching is disabled only for this request (review feedback)
         request_params["cache"] = {"no-cache": True}
-        # Allow passing reasoning effort to Fireworks via LiteLLM using extra_body
-        # Expected: config.input_params may contain {"reasoning": {"effort": "low|medium|high"}}
-        if "reasoning" in config.input_params:
+        # Single-level reasoning effort: expect `reasoning_effort` only
+        effort_val = None
+        if isinstance(config.input_params, dict):
+            if "reasoning_effort" in config.input_params:
+                effort_val = str(config.input_params["reasoning_effort"])  # flat shape
+            elif isinstance(config.input_params.get("extra_body"), dict) and "reasoning_effort" in config.input_params["extra_body"]:
+                # Accept if user passed it directly inside extra_body
+                effort_val = str(config.input_params["extra_body"]["reasoning_effort"])  # already in extra_body
+
+        if effort_val:
+            # Always under extra_body so LiteLLM forwards to provider-specific param set
             request_params.setdefault("extra_body", {})
-            request_params["extra_body"]["reasoning"] = config.input_params["reasoning"]
+            request_params["extra_body"]["reasoning_effort"] = effort_val
+            # Ensure unsupported top-level keys are not present
+            if "reasoning_effort" in request_params:
+                request_params.pop("reasoning_effort", None)
 
         if row.tools is not None:
             request_params["tools"] = row.tools
@@ -88,7 +99,10 @@ async def process_row(row: EvaluationRow) -> EvaluationRow:
 
     async def _sem_wrapper(r: EvaluationRow) -> EvaluationRow:
         async with semaphore:
-            return await process_row(r)
+            try:
+                return await process_row(r)
+            except Exception as e:
+                return r
 
     tasks = [_sem_wrapper(row) for row in rows]
     dataset = list(await asyncio.gather(*tasks))

eval_protocol/pytest/evaluation_test.py

@@ -280,7 +280,9 @@ def wrapper_body(**kwargs):
                     input_params = kwargs.get("input_params") or {}
                     # Optional global overrides via environment for ad-hoc experimentation
                     # EP_INPUT_PARAMS_JSON can contain a JSON object that will be deep-merged
-                    # into input_params (e.g., '{"temperature":0,"extra_body":{"reasoning":{"effort":"low"}}}').
+                    # into input_params (e.g., '{"temperature":0,"reasoning_effort":"low"}' or
+                    # '{"extra_body":{"reasoning_effort":"low"}}'). Legacy nested
+                    # '{"reasoning":{"effort":"low"}}' is also supported by the rollout.
                     try:
                         import json as _json
                         _env_override = os.getenv("EP_INPUT_PARAMS_JSON")
@@ -480,15 +482,11 @@ def _extract_effort_tag(params: dict) -> str | None:
                                 try:
                                     if not isinstance(params, dict):
                                         return None
-                                    # Common locations
-                                    if "extra_body" in params and isinstance(params["extra_body"], dict):
-                                        eb = params["extra_body"]
-                                        if isinstance(eb.get("reasoning"), dict) and "effort" in eb["reasoning"]:
-                                            return str(eb["reasoning"]["effort"]).lower()
-                                        if "reasoning_effort" in eb:
-                                            return str(eb["reasoning_effort"]).lower()
-                                    if "reasoning" in params and isinstance(params["reasoning"], dict) and "effort" in params["reasoning"]:
-                                        return str(params["reasoning"]["effort"]).lower()
+                                    # Preferred locations: top-level or inside extra_body
+                                    if "reasoning_effort" in params:
+                                        return str(params["reasoning_effort"]).lower()
+                                    if isinstance(params.get("extra_body"), dict) and "reasoning_effort" in params["extra_body"]:
+                                        return str(params["extra_body"]["reasoning_effort"]).lower()
                                 except Exception:
                                     return None
                                 return None
@@ -925,14 +923,10 @@ def _extract_effort_tag(params: dict) -> str | None:
                     try:
                         if not isinstance(params, dict):
                             return None
-                        if "extra_body" in params and isinstance(params["extra_body"], dict):
-                            eb = params["extra_body"]
-                            if isinstance(eb.get("reasoning"), dict) and "effort" in eb["reasoning"]:
-                                return str(eb["reasoning"]["effort"]).lower()
-                            if "reasoning_effort" in eb:
-                                return str(eb["reasoning_effort"]).lower()
-                        if "reasoning" in params and isinstance(params["reasoning"], dict) and "effort" in params["reasoning"]:
-                            return str(params["reasoning"]["effort"]).lower()
+                        if "reasoning_effort" in params:
+                            return str(params["reasoning_effort"]).lower()
+                        if isinstance(params.get("extra_body"), dict) and "reasoning_effort" in params["extra_body"]:
+                            return str(params["extra_body"]["reasoning_effort"]).lower()
                     except Exception:
                         return None
                     return None

eval_protocol/pytest/plugin.py

@@ -131,10 +131,9 @@ def pytest_configure(config) -> None:
                         merged[k] = v
         reasoning_effort = config.getoption("--ep-reasoning-effort")
         if reasoning_effort:
-            # Standardize into extra_body.reasoning.effort in EP_INPUT_PARAMS_JSON
+            # Always place under extra_body to avoid LiteLLM rejecting top-level params
             eb = merged.setdefault("extra_body", {})
-            reasoning = eb.setdefault("reasoning", {})
-            reasoning["effort"] = str(reasoning_effort)
+            eb["reasoning_effort"] = str(reasoning_effort)
         if merged:
             os.environ["EP_INPUT_PARAMS_JSON"] = _json.dumps(merged)
     except Exception:

examples/aime2025_chat_completion/tests/test_evaluation.py

@@ -62,7 +62,7 @@ def aime2025_dataset_adapter(rows: List[Dict[str, Any]]) -> List[EvaluationRow]:
         "https://huggingface.co/datasets/opencompass/AIME2025/raw/main/aime2025-II.jsonl",
     ],
     dataset_adapter=aime2025_dataset_adapter,
-    rollout_input_params=[{"extra_body": {"reasoning_effort": "low"}}],
+    rollout_input_params=[{"reasoning_effort": "low"}],
     rollout_processor=default_single_turn_rollout_processor,
     aggregation_method="mean",
     threshold_of_success=None,

examples/gpqa/tests/test_evaluation.py

@@ -65,7 +65,7 @@ def _load_gpqa_messages_from_csv() -> List[List[Message]]:
 @evaluation_test(
     model=["fireworks_ai/accounts/fireworks/models/gpt-oss-120b"],
     input_messages=_GPQA_INPUT_MESSAGES,
-    rollout_input_params=[{"extra_body": {"reasoning_effort": "low"}}], # default to low effort; override via CLI plugin
+    rollout_input_params=[{"reasoning_effort": "low"}], # default to low effort; override via CLI plugin
     rollout_processor=default_single_turn_rollout_processor,
     aggregation_method="mean",
     threshold_of_success=None,