From 70cdf9dcb4797fbce2ad388a1a7a37fc9dc4b136 Mon Sep 17 00:00:00 2001 From: Allen Cheng Date: Thu, 25 Jun 2026 00:07:35 -0700 Subject: [PATCH] feat: emit dataset statistics report alongside exports (#104) Follow-up to #96. Adds an optional `.stats.json` report so a training cut is auditable at a glance without loading the JSONL. - `ExportConfig.emit_stats` gates a stats artifact computed during the existing export pass (no extra dataset re-materialization). - `ExportStats` covers: - counts: examples, by `role`, by `source`, by `tenant`; - token stats (min/median/p95/max/mean) from `state_metadata.tokens_used`, falling back to a whitespace word-count proxy, with a `source` field (`tokens_used` / `length_proxy` / `mixed`); - grouping: `num_groups` + records-per-group distribution; - curation accounting: records excluded by reason (lifecycle, reward threshold, dedup, decontamination), derived from the curation counts; - reward distribution and `reward_source` breakdown (rollouts), plus the preference form. - With a train/eval split, each side gets its own `.stats.json`. Python: `ctx.export_training(..., emit_stats=True)` writes the sibling report; read it from `.stats.json`. Tests: core (role/source/tenant counts, token stats + length-proxy fallback, lifecycle exclusions, rollout reward distribution, flag gating) and python (stats contents + flag gating). README updated. Stacked on #96 (PR #111) and #103 (PR #112). Closes #104 --- README.md | 1 + crates/lance-context-core/src/export.rs | 347 +++++++++++++++++++++++- crates/lance-context-core/src/lib.rs | 7 +- python/python/lance_context/api.py | 9 +- python/src/lib.rs | 7 +- python/tests/test_export_training.py | 27 ++ 6 files changed, 391 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index b06cd82..b802e30 100644 --- a/README.md +++ b/README.md @@ -266,6 +266,7 @@ manifest = ctx.export_training( min_reward=0.7, # keep only high-reward completions dedup_threshold=0.02, # collapse near-duplicates (cosine) version=ctx.version(), # pin for reproducibility + emit_stats=True, # also write sft.jsonl.stats.json (counts, tokens, exclusions) ) print("SFT examples:", manifest["counts"]["examples"]) diff --git a/crates/lance-context-core/src/export.rs b/crates/lance-context-core/src/export.rs index 03d7e10..339053f 100644 --- a/crates/lance-context-core/src/export.rs +++ b/crates/lance-context-core/src/export.rs @@ -29,7 +29,7 @@ //! tombstoned/expired/retired/superseded) and additionally drops //! `contradicted` records. -use std::collections::HashMap; +use std::collections::{BTreeMap, HashMap}; use std::io::{BufWriter, Write}; use chrono::{DateTime, Utc}; @@ -178,6 +178,8 @@ pub struct ExportConfig { pub filters_summary: Option, /// When set, write group-disjoint `train` / `eval` outputs instead of one. pub split: Option, + /// When `true`, also write a `.stats.json` dataset report. + pub emit_stats: bool, } impl Default for ExportConfig { @@ -195,6 +197,7 @@ impl Default for ExportConfig { version: None, filters_summary: None, split: None, + emit_stats: false, } } } @@ -344,6 +347,85 @@ pub struct ExportManifest { pub counts: ExportCounts, } +/// Numeric distribution summary. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct Distribution { + pub count: usize, + pub min: f64, + pub median: f64, + pub p95: f64, + pub max: f64, + pub mean: f64, +} + +impl Distribution { + fn from_sorted(mut values: Vec) -> Option { + if values.is_empty() { + return None; + } + values.sort_by(f64::total_cmp); + let count = values.len(); + let sum: f64 = values.iter().sum(); + let percentile = |p: f64| { + let idx = ((p * (count - 1) as f64).round() as usize).min(count - 1); + values[idx] + }; + Some(Self { + count, + min: values[0], + median: percentile(0.5), + p95: percentile(0.95), + max: values[count - 1], + mean: sum / count as f64, + }) + } +} + +/// Token-length statistics and which source produced them. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct TokenStats { + #[serde(flatten)] + pub distribution: Distribution, + /// `"tokens_used"`, `"length_proxy"`, or `"mixed"`. + pub source: String, +} + +/// Source records excluded during curation, by reason. +#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize)] +pub struct ExcludedCounts { + pub lifecycle: usize, + pub reward_threshold: usize, + pub dedup: usize, + pub decontaminate: usize, +} + +/// Auditable dataset statistics for one export output, written to +/// `.stats.json`. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ExportStats { + pub task: String, + pub examples: usize, + pub num_groups: usize, + /// Record counts by `role`. + pub by_role: BTreeMap, + /// Record counts by `source` (`"__none__"` when absent). + pub by_source: BTreeMap, + /// Record counts by `tenant` (`"__none__"` when absent). + pub by_tenant: BTreeMap, + pub records_per_group: Distribution, + #[serde(skip_serializing_if = "Option::is_none", default)] + pub tokens: Option, + pub excluded: ExcludedCounts, + #[serde(skip_serializing_if = "Option::is_none", default)] + pub preference_form: Option, + /// Reward distribution over records carrying `metadata.reward`. + #[serde(skip_serializing_if = "Option::is_none", default)] + pub reward: Option, + /// Counts by `metadata.reward_source`. + #[serde(skip_serializing_if = "BTreeMap::is_empty", default)] + pub reward_sources: BTreeMap, +} + // ----- metadata helpers ----------------------------------------------------- fn metadata_field<'a>(record: &'a ContextRecord, key: &str) -> Option<&'a Value> { @@ -891,9 +973,126 @@ fn emit_export( let manifest_json = serde_json::to_string_pretty(&manifest).map_err(|err| LanceError::io(err.to_string()))?; std::fs::write(format!("{output_path}.manifest.json"), manifest_json)?; + + if config.emit_stats { + let stats = compute_stats(&groups, &counts, examples, config); + let stats_json = + serde_json::to_string_pretty(&stats).map_err(|err| LanceError::io(err.to_string()))?; + std::fs::write(format!("{output_path}.stats.json"), stats_json)?; + } + Ok(manifest) } +/// Compute the dataset statistics report for one export output. +fn compute_stats( + groups: &[Group], + counts: &ExportCounts, + examples: usize, + config: &ExportConfig, +) -> ExportStats { + let mut by_role: BTreeMap = BTreeMap::new(); + let mut by_source: BTreeMap = BTreeMap::new(); + let mut by_tenant: BTreeMap = BTreeMap::new(); + let mut token_values: Vec = Vec::new(); + let mut used_tokens_used = false; + let mut used_fallback = false; + let mut reward_values: Vec = Vec::new(); + let mut reward_sources: BTreeMap = BTreeMap::new(); + let mut records_per_group: Vec = Vec::new(); + + for group in groups { + records_per_group.push(group.records.len() as f64); + for record in &group.records { + *by_role.entry(record.role.clone()).or_insert(0) += 1; + *by_source + .entry( + record + .source + .clone() + .unwrap_or_else(|| "__none__".to_string()), + ) + .or_insert(0) += 1; + *by_tenant + .entry( + record + .tenant + .clone() + .unwrap_or_else(|| "__none__".to_string()), + ) + .or_insert(0) += 1; + + match record.state_metadata.as_ref().and_then(|m| m.tokens_used) { + Some(tokens) if tokens >= 0 => { + token_values.push(f64::from(tokens)); + used_tokens_used = true; + } + _ => { + // Fallback proxy: whitespace-delimited word count of content. + let proxy = record + .text_payload + .as_deref() + .map_or(0, |text| text.split_whitespace().count()); + token_values.push(proxy as f64); + used_fallback = true; + } + } + + if let Some(reward) = record_reward(record) { + reward_values.push(reward); + } + if let Some(source) = record_reward_source(record) { + *reward_sources.entry(source).or_insert(0) += 1; + } + } + } + + let tokens = Distribution::from_sorted(token_values).map(|distribution| TokenStats { + distribution, + source: match (used_tokens_used, used_fallback) { + (true, true) => "mixed", + (true, false) => "tokens_used", + _ => "length_proxy", + } + .to_string(), + }); + + let excluded = ExcludedCounts { + lifecycle: counts.input_records.saturating_sub(counts.after_lifecycle), + reward_threshold: counts + .after_lifecycle + .saturating_sub(counts.after_reward_filter), + dedup: counts + .after_reward_filter + .saturating_sub(counts.after_dedup), + decontaminate: counts + .after_dedup + .saturating_sub(counts.after_decontaminate), + }; + + ExportStats { + task: config.task.as_str().to_string(), + examples, + num_groups: groups.len(), + by_role, + by_source, + by_tenant, + records_per_group: Distribution::from_sorted(records_per_group).unwrap_or_default(), + tokens, + excluded, + preference_form: matches!(config.task, ExportTask::Preference).then(|| { + match config.preference_form { + PreferenceForm::Paired => "paired", + PreferenceForm::Unpaired => "unpaired", + PreferenceForm::Ranked => "ranked", + } + .to_string() + }), + reward: Distribution::from_sorted(reward_values), + reward_sources, + } +} + impl ContextStore { /// Curate stored records and export them as task-shaped JSONL plus a /// sibling `.manifest.json`, returning the manifest. @@ -1669,4 +1868,150 @@ mod tests { assert_eq!(eval_manifest.split.unwrap().side, "eval"); }); } + + fn read_stats(path: &str) -> ExportStats { + let raw = std::fs::read_to_string(format!("{path}.stats.json")).unwrap(); + serde_json::from_str(&raw).unwrap() + } + + #[test] + fn stats_report_counts_roles_tokens_and_exclusions() { + let dir = TempDir::new().unwrap(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + runtime.block_on(async { + let mut store = open_store(&dir).await; + let mut user = rec("u", "user", "hello there", 1); + user.source = Some("memory".to_string()); + user.tenant = Some("acme".to_string()); + user.state_metadata = Some(crate::record::StateMetadata { + tokens_used: Some(5), + ..Default::default() + }); + let mut asst = rec("a", "assistant", "hi", 2); + asst.source = Some("memory".to_string()); + asst.tenant = Some("acme".to_string()); + asst.state_metadata = Some(crate::record::StateMetadata { + tokens_used: Some(11), + ..Default::default() + }); + // a contradicted record that curation excludes by lifecycle + let mut dropped = rec("d", "user", "nope", 3); + dropped.session_id = Some("other".to_string()); + dropped.lifecycle_status = LIFECYCLE_CONTRADICTED.to_string(); + store.add(&[user, asst, dropped]).await.unwrap(); + + let out = out_path(&dir); + store + .export_training( + &ExportConfig { + task: ExportTask::Sft, + group_by: GroupBy::SessionId, + emit_stats: true, + ..Default::default() + }, + &out, + ) + .await + .unwrap(); + + let stats = read_stats(&out); + assert_eq!(stats.task, "sft"); + assert_eq!(stats.examples, 1); + assert_eq!(stats.num_groups, 1); + assert_eq!(stats.by_role.get("user"), Some(&1)); + assert_eq!(stats.by_role.get("assistant"), Some(&1)); + assert_eq!(stats.by_source.get("memory"), Some(&2)); + assert_eq!(stats.by_tenant.get("acme"), Some(&2)); + assert_eq!(stats.excluded.lifecycle, 1, "contradicted record excluded"); + + let tokens = stats.tokens.unwrap(); + assert_eq!(tokens.source, "tokens_used"); + assert_eq!(tokens.distribution.count, 2); + assert_eq!(tokens.distribution.min, 5.0); + assert_eq!(tokens.distribution.max, 11.0); + assert_eq!(tokens.distribution.mean, 8.0); + }); + } + + #[test] + fn stats_token_fallback_uses_length_proxy() { + let dir = TempDir::new().unwrap(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + runtime.block_on(async { + let mut store = open_store(&dir).await; + // no state_metadata -> fallback to whitespace word count + store + .add(&[rec("u", "user", "one two three four", 1)]) + .await + .unwrap(); + + let out = out_path(&dir); + store + .export_training( + &ExportConfig { + emit_stats: true, + ..Default::default() + }, + &out, + ) + .await + .unwrap(); + + let tokens = read_stats(&out).tokens.unwrap(); + assert_eq!(tokens.source, "length_proxy"); + assert_eq!(tokens.distribution.max, 4.0); + }); + } + + #[test] + fn stats_report_reward_distribution_for_rollout() { + let dir = TempDir::new().unwrap(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + runtime.block_on(async { + let mut store = open_store(&dir).await; + let prompt = rec("p", "user", "solve", 1); + let mut r1 = rec("r1", "assistant", "a1", 2); + r1.metadata = Some(json!({"reward": 1.0, "reward_source": "verifier"})); + let mut r2 = rec("r2", "assistant", "a2", 3); + r2.metadata = Some(json!({"reward": 0.0, "reward_source": "verifier"})); + store.add(&[prompt, r1, r2]).await.unwrap(); + + let out = out_path(&dir); + store + .export_training( + &ExportConfig { + task: ExportTask::Rollout, + group_by: GroupBy::SessionId, + emit_stats: true, + ..Default::default() + }, + &out, + ) + .await + .unwrap(); + + let stats = read_stats(&out); + let reward = stats.reward.unwrap(); + assert_eq!(reward.count, 2); + assert_eq!(reward.min, 0.0); + assert_eq!(reward.max, 1.0); + assert_eq!(stats.reward_sources.get("verifier"), Some(&2)); + }); + } + + #[test] + fn stats_not_written_without_flag() { + let dir = TempDir::new().unwrap(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + runtime.block_on(async { + let mut store = open_store(&dir).await; + store.add(&[rec("u", "user", "hi", 1)]).await.unwrap(); + let out = out_path(&dir); + store + .export_training(&ExportConfig::default(), &out) + .await + .unwrap(); + assert!(std::fs::metadata(format!("{out}.stats.json")).is_err()); + }); + } } diff --git a/crates/lance-context-core/src/lib.rs b/crates/lance-context-core/src/lib.rs index ae79d3c..fe136a6 100644 --- a/crates/lance-context-core/src/lib.rs +++ b/crates/lance-context-core/src/lib.rs @@ -11,9 +11,10 @@ mod store; pub use context::{Context, ContextEntry, Snapshot}; pub use export::{ - ExportConfig, ExportCounts, ExportManifest, ExportTask, GroupBy, Message, PreferenceExample, - PreferenceForm, Provenance, RankedCandidate, RolloutExample, RolloutResponse, SftExample, - SplitConfig, SplitManifest, EXPORT_SCHEMA_VERSION, + Distribution, ExcludedCounts, ExportConfig, ExportCounts, ExportManifest, ExportStats, + ExportTask, GroupBy, Message, PreferenceExample, PreferenceForm, Provenance, RankedCandidate, + RolloutExample, RolloutResponse, SftExample, SplitConfig, SplitManifest, TokenStats, + EXPORT_SCHEMA_VERSION, }; pub use namespace::{ContextNamespace, PartitionInfo, PartitionSelector, PartitionSpec}; pub use record::{ diff --git a/python/python/lance_context/api.py b/python/python/lance_context/api.py index f3718b9..4ce5225 100644 --- a/python/python/lance_context/api.py +++ b/python/python/lance_context/api.py @@ -1260,6 +1260,7 @@ def export_training( include_expired: bool = False, include_retired: bool = False, split: dict[str, Any] | None = None, + emit_stats: bool = False, ) -> dict[str, Any]: """Curate stored records and export a trainable dataset to JSONL. @@ -1280,10 +1281,15 @@ def export_training( ``reward_source``, ``group_id``, ``label``, ``rank``). Pass ``version`` to pin the export to a dataset version for reproducibility. - Pass ``split={"eval_fraction": 0.1, "by": "session_id", "seed": 42}`` + Pass ``split={"eval_fraction": 0.1, "by": "session_id", "seed": 42}`` to write group-disjoint, reproducible ``.train.jsonl`` and ``.eval.jsonl`` outputs (each with its own manifest) instead of a single file; the returned manifest is the train side. + + Pass ``emit_stats=True`` to also write a ``.stats.json`` + dataset report (example/role/source/tenant counts, token-length + distribution, per-group counts, curation exclusions by reason, and + reward distribution). """ if format != "jsonl": raise ValueError("export format currently supports only 'jsonl'") @@ -1312,6 +1318,7 @@ def export_training( split_eval_fraction, split_by, split_seed, + emit_stats, ) return json.loads(manifest_json) diff --git a/python/src/lib.rs b/python/src/lib.rs index 6bedae4..84d20d0 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -294,6 +294,7 @@ fn export_config( version, filters_summary, split, + emit_stats: false, }) } @@ -794,7 +795,7 @@ impl Context { } #[allow(clippy::too_many_arguments)] - #[pyo3(signature = (output_path, task = "sft", group_by = "session_id", preference_form = "paired", filters_json = None, dedup_threshold = None, decontaminate_against = None, decontaminate_threshold = None, min_reward = None, version = None, include_expired = false, include_retired = false, split_eval_fraction = None, split_by = None, split_seed = None))] + #[pyo3(signature = (output_path, task = "sft", group_by = "session_id", preference_form = "paired", filters_json = None, dedup_threshold = None, decontaminate_against = None, decontaminate_threshold = None, min_reward = None, version = None, include_expired = false, include_retired = false, split_eval_fraction = None, split_by = None, split_seed = None, emit_stats = false))] fn export_training( &mut self, py: Python<'_>, @@ -813,8 +814,9 @@ impl Context { split_eval_fraction: Option, split_by: Option, split_seed: Option, + emit_stats: bool, ) -> PyResult { - let config = export_config( + let mut config = export_config( task, group_by, preference_form, @@ -830,6 +832,7 @@ impl Context { split_by, split_seed, )?; + config.emit_stats = emit_stats; let manifest = py .allow_threads(|| { self.runtime diff --git a/python/tests/test_export_training.py b/python/tests/test_export_training.py index f0d0e95..aa6e594 100644 --- a/python/tests/test_export_training.py +++ b/python/tests/test_export_training.py @@ -184,3 +184,30 @@ def test_export_split_is_deterministic(tmp_path: Path) -> None: assert (tmp_path / "a.train.jsonl").read_text() == ( tmp_path / "b.train.jsonl" ).read_text() + + +def test_export_emits_stats_report(tmp_path: Path) -> None: + ctx = Context.create(str(tmp_path / "ctx.lance")) + ctx.add("user", "hello there friend", session_id="s1", source="memory") + ctx.add("assistant", "hi", session_id="s1", source="memory") + + out = tmp_path / "sft.jsonl" + ctx.export_training(str(out), task="sft", emit_stats=True) + + stats_path = tmp_path / "sft.jsonl.stats.json" + assert stats_path.exists() + stats = json.loads(stats_path.read_text()) + assert stats["examples"] == 1 + assert stats["by_role"]["user"] == 1 + assert stats["by_role"]["assistant"] == 1 + assert stats["by_source"]["memory"] == 2 + assert stats["tokens"]["source"] == "length_proxy" + assert stats["tokens"]["max"] == 3.0 # "hello there friend" + + +def test_export_no_stats_without_flag(tmp_path: Path) -> None: + ctx = Context.create(str(tmp_path / "ctx.lance")) + ctx.add("user", "hi", session_id="s1") + out = tmp_path / "sft.jsonl" + ctx.export_training(str(out), task="sft") + assert not (tmp_path / "sft.jsonl.stats.json").exists()