| @@ -118,6 +118,26 @@ mod tests { | |||
| assert_eq!(loaded.cases[0].answer, "4"); | |||
| } | |||
| #[test] | |||
| fn loads_jsonl_normalized_from_hugging_face_rows() { | |||
| let temp_dir = tempfile::tempdir().expect("create temp dir"); | |||
| let dataset_dir = temp_dir.path().join("aime2026"); | |||
| std::fs::create_dir_all(&dataset_dir).expect("create dataset dir"); | |||
| std::fs::write( | |||
| dataset_dir.join("aime2026.jsonl"), | |||
| r#"{"id":"aime2026-7","problem":"Compute 6 times 7.","answer":"42"} | |||
| "#, | |||
| ) | |||
| .expect("write normalized fetch output"); | |||
| let loaded = load_cases(temp_dir.path()).expect("load cases"); | |||
| assert_eq!(loaded.cases.len(), 1); | |||
| assert_eq!(loaded.cases[0].id, "aime2026-7"); | |||
| assert_eq!(loaded.cases[0].problem, "Compute 6 times 7."); | |||
| assert_eq!(loaded.cases[0].answer, "42"); | |||
| } | |||
| #[test] | |||
| fn parquet_only_returns_conversion_hint() { | |||
| let temp_dir = tempfile::tempdir().expect("create temp dir"); | |||
| @@ -54,16 +54,22 @@ pub fn load_cases(data_dir: &Path) -> Result<LoadedGpqaCases> { | |||
| format!("missing Incorrect Answer 3 column in row {}", index + 2) | |||
| })?; | |||
| let id = format!("gpqa-diamond-{}", index + 1); | |||
| let (choices, answer) = rotated_choices( | |||
| [ | |||
| correct.to_string(), | |||
| incorrect_1.to_string(), | |||
| incorrect_2.to_string(), | |||
| incorrect_3.to_string(), | |||
| ], | |||
| index % 4, | |||
| ); | |||
| Ok(GpqaCase { | |||
| id: format!("gpqa-diamond-{}", index + 1), | |||
| id, | |||
| question: question.to_string(), | |||
| choices: [ | |||
| ("A".to_string(), correct.to_string()), | |||
| ("B".to_string(), incorrect_1.to_string()), | |||
| ("C".to_string(), incorrect_2.to_string()), | |||
| ("D".to_string(), incorrect_3.to_string()), | |||
| ], | |||
| answer: 'A', | |||
| choices, | |||
| answer, | |||
| }) | |||
| }) | |||
| .collect::<Result<Vec<_>>>()?; | |||
| @@ -95,6 +101,19 @@ fn normalize_header(header: &str) -> String { | |||
| .to_ascii_lowercase() | |||
| } | |||
| fn rotated_choices(mut values: [String; 4], rotation: usize) -> ([(String, String); 4], char) { | |||
| let rotation = rotation % values.len(); | |||
| values.rotate_left(rotation); | |||
| let answer_index = (4 - rotation) % 4; | |||
| let labels = ["A", "B", "C", "D"]; | |||
| let answer = labels[answer_index] | |||
| .chars() | |||
| .next() | |||
| .expect("labels are non-empty"); | |||
| let choices = std::array::from_fn(|index| (labels[index].to_string(), values[index].clone())); | |||
| (choices, answer) | |||
| } | |||
| fn parse_csv(input: &str) -> Result<Vec<Vec<String>>> { | |||
| let mut rows = Vec::new(); | |||
| let mut row = Vec::new(); | |||
| @@ -187,19 +206,29 @@ mod tests { | |||
| std::fs::create_dir_all(&dataset_dir).expect("create dataset dir"); | |||
| std::fs::write( | |||
| dataset_dir.join("gpqa_diamond.csv"), | |||
| "Question,Correct Answer,Incorrect Answer 1,Incorrect Answer 2,Incorrect Answer 3\n\"Which is correct?\",\"Right\",\"Wrong, with comma\",\"Nope\",\"Never\"\n", | |||
| "Question,Correct Answer,Incorrect Answer 1,Incorrect Answer 2,Incorrect Answer 3\n\"Which is correct?\",\"Right\",\"Wrong, with comma\",\"Nope\",\"Never\"\n\"Which is second?\",\"Second right\",\"Second wrong 1\",\"Second wrong 2\",\"Second wrong 3\"\n", | |||
| ) | |||
| .expect("write csv"); | |||
| let loaded = load_cases(temp_dir.path()).expect("load cases"); | |||
| assert_eq!(loaded.local_path, dataset_dir.join("gpqa_diamond.csv")); | |||
| assert_eq!(loaded.cases.len(), 1); | |||
| assert_eq!(loaded.cases.len(), 2); | |||
| assert_eq!(loaded.cases[0].question, "Which is correct?"); | |||
| assert_eq!( | |||
| loaded.cases[0].choices[0], | |||
| ("A".to_string(), "Right".to_string()) | |||
| ); | |||
| assert_eq!(loaded.cases[0].answer, 'A'); | |||
| assert_eq!(loaded.cases[1].answer, 'D'); | |||
| assert_eq!( | |||
| loaded.cases[1] | |||
| .choices | |||
| .iter() | |||
| .find(|(label, _)| label == "D") | |||
| .map(|(_, text)| text.as_str()), | |||
| Some("Second right") | |||
| ); | |||
| assert!(loaded.cases.iter().any(|case| case.answer != 'A')); | |||
| } | |||
| } | |||
| @@ -4,12 +4,13 @@ pub mod judge; | |||
| use anyhow::{Context, Result, bail}; | |||
| use chrono::Utc; | |||
| use serde::Serialize; | |||
| use serde::{Deserialize, Serialize}; | |||
| use serde_json::{Value, json}; | |||
| 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 AIME_2026_URL: &str = "https://datasets-server.huggingface.co/rows?dataset=MathArena%2Faime_2026&config=default&split=train&offset=0&length=100"; | |||
| const GPQA_DIAMOND_URL: &str = | |||
| "https://huggingface.co/datasets/Idavidrein/gpqa/resolve/main/gpqa_diamond.csv"; | |||
| @@ -58,7 +59,12 @@ pub async fn fetch_dataset(dataset: &str, data_dir: &Path) -> Result<PathBuf> { | |||
| .bytes() | |||
| .await | |||
| .with_context(|| format!("failed to read downloaded bytes for {dataset}"))?; | |||
| fs::write(&spec.file_path, &bytes) | |||
| let file_bytes = match spec.normalize { | |||
| Some(normalize) => normalize(&bytes)?, | |||
| None => bytes.to_vec(), | |||
| }; | |||
| fs::write(&spec.file_path, &file_bytes) | |||
| .await | |||
| .with_context(|| format!("failed to write dataset file {}", spec.file_path.display()))?; | |||
| @@ -66,8 +72,8 @@ pub async fn fetch_dataset(dataset: &str, data_dir: &Path) -> Result<PathBuf> { | |||
| name: dataset, | |||
| source_url: spec.source_url, | |||
| downloaded_at: Utc::now().to_rfc3339(), | |||
| sha256: sha256_hex(&bytes), | |||
| bytes: bytes.len(), | |||
| sha256: sha256_hex(&file_bytes), | |||
| bytes: file_bytes.len(), | |||
| }; | |||
| let metadata_yaml = | |||
| serde_yaml::to_string(&metadata).context("failed to serialize dataset metadata")?; | |||
| @@ -84,6 +90,7 @@ struct DatasetSpec<'a> { | |||
| dir: PathBuf, | |||
| file_path: PathBuf, | |||
| use_hf_token: bool, | |||
| normalize: Option<fn(&[u8]) -> Result<Vec<u8>>>, | |||
| } | |||
| fn dataset_spec<'a>(dataset: &str, data_dir: &'a Path) -> Result<DatasetSpec<'a>> { | |||
| @@ -92,9 +99,10 @@ fn dataset_spec<'a>(dataset: &str, data_dir: &'a Path) -> Result<DatasetSpec<'a> | |||
| let dir = data_dir.join("aime2026"); | |||
| Ok(DatasetSpec { | |||
| source_url: AIME_2026_URL, | |||
| file_path: dir.join("train-00000-of-00001.parquet"), | |||
| file_path: dir.join("aime2026.jsonl"), | |||
| dir, | |||
| use_hf_token: false, | |||
| normalize: Some(normalize_aime_rows_json), | |||
| }) | |||
| } | |||
| "gpqa-diamond" => { | |||
| @@ -104,6 +112,7 @@ fn dataset_spec<'a>(dataset: &str, data_dir: &'a Path) -> Result<DatasetSpec<'a> | |||
| file_path: dir.join("gpqa_diamond.csv"), | |||
| dir, | |||
| use_hf_token: true, | |||
| normalize: None, | |||
| }) | |||
| } | |||
| other => bail!( | |||
| @@ -112,6 +121,45 @@ fn dataset_spec<'a>(dataset: &str, data_dir: &'a Path) -> Result<DatasetSpec<'a> | |||
| } | |||
| } | |||
| #[derive(Debug, Deserialize)] | |||
| struct AimeRowsResponse { | |||
| rows: Vec<AimeRowWrapper>, | |||
| } | |||
| #[derive(Debug, Deserialize)] | |||
| struct AimeRowWrapper { | |||
| row: AimeSourceRow, | |||
| } | |||
| #[derive(Debug, Deserialize)] | |||
| struct AimeSourceRow { | |||
| problem_idx: u64, | |||
| problem: String, | |||
| answer: Value, | |||
| } | |||
| fn normalize_aime_rows_json(bytes: &[u8]) -> Result<Vec<u8>> { | |||
| let response: AimeRowsResponse = | |||
| serde_json::from_slice(bytes).context("failed to parse AIME rows response")?; | |||
| let mut output = Vec::new(); | |||
| for row in response.rows { | |||
| let answer = match row.row.answer { | |||
| Value::String(answer) => answer, | |||
| Value::Number(answer) => answer.to_string(), | |||
| other => bail!("unexpected AIME answer value: {other}"), | |||
| }; | |||
| let line = json!({ | |||
| "id": format!("aime2026-{}", row.row.problem_idx), | |||
| "problem": row.row.problem, | |||
| "answer": answer, | |||
| }); | |||
| serde_json::to_writer(&mut output, &line) | |||
| .context("failed to serialize normalized AIME JSONL")?; | |||
| output.push(b'\n'); | |||
| } | |||
| Ok(output) | |||
| } | |||
| fn sha256_hex(bytes: &[u8]) -> String { | |||
| let digest = Sha256::digest(bytes); | |||
| let mut hex = String::with_capacity(digest.len() * 2); | |||
| @@ -120,3 +168,54 @@ fn sha256_hex(bytes: &[u8]) -> String { | |||
| } | |||
| hex | |||
| } | |||
| #[cfg(test)] | |||
| mod tests { | |||
| use super::*; | |||
| #[test] | |||
| fn normalizes_aime_rows_response_to_loader_jsonl() { | |||
| let rows_response = br#"{ | |||
| "rows": [ | |||
| { | |||
| "row": { | |||
| "problem_idx": 1, | |||
| "problem": "What is 20 + 22?", | |||
| "answer": 42 | |||
| } | |||
| }, | |||
| { | |||
| "row": { | |||
| "problem_idx": 2, | |||
| "problem": "What is 10 squared?", | |||
| "answer": "100" | |||
| } | |||
| } | |||
| ] | |||
| }"#; | |||
| let normalized = normalize_aime_rows_json(rows_response).expect("normalize rows"); | |||
| let lines = std::str::from_utf8(&normalized) | |||
| .expect("utf8 jsonl") | |||
| .lines() | |||
| .map(|line| serde_json::from_str::<Value>(line).expect("json line")) | |||
| .collect::<Vec<_>>(); | |||
| assert_eq!( | |||
| lines, | |||
| vec![ | |||
| json!({ | |||
| "id": "aime2026-1", | |||
| "problem": "What is 20 + 22?", | |||
| "answer": "42", | |||
| }), | |||
| json!({ | |||
| "id": "aime2026-2", | |||
| "problem": "What is 10 squared?", | |||
| "answer": "100", | |||
| }), | |||
| ] | |||
| ); | |||
| } | |||
| } | |||
| @@ -14,7 +14,7 @@ use futures::{StreamExt, stream}; | |||
| use regex::Regex; | |||
| use std::path::{Path, PathBuf}; | |||
| use std::time::{Duration, Instant}; | |||
| use tokio::time::sleep; | |||
| use tokio::time::{Instant as TokioInstant, sleep_until}; | |||
| #[derive(Debug, Parser)] | |||
| #[command( | |||
| @@ -342,22 +342,32 @@ async fn run_rpm( | |||
| let provider_config = config.provider(Some(&provider_name))?; | |||
| let model = model.unwrap_or_else(|| provider_config.default_model.clone()); | |||
| let request = ModelRequest { | |||
| prompt: prompt.clone(), | |||
| prompt, | |||
| ..request_template(provider_config, &model, 0.0, 1024) | |||
| }; | |||
| let delay = Duration::from_secs_f64(60.0 / rpm as f64); | |||
| let schedule = rpm_start_schedule(duration, rpm); | |||
| let started_at = Utc::now(); | |||
| let started = Instant::now(); | |||
| let tokio_started = TokioInstant::now(); | |||
| let in_flight_limit = schedule.len().max(1); | |||
| let mut metrics = Metrics::new(); | |||
| while started.elapsed() < duration { | |||
| match run_model_request(provider_config.protocol, request.clone()).await { | |||
| let results = stream::iter(schedule.into_iter().map(|offset| { | |||
| let request = request.clone(); | |||
| async move { | |||
| sleep_until(tokio_started + offset).await; | |||
| run_model_request(provider_config.protocol, request).await | |||
| } | |||
| })) | |||
| .buffer_unordered(in_flight_limit) | |||
| .collect::<Vec<_>>() | |||
| .await; | |||
| for result in results { | |||
| match result { | |||
| Ok(response) => metrics.record_success(response.status, response.elapsed_ms as u64), | |||
| Err(error) => metrics.record_failure(error_code(&error)), | |||
| } | |||
| if started.elapsed() + delay <= duration { | |||
| sleep(delay).await; | |||
| } | |||
| } | |||
| let summary = metrics.summary(); | |||
| @@ -369,7 +379,6 @@ async fn run_rpm( | |||
| started_at, | |||
| duration_ms: started.elapsed().as_millis(), | |||
| target_rpm: rpm, | |||
| prompt, | |||
| temperature: 0.0, | |||
| max_tokens: 1024, | |||
| }, | |||
| @@ -579,6 +588,22 @@ fn parse_duration(value: &str) -> Result<Duration> { | |||
| Ok(Duration::from_secs(seconds)) | |||
| } | |||
| fn rpm_start_schedule(duration: Duration, rpm: u32) -> Vec<Duration> { | |||
| let interval_nanos = 60_000_000_000u128 / u128::from(rpm); | |||
| if interval_nanos == 0 { | |||
| return Vec::new(); | |||
| } | |||
| let duration_nanos = duration.as_nanos(); | |||
| let mut offsets = Vec::new(); | |||
| let mut offset = 0u128; | |||
| while offset < duration_nanos { | |||
| offsets.push(Duration::from_nanos(offset.min(u128::from(u64::MAX)) as u64)); | |||
| offset += interval_nanos; | |||
| } | |||
| offsets | |||
| } | |||
| fn dataset_data_dir(config_path: &Path) -> Result<PathBuf> { | |||
| if !config_path.exists() { | |||
| return Ok(PathBuf::from("data/benchmarks")); | |||
| @@ -619,6 +644,16 @@ mod tests { | |||
| assert_eq!(parse_duration("5m").expect("minutes").as_secs(), 300); | |||
| } | |||
| #[test] | |||
| fn computes_rpm_start_schedule_from_run_start() { | |||
| let schedule = rpm_start_schedule(Duration::from_secs(60), 120); | |||
| assert_eq!(schedule.len(), 120); | |||
| assert_eq!(schedule[0], Duration::ZERO); | |||
| assert_eq!(schedule[1], Duration::from_millis(500)); | |||
| assert_eq!(schedule[119], Duration::from_millis(59_500)); | |||
| } | |||
| #[test] | |||
| fn rejects_invalid_duration() { | |||
| let error = parse_duration("one hour").expect_err("invalid duration"); | |||
| @@ -66,7 +66,6 @@ pub struct RpmRunReport { | |||
| pub started_at: DateTime<Utc>, | |||
| pub duration_ms: u128, | |||
| pub target_rpm: u32, | |||
| pub prompt: String, | |||
| pub temperature: f32, | |||
| pub max_tokens: u32, | |||
| } | |||
| @@ -221,4 +220,36 @@ mod tests { | |||
| assert!(path.ends_with("reports/aime2026-openai-gpt-test-20260506T010203Z.json")); | |||
| assert!(path.exists()); | |||
| } | |||
| #[test] | |||
| fn rpm_report_does_not_serialize_prompt() { | |||
| let report = RpmReport { | |||
| benchmark: "rpm".to_string(), | |||
| provider: "openai".to_string(), | |||
| model: "gpt/test".to_string(), | |||
| run: RpmRunReport { | |||
| started_at: Utc.with_ymd_and_hms(2026, 5, 6, 1, 2, 3).unwrap(), | |||
| duration_ms: 1000, | |||
| target_rpm: 60, | |||
| temperature: 0.0, | |||
| max_tokens: 1024, | |||
| }, | |||
| summary: RpmSummaryReport { | |||
| actual_requests: 1, | |||
| success: 1, | |||
| failure: 0, | |||
| latency_ms: LatencyReport { | |||
| p50: Some(10), | |||
| p95: Some(10), | |||
| p99: Some(10), | |||
| }, | |||
| }, | |||
| errors: vec![], | |||
| }; | |||
| let json = serde_json::to_string(&report).expect("serialize report"); | |||
| assert!(!json.contains("sensitive prompt")); | |||
| assert!(!json.contains("\"prompt\"")); | |||
| } | |||
| } | |||