Allow using alternate annotator models for AIR-Bench 2024 (#3468)

Author: yifanmai

src/helm/benchmark/annotation/air_bench_annotator.py

@@ -1,7 +1,7 @@
 import datasets
 import os
 import re
-from typing import Any
+from typing import Any, Optional
 
 from helm.common.general import ensure_directory_exists
 from helm.benchmark.adaptation.request_state import RequestState
@@ -15,7 +15,16 @@ class AIRBench2024Annotator(Annotator):
 
     name = "air_bench_2024"
 
-    def __init__(self, auto_client: AutoClient, file_storage_path: str):
+    _DEFAULT_MODEL = "openai/gpt-4o-2024-05-13"
+    _DEFAULT_MODEL_DEPLOYMENT = "openai/gpt-4o-2024-05-13"
+
+    def __init__(
+        self,
+        auto_client: AutoClient,
+        file_storage_path: str,
+        model: Optional[str] = None,
+        model_deployment: Optional[str] = None,
+    ):
         self._auto_client = auto_client
         cache_dir = os.path.join(file_storage_path, "data")
         ensure_directory_exists(cache_dir)
@@ -29,6 +38,8 @@ def __init__(self, auto_client: AutoClient, file_storage_path: str):
         self._category_id_to_judge_prompt = {row["cate-idx"]: row["judge_prompt"] for row in dataset}
         # Regex pattern is lenient to allow for typos e.g. extra whitespace
         self._pattern = re.compile("##\s*short_reasoning\s*:(.*)##\s*the_score\s*:(.*)", re.DOTALL)
+        self._model = model or self._DEFAULT_MODEL
+        self._model_deployment = model_deployment or self._DEFAULT_MODEL_DEPLOYMENT
 
     def annotate(self, request_state: RequestState) -> Any:
         assert request_state.result
@@ -45,8 +56,8 @@ def annotate(self, request_state: RequestState) -> Any:
             "{{ANSWER}}", model_output_text
         )
         annotator_request = Request(
-            model="openai/gpt-4o-2024-05-13",
-            model_deployment="openai/gpt-4o-2024-05-13",
+            model=self._model,
+            model_deployment=self._model_deployment,
             prompt=annotator_prompt,
             temperature=0.0,
             max_tokens=64,

src/helm/benchmark/run_specs/air_bench_run_specs.py

@@ -1,3 +1,5 @@
+from typing import Dict, Optional
+
 from helm.benchmark.adaptation.adapter_spec import ADAPT_GENERATION, AdapterSpec
 from helm.benchmark.annotation.annotator import AnnotatorSpec
 from helm.benchmark.metrics.metric import MetricSpec
@@ -6,7 +8,10 @@
 
 
 @run_spec_function("air_bench_2024")
-def get_air_bench_2024_spec() -> RunSpec:
+def get_air_bench_2024_spec(
+    annotator_model: Optional[str] = None, annotator_model_deployment: Optional[str] = None
+) -> RunSpec:
+    run_spec_name = "air_bench_2024"
     adapter_spec = AdapterSpec(
         method=ADAPT_GENERATION,
         global_prefix="",
@@ -24,14 +29,23 @@ def get_air_bench_2024_spec() -> RunSpec:
         stop_sequences=[],
     )
     scenario_spec = ScenarioSpec(class_name="helm.benchmark.scenarios.air_bench_scenario.AIRBench2024Scenario")
+    annotator_args: Dict[str, str] = {}
+    if annotator_model:
+        annotator_args["model"] = annotator_model
+        annotator_args["model_deployment"] = annotator_model_deployment or annotator_model
+        run_spec_name = (
+            "air_bench_2024:"
+            f"annotator_model={annotator_args['model']},"
+            f"annotator_model_deployment={annotator_args['model_deployment']}"
+        )
     annotator_specs = [AnnotatorSpec(class_name="helm.benchmark.annotation.air_bench_annotator.AIRBench2024Annotator")]
     metric_specs = [
         MetricSpec(class_name="helm.benchmark.metrics.air_bench_metrics.AIRBench2024ScoreMetric"),
         MetricSpec(class_name="helm.benchmark.metrics.air_bench_metrics.AIRBench2024BasicGenerationMetric"),
         MetricSpec(class_name="helm.benchmark.metrics.basic_metrics.InstancesPerSplitMetric"),
     ]
     return RunSpec(
-        name="air_bench_2024",
+        name=run_spec_name,
         scenario_spec=scenario_spec,
         adapter_spec=adapter_spec,
         metric_specs=metric_specs,