Skills › Scientific Agent Skills › pytdc › benchmark_evaluation.py
benchmark_evaluation.py
381 lines · skills/pytdc/scripts/benchmark_evaluation.py at 5f2bc03
1#!/usr/bin/env python3
2"""Plan or explicitly evaluate user-supplied TDC benchmark predictions."""
4from __future__ import annotations
6import argparse
7import importlib
8import math
9import sys
10from pathlib import Path
11from typing import Any
13from _common import (
14 CliError,
15 bounded_int,
16 canonical_name,
17 emit_json,
18 load_pytdc_metadata,
19 read_json_file,
20 safe_directory,
21 safe_input_file,
22)
25# Docking is intentionally excluded: its evaluation can invoke specialized oracles.
26GROUPS: dict[str, tuple[str, str]] = {
27 "admet_group": ("tdc.benchmark_group", "admet_group"),
28 "drugcombo_group": ("tdc.benchmark_group", "drugcombo_group"),
29 "dti_dg_group": ("tdc.benchmark_group", "dti_dg_group"),
30}
32MAX_PREDICTION_FILE_BYTES = 50 * 1024 * 1024
33MAX_PREDICTION_VALUES = 5_000_000
34MAX_RUNS = 100
37def benchmark_names(metadata: Any, group: str) -> list[str]:
38 return [
39 name
40 for names in metadata.benchmark_names[group].values()
41 for name in names
42 ]
45def validate_seeds(seeds: list[int]) -> list[int]:
46 if not seeds:
47 raise CliError("at least one seed is required")
48 if len(seeds) > MAX_RUNS:
49 raise CliError(f"at most {MAX_RUNS} seeds are allowed")
50 if len(set(seeds)) != len(seeds):
51 raise CliError("seeds must be unique")
52 if any(seed < 0 for seed in seeds):
53 raise CliError("seeds must be non-negative")
54 return seeds
57def _validate_values(values: Any, dataset: str) -> list[float]:
58 if not isinstance(values, list) or not values:
59 raise CliError(f"predictions for {dataset} must be a non-empty JSON array")
60 checked: list[float] = []
61 for value in values:
62 if isinstance(value, bool) or not isinstance(value, (int, float)):
63 raise CliError(f"predictions for {dataset} must contain only numbers")
64 number = float(value)
65 if not math.isfinite(number):
66 raise CliError(f"predictions for {dataset} must be finite")
67 checked.append(number)
68 return checked
71def normalize_prediction_mapping(
72 value: Any,
73 *,
74 available: list[str],
75 selected_dataset: str | None,
76) -> tuple[dict[str, list[float]], int]:
77 if not isinstance(value, dict) or not value:
78 raise CliError("each prediction run must be a non-empty JSON object")
79 normalized: dict[str, list[float]] = {}
80 total = 0
81 for query, predictions in value.items():
82 if not isinstance(query, str):
83 raise CliError("benchmark names in prediction objects must be strings")
84 dataset = canonical_name(query, available, "benchmark")
85 if dataset in normalized:
86 raise CliError(f"duplicate benchmark after name normalization: {dataset!r}")
87 if selected_dataset is not None and dataset != selected_dataset:
88 raise CliError(
89 f"prediction object contains {dataset!r}, but --dataset selects "
90 f"{selected_dataset!r}"
91 )
92 checked = _validate_values(predictions, dataset)
93 normalized[dataset] = checked
94 total += len(checked)
95 if selected_dataset is not None and selected_dataset not in normalized:
96 raise CliError(f"prediction object is missing {selected_dataset!r}")
97 return normalized, total
100def normalize_predictions(
101 payload: Any,
102 *,
103 mode: str,
104 available: list[str],
105 selected_dataset: str | None,
106) -> tuple[dict[str, list[float]] | list[dict[str, list[float]]], list[int | None]]:
107 """Validate the documented single-run or evaluate_many JSON shape."""
109 if mode == "single":
110 if isinstance(payload, dict) and "predictions" in payload:
111 payload = payload["predictions"]
112 mapping, total = normalize_prediction_mapping(
113 payload, available=available, selected_dataset=selected_dataset
114 )
115 if total > MAX_PREDICTION_VALUES:
116 raise CliError("prediction value limit exceeded")
117 return mapping, [None]
119 runs = payload.get("runs") if isinstance(payload, dict) else payload
120 if not isinstance(runs, list):
121 raise CliError(
122 "many-run JSON must be a list or an object with a `runs` list"
123 )
124 if not 5 <= len(runs) <= MAX_RUNS:
125 raise CliError(
126 "evaluate_many requires at least five and at most "
127 f"{MAX_RUNS} prediction runs"
128 )
130 normalized_runs: list[dict[str, list[float]]] = []
131 run_seeds: list[int | None] = []
132 total = 0
133 for index, run in enumerate(runs):
134 if not isinstance(run, dict):
135 raise CliError(f"run {index} must be a JSON object")
136 if "predictions" in run:
137 predictions = run["predictions"]
138 seed = run.get("seed")
139 if seed is not None and (
140 isinstance(seed, bool) or not isinstance(seed, int) or seed < 0
141 ):
142 raise CliError(f"run {index} seed must be a non-negative integer")
143 else:
144 predictions = run
145 seed = None
146 mapping, count = normalize_prediction_mapping(
147 predictions, available=available, selected_dataset=selected_dataset
148 )
149 if normalized_runs:
150 if mapping.keys() != normalized_runs[0].keys():
151 raise CliError("every run must contain the same benchmarks")
152 if any(len(mapping[name]) != len(normalized_runs[0][name]) for name in mapping):
153 raise CliError("each benchmark must have the same prediction count in every run")
154 normalized_runs.append(mapping)
155 run_seeds.append(seed)
156 total += count
157 if total > MAX_PREDICTION_VALUES:
158 raise CliError("prediction value limit exceeded")
160 supplied_seeds = [seed for seed in run_seeds if seed is not None]
161 if supplied_seeds and len(supplied_seeds) != len(run_seeds):
162 raise CliError("either provide a seed for every run or for none")
163 if len(set(supplied_seeds)) != len(supplied_seeds):
164 raise CliError("run seeds must be unique")
165 return normalized_runs, run_seeds
168def build_plan(
169 *,
170 group: str,
171 dataset: str | None,
172 seeds: list[int],
173 data_dir: Path,
174 package_version: str,
175 metric: str | None,
176 prediction_summary: dict[str, Any] | None,
177) -> dict[str, Any]:
178 return {
179 "action": "plan",
180 "acknowledgement_required": "--execute",
181 "benchmark_group": group,
182 "data_directory": str(data_dir),
183 "dataset": dataset,
184 "download_performed": False,
185 "metric_from_package_metadata": metric,
186 "network_and_storage": (
187 "Constructing a BenchmarkGroup downloads/extracts the group archive "
188 "under data_directory if no local copy is found."
189 ),
190 "package": "PyTDC",
191 "package_version": package_version,
192 "prediction_input": prediction_summary,
193 "protocol_note": (
194 "TDC documents at least five independent runs for leaderboard "
195 "mean/standard-deviation reporting; these seeds label the plan only."
196 ),
197 "seeds": seeds,
198 }
201def execute_evaluation(
202 *,
203 group: str,
204 dataset: str | None,
205 mode: str,
206 predictions: dict[str, list[float]] | list[dict[str, list[float]]],
207 run_seeds: list[int | None],
208 data_dir: Path,
209 package_version: str,
210) -> dict[str, Any]:
211 module_name, class_name = GROUPS[group]
212 group_class = getattr(importlib.import_module(module_name), class_name)
213 benchmark_group = group_class(path=str(data_dir))
214 runs = [predictions] if mode == "single" else predictions
215 for name, values in runs[0].items():
216 expected = len(benchmark_group.get(name)["test"])
217 if any(len(run[name]) != expected for run in runs):
218 raise CliError(
219 f"{name} requires {expected} test predictions in exact test-row order"
220 )
221 if mode == "single":
222 results = benchmark_group.evaluate(predictions)
223 run_count = 1
224 else:
225 results = benchmark_group.evaluate_many(predictions)
226 run_count = len(predictions)
227 if isinstance(results, Exception):
228 raise CliError(str(results))
229 return {
230 "action": "executed",
231 "benchmark_group": group,
232 "data_directory": str(data_dir),
233 "dataset": dataset,
234 "download_acknowledged": True,
235 "mode": mode,
236 "package": "PyTDC",
237 "package_version": package_version,
238 "results": results,
239 "run_count": run_count,
240 "run_seeds": run_seeds,
241 }
244def build_parser() -> argparse.ArgumentParser:
245 parser = argparse.ArgumentParser(
246 description=(
247 "Plan a non-docking TDC BenchmarkGroup evaluation, or pass --execute "
248 "with bounded JSON predictions to acknowledge group download/cache use. "
249 "No dummy labels or predictions are generated."
250 )
251 )
252 parser.add_argument(
253 "--group",
254 choices=tuple(GROUPS),
255 default="admet_group",
256 help="benchmark group (default: admet_group)",
257 )
258 parser.add_argument("--dataset", help="optional exact benchmark name")
259 parser.add_argument(
260 "--mode",
261 choices=("single", "many"),
262 default="many",
263 help="call evaluate or evaluate_many (default: many)",
264 )
265 parser.add_argument(
266 "--seeds",
267 nargs="+",
268 type=bounded_int(0, 4_294_967_295),
269 default=[1, 2, 3, 4, 5],
270 help="independent run seeds for the plan (default: 1 2 3 4 5)",
271 )
272 parser.add_argument(
273 "--predictions",
274 help=(
275 "bounded JSON input. Single mode: {dataset: [values]}. Many mode: "
276 "[{dataset: [values]}, ...] or {runs: [{seed, predictions}, ...]}"
277 ),
278 )
279 parser.add_argument(
280 "--data-dir",
281 default=".pytdc-benchmarks",
282 help="relative benchmark cache directory (default: .pytdc-benchmarks)",
283 )
284 parser.add_argument(
285 "--execute",
286 action="store_true",
287 help="acknowledge network/storage use and evaluate supplied predictions",
288 )
289 parser.add_argument("--output", help="write JSON to a relative workspace path")
290 parser.add_argument(
291 "--force", action="store_true", help="replace an existing --output file"
292 )
293 return parser
296def main(argv: list[str] | None = None) -> int:
297 parser = build_parser()
298 args = parser.parse_args(argv)
299 if args.execute and not args.predictions:
300 parser.error("--execute requires --predictions")
302 try:
303 seeds = validate_seeds(args.seeds)
304 metadata, package_version = load_pytdc_metadata()
305 group = canonical_name(args.group, GROUPS, "benchmark group")
306 registry_group = canonical_name(
307 group, metadata.benchmark_names, "package benchmark group"
308 )
309 available = benchmark_names(metadata, registry_group)
310 dataset = (
311 canonical_name(args.dataset, available, "benchmark")
312 if args.dataset
313 else None
314 )
315 metric = (
316 metadata.bm_metric_names.get(registry_group, {}).get(dataset)
317 if dataset
318 else None
319 )
320 data_dir = safe_directory(
321 args.data_dir,
322 label="benchmark data directory",
323 create=args.execute,
324 )
326 normalized: (
327 dict[str, list[float]] | list[dict[str, list[float]]] | None
328 ) = None
329 run_seeds: list[int | None] = []
330 prediction_summary: dict[str, Any] | None = None
331 if args.predictions:
332 prediction_path = safe_input_file(
333 args.predictions,
334 max_bytes=MAX_PREDICTION_FILE_BYTES,
335 label="prediction input",
336 )
337 payload = read_json_file(prediction_path)
338 normalized, run_seeds = normalize_predictions(
339 payload,
340 mode=args.mode,
341 available=available,
342 selected_dataset=dataset,
343 )
344 prediction_summary = {
345 "mode": args.mode,
346 "path": str(prediction_path),
347 "run_count": 1 if args.mode == "single" else len(normalized),
348 "validated": True,
349 }
351 if args.execute:
352 if normalized is None:
353 raise CliError("--execute requires validated predictions")
354 result = execute_evaluation(
355 group=group,
356 dataset=dataset,
357 mode=args.mode,
358 predictions=normalized,
359 run_seeds=run_seeds,
360 data_dir=data_dir,
361 package_version=package_version,
362 )
363 else:
364 result = build_plan(
365 group=group,
366 dataset=dataset,
367 seeds=seeds,
368 data_dir=data_dir,
369 package_version=package_version,
370 metric=metric,
371 prediction_summary=prediction_summary,
372 )
373 emit_json(result, args.output, force=args.force)
374 except (CliError, ImportError, OSError, TypeError, ValueError) as exc:
375 print(f"error: {exc}", file=sys.stderr)
376 return 2
377 return 0
380if __name__ == "__main__":
381 raise SystemExit(main())