From 8bac8b4104381ceec7d253ed20596cdeeb0fed4b Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 9 Mar 2026 16:18:25 +0200 Subject: [PATCH] Properly choose training/eval set --- CHANGELOG.md | 5 ++ docs/tuning.md | 10 ++-- evaluations/evaluations/benchmark.py | 10 +++- evaluations/evaluations/optimization.py | 26 +++++++-- evaluations/pyproject.toml | 4 ++ evaluations/tests/test_optimization.py | 75 ++++++++++++++++--------- 6 files changed, 90 insertions(+), 40 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index fcc9fecc..6a4f75a7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,11 @@ - **GEPA prompt optimization**: `evaluations optimize` command for automated QA system prompt improvement using evolutionary optimization with LLM-judged scoring - **Tuning docs**: Added step 7 (Optimize QA Prompts) to the tuning workflow in `docs/tuning.md` +- **Evaluations test coverage**: Tests for evaluators (MAP, MRR), config, benchmark helpers, dataset mappers/builders, and optimization + +### Changed + +- **Optimization train/val split**: Cases are now split 50/50 into train and val sets. GEPA budget is auto-computed from `--iterations` and dataset size, replacing the broken `--max-calls` parameter that counted per-example evaluations ### Fixed diff --git a/docs/tuning.md b/docs/tuning.md index 77cbac6b..00b4824f 100644 --- a/docs/tuning.md +++ b/docs/tuning.md @@ -257,8 +257,8 @@ Once retrieval is tuned (steps 1-6), you can automatically optimize the QA syste # Basic optimization against a dataset evaluations optimize wix -# Limit QA cases and optimization budget -evaluations optimize repliqa --limit 20 --max-calls 30 +# Limit QA cases and iteration count +evaluations optimize repliqa --limit 40 --iterations 30 # Save the optimized prompt to a file evaluations optimize wix --output optimized_prompt.txt @@ -269,13 +269,13 @@ evaluations optimize wix --config haiku.rag.yaml --db /path/to/wix.lancedb | Option | Default | Description | |--------|---------|-------------| -| `--limit` | all cases | Number of QA cases to use for optimization | -| `--max-calls` | 50 | Maximum GEPA metric calls (optimization budget) | +| `--limit` | all cases | Number of QA cases (split 50/50 into train/val) | +| `--iterations` | 50 | Number of optimization iterations | | `--output` | — | Save optimized prompt to a file | | `--config` | auto | Path to haiku.rag YAML config file | | `--db` | auto | Override the database path | -**Cost note:** Each metric call evaluates a minibatch of 3 QA cases, requiring 3 QA calls plus 3 judge calls per batch. With `--max-calls 50`, expect 300+ LLM calls total. Start with `--limit 10 --max-calls 10` to verify your setup before running a full optimization. +**Cost note:** Each iteration evaluates a minibatch of 3 QA cases twice (current + mutant), plus a full valset evaluation on accepted mutations. The GEPA budget is computed automatically from `--iterations` and dataset size. Start with `--limit 20 --iterations 10` to verify your setup before running a full optimization. **Applying the result:** Use `--output` to save the optimized prompt, then set it in your config: diff --git a/evaluations/evaluations/benchmark.py b/evaluations/evaluations/benchmark.py index ab66dd16..76ea0d9b 100644 --- a/evaluations/evaluations/benchmark.py +++ b/evaluations/evaluations/benchmark.py @@ -473,8 +473,12 @@ def optimize( 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."), + limit: int | None = typer.Option( + None, "--limit", help="Limit QA cases (split 50/50 into train/val)." + ), + iterations: int = typer.Option( + 50, "--iterations", help="Number of optimization iterations." + ), output: Path | None = typer.Option( None, "--output", help="Save optimized prompt to file." ), @@ -498,7 +502,7 @@ def optimize( spec=spec, config=app_config, cases=cases, - max_calls=max_calls, + iterations=iterations, db_path=db, output=output, ) diff --git a/evaluations/evaluations/optimization.py b/evaluations/evaluations/optimization.py index cd65e9f4..468b2f04 100644 --- a/evaluations/evaluations/optimization.py +++ b/evaluations/evaluations/optimization.py @@ -207,11 +207,14 @@ class ReflectionLM: return result.output +REFLECTION_MINIBATCH_SIZE = 3 + + def run_optimization( spec: DatasetSpec, config: AppConfig, cases: list[QACase], - max_calls: int, + iterations: int, db_path: Path | None = None, output: Path | None = None, ) -> dict[str, Any]: @@ -237,19 +240,32 @@ def run_optimization( seed_prompt = spec.system_prompt or QA_SYSTEM_PROMPT seed_candidate = {"instructions": seed_prompt} + mid = len(cases) // 2 + trainset = cases[:mid] + valset = cases[mid:] + + # Budget: initial valset eval + worst-case iterations + # (each iteration: 2 minibatch evals + full valset if accepted) + max_metric_calls = len(valset) + iterations * ( + 2 * REFLECTION_MINIBATCH_SIZE + len(valset) + ) + 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"Train: {len(trainset)}, Val: {len(valset)}, " + f"Iterations: {iterations}, GEPA budget: {max_metric_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, + trainset=trainset, + valset=valset, adapter=adapter, reflection_lm=reflection_lm, - max_metric_calls=max_calls, + max_metric_calls=max_metric_calls, display_progress_bar=True, ) diff --git a/evaluations/pyproject.toml b/evaluations/pyproject.toml index ad271ad8..ddb1c655 100644 --- a/evaluations/pyproject.toml +++ b/evaluations/pyproject.toml @@ -27,5 +27,9 @@ build-backend = "hatchling.build" [tool.hatch.build.targets.wheel] packages = ["evaluations"] +[tool.pytest.ini_options] +asyncio_mode = "auto" +testpaths = ["tests"] + [tool.ruff.lint.isort] known-first-party = ["haiku", "evaluations"] diff --git a/evaluations/tests/test_optimization.py b/evaluations/tests/test_optimization.py index c44c23c1..ea53dee1 100644 --- a/evaluations/tests/test_optimization.py +++ b/evaluations/tests/test_optimization.py @@ -294,6 +294,18 @@ class TestProposalAttribute: assert adapter.propose_new_texts is None +def _make_cases(n: int) -> list[Case[str, str, dict[str, str]]]: + return [ + Case( + name=f"q{i}", + inputs=f"Question {i}?", + expected_output=f"Answer {i}.", + metadata={"case_index": str(i)}, + ) + for i in range(1, n + 1) + ] + + class TestRunOptimization: def _make_spec(self, db_path: Path) -> DatasetSpec: return DatasetSpec( @@ -308,14 +320,7 @@ class TestRunOptimization: def test_returns_results(self, tmp_path: Path) -> None: spec = self._make_spec(tmp_path / "test.lancedb") - cases: list[Case[str, str, dict[str, str]]] = [ - Case( - name="q1", - inputs="Q?", - expected_output="A.", - metadata={"case_index": "1"}, - ), - ] + cases = _make_cases(4) mock_result = MagicMock() mock_result.best_idx = 0 @@ -333,7 +338,7 @@ class TestRunOptimization: spec=spec, config=AppConfig(), cases=cases, - max_calls=10, + iterations=10, db_path=tmp_path / "test.lancedb", ) @@ -344,14 +349,7 @@ class TestRunOptimization: def test_saves_output_file(self, tmp_path: Path) -> None: spec = self._make_spec(tmp_path / "test.lancedb") - cases: list[Case[str, str, dict[str, str]]] = [ - Case( - name="q1", - inputs="Q?", - expected_output="A.", - metadata={"case_index": "1"}, - ), - ] + cases = _make_cases(4) output_path = tmp_path / "prompt.txt" mock_result = MagicMock() @@ -370,7 +368,7 @@ class TestRunOptimization: spec=spec, config=AppConfig(), cases=cases, - max_calls=5, + iterations=5, db_path=tmp_path / "test.lancedb", output=output_path, ) @@ -386,14 +384,7 @@ class TestRunOptimization: qa_loader=lambda: None, # type: ignore[return-value] qa_case_builder=lambda idx, doc: None, # type: ignore[return-value] ) - cases: list[Case[str, str, dict[str, str]]] = [ - Case( - name="q1", - inputs="Q?", - expected_output="A.", - metadata={"case_index": "1"}, - ), - ] + cases = _make_cases(4) mock_result = MagicMock() mock_result.best_idx = 0 @@ -411,7 +402,7 @@ class TestRunOptimization: spec=spec, config=AppConfig(), cases=cases, - max_calls=1, + iterations=1, db_path=tmp_path / "test.lancedb", ) @@ -423,3 +414,33 @@ class TestRunOptimization: seed = call_kwargs["seed_candidate"] assert seed["instructions"] is not None assert len(seed["instructions"]) > 0 + + def test_splits_cases_into_train_and_val(self, tmp_path: Path) -> None: + spec = self._make_spec(tmp_path / "test.lancedb") + cases = _make_cases(10) + + mock_result = MagicMock() + mock_result.best_idx = 0 + mock_result.val_aggregate_scores = [0.7] + mock_result.best_candidate = {"instructions": "prompt"} + mock_result.total_metric_calls = 50 + mock_result.num_candidates = 1 + + with ( + patch("evaluations.optimization.get_model"), + patch("evaluations.optimization.ReflectionLM"), + patch("gepa.optimize", return_value=mock_result) as mock_gepa, + ): + run_optimization( + spec=spec, + config=AppConfig(), + cases=cases, + iterations=5, + db_path=tmp_path / "test.lancedb", + ) + + call_kwargs = mock_gepa.call_args[1] + assert len(call_kwargs["trainset"]) == 5 + assert len(call_kwargs["valset"]) == 5 + # Budget = valset_size + iterations * (2*minibatch + valset_size) + assert call_kwargs["max_metric_calls"] == 5 + 5 * (2 * 3 + 5)