| @@ -118,6 +118,26 @@ mod tests { | |||||
| assert_eq!(loaded.cases[0].answer, "4"); | 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] | #[test] | ||||
| fn parquet_only_returns_conversion_hint() { | fn parquet_only_returns_conversion_hint() { | ||||
| let temp_dir = tempfile::tempdir().expect("create temp dir"); | 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) | 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 { | Ok(GpqaCase { | ||||
| id: format!("gpqa-diamond-{}", index + 1), | |||||
| id, | |||||
| question: question.to_string(), | 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<_>>>()?; | .collect::<Result<Vec<_>>>()?; | ||||
| @@ -95,6 +101,19 @@ fn normalize_header(header: &str) -> String { | |||||
| .to_ascii_lowercase() | .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>>> { | fn parse_csv(input: &str) -> Result<Vec<Vec<String>>> { | ||||
| let mut rows = Vec::new(); | let mut rows = Vec::new(); | ||||
| let mut row = 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::create_dir_all(&dataset_dir).expect("create dataset dir"); | ||||
| std::fs::write( | std::fs::write( | ||||
| dataset_dir.join("gpqa_diamond.csv"), | 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"); | .expect("write csv"); | ||||
| let loaded = load_cases(temp_dir.path()).expect("load cases"); | let loaded = load_cases(temp_dir.path()).expect("load cases"); | ||||
| assert_eq!(loaded.local_path, dataset_dir.join("gpqa_diamond.csv")); | 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].question, "Which is correct?"); | ||||
| assert_eq!( | assert_eq!( | ||||
| loaded.cases[0].choices[0], | loaded.cases[0].choices[0], | ||||
| ("A".to_string(), "Right".to_string()) | ("A".to_string(), "Right".to_string()) | ||||
| ); | ); | ||||
| assert_eq!(loaded.cases[0].answer, 'A'); | 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 anyhow::{Context, Result, bail}; | ||||
| use chrono::Utc; | use chrono::Utc; | ||||
| use serde::Serialize; | |||||
| use serde::{Deserialize, Serialize}; | |||||
| use serde_json::{Value, json}; | |||||
| use sha2::{Digest, Sha256}; | use sha2::{Digest, Sha256}; | ||||
| use std::path::{Path, PathBuf}; | use std::path::{Path, PathBuf}; | ||||
| use tokio::fs; | 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 = | const GPQA_DIAMOND_URL: &str = | ||||
| "https://huggingface.co/datasets/Idavidrein/gpqa/resolve/main/gpqa_diamond.csv"; | "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() | .bytes() | ||||
| .await | .await | ||||
| .with_context(|| format!("failed to read downloaded bytes for {dataset}"))?; | .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 | .await | ||||
| .with_context(|| format!("failed to write dataset file {}", spec.file_path.display()))?; | .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, | name: dataset, | ||||
| source_url: spec.source_url, | source_url: spec.source_url, | ||||
| downloaded_at: Utc::now().to_rfc3339(), | 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 = | let metadata_yaml = | ||||
| serde_yaml::to_string(&metadata).context("failed to serialize dataset metadata")?; | serde_yaml::to_string(&metadata).context("failed to serialize dataset metadata")?; | ||||
| @@ -84,6 +90,7 @@ struct DatasetSpec<'a> { | |||||
| dir: PathBuf, | dir: PathBuf, | ||||
| file_path: PathBuf, | file_path: PathBuf, | ||||
| use_hf_token: bool, | use_hf_token: bool, | ||||
| normalize: Option<fn(&[u8]) -> Result<Vec<u8>>>, | |||||
| } | } | ||||
| fn dataset_spec<'a>(dataset: &str, data_dir: &'a Path) -> Result<DatasetSpec<'a>> { | 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"); | let dir = data_dir.join("aime2026"); | ||||
| Ok(DatasetSpec { | Ok(DatasetSpec { | ||||
| source_url: AIME_2026_URL, | source_url: AIME_2026_URL, | ||||
| file_path: dir.join("train-00000-of-00001.parquet"), | |||||
| file_path: dir.join("aime2026.jsonl"), | |||||
| dir, | dir, | ||||
| use_hf_token: false, | use_hf_token: false, | ||||
| normalize: Some(normalize_aime_rows_json), | |||||
| }) | }) | ||||
| } | } | ||||
| "gpqa-diamond" => { | "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"), | file_path: dir.join("gpqa_diamond.csv"), | ||||
| dir, | dir, | ||||
| use_hf_token: true, | use_hf_token: true, | ||||
| normalize: None, | |||||
| }) | }) | ||||
| } | } | ||||
| other => bail!( | 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 { | fn sha256_hex(bytes: &[u8]) -> String { | ||||
| let digest = Sha256::digest(bytes); | let digest = Sha256::digest(bytes); | ||||
| let mut hex = String::with_capacity(digest.len() * 2); | let mut hex = String::with_capacity(digest.len() * 2); | ||||
| @@ -120,3 +168,54 @@ fn sha256_hex(bytes: &[u8]) -> String { | |||||
| } | } | ||||
| hex | 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 regex::Regex; | ||||
| use std::path::{Path, PathBuf}; | use std::path::{Path, PathBuf}; | ||||
| use std::time::{Duration, Instant}; | use std::time::{Duration, Instant}; | ||||
| use tokio::time::sleep; | |||||
| use tokio::time::{Instant as TokioInstant, sleep_until}; | |||||
| #[derive(Debug, Parser)] | #[derive(Debug, Parser)] | ||||
| #[command( | #[command( | ||||
| @@ -342,22 +342,32 @@ async fn run_rpm( | |||||
| let provider_config = config.provider(Some(&provider_name))?; | let provider_config = config.provider(Some(&provider_name))?; | ||||
| let model = model.unwrap_or_else(|| provider_config.default_model.clone()); | let model = model.unwrap_or_else(|| provider_config.default_model.clone()); | ||||
| let request = ModelRequest { | let request = ModelRequest { | ||||
| prompt: prompt.clone(), | |||||
| prompt, | |||||
| ..request_template(provider_config, &model, 0.0, 1024) | ..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_at = Utc::now(); | ||||
| let started = Instant::now(); | let started = Instant::now(); | ||||
| let tokio_started = TokioInstant::now(); | |||||
| let in_flight_limit = schedule.len().max(1); | |||||
| let mut metrics = Metrics::new(); | 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), | Ok(response) => metrics.record_success(response.status, response.elapsed_ms as u64), | ||||
| Err(error) => metrics.record_failure(error_code(&error)), | Err(error) => metrics.record_failure(error_code(&error)), | ||||
| } | } | ||||
| if started.elapsed() + delay <= duration { | |||||
| sleep(delay).await; | |||||
| } | |||||
| } | } | ||||
| let summary = metrics.summary(); | let summary = metrics.summary(); | ||||
| @@ -369,7 +379,6 @@ async fn run_rpm( | |||||
| started_at, | started_at, | ||||
| duration_ms: started.elapsed().as_millis(), | duration_ms: started.elapsed().as_millis(), | ||||
| target_rpm: rpm, | target_rpm: rpm, | ||||
| prompt, | |||||
| temperature: 0.0, | temperature: 0.0, | ||||
| max_tokens: 1024, | max_tokens: 1024, | ||||
| }, | }, | ||||
| @@ -579,6 +588,22 @@ fn parse_duration(value: &str) -> Result<Duration> { | |||||
| Ok(Duration::from_secs(seconds)) | 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> { | fn dataset_data_dir(config_path: &Path) -> Result<PathBuf> { | ||||
| if !config_path.exists() { | if !config_path.exists() { | ||||
| return Ok(PathBuf::from("data/benchmarks")); | return Ok(PathBuf::from("data/benchmarks")); | ||||
| @@ -619,6 +644,16 @@ mod tests { | |||||
| assert_eq!(parse_duration("5m").expect("minutes").as_secs(), 300); | 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] | #[test] | ||||
| fn rejects_invalid_duration() { | fn rejects_invalid_duration() { | ||||
| let error = parse_duration("one hour").expect_err("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 started_at: DateTime<Utc>, | ||||
| pub duration_ms: u128, | pub duration_ms: u128, | ||||
| pub target_rpm: u32, | pub target_rpm: u32, | ||||
| pub prompt: String, | |||||
| pub temperature: f32, | pub temperature: f32, | ||||
| pub max_tokens: u32, | pub max_tokens: u32, | ||||
| } | } | ||||
| @@ -221,4 +220,36 @@ mod tests { | |||||
| assert!(path.ends_with("reports/aime2026-openai-gpt-test-20260506T010203Z.json")); | assert!(path.ends_with("reports/aime2026-openai-gpt-test-20260506T010203Z.json")); | ||||
| assert!(path.exists()); | 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\"")); | |||||
| } | |||||
| } | } | ||||