diff --git a/README.md b/README.md index d693166..d401b77 100644 --- a/README.md +++ b/README.md @@ -232,6 +232,28 @@ ctx.checkout(first_version) print("Entries after checkout:", ctx.entries()) +# Curate stored records into a trainable dataset and export it as JSONL plus a +# reproducible manifest. Curation (lifecycle-correct filtering, semantic dedup, +# decontamination against a holdout set, reward thresholding) runs before +# export; reward/preference signals are read from each record's `metadata` +# (`reward`, `reward_source`, `group_id`, `label`, `rank`). +# +# SFT, grouped by session into conversations (rejection-sampling via min_reward): +manifest = ctx.export_training( + "training/sft.jsonl", + task="sft", + group_by="session_id", + filters={"tenant": "acme", "source": "memory"}, + min_reward=0.7, # keep only high-reward completions + dedup_threshold=0.02, # collapse near-duplicates (cosine) + version=ctx.version(), # pin for reproducibility +) +print("SFT examples:", manifest["counts"]["examples"]) + +# Preference (DPO paired / KTO unpaired / judge-ranked) and RL rollout shapes: +ctx.export_training("training/pref.jsonl", task="preference", preference_form="paired") +ctx.export_training("training/rollout.jsonl", task="rollout") # GRPO/RLVR groups + # Remote persistence on any object_store backend uses a generic `storage_options` # dict, matching the conventions used by `lance` and `lance-graph`. # diff --git a/crates/lance-context-core/src/export.rs b/crates/lance-context-core/src/export.rs new file mode 100644 index 0000000..01f5b66 --- /dev/null +++ b/crates/lance-context-core/src/export.rs @@ -0,0 +1,1364 @@ +//! Curate stored [`ContextRecord`]s into trainable datasets and export them as +//! SFT / preference / RL-rollout JSONL plus a reproducible manifest. +//! +//! This is the downstream half of the post-training pipeline: raw logs are +//! ingested into faithful `ContextRecord`s, then *curated* (lifecycle-correct +//! filtering, semantic dedup, decontamination, reward thresholding), grouped +//! into conversations/trajectories, and *exported* into the shape a trainer +//! expects. +//! +//! ## Field conventions +//! +//! Reward / preference signals live in each record's free-form `metadata` +//! object: +//! - `metadata.reward` (number) — scalar reward / verifier score. +//! - `metadata.reward_source` (string) — e.g. `"verifier"`, `"tests"`, +//! `"judge"`. +//! - `metadata.group_id` (string) — links the N samples generated for one +//! prompt (GRPO/RLVR groups, candidate sets). +//! - `metadata.label` (string `"chosen"`/`"rejected"`) — unpaired (KTO-style) +//! preference label. +//! - `metadata.rank` (integer, 1 = best) — N-way ranking position. +//! +//! Messages are built from each record's `role` + `text_payload`. +//! +//! ## Curation lifecycle +//! +//! Curation is *selection only* — it never deletes or tombstones records. By +//! default it keeps only records visible under default lifecycle (drops +//! tombstoned/expired/retired/superseded) and additionally drops +//! `contradicted` records. + +use std::collections::HashMap; +use std::io::{BufWriter, Write}; + +use chrono::{DateTime, Utc}; +use lance::{Error as LanceError, Result as LanceResult}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +use crate::record::{ContextRecord, LifecycleQueryOptions, RecordFilters, LIFECYCLE_CONTRADICTED}; +use crate::store::ContextStore; + +/// Schema version stamped into every export manifest. +pub const EXPORT_SCHEMA_VERSION: &str = "1"; + +/// Which training shape to export. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)] +#[serde(rename_all = "lowercase")] +pub enum ExportTask { + /// Supervised fine-tuning (also the rejection-sampling / Best-of-N target). + #[default] + Sft, + /// Preference optimization (DPO/SimPO/ORPO paired, KTO unpaired, judge N-way). + Preference, + /// RL rollout (PPO/GRPO/RLVR): prompt + response group(s) with rewards. + Rollout, +} + +impl ExportTask { + #[must_use] + pub fn as_str(self) -> &'static str { + match self { + Self::Sft => "sft", + Self::Preference => "preference", + Self::Rollout => "rollout", + } + } +} + +/// Shape of a preference export (selected by the caller to match their labels). +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)] +#[serde(rename_all = "lowercase")] +pub enum PreferenceForm { + /// Paired `chosen`/`rejected` (DPO/SimPO/ORPO). + #[default] + Paired, + /// One completion with an unpaired binary label (KTO). + Unpaired, + /// N-way ranked candidate list (LLM-judge / RULER). + Ranked, +} + +/// How records are grouped into one training example. +#[derive(Debug, Clone, PartialEq, Eq, Default)] +pub enum GroupBy { + /// One example per record. + None, + #[default] + SessionId, + RunId, + Tenant, + Source, + BotId, + /// Group by the `external_id` prefix up to the first occurrence of the + /// delimiter (e.g. `"doc-7#chunk-1"` -> `"doc-7"` with delimiter `"#"`). + ExternalIdPrefix(String), +} + +impl GroupBy { + fn label(&self) -> String { + match self { + Self::None => "none".to_string(), + Self::SessionId => "session_id".to_string(), + Self::RunId => "run_id".to_string(), + Self::Tenant => "tenant".to_string(), + Self::Source => "source".to_string(), + Self::BotId => "bot_id".to_string(), + Self::ExternalIdPrefix(delim) => format!("external_id_prefix:{delim}"), + } + } + + /// Group key for a record; `None` records become their own singleton group. + fn key(&self, record: &ContextRecord) -> String { + let value = match self { + Self::None => None, + Self::SessionId => record.session_id.clone(), + Self::RunId => Some(record.run_id.clone()), + Self::Tenant => record.tenant.clone(), + Self::Source => record.source.clone(), + Self::BotId => record.bot_id.clone(), + Self::ExternalIdPrefix(delim) => record.external_id.as_ref().map(|external_id| { + external_id + .split_once(delim.as_str()) + .map_or(external_id.as_str(), |(prefix, _)| prefix) + .to_string() + }), + }; + value.unwrap_or_else(|| format!("__rec__{}", record.id)) + } +} + +/// Curation + export configuration. +#[derive(Clone)] +pub struct ExportConfig { + pub task: ExportTask, + pub group_by: GroupBy, + pub preference_form: PreferenceForm, + /// Metadata/quality filters applied before lifecycle curation. + pub filters: Option, + /// Lifecycle visibility; defaults to visible-only. + pub lifecycle: LifecycleQueryOptions, + /// Collapse near-duplicates whose cosine distance is `<=` this threshold. + pub dedup_threshold: Option, + /// Holdout embeddings to decontaminate against. + pub decontaminate_against: Vec>, + /// Drop records within this cosine distance of any holdout embedding. + pub decontaminate_threshold: Option, + /// Drop records whose `metadata.reward` is present and below this value + /// (rejection sampling / Best-of-N / quality threshold). + pub min_reward: Option, + /// Pin the export to a dataset version (time-travel); restored afterward. + pub version: Option, + /// Optional JSON summary of the applied selectors, recorded in the manifest. + pub filters_summary: Option, +} + +impl Default for ExportConfig { + fn default() -> Self { + Self { + task: ExportTask::Sft, + group_by: GroupBy::SessionId, + preference_form: PreferenceForm::Paired, + filters: None, + lifecycle: LifecycleQueryOptions::default(), + dedup_threshold: None, + decontaminate_against: Vec::new(), + decontaminate_threshold: None, + min_reward: None, + version: None, + filters_summary: None, + } + } +} + +/// A single chat message in a training example. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Message { + pub role: String, + pub content: String, +} + +/// Where an exported example came from, for auditability. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Provenance { + pub context_uri: String, + pub version: u64, + pub record_ids: Vec, + #[serde(skip_serializing_if = "Vec::is_empty", default)] + pub external_ids: Vec, + #[serde(skip_serializing_if = "Option::is_none", default)] + pub tenant: Option, + #[serde(skip_serializing_if = "Option::is_none", default)] + pub source: Option, + #[serde(skip_serializing_if = "Option::is_none", default)] + pub bot_id: Option, + #[serde(skip_serializing_if = "Option::is_none", default)] + pub session_id: Option, + #[serde(skip_serializing_if = "Option::is_none", default)] + pub run_id: Option, + pub created_at_start: DateTime, + pub created_at_end: DateTime, +} + +/// SFT example: an ordered message list (doubles as the rejection-sampling / +/// Best-of-N target when filtered by `min_reward`). +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct SftExample { + pub messages: Vec, + #[serde(skip_serializing_if = "Option::is_none", default)] + pub reward: Option, + pub provenance: Provenance, +} + +/// One ranked candidate in an N-way preference example. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct RankedCandidate { + pub messages: Vec, + pub rank: i64, + #[serde(skip_serializing_if = "Option::is_none", default)] + pub reward: Option, +} + +/// Preference example in one of three forms, tagged by `form`. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(tag = "form", rename_all = "lowercase")] +pub enum PreferenceExample { + Paired { + prompt: Vec, + chosen: Vec, + rejected: Vec, + provenance: Provenance, + }, + Unpaired { + prompt: Vec, + completion: Vec, + /// `true` = chosen, `false` = rejected. + label: bool, + provenance: Provenance, + }, + Ranked { + prompt: Vec, + candidates: Vec, + provenance: Provenance, + }, +} + +/// One response within a rollout example. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct RolloutResponse { + pub messages: Vec, + #[serde(skip_serializing_if = "Option::is_none", default)] + pub reward: Option, + #[serde(skip_serializing_if = "Option::is_none", default)] + pub reward_source: Option, +} + +/// RL rollout example: a prompt plus one or more rewarded responses, with an +/// optional `group_id` linking the N samples for one prompt. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct RolloutExample { + pub prompt: Vec, + pub responses: Vec, + #[serde(skip_serializing_if = "Option::is_none", default)] + pub group_id: Option, + pub provenance: Provenance, +} + +/// Record counts at each curation stage. +#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize)] +pub struct ExportCounts { + pub input_records: usize, + pub after_lifecycle: usize, + pub after_dedup: usize, + pub after_decontaminate: usize, + pub after_reward_filter: usize, + pub examples: usize, +} + +/// Reproducible description of an export, written next to the JSONL output. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ExportManifest { + pub context_uri: String, + pub version: u64, + pub task: String, + pub group_by: String, + pub schema_version: String, + #[serde(skip_serializing_if = "Option::is_none", default)] + pub preference_form: Option, + #[serde(skip_serializing_if = "Option::is_none", default)] + pub filters: Option, + #[serde(skip_serializing_if = "Option::is_none", default)] + pub dedup_threshold: Option, + #[serde(skip_serializing_if = "Option::is_none", default)] + pub decontaminate_threshold: Option, + #[serde(skip_serializing_if = "Option::is_none", default)] + pub min_reward: Option, + #[serde(skip_serializing_if = "Option::is_none", default)] + pub created_at_start: Option>, + #[serde(skip_serializing_if = "Option::is_none", default)] + pub created_at_end: Option>, + pub source_record_ids: Vec, + pub counts: ExportCounts, +} + +// ----- metadata helpers ----------------------------------------------------- + +fn metadata_field<'a>(record: &'a ContextRecord, key: &str) -> Option<&'a Value> { + record.metadata.as_ref()?.get(key) +} + +fn record_reward(record: &ContextRecord) -> Option { + metadata_field(record, "reward")?.as_f64() +} + +fn record_reward_source(record: &ContextRecord) -> Option { + Some( + metadata_field(record, "reward_source")? + .as_str()? + .to_string(), + ) +} + +fn record_group_id(record: &ContextRecord) -> Option { + Some(metadata_field(record, "group_id")?.as_str()?.to_string()) +} + +fn record_label(record: &ContextRecord) -> Option { + Some(metadata_field(record, "label")?.as_str()?.to_string()) +} + +fn record_rank(record: &ContextRecord) -> Option { + metadata_field(record, "rank")?.as_i64() +} + +fn message_of(record: &ContextRecord) -> Message { + Message { + role: record.role.clone(), + content: record.text_payload.clone().unwrap_or_default(), + } +} + +fn is_assistant(record: &ContextRecord) -> bool { + record.role.eq_ignore_ascii_case("assistant") +} + +/// Cosine distance in `0.0..=2.0`; `0` = identical direction. Returns `None` +/// for zero-magnitude or mismatched-length vectors. +fn cosine_distance(a: &[f32], b: &[f32]) -> Option { + if a.len() != b.len() || a.is_empty() { + return None; + } + let mut dot = 0.0f32; + let mut na = 0.0f32; + let mut nb = 0.0f32; + for (x, y) in a.iter().zip(b.iter()) { + dot += x * y; + na += x * x; + nb += y * y; + } + if na == 0.0 || nb == 0.0 { + return None; + } + Some(1.0 - dot / (na.sqrt() * nb.sqrt())) +} + +// ----- grouping ------------------------------------------------------------- + +/// Records grouped for one example, already ordered by `(created_at, id)`. +struct Group { + key: String, + records: Vec, +} + +fn group_records(records: Vec, group_by: &GroupBy) -> Vec { + let mut order: Vec = Vec::new(); + let mut groups: HashMap> = HashMap::new(); + for record in records { + let key = group_by.key(&record); + if !groups.contains_key(&key) { + order.push(key.clone()); + } + groups.entry(key).or_default().push(record); + } + + let mut result: Vec = order + .into_iter() + .map(|key| { + let mut records = groups.remove(&key).unwrap_or_default(); + records.sort_by(|a, b| { + a.created_at + .cmp(&b.created_at) + .then_with(|| a.id.cmp(&b.id)) + }); + Group { key, records } + }) + .collect(); + + // Deterministic group order: by each group's earliest (created_at, id). + result.sort_by(|a, b| { + let left = a.records.first(); + let right = b.records.first(); + match (left, right) { + (Some(l), Some(r)) => l + .created_at + .cmp(&r.created_at) + .then_with(|| l.id.cmp(&r.id)), + _ => a.key.cmp(&b.key), + } + }); + result +} + +fn provenance_for(records: &[ContextRecord], context_uri: &str, version: u64) -> Provenance { + let first = records.first(); + let created_at_start = records + .iter() + .map(|r| r.created_at) + .min() + .unwrap_or_else(Utc::now); + let created_at_end = records + .iter() + .map(|r| r.created_at) + .max() + .unwrap_or(created_at_start); + Provenance { + context_uri: context_uri.to_string(), + version, + record_ids: records.iter().map(|r| r.id.clone()).collect(), + external_ids: records + .iter() + .filter_map(|r| r.external_id.clone()) + .collect(), + tenant: first.and_then(|r| r.tenant.clone()), + source: first.and_then(|r| r.source.clone()), + bot_id: first.and_then(|r| r.bot_id.clone()), + session_id: first.and_then(|r| r.session_id.clone()), + run_id: first.map(|r| r.run_id.clone()), + created_at_start, + created_at_end, + } +} + +/// Split a group into leading non-assistant prompt messages and the assistant +/// candidate records that follow. +fn split_prompt_candidates(records: &[ContextRecord]) -> (Vec, Vec<&ContextRecord>) { + let mut prompt = Vec::new(); + let mut candidates = Vec::new(); + for record in records { + if is_assistant(record) { + candidates.push(record); + } else if candidates.is_empty() { + prompt.push(message_of(record)); + } else { + // Non-assistant turn after a candidate (e.g. tool result) — keep it + // attached to the conversation prompt context. + prompt.push(message_of(record)); + } + } + (prompt, candidates) +} + +fn write_line(writer: &mut impl Write, value: &T) -> LanceResult<()> { + let line = serde_json::to_string(value).map_err(|err| LanceError::io(err.to_string()))?; + writeln!(writer, "{line}")?; + Ok(()) +} + +// ----- curation ------------------------------------------------------------- + +/// Greedily collapse near-duplicates: keep a record only if its embedding is +/// not within `threshold` cosine distance of an already-kept record. Records +/// without an embedding are always kept. +fn dedup(records: Vec, threshold: f32) -> Vec { + let mut kept: Vec = Vec::new(); + for record in records { + let is_dup = record.embedding.as_ref().is_some_and(|embedding| { + kept.iter().any(|other| { + other + .embedding + .as_ref() + .and_then(|other_embedding| cosine_distance(embedding, other_embedding)) + .is_some_and(|distance| distance <= threshold) + }) + }); + if !is_dup { + kept.push(record); + } + } + kept +} + +/// Drop records whose embedding is within `threshold` cosine distance of any +/// holdout embedding. Records without an embedding are kept. +fn decontaminate( + records: Vec, + holdout: &[Vec], + threshold: f32, +) -> Vec { + records + .into_iter() + .filter(|record| match &record.embedding { + Some(embedding) => !holdout.iter().any(|held| { + cosine_distance(embedding, held).is_some_and(|distance| distance <= threshold) + }), + None => true, + }) + .collect() +} + +/// Apply the curation pipeline (lifecycle/contradicted, reward threshold, +/// dedup, decontamination) and return the curated records plus stage counts. +fn curate( + records: Vec, + config: &ExportConfig, +) -> (Vec, ExportCounts) { + let mut counts = ExportCounts { + input_records: records.len(), + ..ExportCounts::default() + }; + + // `list` already dropped tombstoned/expired/retired/superseded; also drop + // `contradicted`, which is visible by default. + let lifecycle: Vec = records + .into_iter() + .filter(|record| record.lifecycle_status != LIFECYCLE_CONTRADICTED) + .collect(); + counts.after_lifecycle = lifecycle.len(); + + let rewarded: Vec = match config.min_reward { + Some(min) => lifecycle + .into_iter() + .filter(|record| record_reward(record).is_none_or(|reward| reward >= min)) + .collect(), + None => lifecycle, + }; + counts.after_reward_filter = rewarded.len(); + + let deduped = match config.dedup_threshold { + Some(threshold) => dedup(rewarded, threshold), + None => rewarded, + }; + counts.after_dedup = deduped.len(); + + let clean = match config.decontaminate_threshold { + Some(threshold) if !config.decontaminate_against.is_empty() => { + decontaminate(deduped, &config.decontaminate_against, threshold) + } + _ => deduped, + }; + counts.after_decontaminate = clean.len(); + + (clean, counts) +} + +// ----- exporters ------------------------------------------------------------ + +fn write_sft( + groups: &[Group], + writer: &mut impl Write, + context_uri: &str, + version: u64, +) -> LanceResult { + let mut written = 0; + for group in groups { + if group.records.is_empty() { + continue; + } + let example = SftExample { + messages: group.records.iter().map(message_of).collect(), + reward: group + .records + .iter() + .filter_map(record_reward) + .reduce(f64::max), + provenance: provenance_for(&group.records, context_uri, version), + }; + write_line(writer, &example)?; + written += 1; + } + Ok(written) +} + +/// Preference score used to order candidates: an explicit label dominates a +/// numeric reward, which dominates "unscored". +fn preference_score(record: &ContextRecord) -> f64 { + match record_label(record).as_deref() { + Some("chosen") => f64::INFINITY, + Some("rejected") => f64::NEG_INFINITY, + _ => record_reward(record).unwrap_or(0.0), + } +} + +fn unpaired_label(record: &ContextRecord, min_reward: Option) -> Option { + match record_label(record).as_deref() { + Some("chosen") => Some(true), + Some("rejected") => Some(false), + _ => match (record_reward(record), min_reward) { + (Some(reward), Some(min)) => Some(reward >= min), + _ => None, + }, + } +} + +fn write_preference( + groups: &[Group], + writer: &mut impl Write, + form: PreferenceForm, + min_reward: Option, + context_uri: &str, + version: u64, +) -> LanceResult { + let mut written = 0; + for group in groups { + let (prompt, candidates) = split_prompt_candidates(&group.records); + if candidates.is_empty() { + continue; + } + let provenance = provenance_for(&group.records, context_uri, version); + + match form { + PreferenceForm::Paired => { + if candidates.len() < 2 { + continue; + } + let mut best = candidates[0]; + let mut worst = candidates[0]; + for candidate in &candidates { + if preference_score(candidate) > preference_score(best) { + best = candidate; + } + if preference_score(candidate) < preference_score(worst) { + worst = candidate; + } + } + if best.id == worst.id { + continue; // no signal to separate chosen from rejected + } + let example = PreferenceExample::Paired { + prompt, + chosen: vec![message_of(best)], + rejected: vec![message_of(worst)], + provenance, + }; + write_line(writer, &example)?; + written += 1; + } + PreferenceForm::Unpaired => { + for candidate in &candidates { + let Some(label) = unpaired_label(candidate, min_reward) else { + continue; + }; + let example = PreferenceExample::Unpaired { + prompt: prompt.clone(), + completion: vec![message_of(candidate)], + label, + provenance: provenance.clone(), + }; + write_line(writer, &example)?; + written += 1; + } + } + PreferenceForm::Ranked => { + let mut ordered: Vec<&ContextRecord> = candidates.clone(); + ordered.sort_by(|a, b| { + let rank_a = record_rank(a); + let rank_b = record_rank(b); + match (rank_a, rank_b) { + (Some(ra), Some(rb)) => ra.cmp(&rb), + _ => preference_score(b).total_cmp(&preference_score(a)), + } + }); + let candidates: Vec = ordered + .iter() + .enumerate() + .map(|(index, candidate)| RankedCandidate { + messages: vec![message_of(candidate)], + rank: record_rank(candidate).unwrap_or((index + 1) as i64), + reward: record_reward(candidate), + }) + .collect(); + let example = PreferenceExample::Ranked { + prompt, + candidates, + provenance, + }; + write_line(writer, &example)?; + written += 1; + } + } + } + Ok(written) +} + +fn write_rollout( + groups: &[Group], + writer: &mut impl Write, + context_uri: &str, + version: u64, +) -> LanceResult { + let mut written = 0; + for group in groups { + let (prompt, candidates) = split_prompt_candidates(&group.records); + if candidates.is_empty() { + continue; + } + let responses: Vec = candidates + .iter() + .map(|candidate| RolloutResponse { + messages: vec![message_of(candidate)], + reward: record_reward(candidate), + reward_source: record_reward_source(candidate), + }) + .collect(); + let group_id = candidates + .iter() + .find_map(|candidate| record_group_id(candidate)); + let example = RolloutExample { + prompt, + responses, + group_id, + provenance: provenance_for(&group.records, context_uri, version), + }; + write_line(writer, &example)?; + written += 1; + } + Ok(written) +} + +fn summarize_groups( + groups: &[Group], +) -> (Option>, Option>, Vec) { + let mut start: Option> = None; + let mut end: Option> = None; + let mut ids = Vec::new(); + for group in groups { + for record in &group.records { + start = Some(start.map_or(record.created_at, |s| s.min(record.created_at))); + end = Some(end.map_or(record.created_at, |e| e.max(record.created_at))); + ids.push(record.id.clone()); + } + } + (start, end, ids) +} + +impl ContextStore { + /// Curate stored records and export them as task-shaped JSONL plus a + /// sibling `.manifest.json`, returning the manifest. + /// + /// If `config.version` is set the export is pinned to that dataset version + /// (time-travel) and the store is restored to its current version + /// afterward. Output is written line-by-line, so the full JSONL is never + /// held in memory at once. + pub async fn export_training( + &mut self, + config: &ExportConfig, + output_path: &str, + ) -> LanceResult { + let restore = match config.version { + Some(target) => { + let original = self.version(); + self.checkout(target).await?; + Some(original) + } + None => None, + }; + + let result = self.export_inner(config, output_path).await; + + if let Some(original) = restore { + self.checkout(original).await?; + } + result + } + + async fn export_inner( + &self, + config: &ExportConfig, + output_path: &str, + ) -> LanceResult { + let context_uri = self.uri().to_string(); + let version = self.version(); + + let records = self + .list_filtered_with_options( + None, + None, + config.filters.as_ref(), + config.lifecycle.clone(), + ) + .await?; + let (curated, mut counts) = curate(records, config); + let groups = group_records(curated, &config.group_by); + let (created_at_start, created_at_end, source_record_ids) = summarize_groups(&groups); + + let file = std::fs::File::create(output_path)?; + let mut writer = BufWriter::new(file); + let examples = match config.task { + ExportTask::Sft => write_sft(&groups, &mut writer, &context_uri, version)?, + ExportTask::Preference => write_preference( + &groups, + &mut writer, + config.preference_form, + config.min_reward, + &context_uri, + version, + )?, + ExportTask::Rollout => write_rollout(&groups, &mut writer, &context_uri, version)?, + }; + writer.flush()?; + counts.examples = examples; + + let manifest = ExportManifest { + context_uri, + version, + task: config.task.as_str().to_string(), + group_by: config.group_by.label(), + schema_version: EXPORT_SCHEMA_VERSION.to_string(), + preference_form: matches!(config.task, ExportTask::Preference).then(|| { + match config.preference_form { + PreferenceForm::Paired => "paired", + PreferenceForm::Unpaired => "unpaired", + PreferenceForm::Ranked => "ranked", + } + .to_string() + }), + filters: config.filters_summary.clone(), + dedup_threshold: config.dedup_threshold, + decontaminate_threshold: config.decontaminate_threshold, + min_reward: config.min_reward, + created_at_start, + created_at_end, + source_record_ids, + counts, + }; + + 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)?; + + Ok(manifest) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::record::LIFECYCLE_ACTIVE; + use crate::store::{ContextStore, ContextStoreOptions}; + use chrono::TimeZone; + use serde_json::json; + use tempfile::TempDir; + + const DIM: i32 = 4; + + async fn open_store(dir: &TempDir) -> ContextStore { + let uri = dir.path().join("ctx.lance").to_string_lossy().to_string(); + ContextStore::open_with_options( + &uri, + ContextStoreOptions { + embedding_dim: Some(DIM), + ..Default::default() + }, + ) + .await + .unwrap() + } + + fn rec(id: &str, role: &str, text: &str, secs: i64) -> ContextRecord { + ContextRecord { + id: id.to_string(), + external_id: None, + run_id: "run".to_string(), + bot_id: None, + session_id: Some("s1".to_string()), + tenant: None, + source: None, + created_at: Utc.timestamp_opt(1_700_000_000 + secs, 0).unwrap(), + role: role.to_string(), + state_metadata: None, + metadata: None, + relationships: Vec::new(), + expires_at: None, + retention_policy: None, + lifecycle_status: LIFECYCLE_ACTIVE.to_string(), + retired_at: None, + retired_reason: None, + supersedes_id: None, + superseded_by_id: None, + content_type: "text/plain".to_string(), + text_payload: Some(text.to_string()), + binary_payload: None, + embedding: None, + } + } + + fn emb(lead: &[f32]) -> Vec { + let mut v = vec![0.0f32; DIM as usize]; + for (i, x) in lead.iter().enumerate() { + v[i] = *x; + } + v + } + + fn read_lines(path: &str) -> Vec { + std::fs::read_to_string(path) + .unwrap() + .lines() + .filter(|l| !l.trim().is_empty()) + .map(|l| serde_json::from_str(l).unwrap()) + .collect() + } + + fn read_manifest(path: &str) -> ExportManifest { + let raw = std::fs::read_to_string(format!("{path}.manifest.json")).unwrap(); + serde_json::from_str(&raw).unwrap() + } + + fn out_path(dir: &TempDir) -> String { + dir.path().join("out.jsonl").to_string_lossy().to_string() + } + + #[test] + fn sft_groups_session_into_ordered_conversation() { + let dir = TempDir::new().unwrap(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + runtime.block_on(async { + let mut store = open_store(&dir).await; + // Added out of order; export must order by (created_at, id). + store + .add(&[ + rec("r2", "assistant", "hi there", 2), + rec("r1", "user", "hello", 1), + rec("r3", "user", "bye", 3), + ]) + .await + .unwrap(); + + let out = out_path(&dir); + let manifest = store + .export_training( + &ExportConfig { + task: ExportTask::Sft, + group_by: GroupBy::SessionId, + ..Default::default() + }, + &out, + ) + .await + .unwrap(); + + let lines = read_lines(&out); + assert_eq!(lines.len(), 1); + let messages = lines[0]["messages"].as_array().unwrap(); + let contents: Vec<&str> = messages + .iter() + .map(|m| m["content"].as_str().unwrap()) + .collect(); + assert_eq!(contents, ["hello", "hi there", "bye"]); + assert_eq!(manifest.task, "sft"); + assert_eq!(manifest.counts.examples, 1); + assert_eq!(lines[0]["provenance"]["version"], json!(manifest.version)); + + // The sibling manifest file is written and matches the return value. + let manifest_file = read_manifest(&out); + assert_eq!(manifest_file.task, "sft"); + assert_eq!(manifest_file.counts.examples, 1); + assert_eq!(manifest_file.schema_version, EXPORT_SCHEMA_VERSION); + assert_eq!(manifest_file.source_record_ids.len(), 3); + }); + } + + #[test] + fn sft_rejection_sampling_filters_low_reward() { + 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 good = rec("g", "assistant", "good", 1); + good.session_id = Some("a".to_string()); + good.metadata = Some(json!({"reward": 0.9})); + let mut bad = rec("b", "assistant", "bad", 2); + bad.session_id = Some("b".to_string()); + bad.metadata = Some(json!({"reward": 0.1})); + store.add(&[good, bad]).await.unwrap(); + + let out = out_path(&dir); + let manifest = store + .export_training( + &ExportConfig { + task: ExportTask::Sft, + group_by: GroupBy::SessionId, + min_reward: Some(0.5), + ..Default::default() + }, + &out, + ) + .await + .unwrap(); + + let lines = read_lines(&out); + assert_eq!(lines.len(), 1, "only the high-reward record survives"); + assert_eq!(lines[0]["messages"][0]["content"], "good"); + assert_eq!(manifest.counts.after_reward_filter, 1); + assert_eq!(manifest.min_reward, Some(0.5)); + }); + } + + #[test] + fn preference_paired_uses_reward_for_chosen_rejected() { + 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", "question", 1); + let mut hi = rec("hi", "assistant", "great answer", 2); + hi.metadata = Some(json!({"reward": 0.9})); + let mut lo = rec("lo", "assistant", "poor answer", 3); + lo.metadata = Some(json!({"reward": 0.2})); + store.add(&[prompt, hi, lo]).await.unwrap(); + + let out = out_path(&dir); + store + .export_training( + &ExportConfig { + task: ExportTask::Preference, + preference_form: PreferenceForm::Paired, + group_by: GroupBy::SessionId, + ..Default::default() + }, + &out, + ) + .await + .unwrap(); + + let lines = read_lines(&out); + assert_eq!(lines.len(), 1); + assert_eq!(lines[0]["form"], "paired"); + assert_eq!(lines[0]["prompt"][0]["content"], "question"); + assert_eq!(lines[0]["chosen"][0]["content"], "great answer"); + assert_eq!(lines[0]["rejected"][0]["content"], "poor answer"); + }); + } + + #[test] + fn preference_unpaired_uses_kto_labels() { + 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", "q", 1); + let mut a = rec("a", "assistant", "yes", 2); + a.metadata = Some(json!({"label": "chosen"})); + let mut b = rec("b", "assistant", "no", 3); + b.metadata = Some(json!({"label": "rejected"})); + store.add(&[prompt, a, b]).await.unwrap(); + + let out = out_path(&dir); + store + .export_training( + &ExportConfig { + task: ExportTask::Preference, + preference_form: PreferenceForm::Unpaired, + group_by: GroupBy::SessionId, + ..Default::default() + }, + &out, + ) + .await + .unwrap(); + + let lines = read_lines(&out); + assert_eq!(lines.len(), 2); + assert!(lines.iter().all(|l| l["form"] == "unpaired")); + let chosen = lines + .iter() + .find(|l| l["completion"][0]["content"] == "yes") + .unwrap(); + assert_eq!(chosen["label"], json!(true)); + let rejected = lines + .iter() + .find(|l| l["completion"][0]["content"] == "no") + .unwrap(); + assert_eq!(rejected["label"], json!(false)); + }); + } + + #[test] + fn preference_ranked_orders_by_rank() { + 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", "q", 1); + let mut second = rec("c2", "assistant", "second", 2); + second.metadata = Some(json!({"rank": 2})); + let mut first = rec("c1", "assistant", "first", 3); + first.metadata = Some(json!({"rank": 1})); + store.add(&[prompt, second, first]).await.unwrap(); + + let out = out_path(&dir); + store + .export_training( + &ExportConfig { + task: ExportTask::Preference, + preference_form: PreferenceForm::Ranked, + group_by: GroupBy::SessionId, + ..Default::default() + }, + &out, + ) + .await + .unwrap(); + + let lines = read_lines(&out); + assert_eq!(lines.len(), 1); + assert_eq!(lines[0]["form"], "ranked"); + let cands = lines[0]["candidates"].as_array().unwrap(); + assert_eq!(cands[0]["messages"][0]["content"], "first"); + assert_eq!(cands[0]["rank"], json!(1)); + assert_eq!(cands[1]["messages"][0]["content"], "second"); + }); + } + + #[test] + fn rollout_groups_responses_with_rewards() { + 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 x", 1); + let mut r1 = rec("r1", "assistant", "ans1", 2); + r1.metadata = + Some(json!({"reward": 1.0, "reward_source": "verifier", "group_id": "g1"})); + let mut r2 = rec("r2", "assistant", "ans2", 3); + r2.metadata = + Some(json!({"reward": 0.0, "reward_source": "verifier", "group_id": "g1"})); + store.add(&[prompt, r1, r2]).await.unwrap(); + + let out = out_path(&dir); + store + .export_training( + &ExportConfig { + task: ExportTask::Rollout, + group_by: GroupBy::SessionId, + ..Default::default() + }, + &out, + ) + .await + .unwrap(); + + let lines = read_lines(&out); + assert_eq!(lines.len(), 1); + assert_eq!(lines[0]["group_id"], "g1"); + assert_eq!(lines[0]["prompt"][0]["content"], "solve x"); + let responses = lines[0]["responses"].as_array().unwrap(); + assert_eq!(responses.len(), 2); + assert_eq!(responses[0]["reward"], json!(1.0)); + assert_eq!(responses[0]["reward_source"], "verifier"); + }); + } + + #[test] + fn dedup_collapses_near_duplicates() { + 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 a = rec("a", "user", "dup one", 1); + a.session_id = Some("a".to_string()); + a.embedding = Some(emb(&[1.0, 0.0])); + let mut b = rec("b", "user", "dup two", 2); + b.session_id = Some("b".to_string()); + b.embedding = Some(emb(&[1.0, 0.0])); // identical direction -> distance 0 + store.add(&[a, b]).await.unwrap(); + + let out = out_path(&dir); + let manifest = store + .export_training( + &ExportConfig { + task: ExportTask::Sft, + group_by: GroupBy::SessionId, + dedup_threshold: Some(0.01), + ..Default::default() + }, + &out, + ) + .await + .unwrap(); + + assert_eq!( + manifest.counts.after_dedup, 1, + "one near-duplicate collapsed" + ); + assert_eq!(read_lines(&out).len(), 1); + }); + } + + #[test] + fn decontaminate_drops_holdout_matches() { + 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 keep = rec("k", "user", "keep", 1); + keep.session_id = Some("k".to_string()); + keep.embedding = Some(emb(&[0.0, 1.0])); + let mut leak = rec("l", "user", "leak", 2); + leak.session_id = Some("l".to_string()); + leak.embedding = Some(emb(&[1.0, 0.0])); + store.add(&[keep, leak]).await.unwrap(); + + let out = out_path(&dir); + let manifest = store + .export_training( + &ExportConfig { + task: ExportTask::Sft, + group_by: GroupBy::SessionId, + decontaminate_against: vec![emb(&[1.0, 0.0])], // matches "leak" + decontaminate_threshold: Some(0.01), + ..Default::default() + }, + &out, + ) + .await + .unwrap(); + + assert_eq!(manifest.counts.after_decontaminate, 1); + let lines = read_lines(&out); + assert_eq!(lines.len(), 1); + assert_eq!(lines[0]["messages"][0]["content"], "keep"); + }); + } + + #[test] + fn curation_drops_contradicted_records() { + let dir = TempDir::new().unwrap(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + runtime.block_on(async { + let mut store = open_store(&dir).await; + let good = rec("g", "user", "valid", 1); + let mut bad = rec("c", "user", "contradicted", 2); + bad.session_id = Some("other".to_string()); + bad.lifecycle_status = LIFECYCLE_CONTRADICTED.to_string(); + store.add(&[good, bad]).await.unwrap(); + + let out = out_path(&dir); + let manifest = store + .export_training(&ExportConfig::default(), &out) + .await + .unwrap(); + + assert_eq!(manifest.counts.after_lifecycle, 1); + let lines = read_lines(&out); + assert!(lines + .iter() + .all(|l| l["messages"][0]["content"] != "contradicted")); + }); + } + + #[test] + fn version_pinning_exports_old_state_and_restores() { + 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("r1", "user", "first", 1)]).await.unwrap(); + store.compact(None).await.unwrap(); + let pinned = store.version(); + + store.add(&[rec("r2", "user", "second", 2)]).await.unwrap(); + store.compact(None).await.unwrap(); + let latest = store.version(); + + let out = out_path(&dir); + let manifest = store + .export_training( + &ExportConfig { + task: ExportTask::Sft, + group_by: GroupBy::None, + version: Some(pinned), + ..Default::default() + }, + &out, + ) + .await + .unwrap(); + + assert_eq!(manifest.version, pinned); + assert_eq!( + store.version(), + latest, + "store restored after pinned export" + ); + }); + } + + #[test] + fn export_is_reproducible() { + 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("r1", "user", "a", 1), + rec("r2", "assistant", "b", 2), + rec("r3", "user", "c", 3), + ]) + .await + .unwrap(); + + let config = ExportConfig { + task: ExportTask::Sft, + group_by: GroupBy::SessionId, + ..Default::default() + }; + let first = out_path(&dir); + let second = dir.path().join("out2.jsonl").to_string_lossy().to_string(); + store.export_training(&config, &first).await.unwrap(); + store.export_training(&config, &second).await.unwrap(); + + assert_eq!( + std::fs::read_to_string(&first).unwrap(), + std::fs::read_to_string(&second).unwrap(), + "same version + config produces identical output" + ); + }); + } + + #[test] + fn external_id_prefix_grouping() { + 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 a = rec("a", "user", "doc7 turn1", 1); + a.external_id = Some("doc-7#chunk-1".to_string()); + let mut b = rec("b", "assistant", "doc7 turn2", 2); + b.external_id = Some("doc-7#chunk-2".to_string()); + let mut c = rec("c", "user", "doc8 turn1", 3); + c.external_id = Some("doc-8#chunk-1".to_string()); + store.add(&[a, b, c]).await.unwrap(); + + let out = out_path(&dir); + store + .export_training( + &ExportConfig { + task: ExportTask::Sft, + group_by: GroupBy::ExternalIdPrefix("#".to_string()), + ..Default::default() + }, + &out, + ) + .await + .unwrap(); + + let lines = read_lines(&out); + // doc-7 (2 turns) + doc-8 (1 turn) => 2 examples. + assert_eq!(lines.len(), 2); + assert_eq!(lines[0]["messages"].as_array().unwrap().len(), 2); + }); + } +} diff --git a/crates/lance-context-core/src/lib.rs b/crates/lance-context-core/src/lib.rs index 6024167..4fdbb66 100644 --- a/crates/lance-context-core/src/lib.rs +++ b/crates/lance-context-core/src/lib.rs @@ -3,12 +3,18 @@ mod api_impl; mod context; +mod export; mod namespace; mod record; pub mod serde; mod store; pub use context::{Context, ContextEntry, Snapshot}; +pub use export::{ + ExportConfig, ExportCounts, ExportManifest, ExportTask, GroupBy, Message, PreferenceExample, + PreferenceForm, Provenance, RankedCandidate, RolloutExample, RolloutResponse, SftExample, + EXPORT_SCHEMA_VERSION, +}; pub use namespace::{ContextNamespace, PartitionInfo, PartitionSelector, PartitionSpec}; pub use record::{ ContextRecord, LifecycleQueryOptions, MetadataFilter, RecordFilters, RecordPatch, Relationship, diff --git a/crates/lance-context-core/src/store.rs b/crates/lance-context-core/src/store.rs index 19556fd..9da7632 100644 --- a/crates/lance-context-core/src/store.rs +++ b/crates/lance-context-core/src/store.rs @@ -360,6 +360,12 @@ impl ContextStore { self.embedding_dim } + /// URI of the underlying Lance dataset. + #[must_use] + pub fn uri(&self) -> &str { + self.dataset.uri() + } + /// Append context records to the store and return the new dataset version. pub async fn add(&mut self, entries: &[ContextRecord]) -> LanceResult { if entries.is_empty() { diff --git a/python/python/lance_context/api.py b/python/python/lance_context/api.py index 8bfd43d..1477255 100644 --- a/python/python/lance_context/api.py +++ b/python/python/lance_context/api.py @@ -1141,6 +1141,60 @@ def retrieve( ) return [_normalize_retrieve_hit(item) for item in results] + def export_training( + self, + output_path: str, + *, + task: str = "sft", + format: str = "jsonl", + group_by: str = "session_id", + preference_form: str = "paired", + filters: dict[str, Any] | None = None, + dedup_threshold: float | None = None, + decontaminate_against: list[list[float]] | None = None, + decontaminate_threshold: float | None = None, + min_reward: float | None = None, + version: int | None = None, + include_expired: bool = False, + include_retired: bool = False, + ) -> dict[str, Any]: + """Curate stored records and export a trainable dataset to JSONL. + + Writes one JSON object per line to ``output_path`` (plus a sibling + ``.manifest.json``) and returns the manifest dict. + + ``task`` selects the shape: ``"sft"`` (ordered messages; pass + ``min_reward`` for rejection-sampling / Best-of-N), ``"preference"`` + (set ``preference_form`` to ``"paired"`` for DPO-style chosen/rejected, + ``"unpaired"`` for KTO binary labels, or ``"ranked"`` for N-way judge + rankings), or ``"rollout"`` (RL groups with per-response reward). + + Curation runs before export: lifecycle-correct filtering (drops + tombstoned/expired/retired/superseded/contradicted), optional + ``min_reward`` thresholding, semantic ``dedup_threshold`` (cosine), and + ``decontaminate_against`` a holdout-embedding set. Reward / preference + signals are read from each record's ``metadata`` (``reward``, + ``reward_source``, ``group_id``, ``label``, ``rank``). Pass ``version`` + to pin the export to a dataset version for reproducibility. + """ + if format != "jsonl": + raise ValueError("export format currently supports only 'jsonl'") + manifest_json = self._inner.export_training( + output_path, + task, + group_by, + preference_form, + _json_dumps(filters, "filters"), + dedup_threshold, + decontaminate_against, + decontaminate_threshold, + min_reward, + int(version) if version is not None else None, + include_expired, + include_retired, + ) + return json.loads(manifest_json) + def list( self, limit: int | None = None, diff --git a/python/src/lib.rs b/python/src/lib.rs index c7acaae..6362d4a 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -21,9 +21,9 @@ use lance_context_core::serde::CONTENT_TYPE_TEXT; use lance_context_core::{ CompactionConfig, CompactionMetrics, CompactionStats, Context as RustContext, ContextNamespace as RustContextNamespace, ContextRecord, ContextStore, ContextStoreOptions, - DistanceMetric, IdIndexType, LifecycleQueryOptions, PartitionInfo, PartitionSelector, - PartitionSpec, RecordFilters, RecordPatch, Relationship, RetrieveResult, SearchResult, - StateMetadata, LIFECYCLE_ACTIVE, + DistanceMetric, ExportConfig, ExportTask, GroupBy, IdIndexType, LifecycleQueryOptions, + PartitionInfo, PartitionSelector, PartitionSpec, PreferenceForm, RecordFilters, RecordPatch, + Relationship, RetrieveResult, SearchResult, StateMetadata, LIFECYCLE_ACTIVE, }; const DEFAULT_BINARY_CONTENT_TYPE: &str = "application/octet-stream"; @@ -226,6 +226,78 @@ fn filters_from_json(filters_json: Option) -> PyResult, + dedup_threshold: Option, + decontaminate_against: Option>>, + decontaminate_threshold: Option, + min_reward: Option, + version: Option, + include_expired: bool, + include_retired: bool, +) -> PyResult { + let task = match task { + "sft" => ExportTask::Sft, + "preference" => ExportTask::Preference, + "rollout" => ExportTask::Rollout, + other => { + return Err(PyRuntimeError::new_err(format!( + "invalid export task '{other}'; use 'sft', 'preference', or 'rollout'" + ))); + } + }; + let group_by = match group_by { + "session_id" => GroupBy::SessionId, + "run_id" => GroupBy::RunId, + "tenant" => GroupBy::Tenant, + "source" => GroupBy::Source, + "bot_id" => GroupBy::BotId, + "none" => GroupBy::None, + "external_id_prefix" => GroupBy::ExternalIdPrefix("#".to_string()), + other if other.starts_with("external_id_prefix:") => { + GroupBy::ExternalIdPrefix(other["external_id_prefix:".len()..].to_string()) + } + other => { + return Err(PyRuntimeError::new_err(format!( + "invalid group_by '{other}'" + ))); + } + }; + let preference_form = match preference_form { + "paired" => PreferenceForm::Paired, + "unpaired" => PreferenceForm::Unpaired, + "ranked" => PreferenceForm::Ranked, + other => { + return Err(PyRuntimeError::new_err(format!( + "invalid preference_form '{other}'; use 'paired', 'unpaired', or 'ranked'" + ))); + } + }; + let filters_summary = filters_json + .as_ref() + .map(|raw| serde_json::from_str::(raw)) + .transpose() + .map_err(to_py_err)?; + + Ok(ExportConfig { + task, + group_by, + preference_form, + filters: filters_from_json(filters_json)?, + lifecycle: LifecycleQueryOptions::new(include_expired, include_retired), + dedup_threshold, + decontaminate_against: decontaminate_against.unwrap_or_default(), + decontaminate_threshold, + min_reward, + version, + filters_summary, + }) +} + fn selector_from_dict(dict: &Bound<'_, PyDict>) -> PyResult { let mut selector = BTreeMap::new(); for (key, value) in dict.iter() { @@ -640,6 +712,46 @@ impl Context { .collect() } + #[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))] + fn export_training( + &mut self, + py: Python<'_>, + output_path: &str, + task: &str, + group_by: &str, + preference_form: &str, + filters_json: Option, + dedup_threshold: Option, + decontaminate_against: Option>>, + decontaminate_threshold: Option, + min_reward: Option, + version: Option, + include_expired: bool, + include_retired: bool, + ) -> PyResult { + let config = export_config( + task, + group_by, + preference_form, + filters_json, + dedup_threshold, + decontaminate_against, + decontaminate_threshold, + min_reward, + version, + include_expired, + include_retired, + )?; + let manifest = py + .allow_threads(|| { + self.runtime + .block_on(self.store.export_training(&config, output_path)) + }) + .map_err(to_py_err)?; + serde_json::to_string(&manifest).map_err(to_py_err) + } + #[pyo3(signature = (limit = None, offset = None, filters_json = None, include_expired = false, include_retired = false))] fn list( &self, diff --git a/python/tests/test_export_training.py b/python/tests/test_export_training.py new file mode 100644 index 0000000..76f2ee3 --- /dev/null +++ b/python/tests/test_export_training.py @@ -0,0 +1,135 @@ +"""Tests for curate + export to trainable datasets (export_training).""" + +from __future__ import annotations + +import json +from typing import TYPE_CHECKING, Any + +import pytest + +if TYPE_CHECKING: + from pathlib import Path + +from lance_context.api import Context + + +def _read_jsonl(path: Path) -> list[dict[str, Any]]: + return [json.loads(line) for line in path.read_text().splitlines() if line.strip()] + + +def test_export_sft_groups_by_session(tmp_path: Path) -> None: + ctx = Context.create(str(tmp_path / "ctx.lance")) + ctx.add("user", "hello", session_id="s1") + ctx.add("assistant", "hi there", session_id="s1") + + out = tmp_path / "sft.jsonl" + manifest = ctx.export_training(str(out), task="sft", group_by="session_id") + + rows = _read_jsonl(out) + assert len(rows) == 1 + assert [m["content"] for m in rows[0]["messages"]] == ["hello", "hi there"] + assert manifest["task"] == "sft" + assert manifest["counts"]["examples"] == 1 + # sibling manifest file written + assert (tmp_path / "sft.jsonl.manifest.json").exists() + + +def test_export_sft_rejection_sampling_min_reward(tmp_path: Path) -> None: + ctx = Context.create(str(tmp_path / "ctx.lance")) + ctx.add("assistant", "good", session_id="a", metadata={"reward": 0.9}) + ctx.add("assistant", "bad", session_id="b", metadata={"reward": 0.1}) + + out = tmp_path / "sft.jsonl" + ctx.export_training(str(out), task="sft", min_reward=0.5) + + rows = _read_jsonl(out) + assert len(rows) == 1 + assert rows[0]["messages"][0]["content"] == "good" + + +def test_export_preference_paired(tmp_path: Path) -> None: + ctx = Context.create(str(tmp_path / "ctx.lance")) + ctx.add("user", "q", session_id="s1") + ctx.add("assistant", "great", session_id="s1", metadata={"reward": 0.9}) + ctx.add("assistant", "poor", session_id="s1", metadata={"reward": 0.1}) + + out = tmp_path / "pref.jsonl" + ctx.export_training(str(out), task="preference", preference_form="paired") + + rows = _read_jsonl(out) + assert len(rows) == 1 + assert rows[0]["form"] == "paired" + assert rows[0]["chosen"][0]["content"] == "great" + assert rows[0]["rejected"][0]["content"] == "poor" + + +def test_export_preference_unpaired_kto(tmp_path: Path) -> None: + ctx = Context.create(str(tmp_path / "ctx.lance")) + ctx.add("user", "q", session_id="s1") + ctx.add("assistant", "yes", session_id="s1", metadata={"label": "chosen"}) + ctx.add("assistant", "no", session_id="s1", metadata={"label": "rejected"}) + + out = tmp_path / "pref.jsonl" + ctx.export_training(str(out), task="preference", preference_form="unpaired") + + rows = _read_jsonl(out) + assert len(rows) == 2 + labels = {r["completion"][0]["content"]: r["label"] for r in rows} + assert labels == {"yes": True, "no": False} + + +def test_export_rollout_groups_with_rewards(tmp_path: Path) -> None: + ctx = Context.create(str(tmp_path / "ctx.lance")) + ctx.add("user", "solve", session_id="s1") + ctx.add( + "assistant", + "a1", + session_id="s1", + metadata={"reward": 1.0, "reward_source": "verifier", "group_id": "g1"}, + ) + ctx.add( + "assistant", + "a2", + session_id="s1", + metadata={"reward": 0.0, "reward_source": "verifier", "group_id": "g1"}, + ) + + out = tmp_path / "rollout.jsonl" + ctx.export_training(str(out), task="rollout") + + rows = _read_jsonl(out) + assert len(rows) == 1 + assert rows[0]["group_id"] == "g1" + assert len(rows[0]["responses"]) == 2 + assert rows[0]["responses"][0]["reward_source"] == "verifier" + + +def test_export_dedup_collapses_duplicates(tmp_path: Path) -> None: + ctx = Context.create(str(tmp_path / "ctx.lance"), embedding_dim=4) + ctx.add("user", "one", session_id="a", embedding=[1.0, 0.0, 0.0, 0.0]) + ctx.add("user", "two", session_id="b", embedding=[1.0, 0.0, 0.0, 0.0]) + + out = tmp_path / "sft.jsonl" + manifest = ctx.export_training(str(out), task="sft", dedup_threshold=0.01) + + assert manifest["counts"]["after_dedup"] == 1 + assert len(_read_jsonl(out)) == 1 + + +def test_export_invalid_task_raises(tmp_path: Path) -> None: + ctx = Context.create(str(tmp_path / "ctx.lance")) + ctx.add("user", "hi") + with pytest.raises(RuntimeError, match="invalid export task"): + ctx.export_training(str(tmp_path / "x.jsonl"), task="bogus") + + +def test_export_reproducible(tmp_path: Path) -> None: + ctx = Context.create(str(tmp_path / "ctx.lance")) + ctx.add("user", "a", session_id="s1") + ctx.add("assistant", "b", session_id="s1") + + first = tmp_path / "a.jsonl" + second = tmp_path / "b.jsonl" + ctx.export_training(str(first), task="sft") + ctx.export_training(str(second), task="sft") + assert first.read_text() == second.read_text()