| @@ -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 aime; | ||||
| pub mod gpqa; | pub mod gpqa; | ||||
| pub mod judge; | 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 anyhow::Result; | ||||
| use clap::{Parser, Subcommand}; | use clap::{Parser, Subcommand}; | ||||
| use std::path::PathBuf; | |||||
| use std::path::{Path, PathBuf}; | |||||
| #[derive(Debug, Parser)] | #[derive(Debug, Parser)] | ||||
| #[command( | #[command( | ||||
| @@ -85,8 +87,25 @@ pub enum BenchCommand { | |||||
| pub async fn dispatch(cli: Cli) -> Result<()> { | pub async fn dispatch(cli: Cli) -> Result<()> { | ||||
| match cli.command { | match cli.command { | ||||
| Command::Check { .. } => anyhow::bail!("check is not implemented yet"), | 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::Bench { .. } => anyhow::bail!("bench is not implemented yet"), | ||||
| Command::Rpm { .. } => anyhow::bail!("rpm 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") | |||||
| } | |||||