diff --git a/src/benchmarks/aime.rs b/src/benchmarks/aime.rs index 8b13789..50e9a21 100644 --- a/src/benchmarks/aime.rs +++ b/src/benchmarks/aime.rs @@ -1 +1,34 @@ +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AimeCase { + pub id: String, + pub problem: String, + pub answer: String, +} +impl AimeCase { + pub fn prompt(&self) -> String { + format!( + "Solve the following AIME problem.\n\n{}\n\nRespond with the final integer answer.", + self.problem + ) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn prompt_contains_problem_and_final_integer_instruction() { + let case = AimeCase { + id: "aime-1".to_string(), + problem: "What is 20 + 22?".to_string(), + answer: "42".to_string(), + }; + + let prompt = case.prompt(); + + assert!(prompt.contains("What is 20 + 22?")); + assert!(prompt.contains("final integer answer")); + } +} diff --git a/src/benchmarks/gpqa.rs b/src/benchmarks/gpqa.rs index 8b13789..bc11004 100644 --- a/src/benchmarks/gpqa.rs +++ b/src/benchmarks/gpqa.rs @@ -1 +1,52 @@ +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct GpqaCase { + pub id: String, + pub question: String, + pub choices: [(String, String); 4], + pub answer: char, +} +impl GpqaCase { + pub fn prompt(&self) -> String { + let choices = self + .choices + .iter() + .map(|(label, text)| format!("{label}. {text}")) + .collect::>() + .join("\n"); + + format!( + "Answer the following multiple-choice question.\n\n{}\n\n{}\n\nPlease answer with exactly one letter: A, B, C, or D.", + self.question, choices + ) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn prompt_contains_choices_and_single_letter_instruction() { + let case = GpqaCase { + id: "gpqa-1".to_string(), + question: "Which option is correct?".to_string(), + choices: [ + ("A".to_string(), "Alpha".to_string()), + ("B".to_string(), "Beta".to_string()), + ("C".to_string(), "Gamma".to_string()), + ("D".to_string(), "Delta".to_string()), + ], + answer: 'C', + }; + + let prompt = case.prompt(); + + assert!(prompt.contains("Which option is correct?")); + assert!(prompt.contains("A. Alpha")); + assert!(prompt.contains("B. Beta")); + assert!(prompt.contains("C. Gamma")); + assert!(prompt.contains("D. Delta")); + assert!(prompt.contains("answer with exactly one letter")); + } +} diff --git a/src/benchmarks/judge.rs b/src/benchmarks/judge.rs index 8b13789..0a1e20d 100644 --- a/src/benchmarks/judge.rs +++ b/src/benchmarks/judge.rs @@ -1 +1,83 @@ +use regex::Regex; +pub fn extract_final_integer(text: &str) -> Option { + let boxed = Regex::new(r"\\boxed\{\s*(-?\d+)\s*\}").expect("valid boxed integer regex"); + if let Some(captures) = boxed.captures(text) { + return captures.get(1).map(|value| value.as_str().to_string()); + } + + let integer = Regex::new(r"-?\d+").expect("valid integer regex"); + integer + .find_iter(text) + .last() + .map(|value| value.as_str().to_string()) +} + +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 standalone = Regex::new(r"(?i)\b([A-D])\b").expect("valid standalone choice regex"); + standalone + .captures(text) + .and_then(|captures| captures.get(1)) + .and_then(|value| value.as_str().chars().next()) + .map(|choice| choice.to_ascii_uppercase()) +} + +pub fn judge_integer(output: &str, expected: &str) -> bool { + extract_final_integer(output).is_some_and(|actual| actual == expected.trim()) +} + +pub fn judge_choice(output: &str, expected: char) -> bool { + extract_choice(output).is_some_and(|actual| actual == expected.to_ascii_uppercase()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn extracts_final_integer_from_plain_answer() { + assert_eq!( + extract_final_integer("The answer is 42."), + Some("42".to_string()) + ); + } + + #[test] + fn extracts_final_integer_from_boxed_answer() { + assert_eq!( + extract_final_integer(r"Final: \boxed{17}"), + Some("17".to_string()) + ); + } + + #[test] + fn extracts_choice_from_answer_prefix() { + assert_eq!(extract_choice("Answer: C"), Some('C')); + } + + #[test] + fn extracts_choice_from_parenthesized_lowercase() { + assert_eq!(extract_choice("I choose (b)."), Some('B')); + } + + #[test] + fn judges_integer_by_extracted_value() { + assert!(judge_integer("Final: 42.", "42")); + assert!(!judge_integer("Final: 43.", "42")); + } + + #[test] + fn judges_choice_by_extracted_value() { + assert!(judge_choice("I choose (b).", 'B')); + assert!(!judge_choice("Answer: C", 'B')); + } +} diff --git a/src/benchmarks/mod.rs b/src/benchmarks/mod.rs index 2bffcce..8de2167 100644 --- a/src/benchmarks/mod.rs +++ b/src/benchmarks/mod.rs @@ -1,3 +1,122 @@ pub mod aime; pub mod gpqa; pub mod judge; + +use anyhow::{Context, Result, bail}; +use chrono::Utc; +use serde::Serialize; +use sha2::{Digest, Sha256}; +use std::path::{Path, PathBuf}; +use tokio::fs; + +const AIME_2026_URL: &str = "https://huggingface.co/datasets/MathArena/aime_2026/resolve/main/data/train-00000-of-00001.parquet"; +const GPQA_DIAMOND_URL: &str = + "https://huggingface.co/datasets/Idavidrein/gpqa/resolve/main/gpqa_diamond.csv"; + +#[derive(Debug, Serialize)] +struct DatasetMetadata<'a> { + name: &'a str, + source_url: &'a str, + downloaded_at: String, + sha256: String, + bytes: usize, +} + +pub async fn fetch_dataset(dataset: &str, data_dir: &Path) -> Result { + let spec = dataset_spec(dataset, data_dir)?; + fs::create_dir_all(&spec.dir) + .await + .with_context(|| format!("failed to create dataset directory {}", spec.dir.display()))?; + + let client = reqwest::Client::new(); + let mut request = client.get(spec.source_url); + if spec.use_hf_token { + if let Ok(token) = std::env::var("HF_TOKEN") { + request = request.bearer_auth(token); + } + } + + let response = request + .send() + .await + .with_context(|| format!("failed to download {dataset} from {}", spec.source_url))?; + + let status = response.status(); + if !status.is_success() { + let auth_hint = if spec.use_hf_token && std::env::var("HF_TOKEN").is_err() { + " Set HF_TOKEN if the Hugging Face dataset requires authentication." + } else { + "" + }; + bail!( + "failed to download {dataset} from {}: HTTP {status}.{auth_hint}", + spec.source_url + ); + } + + let bytes = response + .bytes() + .await + .with_context(|| format!("failed to read downloaded bytes for {dataset}"))?; + fs::write(&spec.file_path, &bytes) + .await + .with_context(|| format!("failed to write dataset file {}", spec.file_path.display()))?; + + let metadata = DatasetMetadata { + name: dataset, + source_url: spec.source_url, + downloaded_at: Utc::now().to_rfc3339(), + sha256: sha256_hex(&bytes), + bytes: bytes.len(), + }; + let metadata_yaml = + serde_yaml::to_string(&metadata).context("failed to serialize dataset metadata")?; + let metadata_path = spec.dir.join("metadata.yaml"); + fs::write(&metadata_path, metadata_yaml) + .await + .with_context(|| format!("failed to write metadata file {}", metadata_path.display()))?; + + Ok(spec.file_path) +} + +struct DatasetSpec<'a> { + source_url: &'a str, + dir: PathBuf, + file_path: PathBuf, + use_hf_token: bool, +} + +fn dataset_spec<'a>(dataset: &str, data_dir: &'a Path) -> Result> { + match dataset { + "aime2026" => { + let dir = data_dir.join("aime2026"); + Ok(DatasetSpec { + source_url: AIME_2026_URL, + file_path: dir.join("train-00000-of-00001.parquet"), + dir, + use_hf_token: false, + }) + } + "gpqa-diamond" => { + let dir = data_dir.join("gpqa_diamond"); + Ok(DatasetSpec { + source_url: GPQA_DIAMOND_URL, + file_path: dir.join("gpqa_diamond.csv"), + dir, + use_hf_token: true, + }) + } + other => bail!( + "unknown dataset {other:?}; supported datasets are \"aime2026\" and \"gpqa-diamond\"" + ), + } +} + +fn sha256_hex(bytes: &[u8]) -> String { + let digest = Sha256::digest(bytes); + let mut hex = String::with_capacity(digest.len() * 2); + for byte in digest { + hex.push_str(&format!("{byte:02x}")); + } + hex +} diff --git a/src/cli.rs b/src/cli.rs index 6a841ed..c14eea5 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -1,6 +1,8 @@ +use crate::benchmarks; +use crate::config::AppConfig; use anyhow::Result; use clap::{Parser, Subcommand}; -use std::path::PathBuf; +use std::path::{Path, PathBuf}; #[derive(Debug, Parser)] #[command( @@ -85,8 +87,25 @@ pub enum BenchCommand { pub async fn dispatch(cli: Cli) -> Result<()> { match cli.command { Command::Check { .. } => anyhow::bail!("check is not implemented yet"), - Command::Dataset { .. } => anyhow::bail!("dataset is not implemented yet"), + Command::Dataset { + command: DatasetCommand::Fetch { dataset }, + } => { + let data_dir = dataset_data_dir(Path::new("config.yaml")); + let path = benchmarks::fetch_dataset(&dataset, &data_dir).await?; + println!("{}", path.display()); + Ok(()) + } Command::Bench { .. } => anyhow::bail!("bench is not implemented yet"), Command::Rpm { .. } => anyhow::bail!("rpm is not implemented yet"), } } + +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); + } + } + + PathBuf::from("data/benchmarks") +}