diff --git a/judgearena/cli.py b/judgearena/cli.py index 5006df5..bf45dfe 100644 --- a/judgearena/cli.py +++ b/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, diff --git a/judgearena/cli_common.py b/judgearena/cli_common.py index f22c617..e4e8b65 100644 --- a/judgearena/cli_common.py +++ b/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, diff --git a/judgearena/generate.py b/judgearena/generate.py index 06b381d..2c3c233 100644 --- a/judgearena/generate.py +++ b/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,6 +127,7 @@ 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.""" @@ -120,15 +135,33 @@ def generate_multiturn( 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( diff --git a/judgearena/generate_and_evaluate.py b/judgearena/generate_and_evaluate.py index cba8369..7ed2f6f 100644 --- a/judgearena/generate_and_evaluate.py +++ b/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, ) diff --git a/judgearena/inference.py b/judgearena/inference.py index 358f5fe..4aa4f5a 100644 --- a/judgearena/inference.py +++ b/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: diff --git a/judgearena/mt_bench/fastchat_compat.py b/judgearena/mt_bench/fastchat_compat.py index 3b0e7ec..62868e9 100644 --- a/judgearena/mt_bench/fastchat_compat.py +++ b/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]] = [] diff --git a/judgearena/mt_bench/mt_bench_utils.py b/judgearena/mt_bench/mt_bench_utils.py index b28f859..117246f 100644 --- a/judgearena/mt_bench/mt_bench_utils.py +++ b/judgearena/mt_bench/mt_bench_utils.py @@ -17,6 +17,7 @@ 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 ( @@ -24,7 +25,7 @@ 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, diff --git a/judgearena/utils.py b/judgearena/utils.py index 0e5e247..683c02b 100644 --- a/judgearena/utils.py +++ b/judgearena/utils.py @@ -7,9 +7,7 @@ import pandas as pd from huggingface_hub import snapshot_download -from langchain_community.cache import SQLiteCache from langchain_community.llms import LlamaCpp -from langchain_core.globals import set_llm_cache from langchain_openai import ChatOpenAI from tqdm.asyncio import tqdm from tqdm.contrib.logging import logging_redirect_tqdm @@ -42,10 +40,6 @@ def _data_root_path() -> Path: data_root = _data_root_path() -def set_langchain_cache(): - set_llm_cache(SQLiteCache(database_path=str(data_root / ".langchain.db"))) - - def download_hf(name: str, local_path: Path): local_path.mkdir(exist_ok=True, parents=True) # downloads the model from huggingface into `local_path` folder @@ -305,6 +299,8 @@ def __init__( model: str, max_tokens: int = 8192, chat_template: str | None = None, + temperature: float = VLLM_TEMPERATURE, + top_p: float = VLLM_TOP_P, **vllm_kwargs, ): from vllm import LLM, SamplingParams @@ -340,8 +336,8 @@ def __init__( self.llm = LLM(model=model, trust_remote_code=True, **vllm_kwargs) self.sampling_params = SamplingParams( max_tokens=max_tokens, - temperature=VLLM_TEMPERATURE, - top_p=VLLM_TOP_P, + temperature=temperature, + top_p=top_p, ) # Resolve chat template: @@ -443,6 +439,9 @@ def invoke(self, input_item, **invoke_kwargs) -> str: results = self.batch([input_item], **invoke_kwargs) return results[0] + def set_temperature(self, temperature: float) -> None: + self.sampling_params.temperature = temperature + async def ainvoke(self, input_item, **invoke_kwargs): """Async version - runs sync version in executor for compatibility.""" import asyncio @@ -494,6 +493,8 @@ def _resolve_model_config( key: value for key, value in resolved_kwargs.items() if value is not None } resolved_kwargs["chat_template"] = resolved_kwargs.get("chat_template") + resolved_kwargs.setdefault("temperature", VLLM_TEMPERATURE) + resolved_kwargs.setdefault("top_p", VLLM_TOP_P) elif model_provider == "LlamaCpp": resolved_kwargs["model_path"] = model_name elif model_provider not in {"Dummy", "OpenRouter"}: @@ -506,6 +507,7 @@ def prepare_model( max_tokens: int | None = 8192, *, cache: InferenceCache | None = None, + factory: Callable[[], object] | None = None, **engine_kwargs, ) -> PreparedModel: """Prepare cache identity and a lazy factory without loading the backend.""" @@ -538,10 +540,14 @@ def prepare_model( return PreparedModel( model_spec=model, descriptor=descriptor, - factory=lambda: make_model( - model, - max_tokens=max_tokens, - **factory_kwargs, + factory=( + factory + if factory is not None + else lambda: make_model( + model, + max_tokens=max_tokens, + **factory_kwargs, + ) ), cache=cache, ) @@ -632,92 +638,5 @@ def download_all(): download_mt_bench() -class Timeblock: - """Timer context manager""" - - def __init__(self, name: str | None = None, verbose: bool = True): - self.name = name - self.verbose = verbose - - def __enter__(self): - """Start a new timer as a context manager""" - self.start = time.time() - return self - - def __exit__(self, *args): - """Stop the context manager timer""" - self.end = time.time() - self.duration = self.end - self.start - if self.verbose: - logger.info("%s", self) - - def __str__(self): - name = self.name if self.name else "block" - msg = f"{name} took {self.duration} seconds" - return msg - - -def cache_function_dataframe( - fun: Callable[[], pd.DataFrame], - cache_name: str, - ignore_cache: bool = False, - cache_path: Path | None = None, - parquet: bool = False, -) -> pd.DataFrame: - """ - :param fun: a function whose dataframe result obtained `fun()` will be cached - :param cache_name: the cache of the function result is written into `{cache_path}/{cache_name}.csv.zip` - :param ignore_cache: whether to recompute even if the cache is present - :param cache_path: folder where to write cache files, default to ~/cache-zeroshot/ - :param parquet: whether to store the data in parquet, if not specified use csv.zip - :return: result of fun() - """ - if cache_path is None: - cache_path = data_root / "cache" - - if parquet: - cache_file = cache_path / (cache_name + ".parquet") - else: - cache_file = cache_path / (cache_name + ".csv.zip") - cache_file.parent.mkdir(parents=True, exist_ok=True) - if cache_file.exists() and not ignore_cache: - logger.info("Loading cache %s", cache_file) - if parquet: - return pd.read_parquet(cache_file) - else: - return pd.read_csv(cache_file) - else: - logger.info( - "Cache %s not found or ignore_cache set to True, regenerating the file", - cache_file, - ) - with Timeblock("Evaluate function."): - df = fun() - assert isinstance(df, pd.DataFrame) - if parquet: - # object cols cannot be saved easily in parquet; numpy arrays must be - # deep-converted to plain Python so str() produces ast.literal_eval-safe - # repr (no "array([...])" syntax, which breaks literal_eval) - import numpy as np - - def _to_python(x): - """Recursively convert numpy arrays/scalars to Python lists/dicts.""" - if isinstance(x, np.ndarray): - return [_to_python(i) for i in x] - if isinstance(x, dict): - return {k: _to_python(v) for k, v in x.items()} - if isinstance(x, list): - return [_to_python(i) for i in x] - return x - - for col in df.select_dtypes(include="object").columns: - df[col] = df[col].apply(_to_python).astype(str) - df.to_parquet(cache_file, index=False) - return pd.read_parquet(cache_file) - else: - df.to_csv(cache_file, index=False) - return pd.read_csv(cache_file) - - if __name__ == "__main__": download_all() diff --git a/scripts/fluency/generate_fluency.py b/scripts/fluency/generate_fluency.py index d829d8d..1da9647 100644 --- a/scripts/fluency/generate_fluency.py +++ b/scripts/fluency/generate_fluency.py @@ -3,9 +3,7 @@ import pandas as pd from datasets import Dataset -from judgearena.utils import do_inference, make_model, set_langchain_cache - -set_langchain_cache() +from judgearena.utils import do_inference, make_model dataset_name = "geoalgo/multilingual-fluency" model = "OpenRouter/openai/gpt-5-mini" diff --git a/scripts/multilingual_arena_hard/translate_arena_hard.py b/scripts/multilingual_arena_hard/translate_arena_hard.py index d0849c4..87cd9ef 100644 --- a/scripts/multilingual_arena_hard/translate_arena_hard.py +++ b/scripts/multilingual_arena_hard/translate_arena_hard.py @@ -20,8 +20,6 @@ from judgearena.instruction_dataset import load_instructions from judgearena.utils import do_inference, make_model -# set_langchain_cache() - dataset_name = "openeurollm/ArenaHard-EU-v0-bis" """ diff --git a/tests/test_inference_cache.py b/tests/test_inference_cache.py index 1360307..6e7694d 100644 --- a/tests/test_inference_cache.py +++ b/tests/test_inference_cache.py @@ -1,4 +1,5 @@ import json +from pathlib import Path import pandas as pd import pytest @@ -34,6 +35,13 @@ def batch(self, inputs, **_kwargs): return [f"generated:{item}" for item in inputs] +def test_legacy_runtime_cache_symbols_are_removed(): + package = Path(utils.__file__).parent + source = "\n".join(path.read_text() for path in package.rglob("*.py")) + for symbol in ("cache_function_dataframe", "ignore_cache", "set_langchain_cache"): + assert symbol not in source + + def test_generate_and_judge_full_hits_do_not_materialize_models(tmp_path, monkeypatch): completion_cache = CompletionInferenceCache(tmp_path, "arena-hard") judgement_cache = JudgementInferenceCache(tmp_path, "arena-hard") @@ -122,6 +130,30 @@ def test_mixed_hits_and_misses_preserve_order(tmp_path, monkeypatch): assert backend.calls == [["miss-a", "miss-b"]] +def test_multiturn_temperature_cache_full_hit(tmp_path, monkeypatch): + questions = pd.DataFrame( + { + "turn_1": ["question 1", "question 2"], + "turn_2": ["follow-up 1", "follow-up 2"], + "category": ["writing", "math"], + } + ) + cache = CompletionInferenceCache(tmp_path, "mt-bench") + kwargs = {"temperature_config": {"writing": 0.7, "math": 0.0}, "use_tqdm": False} + first = generate.generate_multiturn( + questions, "Dummy/answer", inference_cache=cache, **kwargs + ) + + def fail_if_materialized(*_args, **_kwargs): + raise AssertionError("cache hit materialized a model") + + monkeypatch.setattr(utils, "make_model", fail_if_materialized) + second = generate.generate_multiturn( + questions, "Dummy/answer", inference_cache=cache, **kwargs + ) + pd.testing.assert_frame_equal(second, first) + + def test_vllm_descriptor_contains_output_configuration(tmp_path, monkeypatch): monkeypatch.setattr(inference.importlib_metadata, "version", lambda _name: "0.10.2") cache = CompletionInferenceCache(tmp_path, "arena-hard") @@ -134,6 +166,7 @@ def test_vllm_descriptor_contains_output_configuration(tmp_path, monkeypatch): gpu_memory_utilization=0.9, max_model_len=4096, tensor_parallel_size=2, + temperature=0.2, ) assert model.descriptor["backend_version"] == "0.10.2" @@ -142,7 +175,7 @@ def test_vllm_descriptor_contains_output_configuration(tmp_path, monkeypatch): "max_model_len": 4096, "chat_template": None, } - assert model.descriptor["sampling"] == {"temperature": 0.6, "top_p": 0.95} + assert model.descriptor["sampling"] == {"temperature": 0.2, "top_p": 0.95} @pytest.mark.parametrize(