(5/6) Adopt inference caching in MT-Bench

judgearena/cli.py

@@ -195,7 +195,6 @@ def _build_elo_args(
         n_instructions=args.n_instructions,
         provide_explanation=args.provide_explanation,
         swap_mode=args.swap_mode,
-        ignore_cache=args.ignore_cache,
         store_root=args.store_root,
         truncate_all_input_chars=args.truncate_all_input_chars,
         max_out_tokens_models=args.max_out_tokens_models,
@@ -224,7 +223,6 @@ def _build_generate_and_evaluate_args(
         n_instructions=args.n_instructions,
         provide_explanation=args.provide_explanation,
         swap_mode=args.swap_mode,
-        ignore_cache=args.ignore_cache,
         store_root=args.store_root,
         truncate_all_input_chars=args.truncate_all_input_chars,
         max_out_tokens_models=args.max_out_tokens_models,

judgearena/cli_common.py

@@ -21,7 +21,6 @@ class BaseCliArgs:
     n_instructions: int | None = None
     provide_explanation: bool = False
     swap_mode: str = "fixed"
-    ignore_cache: bool = False
     store_root: str | None = None
     truncate_all_input_chars: int = 8192
     max_out_tokens_models: int = 32768
@@ -81,11 +80,6 @@ def add_common_arguments(parser: argparse.ArgumentParser) -> None:
             "This helps account for judge position bias. Default is 'fixed'."
         ),
     )
-    parser.add_argument(
-        "--ignore_cache",
-        action="store_true",
-        help="If specified, ignore cache of previous completions.",
-    )
     parser.add_argument(
         "--store_root",
         type=str,

judgearena/generate.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from collections.abc import Callable
+
 import pandas as pd
 from langchain_core.prompts import ChatPromptTemplate
 
@@ -70,11 +72,12 @@ def _set_temperature_on_model(chat_model, temperature: float) -> None:
 def _infer_grouped_by_temperature(
     *,
     model_spec: str,
-    provider: str,
     max_tokens: int | None,
     model_kwargs: dict,
-    base_model,
+    factory_for_temperature: Callable[[float], object] | None,
+    inference_cache: InferenceCache | None,
     inputs: list,
+    cache_metadata: list[dict],
     temperatures: list[float],
     use_tqdm: bool,
 ) -> list[str]:
@@ -87,18 +90,29 @@ def _infer_grouped_by_temperature(
         idxs = groups[temp]
         group_inputs = [inputs[i] for i in idxs]
 
-        if provider in {"VLLM", "LlamaCpp"}:
-            _set_temperature_on_model(base_model, temp)
-            group_model = base_model
+        if factory_for_temperature is not None:
+            group_model = prepare_model(
+                model_spec,
+                max_tokens=max_tokens,
+                cache=inference_cache,
+                factory=lambda temp=temp: factory_for_temperature(temp),
+                temperature=temp,
+                **model_kwargs,
+            )
         else:
-            group_model = make_model(
-                model_spec, max_tokens=max_tokens, temperature=temp, **model_kwargs
+            group_model = prepare_model(
+                model_spec,
+                max_tokens=max_tokens,
+                cache=inference_cache,
+                temperature=temp,
+                **model_kwargs,
             )
 
         group_outs = do_inference(
             chat_model=group_model,
             inputs=group_inputs,
             use_tqdm=use_tqdm,
+            cache_metadata=[cache_metadata[i] for i in idxs],
         )
         for i, out in zip(idxs, group_outs, strict=True):
             outputs[i] = out
@@ -113,22 +127,41 @@ def generate_multiturn(
     max_tokens: int | None = 8192,
     use_tqdm: bool = True,
     temperature_config: dict[str, float] | None = None,
+    inference_cache: InferenceCache | None = None,
     **model_kwargs,
 ) -> pd.DataFrame:
     """Generate two-turn completions for MT-Bench style questions."""
     provider = model.split("/")[0]
     use_category_temperatures = temperature_config is not None
     local_provider = provider in {"VLLM", "LlamaCpp"}
 
-    if use_category_temperatures and local_provider:
-        chat_model = make_model(
-            model, max_tokens=max_tokens, temperature=0.0, **model_kwargs
+    chat_model = None
+    materialized_model = None
+
+    def factory_for_temperature(temperature: float):
+        nonlocal materialized_model
+        if materialized_model is None:
+            materialized_model = make_model(
+                model,
+                max_tokens=max_tokens,
+                temperature=temperature,
+                **model_kwargs,
+            )
+        else:
+            _set_temperature_on_model(materialized_model, temperature)
+        return materialized_model
+
+    if not use_category_temperatures:
+        chat_model = prepare_model(
+            model,
+            max_tokens=max_tokens,
+            cache=inference_cache,
+            **model_kwargs,
         )
-    else:
-        chat_model = make_model(model, max_tokens=max_tokens, **model_kwargs)
 
     system_prompt = "You are a helpful assistant."
     idxs = questions.index.tolist()
+    cache_metadata = [{"instruction_id": index} for index in idxs]
     temperatures: list[float] = []
     if use_category_temperatures:
         temperatures = [
@@ -149,11 +182,14 @@ def generate_multiturn(
     if use_category_temperatures:
         completions_turn_1 = _infer_grouped_by_temperature(
             model_spec=model,
-            provider=provider,
             max_tokens=max_tokens,
             model_kwargs=model_kwargs,
-            base_model=chat_model,
+            factory_for_temperature=(
+                factory_for_temperature if local_provider else None
+            ),
+            inference_cache=inference_cache,
             inputs=turn1_inputs,
+            cache_metadata=cache_metadata,
             temperatures=temperatures,
             use_tqdm=use_tqdm,
         )
@@ -162,6 +198,7 @@ def generate_multiturn(
             chat_model=chat_model,
             inputs=turn1_inputs,
             use_tqdm=use_tqdm,
+            cache_metadata=cache_metadata,
         )
 
     turn2_inputs = []
@@ -196,11 +233,14 @@ def generate_multiturn(
     if use_category_temperatures:
         completions_turn_2 = _infer_grouped_by_temperature(
             model_spec=model,
-            provider=provider,
             max_tokens=max_tokens,
             model_kwargs=model_kwargs,
-            base_model=chat_model,
+            factory_for_temperature=(
+                factory_for_temperature if local_provider else None
+            ),
+            inference_cache=inference_cache,
             inputs=turn2_inputs,
+            cache_metadata=cache_metadata,
             temperatures=temperatures,
             use_tqdm=use_tqdm,
         )
@@ -209,6 +249,7 @@ def generate_multiturn(
             chat_model=chat_model,
             inputs=turn2_inputs,
             use_tqdm=use_tqdm,
+            cache_metadata=cache_metadata,
         )
 
     return pd.DataFrame(

judgearena/generate_and_evaluate.py

@@ -149,7 +149,6 @@ def main(args: CliArgs):
     if args.task == "mt-bench":
         return run_mt_bench(
             args,
-            args.ignore_cache,
             res_folder=res_folder,
             result_name=name,
         )

judgearena/inference.py

@@ -122,7 +122,12 @@ def build_model_descriptor(
         return None
 
     descriptor_kwargs = _without_credentials(resolved_kwargs)
+    sampling = None
     if provider == "VLLM":
+        sampling = {
+            "temperature": descriptor_kwargs.pop("temperature"),
+            "top_p": descriptor_kwargs.pop("top_p"),
+        }
         descriptor_kwargs = {
             key: value
             for key, value in descriptor_kwargs.items()
@@ -138,10 +143,7 @@ def build_model_descriptor(
     }
     if provider == "VLLM":
         descriptor["backend_version"] = importlib_metadata.version("vllm")
-        descriptor["sampling"] = {
-            "temperature": VLLM_TEMPERATURE,
-            "top_p": VLLM_TOP_P,
-        }
+        descriptor["sampling"] = sampling
     elif provider == "LlamaCpp":
         descriptor["backend_version"] = importlib_metadata.version("llama-cpp-python")
     if endpoint is not None:

judgearena/mt_bench/fastchat_compat.py

@@ -267,6 +267,8 @@ def _infer_by_prompt_groups(
     items: list[dict[str, Any]],
     use_tqdm: bool,
     swap_answers: bool,
+    model_a: str,
+    model_b: str,
 ) -> list[str]:
     """Run judge inference, grouping by prompt variant for batching."""
     grouped_indices = _group_indices_by_prompt(items)
@@ -290,6 +292,15 @@ def _infer_by_prompt_groups(
             chat_model=judge_chat_model,
             inputs=prompt_inputs,
             use_tqdm=use_tqdm,
+            cache_metadata=[
+                {
+                    "instruction_id": items[i]["question_id"],
+                    "model_a": model_b if swap_answers else model_a,
+                    "model_b": model_a if swap_answers else model_b,
+                    "orientation": "reversed" if swap_answers else "direct",
+                }
+                for i in idxs
+            ],
         )
         for i, out in zip(idxs, outs, strict=True):
             judgments[i] = str(out)
@@ -456,6 +467,8 @@ def judge_mt_bench_pairwise_fastchat(
         items=items,
         use_tqdm=use_tqdm,
         swap_answers=False,
+        model_a=model_a,
+        model_b=model_b,
     )
 
     g2_judgments: list[str] | None = None
@@ -465,6 +478,8 @@ def judge_mt_bench_pairwise_fastchat(
             items=items,
             use_tqdm=use_tqdm,
             swap_answers=True,
+            model_a=model_a,
+            model_b=model_b,
         )
 
     annotations: list[dict[str, Any]] = []

judgearena/mt_bench/mt_bench_utils.py

@@ -17,14 +17,15 @@
 
 from judgearena.eval_utils import _compute_grouped_stats, print_results
 from judgearena.generate import generate_multiturn
+from judgearena.inference import CompletionInferenceCache, JudgementInferenceCache
 from judgearena.instruction_dataset import load_instructions
 from judgearena.log import get_logger
 from judgearena.mt_bench.fastchat_compat import (
     FASTCHAT_TEMPERATURE_CONFIG,
     judge_mt_bench_pairwise_fastchat,
 )
 from judgearena.repro import _to_jsonable
-from judgearena.utils import cache_function_dataframe, compute_pref_summary, make_model
+from judgearena.utils import compute_pref_summary, prepare_model
 
 logger = get_logger(__name__)
 
@@ -35,9 +36,12 @@
 def _generate_mt_bench_completions(
     args: CliArgs,
     questions_df: pd.DataFrame,
-    ignore_cache: bool,
 ) -> tuple[pd.DataFrame, pd.DataFrame]:
-    cache_prefix = "mt-bench"
+    inference_cache = (
+        CompletionInferenceCache(Path(args.store_root), "mt-bench")
+        if args.store_root is not None
+        else None
+    )
 
     def _run_generation(model_name: str) -> pd.DataFrame:
         return generate_multiturn(
@@ -49,19 +53,12 @@ def _run_generation(model_name: str) -> pd.DataFrame:
             max_model_len=args.max_model_len,
             chat_template=args.chat_template,
             temperature_config=FASTCHAT_TEMPERATURE_CONFIG,
+            inference_cache=inference_cache,
+            **args.engine_kwargs,
         )
 
-    completions_a = cache_function_dataframe(
-        lambda: _run_generation(args.model_A),
-        ignore_cache=ignore_cache,
-        cache_name=f"{cache_prefix}_{args.model_A}_{args.n_instructions}",
-    ).set_index("instruction_index")
-
-    completions_b = cache_function_dataframe(
-        lambda: _run_generation(args.model_B),
-        ignore_cache=ignore_cache,
-        cache_name=f"{cache_prefix}_{args.model_B}_{args.n_instructions}",
-    ).set_index("instruction_index")
+    completions_a = _run_generation(args.model_A).set_index("instruction_index")
+    completions_b = _run_generation(args.model_B).set_index("instruction_index")
     return completions_a, completions_b
 
 
@@ -143,7 +140,6 @@ def _run_mt_bench_fastchat(
 
 def run_mt_bench(
     args: CliArgs,
-    ignore_cache: bool,
     *,
     res_folder: Path,
     result_name: str,
@@ -158,14 +154,20 @@ def run_mt_bench(
     completions_a, completions_b = _generate_mt_bench_completions(
         args=args,
         questions_df=questions_df,
-        ignore_cache=ignore_cache,
     )
-    judge_chat_model = make_model(
+    judgement_cache = (
+        JudgementInferenceCache(Path(args.store_root), "mt-bench")
+        if args.store_root is not None
+        else None
+    )
+    judge_chat_model = prepare_model(
         model=args.judge_model,
         max_tokens=args.max_out_tokens_judge,
+        cache=judgement_cache,
         temperature=0.0,
         max_model_len=args.max_model_len,
         chat_template=args.chat_template,
+        **args.engine_kwargs,
     )
     return _run_mt_bench_fastchat(
         args=args,