diff --git a/CONFIG.md b/CONFIG.md index ab4e093..5ef1bef 100644 --- a/CONFIG.md +++ b/CONFIG.md @@ -26,6 +26,7 @@ For example, if you want to override default parameters for the built-in PubMedQ # Configure the model and sample limit for the built-in PubMedQA benchmark. [benchmarks.pubmedqa] +dataset = "hf://quantiles/PubMedQA" samples = 50 model = "openai:gpt-5.6" ``` @@ -41,6 +42,7 @@ Built-in benchmarks run natively inside the CLI, without any custom code. Below | Field | Type | Required | Description | | ------------- | --------------- | -------- | ---------------------------------------------------------------- | | `type` | string | no | Defaults to `"builtin"`. May be omitted for built-in benchmarks. | +| `dataset` | string | yes | Dataset source. Hugging Face datasets must use `hf://...` or `huggingface://...`. | | `samples` | integer | no | Number of dataset rows to evaluate. | | `model` | string or table | no | Model sampler. See [model format](#model-format). | | `max_workers` | integer | no | Maximum concurrent workers. | @@ -48,10 +50,24 @@ Built-in benchmarks run natively inside the CLI, without any custom code. Below If none of these fields are customized, the built-in benchmark uses the following defaults: - `type`: `builtin` +- `dataset`: The benchmark's default Hugging Face dataset, used when no config section is provided. - `samples`: All samples available in the benchmark's dataset, in order - `model`: The "demo" model, which outputs random values - `max_workers`: The default parallelism provided by the Rust [Tokio runtime](https://tokio.rs/) +#### `dataset` naming + +The `dataset` field accepts a Hugging Face dataset source prefixed with `hf://` +or `huggingface://`. For example: + +```toml +dataset = "hf://quantiles/PubMedQA" +``` + +The CLI resolves `hf://quantiles/PubMedQA` to the Hugging Face dataset ID +`quantiles/PubMedQA` before using the normal Hugging Face download path. +Other URI-style prefixes are rejected. + #### `model` naming The `model` field described above accepts a provider-prefixed string, for example: @@ -156,6 +172,7 @@ When you run `qt resume `, the CLI looks up the stored eval name and inp ```toml [benchmarks.pubmedqa] +dataset = "hf://quantiles/PubMedQA" model = "openai:gpt-5.6" ``` @@ -163,6 +180,7 @@ model = "openai:gpt-5.6" ```toml [benchmarks.simpleqa-verified] +dataset = "hf://quantiles/simpleqa-verified" samples = 10 ``` diff --git a/README.md b/README.md index f0c781f..2da6ce6 100644 --- a/README.md +++ b/README.md @@ -79,6 +79,7 @@ For **built-in benchmarks**, configure settings like `samples`, `model`, and `ma ```toml [benchmarks.pubmedqa] +dataset = "hf://quantiles/PubMedQA" samples = 50 model = "openai:gpt-5.6" max_workers = 100 diff --git a/cli/README.md b/cli/README.md index 8c75c28..934f8cc 100644 --- a/cli/README.md +++ b/cli/README.md @@ -60,10 +60,11 @@ You can customize how the CLI executes built-in benchmarks and custom evaluation ### Built-in benchmarks -For built-in benchmarks, configure settings like `samples`, `model`, and `max_workers`: +For built-in benchmarks, configure settings like `dataset`, `samples`, `model`, and `max_workers`: ```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 7405860..05f6e8a 100644 --- a/cli/examples/configs/anthropic/quantiles.toml +++ b/cli/examples/configs/anthropic/quantiles.toml @@ -5,6 +5,8 @@ # 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 1000 samples. Since Anthropic charges for # usage, we are limiting the number of samples to 50 in this demo, to keep costs lower. samples = 50 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 29bf1cd..ae4f45c 100644 --- a/cli/examples/configs/cloudflare/quantiles.toml +++ b/cli/examples/configs/cloudflare/quantiles.toml @@ -13,6 +13,8 @@ # https://dash.cloudflare.com//ai/ai-gateway [benchmarks.pubmedqa] +dataset = "hf://quantiles/PubMedQA" + # The full expert-labeled PubMedQA benchmark has 1000 samples. Since Cloudflare charges for # usage, we are limiting the number of samples to 50 in this demo, to keep costs lower. samples = 50 diff --git a/cli/examples/configs/gemini/quantiles.toml b/cli/examples/configs/gemini/quantiles.toml index cfc348c..6c5c893 100644 --- a/cli/examples/configs/gemini/quantiles.toml +++ b/cli/examples/configs/gemini/quantiles.toml @@ -7,6 +7,8 @@ # https://aistudio.google.com/prompts/new_chat [benchmarks.pubmedqa] +dataset = "hf://quantiles/PubMedQA" + # The full expert-labeled PubMedQA benchmark has 1000 samples. Since Google charges for # usage, we are limiting the number of samples to 50 in this demo, to keep costs lower. samples = 50 diff --git a/cli/examples/configs/openai/quantiles.toml b/cli/examples/configs/openai/quantiles.toml index db46a7a..4e31d8d 100644 --- a/cli/examples/configs/openai/quantiles.toml +++ b/cli/examples/configs/openai/quantiles.toml @@ -8,6 +8,8 @@ # https://platform.openai.com/home [benchmarks.pubmedqa] +dataset = "hf://quantiles/PubMedQA" + # The full expert-labeled PubMedQA benchmark has 1000 samples. Since OpenAI charges for # usage, we are limiting the number of samples to 50 in this demo, to keep costs lower. samples = 50 diff --git a/cli/examples/configs/simple/quantiles.toml b/cli/examples/configs/simple/quantiles.toml index 06b287e..9f2998f 100644 --- a/cli/examples/configs/simple/quantiles.toml +++ b/cli/examples/configs/simple/quantiles.toml @@ -8,7 +8,9 @@ # number of 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 53ea370..6ed37fa 100644 --- a/cli/src/builtins/common.rs +++ b/cli/src/builtins/common.rs @@ -18,6 +18,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 7ea162c..d6e6cc5 100644 --- a/cli/src/builtins/pubmedqa/mod.rs +++ b/cli/src/builtins/pubmedqa/mod.rs @@ -7,7 +7,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::LLMSampler; use crate::llm::random_label::RandomLabelSampler; @@ -22,6 +22,8 @@ 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 { @@ -49,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?; @@ -62,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 051a260..0fd9083 100644 --- a/cli/src/builtins/similarity.rs +++ b/cli/src/builtins/similarity.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::LLMSampler; use crate::llm::random::RandomSampler; use crate::similarity::{ @@ -45,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, } @@ -53,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", }; @@ -61,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", }; @@ -96,7 +96,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 @@ -109,6 +115,7 @@ impl BuiltinWorkflow for SimilarityBenchmark { set_builtin_run_input( db, run_id, + dataset_source, config.base.model.as_ref(), limit, config.base.max_workers, @@ -120,7 +127,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 91b61af..86e68ed 100644 --- a/cli/src/commands/resume.rs +++ b/cli/src/commands/resume.rs @@ -132,6 +132,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, }); @@ -146,6 +147,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 18fcb06..ab536d8 100644 --- a/cli/src/commands/run.rs +++ b/cli/src/commands/run.rs @@ -103,12 +103,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, }; @@ -382,6 +379,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")] @@ -623,6 +621,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, }; @@ -637,28 +636,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.rs b/cli/src/config.rs index 0da4a18..2afa4b8 100644 --- a/cli/src/config.rs +++ b/cli/src/config.rs @@ -76,6 +76,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 concurrent workers for this benchmark. @@ -159,6 +161,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(); @@ -167,6 +170,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()); } } @@ -177,11 +181,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 b740fd8..5242ff2 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; @@ -145,3 +145,60 @@ 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. +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);