diff --git a/judgearena/cli.py b/judgearena/cli.py index eb94c83..5006df5 100644 --- a/judgearena/cli.py +++ b/judgearena/cli.py @@ -196,6 +196,7 @@ def _build_elo_args( 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, max_out_tokens_judge=args.max_out_tokens_judge, @@ -224,6 +225,7 @@ def _build_generate_and_evaluate_args( 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, max_out_tokens_judge=args.max_out_tokens_judge, diff --git a/judgearena/cli_common.py b/judgearena/cli_common.py index 58ce78b..f22c617 100644 --- a/judgearena/cli_common.py +++ b/judgearena/cli_common.py @@ -22,6 +22,7 @@ class BaseCliArgs: 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 max_out_tokens_judge: int = 32768 @@ -85,6 +86,12 @@ def add_common_arguments(parser: argparse.ArgumentParser) -> None: action="store_true", help="If specified, ignore cache of previous completions.", ) + parser.add_argument( + "--store_root", + type=str, + default=None, + help="Root directory for content-addressed completion and judgement caches.", + ) parser.add_argument( "--result_folder", type=str, diff --git a/judgearena/estimate_elo_ratings.py b/judgearena/estimate_elo_ratings.py index e85840a..dc1628b 100644 --- a/judgearena/estimate_elo_ratings.py +++ b/judgearena/estimate_elo_ratings.py @@ -1,10 +1,7 @@ from __future__ import annotations -import hashlib from dataclasses import dataclass -from functools import partial from pathlib import Path -from typing import TYPE_CHECKING import numpy as np import pandas as pd @@ -12,13 +9,11 @@ from judgearena.arenas_utils import _extract_instruction_text, load_arena_dataframe from judgearena.cli_common import BaseCliArgs -from judgearena.evaluate import PairScore, judge_and_parse_prefs +from judgearena.evaluate import judge_and_parse_prefs from judgearena.generate import generate_instructions +from judgearena.inference import CompletionInferenceCache, JudgementInferenceCache from judgearena.log import get_logger -from judgearena.utils import cache_function_dataframe, compute_pref_summary, make_model - -if TYPE_CHECKING: - pass +from judgearena.utils import compute_pref_summary, prepare_model logger = get_logger(__name__) @@ -40,9 +35,6 @@ class CliEloArgs(BaseCliArgs): n_bootstraps: int = 20 seed: int = 0 baseline_model: str | None = None - store_root: str | None = ( - None # root dir of the SQLite store; enables caching if set - ) def compute_bradley_terry( @@ -157,29 +149,7 @@ def compute_bradley_terry( return dict(pd.Series(elo_scores, index=models.index)) -def _store_folder(store_root: str, kind: str, task: str, model_spec: str) -> Path: - provider, model_path = model_spec.split("/", 1) - model_name = model_path.replace("/", "--") - return Path(store_root) / kind / task / model_name / provider - - def main(args: CliEloArgs) -> dict: - from judgearena.store_sqlite import SQLiteCompletionStore, SQLiteJudgementStore - - if args.store_root is not None: - comp_folder = _store_folder( - args.store_root, "completions", args.arena, args.model - ) - completion_store = SQLiteCompletionStore(comp_folder / "completions.db") - judge_folder = _store_folder( - args.store_root, "judgements", args.arena, args.judge_model - ) - judgement_store = SQLiteJudgementStore(judge_folder / "judgements.db") - logger.info("Using SQLite store at %s", args.store_root) - else: - completion_store = None - judgement_store = None - rng = np.random.default_rng(args.seed) # Step 1: Load arena battles @@ -264,43 +234,19 @@ def main(args: CliEloArgs) -> dict: if args.chat_template is not None: extra_kwargs["chat_template"] = args.chat_template use_tqdm = False - gen_fun = partial( - generate_instructions, + completion_cache = ( + CompletionInferenceCache(Path(args.store_root), args.arena) + if args.store_root is not None + else None + ) + completions_df = generate_instructions( + instructions=instructions, + model=args.model, truncate_input_chars=args.truncate_all_input_chars, max_tokens=args.max_out_tokens_models, use_tqdm=use_tqdm, - completion_store=completion_store, + inference_cache=completion_cache, **extra_kwargs, - ) - - def replace_slash(s: str) -> str: - return s.replace("/", "_") - - languages_str = "-".join(sorted(args.languages)) if args.languages else "all" - extra_kwargs_str = ( - "_".join(f"{k}={v}" for k, v in sorted(extra_kwargs.items())) - if extra_kwargs - else "" - ) - cache_suffix = ( - f"{args.arena}_{replace_slash(args.model)}_" - f"{args.n_instructions}_{args.n_instructions_per_language}_" - f"{languages_str}_{args.truncate_all_input_chars}_{args.max_out_tokens_models}" - + (f"_{extra_kwargs_str}" if extra_kwargs_str else "") - ) - if len(cache_suffix) > 100: - cache_hash = hashlib.sha256(cache_suffix.encode()).hexdigest()[:16] - logger.debug( - "Cache suffix too long (%d chars), using hash: %s (full: %s)", - len(cache_suffix), - cache_hash, - cache_suffix, - ) - cache_suffix = cache_hash - completions_df = cache_function_dataframe( - lambda: gen_fun(instructions=instructions, model=args.model), - ignore_cache=args.ignore_cache, - cache_name=f"elo/{cache_suffix}", ).set_index("instruction_index") completions = completions_df.loc[:, "completion"] @@ -354,116 +300,43 @@ def replace_slash(s: str) -> str: if args.chat_template is not None: judge_extra_kwargs["chat_template"] = args.chat_template - def run_judge() -> pd.DataFrame: - # Determine which indices still need judging - all_indices = list(range(n)) - if judgement_store is not None: - pairs = list( - zip(all_indices, model_A_per_row, model_B_per_row, strict=True) - ) - cached_df = judgement_store.query(model=args.model) - cached_keys = ( - set( - zip( - cached_df["instruction_index"].astype(int), - cached_df["model_A"], - cached_df["model_B"], - strict=True, - ) - ) - if not cached_df.empty - else set() - ) - missing_idx = [i for i, mA, mB in pairs if (i, mA, mB) not in cached_keys] - logger.info( - "Judgement store: %d cached, %d to judge.", - n - len(missing_idx), - len(missing_idx), - ) - else: - missing_idx = all_indices - cached_df = pd.DataFrame() - - new_df = pd.DataFrame() - if missing_idx: - judge_chat_model = make_model( - model=args.judge_model, - max_tokens=args.max_out_tokens_judge, - **judge_extra_kwargs, - ) - annotations, _, prefs = judge_and_parse_prefs( - judge_chat_model=judge_chat_model, - instructions=[instructions[i] for i in missing_idx], - completions_A=[completions_A[i] for i in missing_idx], - completions_B=[completions_B[i] for i in missing_idx], - swap_mode=args.swap_mode, - provide_explanation=args.provide_explanation, - truncate_input_chars=args.truncate_all_input_chars, - use_tqdm=use_tqdm, - ) - new_df = pd.DataFrame( - { - "judge_completion": [a.judge_completion for a in annotations], - "instruction": [a.instruction for a in annotations], - "completion_A": [a.completion_A for a in annotations], - "completion_B": [a.completion_B for a in annotations], - "pref": list(prefs), - "use_model_a_as_opponent": use_model_a_as_opponent[missing_idx], - "our_model_is_position_a": our_model_is_position_a[missing_idx], - "opponent_model": [opponent_models[i] for i in missing_idx], - "instruction_index": missing_idx, - "model_A": [model_A_per_row[i] for i in missing_idx], - "model_B": [model_B_per_row[i] for i in missing_idx], - } - ) - if judgement_store is not None: - judgement_store.save( - new_df.rename(columns={"judge_completion": "judge_output"}), - pushed_by="judgearena", - ) - - if judgement_store is not None and not cached_df.empty: - # Reconstruct full df by merging cached rows back in - score_parser = PairScore() - cached_df = cached_df.copy() - cached_df["pref"] = cached_df["judge_output"].apply( - score_parser.parse_model_raw - ) - cached_df = cached_df.rename(columns={"judge_output": "judge_completion"}) - cached_df["our_model_is_position_a"] = cached_df["model_A"] == args.model - cached_df["opponent_model"] = cached_df.apply( - lambda r: ( - r["model_B"] if r["our_model_is_position_a"] else r["model_A"] - ), - axis=1, - ) - cached_df["use_model_a_as_opponent"] = [ - df_battles.iloc[int(idx)]["model_a"] == opp - for idx, opp in zip( - cached_df["instruction_index"], - cached_df["opponent_model"], - strict=True, - ) - ] - full_df = pd.concat([new_df, cached_df], ignore_index=True) - return full_df.sort_values("instruction_index").reset_index(drop=True) - - return new_df - - judge_cache_suffix = f"judge_{cache_suffix}" - df_judge = cache_function_dataframe( - run_judge, - ignore_cache=args.ignore_cache, - cache_name=f"elo/{judge_cache_suffix}", + judgement_cache = ( + JudgementInferenceCache(Path(args.store_root), args.arena) + 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, + **judge_extra_kwargs, + ) + annotations, annotations_reversed, prefs = judge_and_parse_prefs( + judge_chat_model=judge_chat_model, + instructions=instructions.tolist(), + completions_A=completions_A, + completions_B=completions_B, + swap_mode=args.swap_mode, + provide_explanation=args.provide_explanation, + truncate_input_chars=args.truncate_all_input_chars, + use_tqdm=use_tqdm, + cache_metadata=[ + { + "instruction_id": index, + "model_a": model_A_per_row[index], + "model_b": model_B_per_row[index], + "orientation": "direct", + } + for index in range(n) + ], + ) + prefs = prefs.tolist() + if annotations_reversed is not None: + use_model_a_as_opponent = np.tile(use_model_a_as_opponent, 2) + our_model_is_position_a = np.tile(our_model_is_position_a, 2) + opponent_models *= 2 - # Restore position arrays and prefs from cache (in case loaded from disk) - use_model_a_as_opponent = df_judge["use_model_a_as_opponent"].to_numpy() - our_model_is_position_a = df_judge["our_model_is_position_a"].to_numpy() - opponent_models = df_judge["opponent_model"].tolist() - prefs = df_judge["pref"].tolist() - - logger.debug("First judge output:\n%s", df_judge["judge_completion"].iloc[0][:500]) + logger.debug("First judge output:\n%s", annotations[0].judge_completion[:500]) # Map preferences back to model-name-level battle results model_name = args.model diff --git a/judgearena/evaluate.py b/judgearena/evaluate.py index cb63df7..3dd41c9 100644 --- a/judgearena/evaluate.py +++ b/judgearena/evaluate.py @@ -368,6 +368,7 @@ def judge_and_parse_prefs( user_prompt_template: str | None = None, truncate_input_chars: int = 8192, use_tqdm: bool = False, + cache_metadata: list[dict] | None = None, ) -> tuple[list[JudgeAnnotation], list[JudgeAnnotation] | None, pd.Series]: """Run judge annotation and parse preferences, handling swap_mode='both'. @@ -396,10 +397,22 @@ def judge_and_parse_prefs( user_prompt_template=user_prompt_template, truncate_input_chars=truncate_input_chars, use_tqdm=use_tqdm, + cache_metadata=cache_metadata, ) annotations_reversed = None if swap_mode == "both": + reversed_cache_metadata = None + if cache_metadata is not None: + reversed_cache_metadata = [ + { + **metadata, + "model_a": metadata["model_b"], + "model_b": metadata["model_a"], + "orientation": "reversed", + } + for metadata in cache_metadata + ] annotations_reversed = annotate_battles( judge_chat_model=judge_chat_model, instructions=instructions, @@ -410,6 +423,7 @@ def judge_and_parse_prefs( user_prompt_template=user_prompt_template, truncate_input_chars=truncate_input_chars, use_tqdm=use_tqdm, + cache_metadata=reversed_cache_metadata, ) def _none_to_nan(x): diff --git a/judgearena/generate.py b/judgearena/generate.py index d164ed5..06b381d 100644 --- a/judgearena/generate.py +++ b/judgearena/generate.py @@ -1,7 +1,5 @@ from __future__ import annotations -from typing import TYPE_CHECKING - import pandas as pd from langchain_core.prompts import ChatPromptTemplate @@ -14,9 +12,6 @@ truncate, ) -if TYPE_CHECKING: - from judgearena.store_sqlite import SQLiteCompletionStore - logger = get_logger(__name__) @@ -27,29 +22,9 @@ def generate_instructions( max_tokens: int | None = 32768, use_tqdm: bool = True, system_prompt: str | None = None, - completion_store: SQLiteCompletionStore | None = None, inference_cache: InferenceCache | None = None, - pushed_by: str = "judgearena", **engine_kwargs, ) -> pd.DataFrame: - # Filter to instructions not already in the shared store - if completion_store is not None: - all_indices = instructions.index.tolist() - missing = set(completion_store.missing_indices(all_indices)) - cached_df = completion_store.query([i for i in all_indices if i not in missing]) - instructions_to_run = instructions.loc[sorted(missing)] - logger.info( - "Completion store: %d cached, %d to generate.", - len(cached_df), - len(instructions_to_run), - ) - else: - instructions_to_run = instructions - cached_df = pd.DataFrame() - - if instructions_to_run.empty: - return cached_df[["instruction_index", "completion"]].reset_index(drop=True) - chat_model = prepare_model( model, max_tokens=max_tokens, @@ -67,38 +42,22 @@ def generate_instructions( inputs = prompt_template.batch( [ {"user_prompt": truncate(user_prompt, max_len=truncate_input_chars)} - for user_prompt in instructions_to_run + for user_prompt in instructions ] ) completions = do_inference( chat_model=chat_model, inputs=inputs, use_tqdm=use_tqdm, - cache_metadata=[ - {"instruction_id": index} for index in instructions_to_run.index - ], + cache_metadata=[{"instruction_id": index} for index in instructions.index], ) - df_new = pd.DataFrame( + return pd.DataFrame( { "completion": completions, - "instruction_index": instructions_to_run.index.tolist(), + "instruction_index": instructions.index.tolist(), } ) - if completion_store is not None: - completion_store.save(df_new, pushed_by=pushed_by) - if not cached_df.empty: - df_new = ( - pd.concat( - [cached_df[["instruction_index", "completion"]], df_new], - ignore_index=True, - ) - .sort_values("instruction_index") - .reset_index(drop=True) - ) - - return df_new - def _set_temperature_on_model(chat_model, temperature: float) -> None: if hasattr(chat_model, "set_temperature"): @@ -267,20 +226,27 @@ def generate_base( truncate_input_chars: int | None = 8192, max_tokens: int | None = 32768, use_tqdm: bool = False, + inference_cache: InferenceCache | None = None, **engine_kwargs, ) -> pd.DataFrame: - model = make_model(model, max_tokens=max_tokens, **engine_kwargs) + chat_model = prepare_model( + model, + max_tokens=max_tokens, + cache=inference_cache, + **engine_kwargs, + ) inputs = [ truncate(instruction, max_len=truncate_input_chars) for instruction in instructions ] - completions = model.batch( - inputs=inputs, - max_tokens=max_tokens, + completions = do_inference( + chat_model, + inputs, + use_tqdm=use_tqdm, + cache_metadata=[{"instruction_id": index} for index in instructions.index], ) - completions = [x.content if hasattr(x, "content") else x for x in completions] df_outputs = pd.DataFrame( data={ diff --git a/judgearena/generate_and_evaluate.py b/judgearena/generate_and_evaluate.py index 2919280..cba8369 100644 --- a/judgearena/generate_and_evaluate.py +++ b/judgearena/generate_and_evaluate.py @@ -14,6 +14,7 @@ from judgearena.cli_common import BaseCliArgs from judgearena.evaluate import judge_and_parse_prefs, resolve_judge_prompts from judgearena.generate import generate_base, generate_instructions +from judgearena.inference import CompletionInferenceCache, JudgementInferenceCache from judgearena.instruction_dataset import load_instructions from judgearena.instruction_dataset.arena_hard import ( download_arena_hard, @@ -27,11 +28,10 @@ from judgearena.mt_bench.mt_bench_utils import run_mt_bench from judgearena.repro import _to_jsonable, write_run_metadata from judgearena.utils import ( - cache_function_dataframe, compute_pref_summary, data_root, download_hf, - make_model, + prepare_model, read_df, ) @@ -146,15 +146,10 @@ def main(args: CliArgs): args.model_B, ) - # Not working with vllm, not detecting model changes and serving the same cache for two different models... - # if not args.ignore_cache: - # set_langchain_cache() - ignore_cache = args.ignore_cache - if args.task == "mt-bench": return run_mt_bench( args, - ignore_cache, + args.ignore_cache, res_folder=res_folder, result_name=name, ) @@ -182,6 +177,11 @@ def main(args: CliArgs): args.model_A, args.model_B, ) + completion_cache = ( + CompletionInferenceCache(Path(args.store_root), args.task) + if args.store_root is not None + else None + ) # TODO currently we just support base models for fluency, we could also support instruction-tuned models gen_fun = ( @@ -192,6 +192,7 @@ def main(args: CliArgs): max_model_len=args.max_model_len, chat_template=args.chat_template, use_tqdm=args.use_tqdm, + inference_cache=completion_cache, **args.engine_kwargs, ) if is_fluency_task @@ -202,6 +203,7 @@ def main(args: CliArgs): max_model_len=args.max_model_len, chat_template=args.chat_template, use_tqdm=args.use_tqdm, + inference_cache=completion_cache, **args.engine_kwargs, ) ) @@ -213,14 +215,10 @@ def main(args: CliArgs): :, "completion" ] else: - completions_A = cache_function_dataframe( - lambda: gen_fun( - instructions=instructions, - model=args.model_A, - use_tqdm=args.use_tqdm, - ), - ignore_cache=ignore_cache, - cache_name=f"{args.task}_{args.model_A}_{args.n_instructions}", + completions_A = gen_fun( + instructions=instructions, + model=args.model_A, + use_tqdm=args.use_tqdm, ).set_index("instruction_index") completions_A = completions_A.loc[:, "completion"] @@ -232,14 +230,10 @@ def main(args: CliArgs): :, "completion" ] else: - completions_B = cache_function_dataframe( - lambda: gen_fun( - instructions=instructions, - model=args.model_B, - use_tqdm=args.use_tqdm, - ), - ignore_cache=ignore_cache, - cache_name=f"{args.task}_{args.model_B}_{args.n_instructions}", + completions_B = gen_fun( + instructions=instructions, + model=args.model_B, + use_tqdm=args.use_tqdm, ).set_index("instruction_index") completions_B = completions_B.loc[:, "completion"] logger.debug("First instruction/context: %s", instructions.values[0]) @@ -247,9 +241,15 @@ def main(args: CliArgs): logger.debug("First completion of %s:\n%s", args.model_B, completions_B.values[0]) logger.info("Evaluating completions with judge %s.", args.judge_model) - judge_chat_model = make_model( + judgement_cache = ( + JudgementInferenceCache(Path(args.store_root), args.task) + 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, max_model_len=args.max_model_len, chat_template=args.chat_template, **args.engine_kwargs, @@ -288,6 +288,15 @@ def main(args: CliArgs): user_prompt_template=judge_user_prompt_template, truncate_input_chars=args.truncate_all_input_chars, use_tqdm=args.use_tqdm, + cache_metadata=[ + { + "instruction_id": index, + "model_a": args.model_A, + "model_b": args.model_B, + "orientation": "direct", + } + for index in instructions.head(n_instructions).index + ], ) df = pd.DataFrame(annotations) diff --git a/tests/test_estimate_elo_ratings.py b/tests/test_estimate_elo_ratings.py index 83f9c8a..c000d40 100644 --- a/tests/test_estimate_elo_ratings.py +++ b/tests/test_estimate_elo_ratings.py @@ -69,13 +69,6 @@ def mock_generate(instructions, model, **kwargs): monkeypatch.setattr(estimate_elo_ratings, "generate_instructions", mock_generate) - def _run_without_cache(fun, **_kwargs): - return fun() - - monkeypatch.setattr( - estimate_elo_ratings, "cache_function_dataframe", _run_without_cache - ) - def _default_args(**kwargs) -> CliEloArgs: defaults = dict( diff --git a/tests/test_generate_and_evaluate.py b/tests/test_generate_and_evaluate.py index 03a278f..8edca0b 100644 --- a/tests/test_generate_and_evaluate.py +++ b/tests/test_generate_and_evaluate.py @@ -2,6 +2,7 @@ import pytest import judgearena.generate_and_evaluate as generate_and_evaluate +import judgearena.utils as utils from judgearena.generate_and_evaluate import ( CliArgs, ) @@ -40,13 +41,6 @@ def mock_external_data_and_cache(monkeypatch): lambda dataset, model, n_instructions: None, ) - def _run_without_cache(fun, **_kwargs): - return fun() - - monkeypatch.setattr( - generate_and_evaluate, "cache_function_dataframe", _run_without_cache - ) - @pytest.mark.parametrize( "task", @@ -96,3 +90,22 @@ def test_generate_and_evaluate_correct_order_bias(tmp_path): avg_pref = sum(prefs) / len(prefs) assert avg_pref == 0.5 + + +def test_generate_and_evaluate_reuses_inference_cache(tmp_path, monkeypatch): + args = CliArgs( + task="alpaca-eval", + model_A="Dummy/answer A", + model_B="Dummy/answer B", + judge_model="Dummy/score A: 1 score B: 0", + n_instructions=2, + result_folder=str(tmp_path / "results"), + store_root=str(tmp_path / "cache"), + ) + main_generate_and_eval(args) + monkeypatch.setattr( + utils, + "make_model", + lambda *_args, **_kwargs: pytest.fail("cache hit materialized a model"), + ) + main_generate_and_eval(args)