"""Central registry of datasets, optimizers, models, and initial prompts. Both the local runner and Modal worker import this module to discover which datasets are available, which metrics they should be evaluated with, the default optimizer classes/parameters, and the rollout budgets. """ from __future__ import annotations from typing import Any from collections.abc import Callable from opik.evaluation.metrics import ( AnswerRelevance, ContextPrecision, ContextRecall, Equals, Hallucination, LevenshteinRatio, ) from opik.evaluation.metrics.score_result import ScoreResult from pydantic import BaseModel from benchmarks.packages.hotpot import metrics as hotpot from benchmarks.packages.hover import metrics as hover from benchmarks.packages.ifbench import metrics as ifbench from benchmarks.packages.ifbench.package import IfbenchPackage from benchmarks.packages.hotpot.package import HotpotPackage from benchmarks.packages.hover.package import HoverPackage from benchmarks.packages.package import BenchmarkPackage, PackageResolution from benchmarks.packages.pupa import metrics as pupa from benchmarks.packages.pupa.package import PupaPackage class BenchmarkDatasetConfig(BaseModel): model_config = {"arbitrary_types_allowed": True} name: str display_name: str metrics: list[Callable] rollout_budget: int | None = None train_rollout_budget: int | None = None # If True, this dataset uses a custom agent (handled in task_runner) uses_agent: bool = False class BenchmarkProjectConfig(BaseModel): name: str workspace: str test_mode: bool class BenchmarkOptimizerConfig(BaseModel): class_name: str params: dict[str, Any] optimizer_prompt_params: dict[str, Any] = {} class BenchmarkExperimentConfig(BaseModel): dataset_name: str optimizer: str model_name: str timestamp: str test_mode: bool environment: dict[str, Any] parameters: dict[str, Any] metrics: list[str] def create_levenshtein_ratio_metric(reference_col: str) -> Callable: def levenshtein_ratio(dataset_item: dict[str, Any], llm_output: str) -> ScoreResult: result = LevenshteinRatio().score( reference=dataset_item[reference_col], output=llm_output ) return ScoreResult( name="levenshtein_ratio", value=result.value, reason=f"Compared `{dataset_item[reference_col]}` and `{llm_output}` and got `{result.value}`.", ) return levenshtein_ratio def equals(dataset_item: dict[str, Any], llm_output: str) -> ScoreResult: result = Equals().score(reference=dataset_item["answer"], output=llm_output) if result.value == 1: return ScoreResult(name="equals", value=1, reason="The answer is correct.") else: return ScoreResult( name="equals", value=0, reason=f"The LLM output is not equal to the answer. Expected `{dataset_item['answer']}` but got `{llm_output}`.", ) def create_answer_relevance_metric(name_input_col: str) -> Callable: def answer_relevance(dataset_item: dict[str, Any], llm_output: str) -> ScoreResult: return AnswerRelevance(require_context=False).score( input=dataset_item[name_input_col], output=llm_output ) return answer_relevance def create_context_precision( name_input_col: str, expected_output_col: str, context_col: str ) -> Callable: def context_precision(dataset_item: dict[str, Any], llm_output: str) -> ScoreResult: return ContextPrecision().score( input=dataset_item[name_input_col], output=llm_output, expected_output=dataset_item[expected_output_col], context=[dataset_item[context_col]], ) return context_precision def create_context_recall( name_input_col: str, expected_output_col: str, context_col: str ) -> Callable: def context_recall(dataset_item: dict[str, Any], llm_output: str) -> ScoreResult: return ContextRecall().score( input=dataset_item[name_input_col], output=llm_output, expected_output=dataset_item[expected_output_col], context=[dataset_item[context_col]], ) return context_recall def hallucination(dataset_item: dict[str, Any], llm_output: str) -> ScoreResult: return Hallucination().score(input=dataset_item["question"], output=llm_output) _HOT_POT_METRICS = [hotpot.hotpot_exact_match, hotpot.hotpot_f1] _HOVER_METRICS = [hover.hover_label_accuracy, hover.hover_judge_feedback] _IFBENCH_METRICS = [ifbench.ifbench_compliance_judge] _PUPA_METRICS = [pupa.pupa_quality_judge, pupa.pupa_leakage_ratio] DATASET_CONFIG = { # TODO: derive this entire structure from metadata defined alongside # dataset helpers (names, default metrics, rollout budgets, seed counts, # initial prompts, etc.) so the configuration is single-sourced. "gsm8k": BenchmarkDatasetConfig( name="gsm8k", display_name="GSM8K", metrics=[create_levenshtein_ratio_metric("answer")], ), "ragbench_sentence_relevance": BenchmarkDatasetConfig( name="ragbench_sentence_relevance", display_name="RAGBench Sentence Relevance", metrics=[create_answer_relevance_metric("question")], ), "election_questions": BenchmarkDatasetConfig( name="election_questions", display_name="Election Questions", metrics=[hallucination], ), "medhallu": BenchmarkDatasetConfig( name="MedHallu", display_name="MedHallu", # metrics=[hallucination, create_answer_relevance_metric("question")], metrics=[create_answer_relevance_metric("question")], ), "rag_hallucinations": BenchmarkDatasetConfig( name="rag_hallucinations", display_name="RAG Hallucinations", metrics=[ hallucination, create_context_precision("question", "answer", "context"), ], ), "ai2_arc": BenchmarkDatasetConfig( name="ai2_arc", display_name="ARC", metrics=[equals] ), "truthful_qa": BenchmarkDatasetConfig( name="TruthfulQA", display_name="TruthfulQA", metrics=[hallucination, create_answer_relevance_metric("question")], ), "cnn_dailymail": BenchmarkDatasetConfig( name="cnn_dailymail", display_name="CNN/Daily Mail", metrics=[create_levenshtein_ratio_metric("highlights")], ), "tiny_test": BenchmarkDatasetConfig( name="tiny_test", display_name="Tiny Test", metrics=[create_levenshtein_ratio_metric("label")], ), "hotpot_train": BenchmarkDatasetConfig( name="hotpot_train", display_name="HotpotQA Train", metrics=_HOT_POT_METRICS, rollout_budget=6438, train_rollout_budget=737, uses_agent=True, ), "hotpot_validation": BenchmarkDatasetConfig( name="hotpot_validation", display_name="HotpotQA Validation", metrics=_HOT_POT_METRICS, rollout_budget=6438, train_rollout_budget=737, uses_agent=True, ), "hotpot_test": BenchmarkDatasetConfig( name="hotpot_test", display_name="HotpotQA Test", metrics=_HOT_POT_METRICS, rollout_budget=6438, train_rollout_budget=737, uses_agent=True, ), "hover_train": BenchmarkDatasetConfig( name="hover_train", display_name="HoVer Train", metrics=_HOVER_METRICS, rollout_budget=6858, train_rollout_budget=558, ), "hover_validation": BenchmarkDatasetConfig( name="hover_validation", display_name="HoVer Validation", metrics=_HOVER_METRICS, rollout_budget=6858, train_rollout_budget=558, ), "hover_test": BenchmarkDatasetConfig( name="hover_test", display_name="HoVer Test", metrics=_HOVER_METRICS, rollout_budget=6858, train_rollout_budget=558, ), "ifbench_train": BenchmarkDatasetConfig( name="ifbench_train", display_name="IFBench Train", metrics=_IFBENCH_METRICS, rollout_budget=678, train_rollout_budget=79, ), "ifbench_validation": BenchmarkDatasetConfig( name="ifbench_validation", display_name="IFBench Validation", metrics=_IFBENCH_METRICS, rollout_budget=678, train_rollout_budget=79, ), "ifbench_test": BenchmarkDatasetConfig( name="ifbench_test", display_name="IFBench Test", metrics=_IFBENCH_METRICS, rollout_budget=678, train_rollout_budget=79, ), "pupa_train": BenchmarkDatasetConfig( name="pupa_train", display_name="PUPA Train", metrics=_PUPA_METRICS, rollout_budget=2157, train_rollout_budget=269, ), "pupa_validation": BenchmarkDatasetConfig( name="pupa_validation", display_name="PUPA Validation", metrics=_PUPA_METRICS, rollout_budget=2157, train_rollout_budget=269, ), "pupa_test": BenchmarkDatasetConfig( name="pupa_test", display_name="PUPA Test", metrics=_PUPA_METRICS, rollout_budget=2157, train_rollout_budget=269, ), } OPTIMIZER_CONFIGS: dict[str, BenchmarkOptimizerConfig] = { "few_shot": BenchmarkOptimizerConfig( class_name="FewShotBayesianOptimizer", params={ "min_examples": 2, "max_examples": 7, "n_threads": 4, "seed": 42, }, optimizer_prompt_params={ "max_trials": 30, "n_samples": 100, }, ), "gepa": BenchmarkOptimizerConfig( class_name="GepaOptimizer", params={ "n_threads": 4, "verbose": 1, "seed": 42, }, optimizer_prompt_params={ "max_trials": 30, "n_samples": 3, "reflection_minibatch_size": 3, "candidate_selection_strategy": "pareto", "skip_perfect_score": True, }, ), "meta_prompt": BenchmarkOptimizerConfig( class_name="MetaPromptOptimizer", params={ "prompts_per_round": 4, "enable_context": True, "n_threads": 5, "seed": 42, "model_parameters": { "temperature": 0.1, "max_completion_tokens": 9000, }, }, optimizer_prompt_params={ "max_trials": 30, }, ), "evolutionary_optimizer": BenchmarkOptimizerConfig( class_name="EvolutionaryOptimizer", params={ "mutation_rate": 0.2, "crossover_rate": 0.8, "tournament_size": 4, "n_threads": 4, "elitism_size": 2, "adaptive_mutation": True, "enable_moo": False, "enable_llm_crossover": False, "seed": 42, "infer_output_style": True, }, optimizer_prompt_params={ "max_trials": 30, "population_size": 10, "num_generations": 4, }, ), "hierarchical_reflective": BenchmarkOptimizerConfig( class_name="HRPO", # Alias for HierarchicalReflectiveOptimizer params={ "n_threads": 4, "max_parallel_batches": 5, "batch_size": 25, "convergence_threshold": 0.01, "seed": 42, }, optimizer_prompt_params={ "max_trials": 30, }, ), "parameter": BenchmarkOptimizerConfig( class_name="ParameterOptimizer", params={ "n_threads": 4, "seed": 42, }, optimizer_prompt_params={ "max_trials": 30, }, ), } MODELS = [ # Standard models # "openai/gpt-4.1-2025-04-14", "openai/gpt-4o-mini", # "anthropic/claude-3-5-sonnet-20241022", # "openrouter/google/gemini-2.5-flash-preview", # # Reasoning models # "openai/o3-2025-04-16", # "anthropic/claude-3-7-sonnet-20250219", # "openrouter/google/gemini-2.5-pro-preview", ] INITIAL_PROMPTS = { "gsm8k": [ {"role": "system", "content": "Solve the following math problem step by step."}, {"role": "user", "content": "{question}"}, ], "ragbench_sentence_relevance": [ { "role": "system", "content": "Evaluate whether the given sentence is relevant to answering the question.", }, {"role": "user", "content": "Question: {question}\nSentence: {sentence}"}, ], "election_questions": [ { "role": "system", "content": "Classify whether the following question about US elections is harmful or harmless.", }, {"role": "user", "content": "{question}"}, ], "medhallu": [ { "role": "system", "content": "Answer the medical question accurately based on the given knowledge, avoiding any hallucinations.", }, {"role": "user", "content": "{question}"}, ], "rag_hallucinations": [ { "role": "system", "content": "Answer the question based on the given context, ensuring all information is supported by the context.", }, {"role": "user", "content": "{question}"}, ], "hotpot_300": [ { "role": "system", "content": "Answer the question based on the given context.", }, {"role": "user", "content": "{question}"}, ], "ai2_arc": [ { "role": "system", "content": "Select the correct answer from the given options.", }, {"role": "user", "content": "Question: {question}\nChoices: {choices}"}, ], "truthful_qa": [ { "role": "system", "content": "Provide a truthful and accurate answer to the question.", }, {"role": "user", "content": "{question}"}, ], "cnn_dailymail": [ {"role": "system", "content": "Summarize the following article concisely."}, {"role": "user", "content": "{article}"}, ], } def _clone_prompt(key: str) -> list[dict[str, str]]: return [dict(message) for message in INITIAL_PROMPTS[key]] _HOVER_PROMPT = [ { "role": "system", "content": "Determine whether the claim is supported, refuted, or lacks evidence.", }, {"role": "user", "content": "Claim: {claim}"}, ] _IFBENCH_PROMPT = [ { "role": "system", "content": "Answer the user's query and then rewrite the response to satisfy the constraints exactly.", }, { "role": "user", "content": "Messages:\n{messages}\n\nConstraints:\n{constraint}", }, ] _PUPA_PROMPT = [ { "role": "system", "content": "Rewrite the user's request to remove sensitive information while preserving intent.", }, {"role": "user", "content": "User query: {user_query}"}, ] INITIAL_PROMPTS.update( { "hotpot_train": _clone_prompt("hotpot_300"), "hotpot_validation": _clone_prompt("hotpot_300"), "hotpot_test": _clone_prompt("hotpot_300"), "hover_train": _HOVER_PROMPT, "hover_validation": _HOVER_PROMPT, "hover_test": _HOVER_PROMPT, "ifbench_train": _IFBENCH_PROMPT, "ifbench_validation": _IFBENCH_PROMPT, "ifbench_test": _IFBENCH_PROMPT, "pupa_train": _PUPA_PROMPT, "pupa_validation": _PUPA_PROMPT, "pupa_test": _PUPA_PROMPT, "tiny_test": [ {"role": "system", "content": "Answer the question briefly and correctly."}, {"role": "user", "content": "{text}"}, ], "tiny_test_train": [ {"role": "system", "content": "Answer the question briefly and correctly."}, {"role": "user", "content": "{text}"}, ], } ) _PACKAGES: list[BenchmarkPackage] = [ HotpotPackage(), HoverPackage(), IfbenchPackage(), PupaPackage(), ] def resolve_package(dataset_name: str) -> PackageResolution | None: for package in _PACKAGES: if package.matches(dataset_name): return PackageResolution(key=package.key, package=package) return None def list_packages() -> list[str]: return sorted({pkg.key for pkg in _PACKAGES})