Просмотр исходного кода

feat: add benchmark dataset fetch and prompts

main
orangels 1 неделю назад
Родитель
Сommit
67a4ea8af4
5 измененных файлов: 306 добавлений и 2 удалений
  1. +33
    -0
      src/benchmarks/aime.rs
  2. +51
    -0
      src/benchmarks/gpqa.rs
  3. +82
    -0
      src/benchmarks/judge.rs
  4. +119
    -0
      src/benchmarks/mod.rs
  5. +21
    -2
      src/cli.rs

+ 33
- 0
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"));
}
}

+ 51
- 0
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::<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"));
}
}

+ 82
- 0
src/benchmarks/judge.rs Просмотреть файл

@@ -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'));
}
}

+ 119
- 0
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<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
}

+ 21
- 2
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")
}

Загрузка…
Отмена
Сохранить