diff --git a/cli/README.md b/cli/README.md index e9c1f22..0343732 100644 --- a/cli/README.md +++ b/cli/README.md @@ -43,6 +43,7 @@ Add a `quantiles.toml` or `.quantiles.toml` file to configure an evaluation. For ```toml [benchmarks.pubmedqa] +dataset = "hf://quantiles/PubMedQA" samples = 50 model = "openai:gpt-5.6" max_workers = 100 diff --git a/cli/examples/configs/anthropic/quantiles.toml b/cli/examples/configs/anthropic/quantiles.toml index 66d0544..950c435 100644 --- a/cli/examples/configs/anthropic/quantiles.toml +++ b/cli/examples/configs/anthropic/quantiles.toml @@ -4,9 +4,12 @@ # variable. To get an API key, go to https://platform.claude.com/settings/workspaces/default/keys [benchmarks.pubmedqa] +dataset = "hf://quantiles/PubMedQA" + # The full expert-labeled PubMedQA benchmark has 1,000 samples. Because Anthropic charges for -# usage, this example limits the run to 50 samples to control costs. Omit `samples` to run the -# full benchmark. +# usage, this configuration entry limits the run to 50 samples to control costs. +# +# Omit the below `samples` key to run the full benchmark. samples = 50 # This example uses a relatively low-cost, fast model. Model details are available at: diff --git a/cli/examples/configs/both_error/.quantiles.toml b/cli/examples/configs/both_error/.quantiles.toml index e2a96a1..040121c 100644 --- a/cli/examples/configs/both_error/.quantiles.toml +++ b/cli/examples/configs/both_error/.quantiles.toml @@ -1,8 +1,11 @@ [benchmarks.pubmedqa] +dataset = "hf://quantiles/PubMedQA" samples = 10 [benchmarks.simpleqa-verified] +dataset = "hf://quantiles/simpleqa-verified" samples = 100 [benchmarks.financebench] +dataset = "hf://quantiles/financebench" samples = 50 diff --git a/cli/examples/configs/cloudflare/quantiles.toml b/cli/examples/configs/cloudflare/quantiles.toml index 42dee7e..e017700 100644 --- a/cli/examples/configs/cloudflare/quantiles.toml +++ b/cli/examples/configs/cloudflare/quantiles.toml @@ -12,6 +12,8 @@ # https://dash.cloudflare.com//ai/ai-gateway [benchmarks.pubmedqa] +dataset = "hf://quantiles/PubMedQA" + # The full expert-labeled PubMedQA benchmark has 1,000 samples. Because Cloudflare charges for # usage, this example limits the run to 50 samples to control costs. Omit `samples` to run the # full benchmark. diff --git a/cli/examples/configs/gemini/quantiles.toml b/cli/examples/configs/gemini/quantiles.toml index d10b1fd..82f851a 100644 --- a/cli/examples/configs/gemini/quantiles.toml +++ b/cli/examples/configs/gemini/quantiles.toml @@ -6,6 +6,8 @@ # https://aistudio.google.com/prompts/new_chat [benchmarks.pubmedqa] +dataset = "hf://quantiles/PubMedQA" + # The full expert-labeled PubMedQA benchmark has 1,000 samples. Because Google charges for # usage, this example limits the run to 50 samples to control costs. Omit `samples` to run the # full benchmark. diff --git a/cli/examples/configs/openai/quantiles.toml b/cli/examples/configs/openai/quantiles.toml index eef0f56..ba41ed6 100644 --- a/cli/examples/configs/openai/quantiles.toml +++ b/cli/examples/configs/openai/quantiles.toml @@ -7,6 +7,8 @@ # https://platform.openai.com/home [benchmarks.pubmedqa] +dataset = "hf://quantiles/PubMedQA" + # The full expert-labeled PubMedQA benchmark has 1,000 samples. Because OpenAI charges for # usage, this example limits the run to 50 samples to control costs. Omit `samples` to run the # full benchmark. diff --git a/cli/examples/configs/simple/quantiles.toml b/cli/examples/configs/simple/quantiles.toml index f6ff43e..1a0b772 100644 --- a/cli/examples/configs/simple/quantiles.toml +++ b/cli/examples/configs/simple/quantiles.toml @@ -8,7 +8,9 @@ # dataset rows processed. [benchmarks.pubmedqa] +dataset = "hf://quantiles/PubMedQA" samples = 10 [benchmarks.simpleqa-verified] +dataset = "hf://quantiles/simpleqa-verified" samples = 100 diff --git a/cli/src/builtins/common.rs b/cli/src/builtins/common.rs index 374fc26..9e2e22e 100644 --- a/cli/src/builtins/common.rs +++ b/cli/src/builtins/common.rs @@ -20,6 +20,9 @@ pub(crate) struct BuiltinConfig { /// Number of dataset rows to evaluate. If omitted, the entire dataset is used. #[serde(default)] pub(crate) limit: Option, + /// Dataset source to evaluate. Builtins currently support Hugging Face via `hf://...`. + #[serde(default)] + pub(crate) dataset: Option, /// Which model sampler to use. If omitted, the builtin chooses a sensible default. #[serde(default)] pub(crate) model: Option, diff --git a/cli/src/builtins/input.rs b/cli/src/builtins/input.rs index 7d443b5..437a124 100644 --- a/cli/src/builtins/input.rs +++ b/cli/src/builtins/input.rs @@ -5,6 +5,7 @@ use crate::llm::Sampler; /// Normalized run input schema for all builtins. #[derive(Serialize)] pub(crate) struct BuiltinRunInput { + pub(crate) dataset: String, pub(crate) model: String, pub(crate) num_samples: usize, #[serde(skip_serializing_if = "Option::is_none")] @@ -15,11 +16,13 @@ pub(crate) struct BuiltinRunInput { pub(crate) async fn set_builtin_run_input( db: &sea_orm::DatabaseConnection, run_id: i64, + dataset: &str, model: Option<&Sampler>, num_samples: usize, max_workers: Option, ) -> anyhow::Result<()> { let input = serde_json::to_string(&BuiltinRunInput { + dataset: dataset.to_owned(), model: builtin_model_name(model), num_samples, max_workers, diff --git a/cli/src/builtins/pubmedqa/mod.rs b/cli/src/builtins/pubmedqa/mod.rs index 4e0df77..4385097 100644 --- a/cli/src/builtins/pubmedqa/mod.rs +++ b/cli/src/builtins/pubmedqa/mod.rs @@ -9,7 +9,7 @@ use crate::builtins::dataset_runner::DatasetRunner; use crate::builtins::input::set_builtin_run_input; use crate::builtins::output::set_builtin_run_output; use crate::builtins::{BuiltinContext, BuiltinWorkflow}; -use crate::dataset::DatasetManager; +use crate::dataset::{DatasetManager, resolve_hf_dataset_source}; use crate::llm::random_label::RandomLabelSampler; use config::{PubMedQAConfig, RowOutput}; @@ -23,6 +23,9 @@ mod eval; /// `PubMedQA` builtin using the quantiles/PubMedQA dataset. pub struct PubmedqaBuiltin; +const DEFAULT_DATASET_SOURCE: &str = "hf://quantiles/PubMedQA"; + +#[expect(clippy::too_many_lines)] #[async_trait::async_trait] impl BuiltinWorkflow for PubmedqaBuiltin { fn name(&self) -> String { @@ -48,7 +51,12 @@ impl BuiltinWorkflow for PubmedqaBuiltin { })?; let manager = DatasetManager::new()?; - let dataset_id = "quantiles/PubMedQA"; + let dataset_source = config + .base + .dataset + .as_deref() + .unwrap_or(DEFAULT_DATASET_SOURCE); + let dataset_id = resolve_hf_dataset_source(dataset_source)?; let info = manager .init(dataset_id, Some("pqa_labeled"), Some("train"), None) .await?; @@ -61,6 +69,7 @@ impl BuiltinWorkflow for PubmedqaBuiltin { set_builtin_run_input( ctx.db, ctx.run_id, + dataset_source, config.base.model.as_ref(), limit, config.base.max_workers, diff --git a/cli/src/builtins/similarity.rs b/cli/src/builtins/similarity.rs index 1c8fe0f..b4e3a01 100644 --- a/cli/src/builtins/similarity.rs +++ b/cli/src/builtins/similarity.rs @@ -9,7 +9,8 @@ use crate::builtins::dataset_runner::DatasetRunner; use crate::builtins::input::set_builtin_run_input; use crate::builtins::output::set_builtin_run_output; use crate::builtins::{BuiltinContext, BuiltinWorkflow}; -use crate::dataset::DatasetManager; + +use crate::dataset::{DatasetManager, resolve_hf_dataset_source}; use crate::llm::random::RandomSampler; use crate::similarity::{ SimilarityMetric, SimilarityMetricName, levenshtein::LevenshteinSimilarity, @@ -44,7 +45,7 @@ struct RowOutput { #[derive(Clone, Copy)] pub struct SimilarityBenchmark { name: &'static str, - dataset_id: &'static str, + dataset_source: &'static str, input_field: &'static str, target_field: &'static str, } @@ -52,7 +53,7 @@ pub struct SimilarityBenchmark { /// `simpleqa-verified` builtin. pub const SIMPLEQA: SimilarityBenchmark = SimilarityBenchmark { name: "simpleqa-verified", - dataset_id: "quantiles/simpleqa-verified", + dataset_source: "hf://quantiles/simpleqa-verified", input_field: "problem", target_field: "answer", }; @@ -60,7 +61,7 @@ pub const SIMPLEQA: SimilarityBenchmark = SimilarityBenchmark { /// `financebench` builtin. pub const FINANCEBENCH: SimilarityBenchmark = SimilarityBenchmark { name: "financebench", - dataset_id: "quantiles/financebench", + dataset_source: "hf://quantiles/financebench", input_field: "question", target_field: "answer", }; @@ -94,7 +95,13 @@ impl BuiltinWorkflow for SimilarityBenchmark { })?; let manager = DatasetManager::new()?; - let info = manager.init(self.dataset_id, None, None, None).await?; + let dataset_source = config + .base + .dataset + .as_deref() + .unwrap_or(self.dataset_source); + let dataset_id = resolve_hf_dataset_source(dataset_source)?; + let info = manager.init(dataset_id, None, None, None).await?; let total = info .total_rows @@ -107,6 +114,7 @@ impl BuiltinWorkflow for SimilarityBenchmark { set_builtin_run_input( db, run_id, + dataset_source, config.base.model.as_ref(), limit, config.base.max_workers, @@ -118,7 +126,7 @@ impl BuiltinWorkflow for SimilarityBenchmark { let metric = &metric; let max_workers = config.base.max_workers.unwrap_or_else(get_max_workers); - let scores = DatasetRunner::new(&manager, self.dataset_id, &info, limit) + let scores = DatasetRunner::new(&manager, dataset_id, &info, limit) .desc(self.name) .set_quiet(ctx.quiet) .for_each_concurrent(max_workers, |i, row| { diff --git a/cli/src/commands/resume.rs b/cli/src/commands/resume.rs index ef8033b..710e996 100644 --- a/cli/src/commands/resume.rs +++ b/cli/src/commands/resume.rs @@ -147,6 +147,7 @@ mod tests { let bench = qt::config::BenchmarkConfig::Builtin(qt::config::BuiltinBenchmarkConfig { type_: "builtin".to_owned(), samples: None, + dataset: "hf://quantiles/PubMedQA".to_owned(), model: None, max_workers: None, }); @@ -161,6 +162,7 @@ mod tests { let bench = qt::config::BenchmarkConfig::Builtin(qt::config::BuiltinBenchmarkConfig { type_: "builtin".to_owned(), samples: Some(10), + dataset: "hf://quantiles/PubMedQA".to_owned(), model: None, max_workers: None, }); diff --git a/cli/src/commands/run.rs b/cli/src/commands/run.rs index aeb0f07..c64acaa 100644 --- a/cli/src/commands/run.rs +++ b/cli/src/commands/run.rs @@ -131,12 +131,9 @@ fn assemble_builtin_input( } if let Some(bench) = bench { - if bench.samples.is_none() && bench.model.is_none() && bench.max_workers.is_none() { - return (None, Vec::new()); - } - let input = BuiltinConfigInput { limit: bench.samples, + dataset: bench.dataset.clone(), model: bench.model.clone(), max_workers: bench.max_workers, }; @@ -480,6 +477,7 @@ struct BuiltinRunJsonOutput { struct BuiltinConfigInput { #[serde(skip_serializing_if = "Option::is_none")] limit: Option, + dataset: String, #[serde(skip_serializing_if = "Option::is_none")] model: Option, #[serde(skip_serializing_if = "Option::is_none")] @@ -721,6 +719,7 @@ mod tests { let bench = qt::config::BuiltinBenchmarkConfig { type_: "builtin".to_owned(), samples: Some(10), + dataset: "hf://quantiles/PubMedQA".to_owned(), model: None, max_workers: None, }; @@ -735,28 +734,32 @@ mod tests { let bench = qt::config::BuiltinBenchmarkConfig { type_: "builtin".to_owned(), samples: Some(5), + dataset: "hf://quantiles/PubMedQA".to_owned(), model: Some(qt::llm::Sampler::Random {}), max_workers: Some(8), }; let (input, _) = super::assemble_builtin_input(Some(&bench), None); let parsed: serde_json::Value = serde_json::from_str(&input.unwrap()).unwrap(); assert_eq!(parsed["limit"], 5); + assert_eq!(parsed["dataset"], "hf://quantiles/PubMedQA"); assert_eq!(parsed["model"], "random"); assert_eq!(parsed["max_workers"], 8); } - /// When the builtin config section exists but has no runtime-relevant fields, the - /// input should be `None` rather than an empty JSON object. + /// When the builtin config section only has the required dataset, the input should + /// still carry that dataset source into builtin execution. #[test] - fn assemble_builtin_input_none_when_no_config_fields() { + fn assemble_builtin_input_with_dataset_only_config() { let bench = qt::config::BuiltinBenchmarkConfig { type_: "builtin".to_owned(), samples: None, + dataset: "hf://quantiles/PubMedQA".to_owned(), model: None, max_workers: None, }; let (input, _) = super::assemble_builtin_input(Some(&bench), None); - assert!(input.is_none()); + let parsed: serde_json::Value = serde_json::from_str(&input.unwrap()).unwrap(); + assert_eq!(parsed["dataset"], "hf://quantiles/PubMedQA"); } /// When there is no config section at all and no CLI `--input`, builtin runs should diff --git a/cli/src/config/mod.rs b/cli/src/config/mod.rs index f51de57..686df37 100644 --- a/cli/src/config/mod.rs +++ b/cli/src/config/mod.rs @@ -96,6 +96,8 @@ pub struct BuiltinBenchmarkConfig { pub type_: String, /// Number of samples (rows) to evaluate. pub samples: Option, + /// Dataset source for this benchmark. + pub dataset: String, /// Which model sampler to use for this benchmark. pub model: Option, /// Maximum number of concurrent workers for this benchmark. @@ -182,6 +184,7 @@ mod tests { fn deserialize_builtin_without_type() { let toml = r#" [benchmarks.demo] + dataset = "hf://quantiles/demo" samples = 10 "#; let config: WorkspaceConfig = toml::from_str(toml).unwrap(); @@ -190,6 +193,7 @@ mod tests { if let BenchmarkConfig::Builtin(b) = bench { assert_eq!(b.type_, "builtin"); assert_eq!(b.samples, Some(10)); + assert_eq!(b.dataset, "hf://quantiles/demo"); assert!(b.model.is_none()); } } @@ -200,11 +204,25 @@ mod tests { [benchmarks.demo] type = "builtin" samples = 5 + dataset = "hf://quantiles/demo" model = "openai:gpt-4" "#; let config: WorkspaceConfig = toml::from_str(toml).unwrap(); let bench = config.benchmarks.get("demo").unwrap(); assert!(matches!(bench, BenchmarkConfig::Builtin(_))); + if let BenchmarkConfig::Builtin(b) = bench { + assert_eq!(b.dataset, "hf://quantiles/demo"); + } + } + + #[test] + fn builtin_requires_dataset_field() { + let toml = r#" + [benchmarks.demo] + samples = 5 + "#; + let result: Result = toml::from_str(toml); + assert!(result.is_err(), "builtin should require dataset field"); } #[test] diff --git a/cli/src/dataset/mod.rs b/cli/src/dataset/mod.rs index e726d87..d144681 100644 --- a/cli/src/dataset/mod.rs +++ b/cli/src/dataset/mod.rs @@ -3,7 +3,7 @@ pub mod hf_client; use crate::dataset::cache::DatasetCache; use crate::dataset::hf_client::HuggingFaceClient; -use anyhow::{Context, Result}; +use anyhow::{Context, Result, bail}; use serde::Serialize; use serde_json::Value; @@ -156,3 +156,65 @@ impl DatasetManager { splits.first().cloned().context("dataset has no splits") } } + +/// Resolve a configured Hugging Face dataset source to the dataset ID expected +/// by the existing Hugging Face download client. +/// +/// # Errors +/// +/// If the given `source` is missing a `hf://` / `huggingface://` prefix, or the +/// source is otherwise unresolvable, this function returns a descriptive error. +pub fn resolve_hf_dataset_source(source: &str) -> Result<&str> { + if let Some(dataset_id) = source + .strip_prefix("hf://") + .or_else(|| source.strip_prefix("huggingface://")) + { + if dataset_id.is_empty() { + bail!("dataset source `{source}` is missing a Hugging Face dataset id"); + } + Ok(dataset_id) + } else if source.contains("://") { + bail!("unsupported dataset source `{source}`; expected `hf://...` or `huggingface://...`"); + } else { + bail!("dataset source `{source}` is missing required `hf://` or `huggingface://` prefix"); + } +} + +#[cfg(test)] +mod tests { + use super::resolve_hf_dataset_source; + + #[test] + fn resolve_hf_dataset_source_strips_hf_prefix() { + assert_eq!( + resolve_hf_dataset_source("hf://quantiles/PubMedQA").unwrap(), + "quantiles/PubMedQA" + ); + } + + #[test] + fn resolve_hf_dataset_source_strips_huggingface_prefix() { + assert_eq!( + resolve_hf_dataset_source("huggingface://quantiles/PubMedQA").unwrap(), + "quantiles/PubMedQA" + ); + } + + #[test] + fn resolve_hf_dataset_source_rejects_other_prefixes() { + let err = resolve_hf_dataset_source("s3://bucket/dataset").unwrap_err(); + assert!( + err.to_string() + .contains("unsupported dataset source `s3://bucket/dataset`") + ); + } + + #[test] + fn resolve_hf_dataset_source_requires_prefix() { + let err = resolve_hf_dataset_source("quantiles/PubMedQA").unwrap_err(); + assert!( + err.to_string() + .contains("missing required `hf://` or `huggingface://` prefix") + ); + } +} diff --git a/cli/src/server/mod.rs b/cli/src/server/mod.rs index 40b9631..94d9fbd 100644 --- a/cli/src/server/mod.rs +++ b/cli/src/server/mod.rs @@ -15,7 +15,7 @@ use sea_orm::DatabaseConnection; use serde_json::json; use tokio::net::TcpListener; -use crate::dataset::DatasetManager; +use crate::dataset::{DatasetManager, resolve_hf_dataset_source}; use crate::db::steps::{self, StepDecision}; use crate::db::{self, DBUrl, SQLitePathURL}; use crate::metrics_store::MetricsStore; @@ -321,14 +321,7 @@ async fn dataset_batch( } fn parse_hf_source(source: &str) -> Result<&str, ApiError> { - source - .strip_prefix("huggingface://") - .or_else(|| source.strip_prefix("hf://")) - .ok_or_else(|| { - ApiError(anyhow::anyhow!( - "unsupported dataset source `{source}`; expected `huggingface://...` or `hf://...`" - )) - }) + resolve_hf_dataset_source(source).map_err(ApiError) } struct ApiError(anyhow::Error);