| @@ -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")); | |||
| } | |||
| } | |||
| @@ -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::<Vec<_>>() | |||
| .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")); | |||
| } | |||
| } | |||
| @@ -1 +1,83 @@ | |||
| use regex::Regex; | |||
| pub fn extract_final_integer(text: &str) -> Option<String> { | |||
| 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<char> { | |||
| 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')); | |||
| } | |||
| } | |||
| @@ -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<PathBuf> { | |||
| 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<DatasetSpec<'a>> { | |||
| 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 | |||
| } | |||
| @@ -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") | |||
| } | |||