From 5124edc599c93de8d7b1761a1ca20d56241d5263 Mon Sep 17 00:00:00 2001 From: Allen Cheng Date: Wed, 24 Jun 2026 23:55:34 -0700 Subject: [PATCH] feat: reproducible train/eval split for post-training exports (#103) Follow-up to #96. Adds a group-disjoint, reproducible train/eval split to the export path so eval scores aren't inflated by conversation leakage. - `SplitConfig { eval_fraction, by, seed }` on `ExportConfig`. Each record is assigned to train or eval by a stable hash (FNV-1a + MurmurHash3 fmix64 finalizer) of its `by` grouping key plus `seed`, so: - no group (default `session_id`, configurable to run_id / tenant / source / bot_id / external-id prefix) spans both sides, and - the same seed reproduces the identical partition across runs and platforms (the hash is fully defined, not `DefaultHasher`). - When `split` is set, `export_training` writes `.train.` and `.eval.`, each with its own manifest recording the split params (`eval_fraction`, `by`, `seed`, `side`) and the complementary output path. The returned manifest is the train side. - Curation (dedup/decontam/lifecycle) runs once before the split, so the two compose: decontamination drops train rows near eval, the split keeps groups disjoint. Python: `ctx.export_training(..., split={"eval_fraction": 0.1, "by": "session_id", "seed": 42})`. Tests: core (determinism, group-disjointness, fraction-within-tolerance, split manifests + complement path) and python (split disjointness + determinism). README updated. Stacked on #96 (PR #111). Closes #103 --- README.md | 9 + crates/lance-context-core/src/export.rs | 394 +++++++++++++++++++++--- crates/lance-context-core/src/lib.rs | 2 +- python/python/lance_context/api.py | 18 ++ python/src/lib.rs | 63 ++-- python/tests/test_export_training.py | 51 +++ 6 files changed, 474 insertions(+), 63 deletions(-) diff --git a/README.md b/README.md index 5045ead..b06cd82 100644 --- a/README.md +++ b/README.md @@ -273,6 +273,15 @@ print("SFT examples:", manifest["counts"]["examples"]) ctx.export_training("training/pref.jsonl", task="preference", preference_form="paired") ctx.export_training("training/rollout.jsonl", task="rollout") # GRPO/RLVR groups +# Reproducible, group-disjoint train/eval split (no session leaks across the +# boundary; same seed reproduces the partition). Writes sft.train.jsonl + +# sft.eval.jsonl, each with its own manifest. +ctx.export_training( + "training/sft.jsonl", + task="sft", + split={"eval_fraction": 0.1, "by": "session_id", "seed": 42}, +) + # 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 index 01f5b66..03d7e10 100644 --- a/crates/lance-context-core/src/export.rs +++ b/crates/lance-context-core/src/export.rs @@ -129,6 +129,30 @@ impl GroupBy { } } +/// Reproducible train/eval split configuration. +/// +/// Each record is assigned to train or eval by a stable hash of its `by` +/// grouping key plus `seed`, so no group (e.g. session) spans both sides and +/// the same `seed` reproduces the identical partition. +#[derive(Clone)] +pub struct SplitConfig { + /// Fraction of groups assigned to the eval side, in `0.0..=1.0`. + pub eval_fraction: f64, + /// Grouping key the split is disjoint on (defaults to `session_id`). + pub by: GroupBy, + pub seed: u64, +} + +impl Default for SplitConfig { + fn default() -> Self { + Self { + eval_fraction: 0.1, + by: GroupBy::SessionId, + seed: 0, + } + } +} + /// Curation + export configuration. #[derive(Clone)] pub struct ExportConfig { @@ -152,6 +176,8 @@ pub struct ExportConfig { pub version: Option, /// Optional JSON summary of the applied selectors, recorded in the manifest. pub filters_summary: Option, + /// When set, write group-disjoint `train` / `eval` outputs instead of one. + pub split: Option, } impl Default for ExportConfig { @@ -168,6 +194,7 @@ impl Default for ExportConfig { min_reward: None, version: None, filters_summary: None, + split: None, } } } @@ -276,6 +303,19 @@ pub struct ExportCounts { pub examples: usize, } +/// Records which side of a train/eval split an output is, and the parameters +/// needed to reproduce the partition. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct SplitManifest { + /// `"train"` or `"eval"`. + pub side: String, + pub eval_fraction: f64, + pub by: String, + pub seed: u64, + /// Path of the complementary (other side's) output. + pub complement_path: String, +} + /// Reproducible description of an export, written next to the JSONL output. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ExportManifest { @@ -295,6 +335,8 @@ pub struct ExportManifest { #[serde(skip_serializing_if = "Option::is_none", default)] pub min_reward: Option, #[serde(skip_serializing_if = "Option::is_none", default)] + pub split: 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>, @@ -742,6 +784,116 @@ fn summarize_groups( (start, end, ids) } +/// Stable 64-bit FNV-1a hash of `seed` + `key`, used for reproducible, +/// platform-independent split assignment. +fn stable_hash(seed: u64, key: &str) -> u64 { + let mut hash: u64 = 0xcbf2_9ce4_8422_2325; + let mix = |hash: &mut u64, byte: u8| { + *hash ^= u64::from(byte); + *hash = hash.wrapping_mul(0x0000_0100_0000_01b3); + }; + for byte in seed.to_le_bytes() { + mix(&mut hash, byte); + } + for byte in key.as_bytes() { + mix(&mut hash, *byte); + } + hash +} + +/// Map a group key to a stable fraction in `0.0..1.0`. +fn split_fraction(seed: u64, key: &str) -> f64 { + // FNV-1a alone has biased high bits; run the MurmurHash3 fmix64 finalizer + // for good avalanche before taking the top 53 bits. + let mut hash = stable_hash(seed, key); + hash ^= hash >> 33; + hash = hash.wrapping_mul(0xff51_afd7_ed55_8ccd); + hash ^= hash >> 33; + hash = hash.wrapping_mul(0xc4ce_b9fe_1a85_ec53); + hash ^= hash >> 33; + (hash >> 11) as f64 / (1u64 << 53) as f64 +} + +/// Derive `train` / `eval` output paths by inserting the side before the final +/// extension (e.g. `cut.jsonl` -> `cut.train.jsonl`). +fn split_paths(output_path: &str) -> (String, String) { + let slash = output_path.rfind('/').map_or(0, |i| i + 1); + match output_path[slash..].rfind('.') { + Some(rel_dot) => { + let dot = slash + rel_dot; + let (stem, ext) = output_path.split_at(dot); + (format!("{stem}.train{ext}"), format!("{stem}.eval{ext}")) + } + None => ( + format!("{output_path}.train"), + format!("{output_path}.eval"), + ), + } +} + +/// Group `records`, write the task-shaped JSONL to `output_path`, write the +/// sibling manifest, and return it. +#[allow(clippy::too_many_arguments)] +fn emit_export( + records: Vec, + config: &ExportConfig, + context_uri: &str, + version: u64, + mut counts: ExportCounts, + output_path: &str, + split: Option, +) -> LanceResult { + let groups = group_records(records, &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: context_uri.to_string(), + 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, + split, + 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) +} + impl ContextStore { /// Curate stored records and export them as task-shaped JSONL plus a /// sibling `.manifest.json`, returning the manifest. @@ -788,56 +940,65 @@ impl ContextStore { 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, + let (curated, counts) = curate(records, config); + + let Some(split) = &config.split else { + return emit_export( + curated, + config, &context_uri, version, - )?, - ExportTask::Rollout => write_rollout(&groups, &mut writer, &context_uri, version)?, + counts, + output_path, + None, + ); }; - writer.flush()?; - counts.examples = examples; - let manifest = ExportManifest { - context_uri, + // Group-disjoint, reproducible partition: a record's side is decided by + // a stable hash of its `split.by` key, so no group spans both sides. + let (train_path, eval_path) = split_paths(output_path); + let mut train_records = Vec::new(); + let mut eval_records = Vec::new(); + for record in curated { + let key = split.by.key(&record); + if split_fraction(split.seed, &key) < split.eval_fraction { + eval_records.push(record); + } else { + train_records.push(record); + } + } + + let train_manifest = emit_export( + train_records, + config, + &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() + counts, + &train_path, + Some(SplitManifest { + side: "train".to_string(), + eval_fraction: split.eval_fraction, + by: split.by.label(), + seed: split.seed, + complement_path: eval_path.clone(), }), - 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, + )?; + emit_export( + eval_records, + config, + &context_uri, + version, 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) + &eval_path, + Some(SplitManifest { + side: "eval".to_string(), + eval_fraction: split.eval_fraction, + by: split.by.label(), + seed: split.seed, + complement_path: train_path.clone(), + }), + )?; + Ok(train_manifest) } } @@ -1361,4 +1522,151 @@ mod tests { assert_eq!(lines[0]["messages"].as_array().unwrap().len(), 2); }); } + + fn session_ids(path: &str) -> std::collections::HashSet { + read_lines(path) + .iter() + .filter_map(|line| { + line["provenance"]["session_id"] + .as_str() + .map(str::to_string) + }) + .collect() + } + + #[test] + fn split_is_deterministic_and_group_disjoint() { + 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 records = Vec::new(); + for s in 0..10 { + let mut user = rec(&format!("u{s}"), "user", "q", s * 2); + user.session_id = Some(format!("s{s}")); + let mut asst = rec(&format!("a{s}"), "assistant", "r", s * 2 + 1); + asst.session_id = Some(format!("s{s}")); + records.push(user); + records.push(asst); + } + store.add(&records).await.unwrap(); + + let config = ExportConfig { + task: ExportTask::Sft, + group_by: GroupBy::SessionId, + split: Some(SplitConfig { + eval_fraction: 0.5, + by: GroupBy::SessionId, + seed: 42, + }), + ..Default::default() + }; + + let base1 = dir.path().join("a.jsonl").to_string_lossy().to_string(); + let base2 = dir.path().join("b.jsonl").to_string_lossy().to_string(); + store.export_training(&config, &base1).await.unwrap(); + store.export_training(&config, &base2).await.unwrap(); + + let (train1, eval1) = split_paths(&base1); + let (train2, eval2) = split_paths(&base2); + + // Determinism: identical partition across runs with the same seed. + assert_eq!( + std::fs::read_to_string(&train1).unwrap(), + std::fs::read_to_string(&train2).unwrap() + ); + assert_eq!( + std::fs::read_to_string(&eval1).unwrap(), + std::fs::read_to_string(&eval2).unwrap() + ); + + // Group-disjoint: no session appears on both sides. + let train_sessions = session_ids(&train1); + let eval_sessions = session_ids(&eval1); + assert!(!train_sessions.is_empty() && !eval_sessions.is_empty()); + assert!(train_sessions.is_disjoint(&eval_sessions)); + assert_eq!(train_sessions.len() + eval_sessions.len(), 10); + }); + } + + #[test] + fn split_fraction_is_approximately_respected() { + 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 records = Vec::new(); + for s in 0..200 { + let mut r = rec(&format!("r{s}"), "user", "q", s); + r.session_id = Some(format!("s{s}")); + records.push(r); + } + store.add(&records).await.unwrap(); + + let config = ExportConfig { + task: ExportTask::Sft, + group_by: GroupBy::SessionId, + split: Some(SplitConfig { + eval_fraction: 0.25, + by: GroupBy::SessionId, + seed: 7, + }), + ..Default::default() + }; + let base = dir.path().join("c.jsonl").to_string_lossy().to_string(); + store.export_training(&config, &base).await.unwrap(); + + let (train, eval) = split_paths(&base); + let eval_count = read_lines(&eval).len(); + let train_count = read_lines(&train).len(); + assert_eq!(train_count + eval_count, 200); + let fraction = eval_count as f64 / 200.0; + assert!( + (fraction - 0.25).abs() < 0.1, + "eval fraction {fraction} too far from 0.25" + ); + }); + } + + #[test] + fn split_manifests_record_params_and_complement() { + 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", "x", 1); + a.session_id = Some("s1".to_string()); + let mut b = rec("b", "user", "y", 2); + b.session_id = Some("s2".to_string()); + store.add(&[a, b]).await.unwrap(); + + let config = ExportConfig { + task: ExportTask::Sft, + group_by: GroupBy::SessionId, + split: Some(SplitConfig { + eval_fraction: 0.5, + by: GroupBy::SessionId, + seed: 99, + }), + ..Default::default() + }; + let base = dir.path().join("d.jsonl").to_string_lossy().to_string(); + store.export_training(&config, &base).await.unwrap(); + + let (train, eval) = split_paths(&base); + assert!(std::fs::metadata(&train).is_ok()); + assert!(std::fs::metadata(&eval).is_ok()); + + let train_manifest = read_manifest(&train); + let split = train_manifest.split.unwrap(); + assert_eq!(split.side, "train"); + assert_eq!(split.seed, 99); + assert_eq!(split.eval_fraction, 0.5); + assert_eq!(split.by, "session_id"); + assert_eq!(split.complement_path, eval); + + let eval_manifest = read_manifest(&eval); + assert_eq!(eval_manifest.split.unwrap().side, "eval"); + }); + } } diff --git a/crates/lance-context-core/src/lib.rs b/crates/lance-context-core/src/lib.rs index 4fdbb66..ae79d3c 100644 --- a/crates/lance-context-core/src/lib.rs +++ b/crates/lance-context-core/src/lib.rs @@ -13,7 +13,7 @@ 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, + SplitConfig, SplitManifest, 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 2e8b027..f3718b9 100644 --- a/python/python/lance_context/api.py +++ b/python/python/lance_context/api.py @@ -1259,6 +1259,7 @@ def export_training( version: int | None = None, include_expired: bool = False, include_retired: bool = False, + split: dict[str, Any] | None = None, ) -> dict[str, Any]: """Curate stored records and export a trainable dataset to JSONL. @@ -1278,9 +1279,23 @@ def export_training( 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. + + 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. """ if format != "jsonl": raise ValueError("export format currently supports only 'jsonl'") + split_eval_fraction = None + split_by = None + split_seed = None + if split is not None: + if "eval_fraction" not in split: + raise ValueError("split requires 'eval_fraction'") + split_eval_fraction = float(split["eval_fraction"]) + split_by = split.get("by") + split_seed = int(split["seed"]) if split.get("seed") is not None else None manifest_json = self._inner.export_training( output_path, task, @@ -1294,6 +1309,9 @@ def export_training( int(version) if version is not None else None, include_expired, include_retired, + split_eval_fraction, + split_by, + split_seed, ) return json.loads(manifest_json) diff --git a/python/src/lib.rs b/python/src/lib.rs index 6fb6b1c..6bedae4 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -23,7 +23,7 @@ use lance_context_core::{ ContextNamespace as RustContextNamespace, ContextRecord, ContextStore, ContextStoreOptions, DistanceMetric, ExportConfig, ExportTask, GroupBy, IdIndexType, LifecycleQueryOptions, PartitionInfo, PartitionSelector, PartitionSpec, PreferenceForm, RecordFilters, RecordPatch, - Relationship, RetrieveResult, SearchResult, StateMetadata, LIFECYCLE_ACTIVE, + Relationship, RetrieveResult, SearchResult, SplitConfig, StateMetadata, LIFECYCLE_ACTIVE, }; const DEFAULT_BINARY_CONTENT_TYPE: &str = "application/octet-stream"; @@ -239,6 +239,9 @@ fn export_config( version: Option, include_expired: bool, include_retired: bool, + split_eval_fraction: Option, + split_by: Option, + split_seed: Option, ) -> PyResult { let task = match task { "sft" => ExportTask::Sft, @@ -250,23 +253,7 @@ fn export_config( ))); } }; - 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 group_by = parse_group_by(group_by)?; let preference_form = match preference_form { "paired" => PreferenceForm::Paired, "unpaired" => PreferenceForm::Unpaired, @@ -282,6 +269,17 @@ fn export_config( .map(|raw| serde_json::from_str::(raw)) .transpose() .map_err(to_py_err)?; + let split = match split_eval_fraction { + Some(eval_fraction) => Some(SplitConfig { + eval_fraction, + by: match split_by { + Some(by) => parse_group_by(&by)?, + None => GroupBy::SessionId, + }, + seed: split_seed.unwrap_or(0), + }), + None => None, + }; Ok(ExportConfig { task, @@ -295,6 +293,27 @@ fn export_config( min_reward, version, filters_summary, + split, + }) +} + +fn parse_group_by(group_by: &str) -> PyResult { + Ok(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}'" + ))); + } }) } @@ -775,7 +794,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))] + #[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))] fn export_training( &mut self, py: Python<'_>, @@ -791,6 +810,9 @@ impl Context { version: Option, include_expired: bool, include_retired: bool, + split_eval_fraction: Option, + split_by: Option, + split_seed: Option, ) -> PyResult { let config = export_config( task, @@ -804,6 +826,9 @@ impl Context { version, include_expired, include_retired, + split_eval_fraction, + split_by, + split_seed, )?; let manifest = py .allow_threads(|| { diff --git a/python/tests/test_export_training.py b/python/tests/test_export_training.py index 76f2ee3..f0d0e95 100644 --- a/python/tests/test_export_training.py +++ b/python/tests/test_export_training.py @@ -133,3 +133,54 @@ def test_export_reproducible(tmp_path: Path) -> None: ctx.export_training(str(first), task="sft") ctx.export_training(str(second), task="sft") assert first.read_text() == second.read_text() + + +def test_export_train_eval_split(tmp_path: Path) -> None: + ctx = Context.create(str(tmp_path / "ctx.lance")) + for s in range(20): + ctx.add("user", f"q{s}", session_id=f"s{s}") + ctx.add("assistant", f"a{s}", session_id=f"s{s}") + + base = tmp_path / "cut.jsonl" + manifest = ctx.export_training( + str(base), + task="sft", + group_by="session_id", + split={"eval_fraction": 0.3, "by": "session_id", "seed": 42}, + ) + + train = tmp_path / "cut.train.jsonl" + eval_ = tmp_path / "cut.eval.jsonl" + assert train.exists() and eval_.exists() + assert manifest["split"]["side"] == "train" + assert manifest["split"]["seed"] == 42 + + def sessions(path: Path) -> set[str]: + return { + json.loads(line)["provenance"]["session_id"] + for line in path.read_text().splitlines() + if line.strip() + } + + train_sessions = sessions(train) + eval_sessions = sessions(eval_) + assert train_sessions and eval_sessions + assert train_sessions.isdisjoint(eval_sessions) # group-disjoint + assert len(train_sessions) + len(eval_sessions) == 20 + + +def test_export_split_is_deterministic(tmp_path: Path) -> None: + ctx = Context.create(str(tmp_path / "ctx.lance")) + for s in range(10): + ctx.add("user", f"q{s}", session_id=f"s{s}") + + split = {"eval_fraction": 0.5, "by": "session_id", "seed": 7} + ctx.export_training(str(tmp_path / "a.jsonl"), task="sft", split=split) + ctx.export_training(str(tmp_path / "b.jsonl"), task="sft", split=split) + + assert (tmp_path / "a.eval.jsonl").read_text() == ( + tmp_path / "b.eval.jsonl" + ).read_text() + assert (tmp_path / "a.train.jsonl").read_text() == ( + tmp_path / "b.train.jsonl" + ).read_text()