From f21ab8bd4c30e6ddea35b09fee0f388847b9be92 Mon Sep 17 00:00:00 2001 From: orangels Date: Thu, 7 May 2026 02:17:36 +0800 Subject: [PATCH] feat: add --max-tokens CLI param and correct_samples to bench reports - Add --max-tokens parameter to bench subcommands (default 32768) - Record up to 5 correct case samples in report JSON for output inspection - Report max_tokens reflects actual value used instead of hardcoded 1024 --- src/cli.rs | 58 +++++++++++++++++++++++++++++++++++++++++---------- src/report.rs | 10 +++++++++ 2 files changed, 57 insertions(+), 11 deletions(-) diff --git a/src/cli.rs b/src/cli.rs index 195ed48..5ada6dc 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -3,10 +3,11 @@ use crate::benchmarks::judge; use crate::config::AppConfig; use crate::metrics::{LatencySummary, Metrics, MetricsSummary}; use crate::report::{ - BenchmarkParamsReport, BenchmarkReport, BenchmarkSummaryReport, DatasetReport, LatencyReport, - LimiterInferenceKind, LimiterInferenceReport, PhaseSummaryReport, ProbeSecondReport, - RpmModeDetailReport, RpmParamsReport, RpmReport, RpmRunReport, RpmSummaryReport, RunReport, - WindowBoundaryReport, WrongCaseReport, write_benchmark_report, write_rpm_report, + BenchmarkParamsReport, BenchmarkReport, BenchmarkSummaryReport, CorrectCaseReport, + DatasetReport, LatencyReport, LimiterInferenceKind, LimiterInferenceReport, + PhaseSummaryReport, ProbeSecondReport, RpmModeDetailReport, RpmParamsReport, RpmReport, + RpmRunReport, RpmSummaryReport, RunReport, WindowBoundaryReport, WrongCaseReport, + write_benchmark_report, write_rpm_report, }; use crate::rpm_modes::{ ProbePhase, RpmMode, ScheduledProbe, burst_schedule, sliding_window_schedule, @@ -105,6 +106,8 @@ pub enum BenchCommand { limit: Option, #[arg(long)] stream: Option, + #[arg(long, default_value_t = 32768)] + max_tokens: u32, }, GpqaDiamond { #[arg(long, default_value = "config.yaml")] @@ -119,6 +122,8 @@ pub enum BenchCommand { limit: Option, #[arg(long)] stream: Option, + #[arg(long, default_value_t = 32768)] + max_tokens: u32, }, } @@ -206,7 +211,8 @@ async fn dispatch_bench(command: BenchCommand) -> Result<()> { concurrency, limit, stream, - } => run_aime_benchmark(config, provider, model, concurrency, limit, stream).await, + max_tokens, + } => run_aime_benchmark(config, provider, model, concurrency, limit, stream, max_tokens).await, BenchCommand::GpqaDiamond { config, provider, @@ -214,7 +220,8 @@ async fn dispatch_bench(command: BenchCommand) -> Result<()> { concurrency, limit, stream, - } => run_gpqa_benchmark(config, provider, model, concurrency, limit, stream).await, + max_tokens, + } => run_gpqa_benchmark(config, provider, model, concurrency, limit, stream, max_tokens).await, } } @@ -225,6 +232,7 @@ async fn run_aime_benchmark( concurrency: usize, limit: Option, stream: Option, + max_tokens: u32, ) -> Result<()> { let config = AppConfig::load(&config_path)?; let provider_name = provider_name(&config, provider.as_deref())?; @@ -243,7 +251,7 @@ async fn run_aime_benchmark( let total = cases.len() as u64; let started_at = Utc::now(); let started = Instant::now(); - let mut base_request = request_template(&provider_config, &model, 0.0, 1024); + let mut base_request = request_template(&provider_config, &model, 0.0, max_tokens); base_request.stream = stream.unwrap_or(provider_config.stream); let protocol = provider_config.protocol; @@ -267,6 +275,7 @@ async fn run_aime_benchmark( let mut metrics = Metrics::new(); let mut wrong_cases = Vec::new(); + let mut correct_samples = Vec::new(); while let Some((case, result)) = results.next().await { pb.inc(1); match result { @@ -276,7 +285,16 @@ async fn run_aime_benchmark( .unwrap_or_else(|| "no_answer".to_string()); let correct = judge::judge_integer(&response.text, &case.answer); metrics.record_judgement(correct); - if !correct { + if correct { + if correct_samples.len() < 5 { + correct_samples.push(CorrectCaseReport { + id: case.id, + question: case.problem, + expected: case.answer, + raw_output: response.text, + }); + } + } else { wrong_cases.push(WrongCaseReport { id: case.id, question: case.problem, @@ -302,7 +320,9 @@ async fn run_aime_benchmark( duration_ms: started.elapsed().as_millis(), concurrency, limit, + max_tokens, summary, + correct_samples, wrong_cases, }); let report_path = write_benchmark_report(Path::new("."), &report)?; @@ -317,6 +337,7 @@ async fn run_gpqa_benchmark( concurrency: usize, limit: Option, stream: Option, + max_tokens: u32, ) -> Result<()> { let config = AppConfig::load(&config_path)?; let provider_name = provider_name(&config, provider.as_deref())?; @@ -335,7 +356,7 @@ async fn run_gpqa_benchmark( let total = cases.len() as u64; let started_at = Utc::now(); let started = Instant::now(); - let mut base_request = request_template(&provider_config, &model, 0.0, 1024); + let mut base_request = request_template(&provider_config, &model, 0.0, max_tokens); base_request.stream = stream.unwrap_or(provider_config.stream); let protocol = provider_config.protocol; @@ -359,6 +380,7 @@ async fn run_gpqa_benchmark( let mut metrics = Metrics::new(); let mut wrong_cases = Vec::new(); + let mut correct_samples = Vec::new(); while let Some((case, result)) = results.next().await { pb.inc(1); match result { @@ -370,7 +392,16 @@ async fn run_gpqa_benchmark( let expected = case.answer.to_string(); let correct = judge::judge_choice(&response.text, case.answer); metrics.record_judgement(correct); - if !correct { + if correct { + if correct_samples.len() < 5 { + correct_samples.push(CorrectCaseReport { + id: case.id, + question: case.question, + expected, + raw_output: response.text, + }); + } + } else { wrong_cases.push(WrongCaseReport { id: case.id, question: case.question, @@ -396,7 +427,9 @@ async fn run_gpqa_benchmark( duration_ms: started.elapsed().as_millis(), concurrency, limit, + max_tokens, summary, + correct_samples, wrong_cases, }); let report_path = write_benchmark_report(Path::new("."), &report)?; @@ -898,7 +931,9 @@ struct BenchmarkReportInput { duration_ms: u128, concurrency: usize, limit: Option, + max_tokens: u32, summary: MetricsSummary, + correct_samples: Vec, wrong_cases: Vec, } @@ -917,7 +952,7 @@ fn benchmark_report(input: BenchmarkReportInput) -> BenchmarkReport { concurrency: input.concurrency, limit: input.limit, temperature: 0.0, - max_tokens: 1024, + max_tokens: input.max_tokens, }, summary: BenchmarkSummaryReport { accuracy: input.summary.accuracy, @@ -930,6 +965,7 @@ fn benchmark_report(input: BenchmarkReportInput) -> BenchmarkReport { ttft_ms: latency_report(&input.summary.ttft_ms), }, errors: input.summary.errors, + correct_samples: input.correct_samples, wrong_cases: input.wrong_cases, } } diff --git a/src/report.rs b/src/report.rs index 97d235a..9552557 100644 --- a/src/report.rs +++ b/src/report.rs @@ -50,6 +50,14 @@ pub struct WrongCaseReport { pub raw_output: String, } +#[derive(Debug, Clone, Serialize)] +pub struct CorrectCaseReport { + pub id: String, + pub question: String, + pub expected: String, + pub raw_output: String, +} + #[derive(Debug, Clone, Serialize)] pub struct BenchmarkReport { pub benchmark: String, @@ -60,6 +68,7 @@ pub struct BenchmarkReport { pub run: RunReport, pub summary: BenchmarkSummaryReport, pub errors: Vec, + pub correct_samples: Vec, pub wrong_cases: Vec, } @@ -287,6 +296,7 @@ mod tests { }, }, errors: vec![], + correct_samples: vec![], wrong_cases: vec![], };