Oleksiy Syvokon
095b57bb26
`ep split --stratify=project` will use the first component of `cursor_path`, which is a worktree name. This is an alternative to `--stratify=repo` in absence of `repository_url` Release Notes: - N/A
680 lines
22 KiB
Rust
680 lines
22 KiB
Rust
//! `ep split` implementation.
|
||
//!
|
||
//! This command splits a JSONL dataset into multiple files based on size specifications,
|
||
//! with optional stratification by a JSON field.
|
||
//!
|
||
//! # Usage
|
||
//!
|
||
//! ```text
|
||
//! ep split [--stratify=<field>] [input.jsonl] <out1>=<size1> <out2>=<size2> ...
|
||
//! ```
|
||
//!
|
||
//! If `input.jsonl` is not provided or is `-`, reads from stdin.
|
||
//!
|
||
//! # Size specifications
|
||
//!
|
||
//! - `80%` - percentage of total examples (lines)
|
||
//! - `100` - approximate absolute count of examples (lines)
|
||
//! - `rest` - all remaining items (only one split can use this)
|
||
//!
|
||
//! # Stratification
|
||
//!
|
||
//! The `--stratify` flag controls how examples are grouped before splitting:
|
||
//!
|
||
//! - `cursor-path` (default): group by the `cursor_path` JSON field
|
||
//! - `project`: group by the first component of the `cursor_path` JSON field
|
||
//! - `repo`: group by the `repository_url` JSON field
|
||
//! - `none`: no grouping, split individual examples
|
||
//!
|
||
//! When stratifying, the split ensures each output file contains examples from
|
||
//! non-overlapping groups. Size specifications always apply to the number of
|
||
//! examples (lines), with whole groups assigned greedily to meet the target.
|
||
//! Examples missing the stratification field are treated as individual groups.
|
||
|
||
use anyhow::{Context as _, Result, bail};
|
||
use clap::Args;
|
||
use rand::SeedableRng;
|
||
use rand::seq::SliceRandom;
|
||
use serde_json::Value;
|
||
use std::collections::BTreeMap;
|
||
use std::fs::File;
|
||
use std::io::{self, BufRead, BufReader, BufWriter, Write};
|
||
use std::path::{Path, PathBuf};
|
||
|
||
/// `ep split` CLI args.
|
||
#[derive(Debug, Args, Clone)]
|
||
#[command(
|
||
about = "Split a JSONL dataset into multiple files with optional stratification",
|
||
after_help = r#"SIZE SPECIFICATIONS:
|
||
<percentage>% Percentage of total (e.g., 80%)
|
||
<count> Absolute number (e.g., 100)
|
||
rest All remaining items (only one output can use this)
|
||
|
||
Sizes always apply to examples (lines). When stratifying, whole groups
|
||
are assigned greedily to approximate the target count.
|
||
|
||
EXAMPLES:
|
||
# Split 80% train, 20% validation (default: stratify by cursor_path)
|
||
ep split input.jsonl train.jsonl=80% valid.jsonl=rest
|
||
|
||
# Split into train/valid/test
|
||
ep split input.jsonl train.jsonl=80% valid.jsonl=10% test.jsonl=rest
|
||
|
||
# Stratify by repository_url instead of cursor_path
|
||
ep split --stratify=repo input.jsonl train.jsonl=80% valid.jsonl=rest
|
||
|
||
# No stratification (split by individual examples)
|
||
ep split --stratify=none input.jsonl train.jsonl=80% valid.jsonl=rest
|
||
|
||
# Read from stdin
|
||
cat input.jsonl | ep split train.jsonl=80% valid.jsonl=rest
|
||
|
||
# Reproducible split with seed
|
||
ep split --seed 42 input.jsonl train.jsonl=80% valid.jsonl=rest
|
||
|
||
STRATIFICATION:
|
||
Controls how examples are grouped before splitting:
|
||
cursor-path Group by "cursor_path" field (default)
|
||
project Group by the first component of the "cursor_path" field
|
||
repo Group by "repository_url" field
|
||
none No grouping, split individual examples
|
||
|
||
When stratifying, the split ensures each output file contains examples
|
||
from non-overlapping groups. This prevents data leakage between
|
||
train/test splits.
|
||
"#
|
||
)]
|
||
pub struct SplitArgs {
|
||
/// Random seed for reproducibility
|
||
#[arg(long)]
|
||
pub seed: Option<u64>,
|
||
|
||
/// Stratification field for splitting the dataset
|
||
#[arg(long, default_value = "cursor-path")]
|
||
pub stratify: Stratify,
|
||
}
|
||
|
||
#[derive(Debug, Clone, Copy, PartialEq, Eq, clap::ValueEnum, strum::Display)]
|
||
pub enum Stratify {
|
||
#[strum(serialize = "cursor_path")]
|
||
CursorPath,
|
||
#[strum(serialize = "project")]
|
||
Project,
|
||
#[strum(serialize = "repo")]
|
||
Repo,
|
||
#[strum(serialize = "none")]
|
||
None,
|
||
}
|
||
|
||
#[derive(Debug, Clone)]
|
||
pub enum SplitSize {
|
||
Percentage(f64),
|
||
Absolute(usize),
|
||
Rest,
|
||
}
|
||
|
||
#[derive(Debug, Clone)]
|
||
pub struct SplitSpec {
|
||
pub path: PathBuf,
|
||
pub size: SplitSize,
|
||
}
|
||
|
||
fn parse_split_spec(spec: &str) -> Result<SplitSpec> {
|
||
let (path, size_str) = spec
|
||
.rsplit_once('=')
|
||
.with_context(|| format!("invalid split spec '{}': expected <path>=<size>", spec))?;
|
||
|
||
let size = if size_str == "rest" {
|
||
SplitSize::Rest
|
||
} else if size_str.ends_with('%') {
|
||
let pct_str = size_str.trim_end_matches('%');
|
||
let pct: f64 = pct_str
|
||
.parse()
|
||
.with_context(|| format!("invalid percentage '{}' in '{}'", pct_str, spec))?;
|
||
if !(0.0..=100.0).contains(&pct) {
|
||
bail!("percentage must be between 0 and 100, got {}", pct);
|
||
}
|
||
SplitSize::Percentage(pct / 100.0)
|
||
} else {
|
||
let count: usize = size_str
|
||
.parse()
|
||
.with_context(|| format!("invalid count '{}' in '{}'", size_str, spec))?;
|
||
SplitSize::Absolute(count)
|
||
};
|
||
|
||
Ok(SplitSpec {
|
||
path: PathBuf::from(path),
|
||
size,
|
||
})
|
||
}
|
||
|
||
fn read_lines_from_input(input: Option<&Path>) -> Result<Vec<String>> {
|
||
let reader: Box<dyn BufRead> = match input {
|
||
Some(path) => {
|
||
let file =
|
||
File::open(path).with_context(|| format!("failed to open '{}'", path.display()))?;
|
||
Box::new(BufReader::new(file))
|
||
}
|
||
None => Box::new(BufReader::new(io::stdin())),
|
||
};
|
||
|
||
let lines: Vec<String> = reader
|
||
.lines()
|
||
.collect::<io::Result<Vec<_>>>()
|
||
.context("failed to read input lines")?;
|
||
|
||
Ok(lines)
|
||
}
|
||
|
||
fn compute_split_counts(specs: &[SplitSpec], total: usize) -> Result<Vec<usize>> {
|
||
let mut counts = vec![0usize; specs.len()];
|
||
let mut remaining = total;
|
||
let mut rest_index: Option<usize> = None;
|
||
|
||
for (i, spec) in specs.iter().enumerate() {
|
||
match &spec.size {
|
||
SplitSize::Percentage(pct) => {
|
||
let count = (total as f64 * pct).round() as usize;
|
||
counts[i] = count.min(remaining);
|
||
remaining = remaining.saturating_sub(counts[i]);
|
||
}
|
||
SplitSize::Absolute(count) => {
|
||
counts[i] = (*count).min(remaining);
|
||
remaining = remaining.saturating_sub(counts[i]);
|
||
}
|
||
SplitSize::Rest => {
|
||
if rest_index.is_some() {
|
||
bail!("only one split can use 'rest'");
|
||
}
|
||
rest_index = Some(i);
|
||
}
|
||
}
|
||
}
|
||
|
||
if let Some(idx) = rest_index {
|
||
counts[idx] = remaining;
|
||
}
|
||
|
||
Ok(counts)
|
||
}
|
||
|
||
fn write_lines_to_file(path: &Path, lines: &[String]) -> Result<()> {
|
||
if let Some(parent) = path.parent() {
|
||
if !parent.as_os_str().is_empty() {
|
||
std::fs::create_dir_all(parent)
|
||
.with_context(|| format!("failed to create directory '{}'", parent.display()))?;
|
||
}
|
||
}
|
||
|
||
let file =
|
||
File::create(path).with_context(|| format!("failed to create '{}'", path.display()))?;
|
||
let mut writer = BufWriter::new(file);
|
||
|
||
for line in lines {
|
||
writeln!(writer, "{}", line)
|
||
.with_context(|| format!("failed to write to '{}'", path.display()))?;
|
||
}
|
||
|
||
writer
|
||
.flush()
|
||
.with_context(|| format!("failed to flush '{}'", path.display()))?;
|
||
|
||
Ok(())
|
||
}
|
||
|
||
pub fn run_split(args: &SplitArgs, inputs: &[PathBuf]) -> Result<()> {
|
||
if inputs.is_empty() {
|
||
bail!("usage: ep split [input.jsonl] train.jsonl=80% valid.jsonl=rest");
|
||
}
|
||
|
||
let (input_path, split_specs_raw): (Option<&Path>, &[PathBuf]) =
|
||
if inputs.first().is_some_and(|p| {
|
||
let s = p.to_string_lossy();
|
||
!s.contains('=')
|
||
}) {
|
||
let first = inputs.first().map(|p| p.as_path());
|
||
let first = if first == Some(Path::new("-")) {
|
||
None
|
||
} else {
|
||
first
|
||
};
|
||
(first, &inputs[1..])
|
||
} else {
|
||
(None, inputs)
|
||
};
|
||
|
||
if split_specs_raw.is_empty() {
|
||
bail!("no split specifications provided");
|
||
}
|
||
|
||
let specs: Vec<SplitSpec> = split_specs_raw
|
||
.iter()
|
||
.map(|p| parse_split_spec(&p.to_string_lossy()))
|
||
.collect::<Result<Vec<_>>>()?;
|
||
|
||
let lines = read_lines_from_input(input_path)?;
|
||
let total_lines = lines.len();
|
||
|
||
if total_lines == 0 {
|
||
for spec in &specs {
|
||
write_lines_to_file(&spec.path, &[])?;
|
||
}
|
||
return Ok(());
|
||
}
|
||
|
||
let mut grouped_lines = group_lines(&lines, args.stratify);
|
||
|
||
if args.stratify != Stratify::None {
|
||
eprintln!(
|
||
"Stratifying by {} ({} unique groups, {} examples)",
|
||
args.stratify,
|
||
grouped_lines.len(),
|
||
total_lines
|
||
);
|
||
} else {
|
||
eprintln!(
|
||
"No stratification, splitting {} examples by line",
|
||
total_lines
|
||
);
|
||
}
|
||
|
||
let mut rng = match args.seed {
|
||
Some(seed) => rand::rngs::StdRng::seed_from_u64(seed),
|
||
None => rand::rngs::StdRng::from_os_rng(),
|
||
};
|
||
|
||
grouped_lines.shuffle(&mut rng);
|
||
|
||
let line_targets = compute_split_counts(&specs, total_lines)?;
|
||
let rest_index = specs.iter().position(|s| matches!(s.size, SplitSize::Rest));
|
||
let mut split_outputs: Vec<Vec<String>> = vec![Vec::new(); specs.len()];
|
||
let mut group_iter = grouped_lines.into_iter();
|
||
|
||
for (split_idx, &target) in line_targets.iter().enumerate() {
|
||
if Some(split_idx) == rest_index {
|
||
continue;
|
||
}
|
||
let mut accumulated = 0;
|
||
while accumulated < target {
|
||
if let Some(group) = group_iter.next() {
|
||
accumulated += group.len();
|
||
split_outputs[split_idx].extend(group);
|
||
} else {
|
||
break;
|
||
}
|
||
}
|
||
}
|
||
|
||
if let Some(idx) = rest_index {
|
||
for group in group_iter {
|
||
split_outputs[idx].extend(group);
|
||
}
|
||
}
|
||
|
||
for (spec, output_lines) in specs.iter().zip(split_outputs.iter()) {
|
||
write_lines_to_file(&spec.path, output_lines)?;
|
||
eprintln!("{}: {} examples", spec.path.display(), output_lines.len());
|
||
}
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// Groups lines by the specified stratification field.
|
||
///
|
||
/// When `stratify` is `None`, each line becomes its own group.
|
||
/// When a line is missing the stratification field, it is also placed in its own group.
|
||
fn group_lines(lines: &[String], stratify: Stratify) -> Vec<Vec<String>> {
|
||
if stratify == Stratify::None {
|
||
return lines.iter().map(|line| vec![line.clone()]).collect();
|
||
}
|
||
|
||
let get_key = |line: &str| {
|
||
let json: Value = serde_json::from_str(line).unwrap_or_default();
|
||
match stratify {
|
||
Stratify::Repo => json
|
||
.get("repository_url")
|
||
.and_then(|v| v.as_str())
|
||
.map(|s| s.to_string()),
|
||
Stratify::CursorPath => json
|
||
.get("cursor_path")
|
||
.and_then(|v| v.as_str())
|
||
.map(|s| s.to_string()),
|
||
Stratify::Project => json
|
||
.get("cursor_path")
|
||
.and_then(|v| v.as_str())
|
||
.and_then(|s| s.split(['/', '\\']).next())
|
||
.map(|s| s.to_string()),
|
||
Stratify::None => unreachable!(),
|
||
}
|
||
};
|
||
|
||
let mut groups: BTreeMap<String, Vec<String>> = BTreeMap::new();
|
||
let mut ungrouped: Vec<Vec<String>> = Vec::new();
|
||
|
||
for line in lines {
|
||
let key = get_key(line);
|
||
match key {
|
||
Some(key) => groups.entry(key).or_default().push(line.clone()),
|
||
None => ungrouped.push(vec![line.clone()]),
|
||
}
|
||
}
|
||
|
||
let mut result: Vec<Vec<String>> = groups.into_values().collect();
|
||
result.extend(ungrouped);
|
||
result
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
use std::io::Write;
|
||
use tempfile::NamedTempFile;
|
||
|
||
fn create_temp_jsonl(lines: &[&str]) -> NamedTempFile {
|
||
let mut file = NamedTempFile::new().unwrap();
|
||
for line in lines {
|
||
writeln!(file, "{}", line).unwrap();
|
||
}
|
||
file.flush().unwrap();
|
||
file
|
||
}
|
||
|
||
#[test]
|
||
fn test_parse_split_spec_percentage() {
|
||
let spec = parse_split_spec("train.jsonl=80%").unwrap();
|
||
assert_eq!(spec.path, PathBuf::from("train.jsonl"));
|
||
match spec.size {
|
||
SplitSize::Percentage(p) => assert!((p - 0.8).abs() < 0.001),
|
||
_ => panic!("expected percentage"),
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn test_parse_split_spec_absolute() {
|
||
let spec = parse_split_spec("test.jsonl=100").unwrap();
|
||
assert_eq!(spec.path, PathBuf::from("test.jsonl"));
|
||
match spec.size {
|
||
SplitSize::Absolute(n) => assert_eq!(n, 100),
|
||
_ => panic!("expected absolute"),
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn test_parse_split_spec_rest() {
|
||
let spec = parse_split_spec("valid.jsonl=rest").unwrap();
|
||
assert_eq!(spec.path, PathBuf::from("valid.jsonl"));
|
||
assert!(matches!(spec.size, SplitSize::Rest));
|
||
}
|
||
|
||
#[test]
|
||
fn test_group_lines_none() {
|
||
let lines = vec!["a".to_string(), "b".to_string(), "c".to_string()];
|
||
let groups = group_lines(&lines, Stratify::None);
|
||
assert_eq!(groups.len(), 3);
|
||
assert!(groups.iter().all(|g| g.len() == 1));
|
||
}
|
||
|
||
#[test]
|
||
fn test_compute_split_counts_percentage() {
|
||
let specs = vec![
|
||
SplitSpec {
|
||
path: PathBuf::from("a"),
|
||
size: SplitSize::Percentage(0.8),
|
||
},
|
||
SplitSpec {
|
||
path: PathBuf::from("b"),
|
||
size: SplitSize::Percentage(0.2),
|
||
},
|
||
];
|
||
let counts = compute_split_counts(&specs, 100).unwrap();
|
||
assert_eq!(counts, vec![80, 20]);
|
||
}
|
||
|
||
#[test]
|
||
fn test_compute_split_counts_with_rest() {
|
||
let specs = vec![
|
||
SplitSpec {
|
||
path: PathBuf::from("a"),
|
||
size: SplitSize::Percentage(0.8),
|
||
},
|
||
SplitSpec {
|
||
path: PathBuf::from("b"),
|
||
size: SplitSize::Rest,
|
||
},
|
||
];
|
||
let counts = compute_split_counts(&specs, 100).unwrap();
|
||
assert_eq!(counts, vec![80, 20]);
|
||
}
|
||
|
||
#[test]
|
||
fn test_compute_split_counts_absolute() {
|
||
let specs = vec![
|
||
SplitSpec {
|
||
path: PathBuf::from("a"),
|
||
size: SplitSize::Absolute(50),
|
||
},
|
||
SplitSpec {
|
||
path: PathBuf::from("b"),
|
||
size: SplitSize::Rest,
|
||
},
|
||
];
|
||
let counts = compute_split_counts(&specs, 100).unwrap();
|
||
assert_eq!(counts, vec![50, 50]);
|
||
}
|
||
|
||
#[test]
|
||
fn test_group_lines_by_repo() {
|
||
let lines = vec![
|
||
r#"{"repository_url": "repo1", "id": 1}"#.to_string(),
|
||
r#"{"repository_url": "repo1", "id": 2}"#.to_string(),
|
||
r#"{"repository_url": "repo2", "id": 3}"#.to_string(),
|
||
r#"{"id": 4}"#.to_string(),
|
||
];
|
||
|
||
let groups = group_lines(&lines, Stratify::Repo);
|
||
|
||
let grouped_count: usize = groups.iter().filter(|g| g.len() > 1).count();
|
||
let ungrouped_count: usize = groups.iter().filter(|g| g.len() == 1).count();
|
||
let total_lines: usize = groups.iter().map(|g| g.len()).sum();
|
||
|
||
assert_eq!(grouped_count, 1); // repo1 has 2 lines
|
||
assert_eq!(ungrouped_count, 2); // repo2 (1 line) + line without repo
|
||
assert_eq!(total_lines, 4);
|
||
}
|
||
|
||
#[test]
|
||
fn test_group_lines_by_cursor_path() {
|
||
let lines = vec![
|
||
r#"{"cursor_path": "src/main.rs", "id": 1}"#.to_string(),
|
||
r#"{"cursor_path": "src/main.rs", "id": 2}"#.to_string(),
|
||
r#"{"cursor_path": "src/lib.rs", "id": 3}"#.to_string(),
|
||
];
|
||
|
||
let groups = group_lines(&lines, Stratify::CursorPath);
|
||
|
||
let total_lines: usize = groups.iter().map(|g| g.len()).sum();
|
||
assert_eq!(groups.len(), 2);
|
||
assert_eq!(total_lines, 3);
|
||
}
|
||
|
||
#[test]
|
||
fn test_run_split_basic() {
|
||
let input = create_temp_jsonl(&[
|
||
r#"{"repository_url": "repo1", "id": 1}"#,
|
||
r#"{"repository_url": "repo1", "id": 2}"#,
|
||
r#"{"repository_url": "repo2", "id": 3}"#,
|
||
r#"{"repository_url": "repo2", "id": 4}"#,
|
||
r#"{"repository_url": "repo3", "id": 5}"#,
|
||
r#"{"repository_url": "repo3", "id": 6}"#,
|
||
r#"{"repository_url": "repo4", "id": 7}"#,
|
||
r#"{"repository_url": "repo4", "id": 8}"#,
|
||
]);
|
||
|
||
let temp_dir = tempfile::tempdir().unwrap();
|
||
let train_path = temp_dir.path().join("train.jsonl");
|
||
let valid_path = temp_dir.path().join("valid.jsonl");
|
||
|
||
let args = SplitArgs {
|
||
seed: Some(42),
|
||
stratify: Stratify::Repo,
|
||
};
|
||
let inputs = vec![
|
||
input.path().to_path_buf(),
|
||
PathBuf::from(format!("{}=50%", train_path.display())),
|
||
PathBuf::from(format!("{}=rest", valid_path.display())),
|
||
];
|
||
|
||
run_split(&args, &inputs).unwrap();
|
||
|
||
let train_content = std::fs::read_to_string(&train_path).unwrap();
|
||
let valid_content = std::fs::read_to_string(&valid_path).unwrap();
|
||
|
||
let train_lines: Vec<&str> = train_content.lines().collect();
|
||
let valid_lines: Vec<&str> = valid_content.lines().collect();
|
||
|
||
assert_eq!(train_lines.len() + valid_lines.len(), 8);
|
||
|
||
let get_repo = |line: &str| -> Option<String> {
|
||
let value: Value = serde_json::from_str(line).ok()?;
|
||
value
|
||
.get("repository_url")
|
||
.and_then(|v| v.as_str())
|
||
.map(|s| s.to_string())
|
||
};
|
||
|
||
let train_repos: std::collections::HashSet<_> =
|
||
train_lines.iter().filter_map(|l| get_repo(l)).collect();
|
||
let valid_repos: std::collections::HashSet<_> =
|
||
valid_lines.iter().filter_map(|l| get_repo(l)).collect();
|
||
|
||
assert!(
|
||
train_repos.is_disjoint(&valid_repos),
|
||
"train and valid should have non-overlapping repos"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn test_multiple_rest_fails() {
|
||
let specs = vec![
|
||
SplitSpec {
|
||
path: PathBuf::from("a"),
|
||
size: SplitSize::Rest,
|
||
},
|
||
SplitSpec {
|
||
path: PathBuf::from("b"),
|
||
size: SplitSize::Rest,
|
||
},
|
||
];
|
||
assert!(compute_split_counts(&specs, 100).is_err());
|
||
}
|
||
|
||
#[test]
|
||
fn test_absolute_targets_lines_not_groups() {
|
||
// 5 repos × 3 lines each = 15 total lines.
|
||
// `train=6` should target ~6 lines (2 groups), NOT 6 groups (all 15 lines).
|
||
let input = create_temp_jsonl(&[
|
||
r#"{"repository_url": "r1", "id": 1}"#,
|
||
r#"{"repository_url": "r1", "id": 2}"#,
|
||
r#"{"repository_url": "r1", "id": 3}"#,
|
||
r#"{"repository_url": "r2", "id": 4}"#,
|
||
r#"{"repository_url": "r2", "id": 5}"#,
|
||
r#"{"repository_url": "r2", "id": 6}"#,
|
||
r#"{"repository_url": "r3", "id": 7}"#,
|
||
r#"{"repository_url": "r3", "id": 8}"#,
|
||
r#"{"repository_url": "r3", "id": 9}"#,
|
||
r#"{"repository_url": "r4", "id": 10}"#,
|
||
r#"{"repository_url": "r4", "id": 11}"#,
|
||
r#"{"repository_url": "r4", "id": 12}"#,
|
||
r#"{"repository_url": "r5", "id": 13}"#,
|
||
r#"{"repository_url": "r5", "id": 14}"#,
|
||
r#"{"repository_url": "r5", "id": 15}"#,
|
||
]);
|
||
|
||
let temp_dir = tempfile::tempdir().unwrap();
|
||
let train_path = temp_dir.path().join("train.jsonl");
|
||
let valid_path = temp_dir.path().join("valid.jsonl");
|
||
|
||
let args = SplitArgs {
|
||
seed: Some(42),
|
||
stratify: Stratify::Repo,
|
||
};
|
||
let inputs = vec![
|
||
input.path().to_path_buf(),
|
||
PathBuf::from(format!("{}=6", train_path.display())),
|
||
PathBuf::from(format!("{}=rest", valid_path.display())),
|
||
];
|
||
|
||
run_split(&args, &inputs).unwrap();
|
||
|
||
let train_content = std::fs::read_to_string(&train_path).unwrap();
|
||
let valid_content = std::fs::read_to_string(&valid_path).unwrap();
|
||
|
||
let train_lines: Vec<&str> = train_content.lines().collect();
|
||
let valid_lines: Vec<&str> = valid_content.lines().collect();
|
||
|
||
// With 3-line groups, train should get 2 groups (6 lines) to meet the
|
||
// target of 6, NOT 6 groups (which don't even exist). Valid gets the rest.
|
||
assert_eq!(train_lines.len(), 6);
|
||
assert_eq!(valid_lines.len(), 9);
|
||
}
|
||
|
||
#[test]
|
||
fn test_stratify_by_project() {
|
||
// 5 repos × 3 lines each = 15 total lines.
|
||
// `train=6` should target ~6 lines (2 groups), NOT 6 groups (all 15 lines).
|
||
let input = create_temp_jsonl(&[
|
||
r#"{"cursor_path": "project1/some/file.rs", "id": 1}"#,
|
||
r#"{"cursor_path": "project2/some/file.rs", "id": 2}"#,
|
||
r#"{"cursor_path": "project3/some/file.rs", "id": 3}"#,
|
||
r#"{"cursor_path": "project1/other/file.rs", "id": 4}"#,
|
||
r#"{"cursor_path": "project2/other/file.rs", "id": 5}"#,
|
||
r#"{"cursor_path": "project3/other/file.rs", "id": 6}"#,
|
||
r#"{"cursor_path": "project3/another/file.rs", "id": 7}"#,
|
||
r#"{"cursor_path": "project3/even/more.rs", "id": 8}"#,
|
||
]);
|
||
|
||
let temp_dir = tempfile::tempdir().unwrap();
|
||
let train_path = temp_dir.path().join("train.jsonl");
|
||
let valid_path = temp_dir.path().join("valid.jsonl");
|
||
|
||
let args = SplitArgs {
|
||
seed: Some(1),
|
||
stratify: Stratify::Project,
|
||
};
|
||
let inputs = vec![
|
||
input.path().to_path_buf(),
|
||
PathBuf::from(format!("{}=4", train_path.display())),
|
||
PathBuf::from(format!("{}=rest", valid_path.display())),
|
||
];
|
||
|
||
run_split(&args, &inputs).unwrap();
|
||
|
||
let train_content = std::fs::read_to_string(&train_path).unwrap();
|
||
let valid_content = std::fs::read_to_string(&valid_path).unwrap();
|
||
|
||
// Make sure project 1 and project 2 are in the train set, and project 3 is in the valid set.
|
||
let mut train_ids: Vec<u64> = train_content
|
||
.lines()
|
||
.map(|l| {
|
||
serde_json::from_str::<serde_json::Value>(l).unwrap()["id"]
|
||
.as_u64()
|
||
.unwrap()
|
||
})
|
||
.collect();
|
||
let mut valid_ids: Vec<u64> = valid_content
|
||
.lines()
|
||
.map(|l| {
|
||
serde_json::from_str::<serde_json::Value>(l).unwrap()["id"]
|
||
.as_u64()
|
||
.unwrap()
|
||
})
|
||
.collect();
|
||
|
||
train_ids.sort();
|
||
valid_ids.sort();
|
||
|
||
assert_eq!(train_ids, vec![1, 2, 4, 5]);
|
||
assert_eq!(valid_ids, vec![3, 6, 7, 8]);
|
||
}
|
||
}
|