From 2064c699243311b9625ef3b400ee6fbd9d05145b Mon Sep 17 00:00:00 2001 From: orangels Date: Wed, 6 May 2026 16:59:24 +0800 Subject: [PATCH] fix: tighten benchmark dataset helpers --- src/benchmarks/judge.rs | 61 +++++++++++++++++++++++++++++++++-------- src/cli.rs | 39 +++++++++++++++++++++----- 2 files changed, 82 insertions(+), 18 deletions(-) diff --git a/src/benchmarks/judge.rs b/src/benchmarks/judge.rs index 0a1e20d..a079fad 100644 --- a/src/benchmarks/judge.rs +++ b/src/benchmarks/judge.rs @@ -14,19 +14,24 @@ pub fn extract_final_integer(text: &str) -> Option { } pub fn extract_choice(text: &str) -> Option { - let answer = Regex::new(r"(?i)\b(?:answer|choose|choice)\s*(?:is|:|-)?\s*\(?([A-D])\)?\b") - .expect("valid answer choice regex"); - if let Some(captures) = answer.captures(text) { - return captures - .get(1) - .and_then(|value| value.as_str().chars().next()) - .map(|choice| choice.to_ascii_uppercase()); + let explicit = Regex::new( + r"(?i)\b(?:final\s+answer|answer|choose|choice)\s*(?:is|:|-)?\s*\(?([A-D])\)?\b", + ) + .expect("valid explicit choice regex"); + if let Some(captures) = explicit.captures_iter(text).last() { + return choice_from_capture(&captures, 1); } - let standalone = Regex::new(r"(?i)\b([A-D])\b").expect("valid standalone choice regex"); - standalone - .captures(text) - .and_then(|captures| captures.get(1)) + let whole_output = + Regex::new(r"(?i)^(?:([A-D])|\(([A-D])\))$").expect("valid whole output choice regex"); + whole_output.captures(text.trim()).and_then(|captures| { + choice_from_capture(&captures, 1).or_else(|| choice_from_capture(&captures, 2)) + }) +} + +fn choice_from_capture(captures: ®ex::Captures<'_>, index: usize) -> Option { + captures + .get(index) .and_then(|value| value.as_str().chars().next()) .map(|choice| choice.to_ascii_uppercase()) } @@ -69,6 +74,40 @@ mod tests { assert_eq!(extract_choice("I choose (b)."), Some('B')); } + #[test] + fn extracts_choice_from_final_answer_instead_of_earlier_prose_option() { + assert_eq!( + extract_choice( + "Option A is tempting, but after checking the calculation the answer is D" + ), + Some('D') + ); + } + + #[test] + fn extracts_choice_from_final_answer_after_considered_options() { + assert_eq!( + extract_choice("I considered A and B before deciding. Final answer: C"), + Some('C') + ); + } + + #[test] + fn extracts_last_explicit_choice() { + assert_eq!( + extract_choice("Answer: A was my first thought. Final answer: D"), + Some('D') + ); + } + + #[test] + fn ignores_prose_without_final_choice() { + assert_eq!( + extract_choice("This explanation has no final option."), + None + ); + } + #[test] fn judges_integer_by_extracted_value() { assert!(judge_integer("Final: 42.", "42")); diff --git a/src/cli.rs b/src/cli.rs index c14eea5..a58cdb2 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -90,7 +90,7 @@ pub async fn dispatch(cli: Cli) -> Result<()> { Command::Dataset { command: DatasetCommand::Fetch { dataset }, } => { - let data_dir = dataset_data_dir(Path::new("config.yaml")); + let data_dir = dataset_data_dir(Path::new("config.yaml"))?; let path = benchmarks::fetch_dataset(&dataset, &data_dir).await?; println!("{}", path.display()); Ok(()) @@ -100,12 +100,37 @@ pub async fn dispatch(cli: Cli) -> Result<()> { } } -fn dataset_data_dir(config_path: &Path) -> PathBuf { - if config_path.exists() { - if let Ok(config) = AppConfig::load(config_path) { - return PathBuf::from(config.benchmarks.data_dir); - } +fn dataset_data_dir(config_path: &Path) -> Result { + if !config_path.exists() { + return Ok(PathBuf::from("data/benchmarks")); } - PathBuf::from("data/benchmarks") + let config = AppConfig::load(config_path)?; + Ok(PathBuf::from(config.benchmarks.data_dir)) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn dataset_data_dir_defaults_when_config_is_missing() { + let temp_dir = tempfile::tempdir().expect("create temp dir"); + let missing_config = temp_dir.path().join("missing-config.yaml"); + + let data_dir = dataset_data_dir(&missing_config).expect("default data dir"); + + assert_eq!(data_dir, PathBuf::from("data/benchmarks")); + } + + #[test] + fn dataset_data_dir_propagates_invalid_existing_config() { + let temp_dir = tempfile::tempdir().expect("create temp dir"); + let config_path = temp_dir.path().join("config.yaml"); + std::fs::write(&config_path, "providers: [").expect("write invalid config"); + + let error = dataset_data_dir(&config_path).expect_err("invalid config should fail"); + + assert!(error.to_string().contains("failed to parse config")); + } }