//! `ep split` implementation.//!//! This command splits a JSONL dataset into multiple files based on size specifications,//! with stratification by repository URL (if the field is present).//!//! # Usage//!//! ```text//! ep split [input.jsonl] <out1>=<size1> <out2>=<size2> ...//! ```//!//! If `input.jsonl` is not provided or is `-`, reads from stdin.//!//! # Size specifications//!//! - `80%` - percentage of total (repositories if stratified, examples otherwise)//! - `100` - absolute count of repositories (if stratified) or examples//! - `rest` - all remaining items (only one split can use this)//!//! # Stratification//!//! When examples have a `repository_url` field, the split is stratified by repository.//! This ensures each output file contains examples from non-overlapping repositories.//! Size specifications apply to the number of repositories, not individual examples.//!//! Examples without `repository_url` are distributed proportionally across all outputs. use anyhow::{Context as _, Result, bail};use clap::Args;use rand::SeedableRng;use rand::seq::SliceRandom;use serde_json::Value;use std::collections::HashMap;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 (stratified by repository_url if present)",    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)   When stratifying by repository_url, sizes apply to repositories, not examples. EXAMPLES:  # Split 80% train, 20% validation  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   # Use absolute counts (100 repos to train, rest to valid)  ep split input.jsonl train.jsonl=100 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:  When examples have a "repository_url" field, the split ensures each output  file contains examples from non-overlapping repositories. This prevents  data leakage between train/test splits."#)]pub struct SplitArgs {    /// Random seed for reproducibility    #[arg(long)]    pub seed: Option<u64>,} #[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 get_repository_url(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())} fn group_lines_by_repo(lines: Vec<String>) -> (HashMap<String, Vec<String>>, Vec<String>) {    let mut by_repo: HashMap<String, Vec<String>> = HashMap::new();    let mut without_repo: Vec<String> = Vec::new();     for line in lines {        if let Some(repo_url) = get_repository_url(&line) {            by_repo.entry(repo_url).or_default().push(line);        } else {            without_repo.push(line);        }    }     (by_repo, without_repo)} 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 (by_repo, without_repo) = group_lines_by_repo(lines);    let has_repos = !by_repo.is_empty();     if has_repos {        eprintln!(            "Stratifying by repository_url ({} unique repositories, {} examples)",            by_repo.len(),            total_lines - without_repo.len()        );        if !without_repo.is_empty() {            eprintln!(                "  + {} examples without repository_url (distributed proportionally)",                without_repo.len()            );        }    }     let mut rng = match args.seed {        Some(seed) => rand::rngs::StdRng::seed_from_u64(seed),        None => rand::rngs::StdRng::from_os_rng(),    };     let mut split_outputs: Vec<Vec<String>> = vec![Vec::new(); specs.len()];     if has_repos {        let mut repos: Vec<String> = by_repo.keys().cloned().collect();        repos.shuffle(&mut rng);         let repo_counts = compute_split_counts(&specs, repos.len())?;         let mut repo_iter = repos.into_iter();        for (split_idx, &count) in repo_counts.iter().enumerate() {            for _ in 0..count {                if let Some(repo) = repo_iter.next() {                    if let Some(repo_lines) = by_repo.get(&repo) {                        split_outputs[split_idx].extend(repo_lines.iter().cloned());                    }                }            }        }         if !without_repo.is_empty() {            let no_repo_counts = compute_split_counts(&specs, without_repo.len())?;            let mut no_repo_shuffled = without_repo;            no_repo_shuffled.shuffle(&mut rng);             let mut line_iter = no_repo_shuffled.into_iter();            for (split_idx, &count) in no_repo_counts.iter().enumerate() {                for _ in 0..count {                    if let Some(line) = line_iter.next() {                        split_outputs[split_idx].push(line);                    }                }            }        }    } else {        let line_counts = compute_split_counts(&specs, total_lines)?;        let mut shuffled_lines = without_repo;        shuffled_lines.shuffle(&mut rng);         let mut line_iter = shuffled_lines.into_iter();        for (split_idx, &count) in line_counts.iter().enumerate() {            for _ in 0..count {                if let Some(line) = line_iter.next() {                    split_outputs[split_idx].push(line);                }            }        }    }     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(())} #[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_get_repository_url() {        let line = r#"{"repository_url": "https://github.com/example/repo", "data": 123}"#;        assert_eq!(            get_repository_url(line),            Some("https://github.com/example/repo".to_string())        );         let line_no_repo = r#"{"data": 123}"#;        assert_eq!(get_repository_url(line_no_repo), None);    }     #[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 (by_repo, without_repo) = group_lines_by_repo(lines);         assert_eq!(by_repo.len(), 2);        assert_eq!(by_repo.get("repo1").unwrap().len(), 2);        assert_eq!(by_repo.get("repo2").unwrap().len(), 1);        assert_eq!(without_repo.len(), 1);    }     #[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) };        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 train_repos: std::collections::HashSet<_> = train_lines            .iter()            .filter_map(|l| get_repository_url(l))            .collect();        let valid_repos: std::collections::HashSet<_> = valid_lines            .iter()            .filter_map(|l| get_repository_url(l))            .collect();         assert!(            train_repos.is_disjoint(&valid_repos),            "train and valid should have non-overlapping repos"        );    }

Lines 1–500 of 516. Load more