Add GEPA prompt optimization for QA evaluations
This commit is contained in:
Yiorgis Gozadinos
parent
e31b692124
commit
60b1d3c013
4 changed files with 361 additions and 24 deletions
|
|
@ -390,6 +390,36 @@ async def evaluate_dataset(
|
|||
app = typer.Typer(help="Run retrieval and QA benchmarks for configured datasets.")
|
||||
|
||||
|
||||
def _load_config(config_path: Path | None) -> AppConfig:
|
||||
"""Load AppConfig from a file path or standard search path."""
|
||||
if config_path:
|
||||
if not config_path.exists():
|
||||
raise typer.BadParameter(f"Config file not found: {config_path}")
|
||||
console.print(f"Loading config from: {config_path}", style="dim")
|
||||
yaml_data = load_yaml_config(config_path)
|
||||
return AppConfig.model_validate(yaml_data)
|
||||
|
||||
found = find_config_file(None)
|
||||
if found:
|
||||
console.print(f"Loading config from: {found}", style="dim")
|
||||
yaml_data = load_yaml_config(found)
|
||||
return AppConfig.model_validate(yaml_data)
|
||||
|
||||
console.print("No config file found, using defaults", style="dim")
|
||||
return AppConfig()
|
||||
|
||||
|
||||
def _resolve_dataset(dataset: str) -> DatasetSpec:
|
||||
"""Resolve a dataset key to a DatasetSpec or raise BadParameter."""
|
||||
spec = DATASETS.get(dataset.lower())
|
||||
if spec is None:
|
||||
valid_datasets = ", ".join(sorted(DATASETS))
|
||||
raise typer.BadParameter(
|
||||
f"Unknown dataset '{dataset}'. Choose from: {valid_datasets}"
|
||||
)
|
||||
return spec
|
||||
|
||||
|
||||
@app.command()
|
||||
def run(
|
||||
dataset: str = typer.Argument(..., help="Dataset key to evaluate."),
|
||||
|
|
@ -417,30 +447,8 @@ def run(
|
|||
help="Only evaluate queries requiring image understanding.",
|
||||
),
|
||||
) -> None:
|
||||
spec = DATASETS.get(dataset.lower())
|
||||
if spec is None:
|
||||
valid_datasets = ", ".join(sorted(DATASETS))
|
||||
raise typer.BadParameter(
|
||||
f"Unknown dataset '{dataset}'. Choose from: {valid_datasets}"
|
||||
)
|
||||
|
||||
# Load config from file or use defaults
|
||||
if config:
|
||||
if not config.exists():
|
||||
raise typer.BadParameter(f"Config file not found: {config}")
|
||||
console.print(f"Loading config from: {config}", style="dim")
|
||||
yaml_data = load_yaml_config(config)
|
||||
app_config = AppConfig.model_validate(yaml_data)
|
||||
else:
|
||||
# Try to find config file using standard search path
|
||||
config_path = find_config_file(None)
|
||||
if config_path:
|
||||
console.print(f"Loading config from: {config_path}", style="dim")
|
||||
yaml_data = load_yaml_config(config_path)
|
||||
app_config = AppConfig.model_validate(yaml_data)
|
||||
else:
|
||||
console.print("No config file found, using defaults", style="dim")
|
||||
app_config = AppConfig()
|
||||
spec = _resolve_dataset(dataset)
|
||||
app_config = _load_config(config)
|
||||
|
||||
asyncio.run(
|
||||
evaluate_dataset(
|
||||
|
|
@ -458,6 +466,46 @@ def run(
|
|||
)
|
||||
|
||||
|
||||
@app.command()
|
||||
def optimize(
|
||||
dataset: str = typer.Argument(..., help="Dataset key to optimize prompt for."),
|
||||
config: Path | None = typer.Option(
|
||||
None, "--config", help="Path to haiku.rag YAML config file."
|
||||
),
|
||||
db: Path | None = typer.Option(None, "--db", help="Override the database path."),
|
||||
limit: int | None = typer.Option(None, "--limit", help="Limit number of QA cases."),
|
||||
max_calls: int = typer.Option(50, "--max-calls", help="Maximum GEPA metric calls."),
|
||||
output: Path | None = typer.Option(
|
||||
None, "--output", help="Save optimized prompt to file."
|
||||
),
|
||||
) -> None:
|
||||
"""Optimize QA system prompt using GEPA evolutionary optimization."""
|
||||
from evaluations.optimization import run_optimization
|
||||
|
||||
spec = _resolve_dataset(dataset)
|
||||
app_config = _load_config(config)
|
||||
|
||||
corpus = spec.qa_loader()
|
||||
if limit is not None:
|
||||
corpus = corpus.select(range(min(limit, len(corpus))))
|
||||
|
||||
cases: list[Case[str, str, dict[str, str]]] = [
|
||||
spec.qa_case_builder(index, cast(Mapping[str, Any], doc))
|
||||
for index, doc in enumerate(corpus, start=1)
|
||||
]
|
||||
|
||||
asyncio.run(
|
||||
run_optimization(
|
||||
spec=spec,
|
||||
config=app_config,
|
||||
cases=cases,
|
||||
max_calls=max_calls,
|
||||
db_path=db,
|
||||
output=output,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@app.command()
|
||||
def download(
|
||||
dataset: str = typer.Argument(..., help="Dataset key or 'all' to download all."),
|
||||
|
|
|
|||
277
evaluations/evaluations/optimization.py
Normal file
277
evaluations/evaluations/optimization.py
Normal file
|
|
@ -0,0 +1,277 @@
|
|||
import asyncio
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from pydantic_evals import Case, Dataset as EvalDataset
|
||||
from pydantic_evals.evaluators import LLMJudge
|
||||
|
||||
from gepa.core.adapter import EvaluationBatch
|
||||
|
||||
from evaluations.config import DatasetSpec
|
||||
from haiku.rag.agents.qa import get_qa_agent
|
||||
from haiku.rag.agents.qa.prompts import QA_SYSTEM_PROMPT
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.config.models import AppConfig, ModelConfig
|
||||
from haiku.rag.utils import get_model
|
||||
|
||||
|
||||
OPTIMIZATION_SCORING_RUBRIC = """You are evaluating the quality of an answer to a question,
|
||||
comparing it against a reference answer.
|
||||
|
||||
Score on a scale of 0.0 to 1.0:
|
||||
- 1.0: The answer is factually correct, complete, and concise. It covers all key points
|
||||
from the reference answer without contradictions or significant omissions.
|
||||
- 0.7-0.9: The answer is mostly correct and addresses the core question, but may miss
|
||||
some secondary details or include minor inaccuracies.
|
||||
- 0.4-0.6: The answer is partially correct — it addresses some aspects of the question
|
||||
but misses key information or contains notable inaccuracies.
|
||||
- 0.1-0.3: The answer is mostly incorrect or fails to address the core question,
|
||||
though it may contain some tangentially relevant information.
|
||||
- 0.0: The answer is completely wrong, irrelevant, or empty.
|
||||
|
||||
GUIDELINES:
|
||||
- Focus on factual correctness relative to the reference answer
|
||||
- Ignore differences in phrasing, style, or formatting
|
||||
- A concise correct answer scores higher than a verbose partially correct one
|
||||
- "I cannot find enough information" when the reference has an answer scores 0.0
|
||||
"""
|
||||
|
||||
|
||||
@dataclass
|
||||
class EvalTrajectory:
|
||||
"""Per-case evaluation result for GEPA reflection."""
|
||||
|
||||
question: str
|
||||
expected_answer: str
|
||||
actual_answer: str | None
|
||||
score: float
|
||||
judge_reason: str | None = None
|
||||
|
||||
|
||||
QACase = Case[str, str, dict[str, str]]
|
||||
|
||||
|
||||
@dataclass
|
||||
class QAPromptAdapter:
|
||||
"""GEPA adapter that evaluates QA prompt candidates against a dataset.
|
||||
|
||||
Implements the GEPAAdapter protocol:
|
||||
- evaluate(): Run QA agent with candidate prompt, score with LLMJudge
|
||||
- make_reflective_dataset(): Build failure records for the GEPA proposer
|
||||
"""
|
||||
|
||||
config: AppConfig
|
||||
db_path: Path
|
||||
judge_model: Any
|
||||
|
||||
def evaluate(
|
||||
self,
|
||||
batch: list[QACase],
|
||||
candidate: dict[str, str],
|
||||
capture_traces: bool = False,
|
||||
) -> EvaluationBatch[EvalTrajectory, str | None]:
|
||||
instructions = candidate["instructions"]
|
||||
return asyncio.run(self._evaluate_async(batch, instructions, capture_traces))
|
||||
|
||||
async def _evaluate_async(
|
||||
self,
|
||||
batch: list[QACase],
|
||||
instructions: str,
|
||||
capture_traces: bool,
|
||||
) -> EvaluationBatch[EvalTrajectory, str | None]:
|
||||
outputs: list[str | None] = []
|
||||
scores: list[float] = []
|
||||
trajectories: list[EvalTrajectory] | None = [] if capture_traces else None
|
||||
|
||||
async with HaikuRAG(self.db_path, config=self.config) as rag:
|
||||
qa = get_qa_agent(rag, self.config, system_prompt=instructions)
|
||||
|
||||
for case in batch:
|
||||
question = case.inputs
|
||||
expected = case.expected_output or ""
|
||||
|
||||
try:
|
||||
answer, _ = await qa.answer(question)
|
||||
except Exception:
|
||||
answer = None
|
||||
|
||||
if answer is not None:
|
||||
score, reason = await self._judge(question, answer, expected)
|
||||
else:
|
||||
score, reason = 0.0, "QA agent failed to produce an answer"
|
||||
|
||||
outputs.append(answer)
|
||||
scores.append(score)
|
||||
|
||||
if capture_traces and trajectories is not None:
|
||||
trajectories.append(
|
||||
EvalTrajectory(
|
||||
question=question,
|
||||
expected_answer=expected,
|
||||
actual_answer=answer,
|
||||
score=score,
|
||||
judge_reason=reason,
|
||||
)
|
||||
)
|
||||
|
||||
return EvaluationBatch(
|
||||
outputs=outputs,
|
||||
scores=scores,
|
||||
trajectories=trajectories,
|
||||
)
|
||||
|
||||
async def _judge(
|
||||
self, question: str, answer: str, expected: str
|
||||
) -> tuple[float, str | None]:
|
||||
"""Score an answer using pydantic-evals LLMJudge with float scoring."""
|
||||
judge = LLMJudge(
|
||||
rubric=OPTIMIZATION_SCORING_RUBRIC,
|
||||
include_input=True,
|
||||
include_expected_output=True,
|
||||
model=self.judge_model,
|
||||
score={"evaluation_name": "accuracy", "include_reason": True},
|
||||
assertion=False,
|
||||
)
|
||||
|
||||
dataset = EvalDataset(
|
||||
cases=[Case(inputs=question, expected_output=expected)],
|
||||
evaluators=[judge],
|
||||
)
|
||||
|
||||
async def identity(q: str) -> str:
|
||||
return answer
|
||||
|
||||
report = await dataset.evaluate(identity, max_concurrency=1, progress=False)
|
||||
case_report = report.cases[0]
|
||||
|
||||
score_result = case_report.scores.get("accuracy")
|
||||
if score_result is not None:
|
||||
score = float(score_result.value)
|
||||
reason = getattr(score_result, "reason", None)
|
||||
return score, reason
|
||||
|
||||
return 0.0, None
|
||||
|
||||
def make_reflective_dataset(
|
||||
self,
|
||||
candidate: dict[str, str],
|
||||
eval_batch: EvaluationBatch[EvalTrajectory, str | None],
|
||||
components_to_update: list[str],
|
||||
) -> Mapping[str, Sequence[Mapping[str, Any]]]:
|
||||
if eval_batch.trajectories is None:
|
||||
return {}
|
||||
|
||||
records: list[dict[str, Any]] = []
|
||||
for traj in eval_batch.trajectories:
|
||||
records.append(
|
||||
{
|
||||
"Inputs": {"question": traj.question},
|
||||
"Generated Outputs": {
|
||||
"answer": traj.actual_answer or "(no answer)"
|
||||
},
|
||||
"Feedback": (
|
||||
f"Expected answer: {traj.expected_answer}\n"
|
||||
f"Score: {traj.score:.2f}\n"
|
||||
f"Judge reasoning: {traj.judge_reason or 'N/A'}"
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
return {"instructions": records}
|
||||
|
||||
propose_new_texts = None
|
||||
|
||||
|
||||
class ReflectionLM:
|
||||
"""LanguageModel implementation for GEPA's ReflectiveMutationProposer.
|
||||
|
||||
Wraps a pydantic-ai Agent to satisfy GEPA's LanguageModel protocol.
|
||||
"""
|
||||
|
||||
def __init__(self, model_config: ModelConfig, config: AppConfig) -> None:
|
||||
from pydantic_ai import Agent
|
||||
|
||||
model = get_model(model_config, config)
|
||||
self._agent: Agent[None, str] = Agent(model=model, output_type=str)
|
||||
|
||||
def __call__(self, prompt: str | list[dict[str, Any]]) -> str:
|
||||
if isinstance(prompt, list):
|
||||
text = "\n".join(
|
||||
f"{msg.get('role', 'user')}: {msg.get('content', '')}" for msg in prompt
|
||||
)
|
||||
else:
|
||||
text = prompt
|
||||
result = self._agent.run_sync(text)
|
||||
return result.output
|
||||
|
||||
|
||||
async def run_optimization(
|
||||
spec: DatasetSpec,
|
||||
config: AppConfig,
|
||||
cases: list[QACase],
|
||||
max_calls: int,
|
||||
db_path: Path | None = None,
|
||||
output: Path | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Run GEPA optimization and return results summary."""
|
||||
from rich.console import Console
|
||||
|
||||
console = Console()
|
||||
|
||||
judge_config = ModelConfig(
|
||||
provider="ollama", name="gpt-oss", enable_thinking=False, temperature=0.0
|
||||
)
|
||||
judge_model = get_model(judge_config, config)
|
||||
|
||||
db = spec.db_path(db_path)
|
||||
adapter = QAPromptAdapter(
|
||||
config=config,
|
||||
db_path=db,
|
||||
judge_model=judge_model,
|
||||
)
|
||||
|
||||
reflection_lm = ReflectionLM(config.qa.model, config)
|
||||
|
||||
seed_prompt = spec.system_prompt or QA_SYSTEM_PROMPT
|
||||
seed_candidate = {"instructions": seed_prompt}
|
||||
|
||||
console.print(f"Optimizing prompt for dataset: {spec.key}", style="bold magenta")
|
||||
console.print(f"QA cases: {len(cases)}, Max metric calls: {max_calls}")
|
||||
console.print(f"Seed prompt length: {len(seed_prompt)} chars")
|
||||
|
||||
from gepa import optimize as gepa_optimize
|
||||
|
||||
result = gepa_optimize(
|
||||
seed_candidate=seed_candidate,
|
||||
trainset=cases,
|
||||
valset=cases,
|
||||
adapter=adapter,
|
||||
reflection_lm=reflection_lm,
|
||||
max_metric_calls=max_calls,
|
||||
display_progress_bar=True,
|
||||
)
|
||||
|
||||
best_score = result.val_aggregate_scores[result.best_idx]
|
||||
best_prompt = result.best_candidate
|
||||
if isinstance(best_prompt, dict):
|
||||
best_prompt = best_prompt["instructions"]
|
||||
total_calls = result.total_metric_calls or "unknown"
|
||||
|
||||
console.print("\n=== Optimization Results ===", style="bold cyan")
|
||||
console.print(f"Total metric calls: {total_calls}")
|
||||
console.print(f"Candidates explored: {result.num_candidates}")
|
||||
console.print(f"Best score: {best_score:.4f}")
|
||||
console.print(f"\nOptimized prompt:\n{best_prompt}")
|
||||
|
||||
if output:
|
||||
output.write_text(best_prompt)
|
||||
console.print(f"\nSaved to: {output}", style="green")
|
||||
|
||||
return {
|
||||
"best_score": best_score,
|
||||
"best_prompt": best_prompt,
|
||||
"total_calls": total_calls,
|
||||
"num_candidates": result.num_candidates,
|
||||
}
|
||||
|
|
@ -14,6 +14,7 @@ dependencies = [
|
|||
"huggingface_hub>=0.20.0",
|
||||
"typer>=0.21.0,<0.22.0",
|
||||
"python-dotenv>=1.2.2",
|
||||
"gepa>=0.1.0",
|
||||
]
|
||||
|
||||
[project.scripts]
|
||||
|
|
|
|||
11
uv.lock
11
uv.lock
|
|
@ -1252,6 +1252,15 @@ wheels = [
|
|||
{ url = "https://files.pythonhosted.org/packages/ae/1c/d544657381270fe21afc2a4dc560e1607e28d7115555cef148271cb25522/genai_prices-0.0.53-py3-none-any.whl", hash = "sha256:5a5dfd92089e9e8a174f7097a1521e36f4e75c74cfbdfb1ec56283bae3c0c96e", size = 61850, upload-time = "2026-02-11T20:47:16.774Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "gepa"
|
||||
version = "0.1.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/f5/30/511e52916956508f56eca721260fcd524cfffd580e57782dd471be925f7e/gepa-0.1.0.tar.gz", hash = "sha256:f8b3d7918d4cdcf8593f39ef1cc757c4ba1a4e6793e3ffb622e6c0bc60a1efd9", size = 226064, upload-time = "2026-02-19T19:43:08.272Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/1f/32/fe8afb3d2a6605a6bcbc8f119f0a2adae96e9e5d57ebed055490219956a8/gepa-0.1.0-py3-none-any.whl", hash = "sha256:4e3f8fe8ca20169e60518b2e9d416e8c4a579459848adffdcad12223fbf9643e", size = 191392, upload-time = "2026-02-19T19:43:07.065Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ghp-import"
|
||||
version = "2.1.0"
|
||||
|
|
@ -1413,6 +1422,7 @@ version = "0.33.2"
|
|||
source = { editable = "evaluations" }
|
||||
dependencies = [
|
||||
{ name = "datasets" },
|
||||
{ name = "gepa" },
|
||||
{ name = "haiku-rag-slim" },
|
||||
{ name = "huggingface-hub" },
|
||||
{ name = "pydantic-ai-slim", extra = ["evals", "logfire"] },
|
||||
|
|
@ -1423,6 +1433,7 @@ dependencies = [
|
|||
[package.metadata]
|
||||
requires-dist = [
|
||||
{ name = "datasets", specifier = ">=4.6.1" },
|
||||
{ name = "gepa", specifier = ">=0.1.0" },
|
||||
{ name = "haiku-rag-slim", editable = "haiku_rag_slim" },
|
||||
{ name = "huggingface-hub", specifier = ">=0.20.0" },
|
||||
{ name = "pydantic-ai-slim", extras = ["evals", "logfire"], specifier = ">=1.66.0" },
|
||||
|
|
|
|||
Loading…
Reference in a new issue