tenant.openagents/omega

1//! `ep split` implementation.
2//!
3//! This command splits a JSONL dataset into multiple files based on size specifications,
4//! with optional stratification by a JSON field.
5//!
6//! # Usage
7//!
8//! ```text
9//! ep split [--stratify=<field>] [input.jsonl] <out1>=<size1> <out2>=<size2> ...
10//! ```
11//!
12//! If `input.jsonl` is not provided or is `-`, reads from stdin.
13//!
14//! # Size specifications
15//!
16//! - `80%` - percentage of total examples (lines)
17//! - `100` - approximate absolute count of examples (lines)
18//! - `rest` - all remaining items (only one split can use this)
19//!
20//! # Stratification
21//!
22//! The `--stratify` flag controls how examples are grouped before splitting:
23//!
24//! - `cursor-path` (default): group by the `cursor_path` JSON field
25//! - `project`: group by the first component of the `cursor_path` JSON field
26//! - `repo`: group by the `repository_url` JSON field
27//! - `none`: no grouping, split individual examples
28//!
29//! When stratifying, the split ensures each output file contains examples from
30//! non-overlapping groups. Size specifications always apply to the number of
31//! examples (lines), with whole groups assigned greedily to meet the target.
32//! Examples missing the stratification field are treated as individual groups.
33
34use anyhow::{Context as _, Result, bail};
35use clap::Args;
36use rand::SeedableRng;
37use rand::seq::SliceRandom;
38use serde_json::Value;
39use std::collections::BTreeMap;
40use std::fs::File;
41use std::io::{self, BufRead, BufReader, BufWriter, Write};
42use std::path::{Path, PathBuf};
43
44/// `ep split` CLI args.
45#[derive(Debug, Args, Clone)]
46#[command(
47    about = "Split a JSONL dataset into multiple files with optional stratification",
48    after_help = r#"SIZE SPECIFICATIONS:
49  <percentage>%    Percentage of total (e.g., 80%)
50  <count>          Absolute number (e.g., 100)
51  rest             All remaining items (only one output can use this)
52
53  Sizes always apply to examples (lines). When stratifying, whole groups
54  are assigned greedily to approximate the target count.
55
56EXAMPLES:
57  # Split 80% train, 20% validation (default: stratify by cursor_path)
58  ep split input.jsonl train.jsonl=80% valid.jsonl=rest
59
60  # Split into train/valid/test
61  ep split input.jsonl train.jsonl=80% valid.jsonl=10% test.jsonl=rest
62
63  # Stratify by repository_url instead of cursor_path
64  ep split --stratify=repo input.jsonl train.jsonl=80% valid.jsonl=rest
65
66  # No stratification (split by individual examples)
67  ep split --stratify=none input.jsonl train.jsonl=80% valid.jsonl=rest
68
69  # Read from stdin
70  cat input.jsonl | ep split train.jsonl=80% valid.jsonl=rest
71
72  # Reproducible split with seed
73  ep split --seed 42 input.jsonl train.jsonl=80% valid.jsonl=rest
74
75STRATIFICATION:
76  Controls how examples are grouped before splitting:
77    cursor-path  Group by "cursor_path" field (default)
78    project      Group by the first component of the "cursor_path" field
79    repo         Group by "repository_url" field
80    none         No grouping, split individual examples
81
82  When stratifying, the split ensures each output file contains examples
83  from non-overlapping groups. This prevents data leakage between
84  train/test splits.
85"#
86)]
87pub struct SplitArgs {
88    /// Random seed for reproducibility
89    #[arg(long)]
90    pub seed: Option<u64>,
91
92    /// Stratification field for splitting the dataset
93    #[arg(long, default_value = "cursor-path")]
94    pub stratify: Stratify,
95}
96
97#[derive(Debug, Clone, Copy, PartialEq, Eq, clap::ValueEnum, strum::Display)]
98pub enum Stratify {
99    #[strum(serialize = "cursor_path")]
100    CursorPath,
101    #[strum(serialize = "project")]
102    Project,
103    #[strum(serialize = "repo")]
104    Repo,
105    #[strum(serialize = "none")]
106    None,
107}
108
109#[derive(Debug, Clone)]
110pub enum SplitSize {
111    Percentage(f64),
112    Absolute(usize),
113    Rest,
114}
115
116#[derive(Debug, Clone)]
117pub struct SplitSpec {
118    pub path: PathBuf,
119    pub size: SplitSize,
120}
121
122fn parse_split_spec(spec: &str) -> Result<SplitSpec> {
123    let (path, size_str) = spec
124        .rsplit_once('=')
125        .with_context(|| format!("invalid split spec '{}': expected <path>=<size>", spec))?;
126
127    let size = if size_str == "rest" {
128        SplitSize::Rest
129    } else if size_str.ends_with('%') {
130        let pct_str = size_str.trim_end_matches('%');
131        let pct: f64 = pct_str
132            .parse()
133            .with_context(|| format!("invalid percentage '{}' in '{}'", pct_str, spec))?;
134        if !(0.0..=100.0).contains(&pct) {
135            bail!("percentage must be between 0 and 100, got {}", pct);
136        }
137        SplitSize::Percentage(pct / 100.0)
138    } else {
139        let count: usize = size_str
140            .parse()
141            .with_context(|| format!("invalid count '{}' in '{}'", size_str, spec))?;
142        SplitSize::Absolute(count)
143    };
144
145    Ok(SplitSpec {
146        path: PathBuf::from(path),
147        size,
148    })
149}
150
151fn read_lines_from_input(input: Option<&Path>) -> Result<Vec<String>> {
152    let reader: Box<dyn BufRead> = match input {
153        Some(path) => {
154            let file =
155                File::open(path).with_context(|| format!("failed to open '{}'", path.display()))?;
156            Box::new(BufReader::new(file))
157        }
158        None => Box::new(BufReader::new(io::stdin())),
159    };
160
161    let lines: Vec<String> = reader
162        .lines()
163        .collect::<io::Result<Vec<_>>>()
164        .context("failed to read input lines")?;
165
166    Ok(lines)
167}
168
169fn compute_split_counts(specs: &[SplitSpec], total: usize) -> Result<Vec<usize>> {
170    let mut counts = vec![0usize; specs.len()];
171    let mut remaining = total;
172    let mut rest_index: Option<usize> = None;
173
174    for (i, spec) in specs.iter().enumerate() {
175        match &spec.size {
176            SplitSize::Percentage(pct) => {
177                let count = (total as f64 * pct).round() as usize;
178                counts[i] = count.min(remaining);
179                remaining = remaining.saturating_sub(counts[i]);
180            }
181            SplitSize::Absolute(count) => {
182                counts[i] = (*count).min(remaining);
183                remaining = remaining.saturating_sub(counts[i]);
184            }
185            SplitSize::Rest => {
186                if rest_index.is_some() {
187                    bail!("only one split can use 'rest'");
188                }
189                rest_index = Some(i);
190            }
191        }
192    }
193
194    if let Some(idx) = rest_index {
195        counts[idx] = remaining;
196    }
197
198    Ok(counts)
199}
200
201fn write_lines_to_file(path: &Path, lines: &[String]) -> Result<()> {
202    if let Some(parent) = path.parent() {
203        if !parent.as_os_str().is_empty() {
204            std::fs::create_dir_all(parent)
205                .with_context(|| format!("failed to create directory '{}'", parent.display()))?;
206        }
207    }
208
209    let file =
210        File::create(path).with_context(|| format!("failed to create '{}'", path.display()))?;
211    let mut writer = BufWriter::new(file);
212
213    for line in lines {
214        writeln!(writer, "{}", line)
215            .with_context(|| format!("failed to write to '{}'", path.display()))?;
216    }
217
218    writer
219        .flush()
220        .with_context(|| format!("failed to flush '{}'", path.display()))?;
221
222    Ok(())
223}
224
225pub fn run_split(args: &SplitArgs, inputs: &[PathBuf]) -> Result<()> {
226    if inputs.is_empty() {
227        bail!("usage: ep split [input.jsonl] train.jsonl=80% valid.jsonl=rest");
228    }
229
230    let (input_path, split_specs_raw): (Option<&Path>, &[PathBuf]) =
231        if inputs.first().is_some_and(|p| {
232            let s = p.to_string_lossy();
233            !s.contains('=')
234        }) {
235            let first = inputs.first().map(|p| p.as_path());
236            let first = if first == Some(Path::new("-")) {
237                None
238            } else {
239                first
240            };
241            (first, &inputs[1..])
242        } else {
243            (None, inputs)
244        };
245
246    if split_specs_raw.is_empty() {
247        bail!("no split specifications provided");
248    }
249
250    let specs: Vec<SplitSpec> = split_specs_raw
251        .iter()
252        .map(|p| parse_split_spec(&p.to_string_lossy()))
253        .collect::<Result<Vec<_>>>()?;
254
255    let lines = read_lines_from_input(input_path)?;
256    let total_lines = lines.len();
257
258    if total_lines == 0 {
259        for spec in &specs {
260            write_lines_to_file(&spec.path, &[])?;
261        }
262        return Ok(());
263    }
264
265    let mut grouped_lines = group_lines(&lines, args.stratify);
266
267    if args.stratify != Stratify::None {
268        eprintln!(
269            "Stratifying by {} ({} unique groups, {} examples)",
270            args.stratify,
271            grouped_lines.len(),
272            total_lines
273        );
274    } else {
275        eprintln!(
276            "No stratification, splitting {} examples by line",
277            total_lines
278        );
279    }
280
281    let mut rng = match args.seed {
282        Some(seed) => rand::rngs::StdRng::seed_from_u64(seed),
283        None => rand::rngs::StdRng::from_os_rng(),
284    };
285
286    grouped_lines.shuffle(&mut rng);
287
288    let line_targets = compute_split_counts(&specs, total_lines)?;
289    let rest_index = specs.iter().position(|s| matches!(s.size, SplitSize::Rest));
290    let mut split_outputs: Vec<Vec<String>> = vec![Vec::new(); specs.len()];
291    let mut group_iter = grouped_lines.into_iter();
292
293    for (split_idx, &target) in line_targets.iter().enumerate() {
294        if Some(split_idx) == rest_index {
295            continue;
296        }
297        let mut accumulated = 0;
298        while accumulated < target {
299            if let Some(group) = group_iter.next() {
300                accumulated += group.len();
301                split_outputs[split_idx].extend(group);
302            } else {
303                break;
304            }
305        }
306    }
307
308    if let Some(idx) = rest_index {
309        for group in group_iter {
310            split_outputs[idx].extend(group);
311        }
312    }
313
314    for (spec, output_lines) in specs.iter().zip(split_outputs.iter()) {
315        write_lines_to_file(&spec.path, output_lines)?;
316        eprintln!("{}: {} examples", spec.path.display(), output_lines.len());
317    }
318
319    Ok(())
320}
321
322/// Groups lines by the specified stratification field.
323///
324/// When `stratify` is `None`, each line becomes its own group.
325/// When a line is missing the stratification field, it is also placed in its own group.
326fn group_lines(lines: &[String], stratify: Stratify) -> Vec<Vec<String>> {
327    if stratify == Stratify::None {
328        return lines.iter().map(|line| vec![line.clone()]).collect();
329    }
330
331    let get_key = |line: &str| {
332        let json: Value = serde_json::from_str(line).unwrap_or_default();
333        match stratify {
334            Stratify::Repo => json
335                .get("repository_url")
336                .and_then(|v| v.as_str())
337                .map(|s| s.to_string()),
338            Stratify::CursorPath => json
339                .get("cursor_path")
340                .and_then(|v| v.as_str())
341                .map(|s| s.to_string()),
342            Stratify::Project => json
343                .get("cursor_path")
344                .and_then(|v| v.as_str())
345                .and_then(|s| s.split(['/', '\\']).next())
346                .map(|s| s.to_string()),
347            Stratify::None => unreachable!(),
348        }
349    };
350
351    let mut groups: BTreeMap<String, Vec<String>> = BTreeMap::new();
352    let mut ungrouped: Vec<Vec<String>> = Vec::new();
353
354    for line in lines {
355        let key = get_key(line);
356        match key {
357            Some(key) => groups.entry(key).or_default().push(line.clone()),
358            None => ungrouped.push(vec![line.clone()]),
359        }
360    }
361
362    let mut result: Vec<Vec<String>> = groups.into_values().collect();
363    result.extend(ungrouped);
364    result
365}
366
367#[cfg(test)]
368mod tests {
369    use super::*;
370    use std::io::Write;
371    use tempfile::NamedTempFile;
372
373    fn create_temp_jsonl(lines: &[&str]) -> NamedTempFile {
374        let mut file = NamedTempFile::new().unwrap();
375        for line in lines {
376            writeln!(file, "{}", line).unwrap();
377        }
378        file.flush().unwrap();
379        file
380    }
381
382    #[test]
383    fn test_parse_split_spec_percentage() {
384        let spec = parse_split_spec("train.jsonl=80%").unwrap();
385        assert_eq!(spec.path, PathBuf::from("train.jsonl"));
386        match spec.size {
387            SplitSize::Percentage(p) => assert!((p - 0.8).abs() < 0.001),
388            _ => panic!("expected percentage"),
389        }
390    }
391
392    #[test]
393    fn test_parse_split_spec_absolute() {
394        let spec = parse_split_spec("test.jsonl=100").unwrap();
395        assert_eq!(spec.path, PathBuf::from("test.jsonl"));
396        match spec.size {
397            SplitSize::Absolute(n) => assert_eq!(n, 100),
398            _ => panic!("expected absolute"),
399        }
400    }
401
402    #[test]
403    fn test_parse_split_spec_rest() {
404        let spec = parse_split_spec("valid.jsonl=rest").unwrap();
405        assert_eq!(spec.path, PathBuf::from("valid.jsonl"));
406        assert!(matches!(spec.size, SplitSize::Rest));
407    }
408
409    #[test]
410    fn test_group_lines_none() {
411        let lines = vec!["a".to_string(), "b".to_string(), "c".to_string()];
412        let groups = group_lines(&lines, Stratify::None);
413        assert_eq!(groups.len(), 3);
414        assert!(groups.iter().all(|g| g.len() == 1));
415    }
416
417    #[test]
418    fn test_compute_split_counts_percentage() {
419        let specs = vec![
420            SplitSpec {
421                path: PathBuf::from("a"),
422                size: SplitSize::Percentage(0.8),
423            },
424            SplitSpec {
425                path: PathBuf::from("b"),
426                size: SplitSize::Percentage(0.2),
427            },
428        ];
429        let counts = compute_split_counts(&specs, 100).unwrap();
430        assert_eq!(counts, vec![80, 20]);
431    }
432
433    #[test]
434    fn test_compute_split_counts_with_rest() {
435        let specs = vec![
436            SplitSpec {
437                path: PathBuf::from("a"),
438                size: SplitSize::Percentage(0.8),
439            },
440            SplitSpec {
441                path: PathBuf::from("b"),
442                size: SplitSize::Rest,
443            },
444        ];
445        let counts = compute_split_counts(&specs, 100).unwrap();
446        assert_eq!(counts, vec![80, 20]);
447    }
448
449    #[test]
450    fn test_compute_split_counts_absolute() {
451        let specs = vec![
452            SplitSpec {
453                path: PathBuf::from("a"),
454                size: SplitSize::Absolute(50),
455            },
456            SplitSpec {
457                path: PathBuf::from("b"),
458                size: SplitSize::Rest,
459            },
460        ];
461        let counts = compute_split_counts(&specs, 100).unwrap();
462        assert_eq!(counts, vec![50, 50]);
463    }
464
465    #[test]
466    fn test_group_lines_by_repo() {
467        let lines = vec![
468            r#"{"repository_url": "repo1", "id": 1}"#.to_string(),
469            r#"{"repository_url": "repo1", "id": 2}"#.to_string(),
470            r#"{"repository_url": "repo2", "id": 3}"#.to_string(),
471            r#"{"id": 4}"#.to_string(),
472        ];
473
474        let groups = group_lines(&lines, Stratify::Repo);
475
476        let grouped_count: usize = groups.iter().filter(|g| g.len() > 1).count();
477        let ungrouped_count: usize = groups.iter().filter(|g| g.len() == 1).count();
478        let total_lines: usize = groups.iter().map(|g| g.len()).sum();
479
480        assert_eq!(grouped_count, 1); // repo1 has 2 lines
481        assert_eq!(ungrouped_count, 2); // repo2 (1 line) + line without repo
482        assert_eq!(total_lines, 4);
483    }
484
485    #[test]
486    fn test_group_lines_by_cursor_path() {
487        let lines = vec![
488            r#"{"cursor_path": "src/main.rs", "id": 1}"#.to_string(),
489            r#"{"cursor_path": "src/main.rs", "id": 2}"#.to_string(),
490            r#"{"cursor_path": "src/lib.rs", "id": 3}"#.to_string(),
491        ];
492
493        let groups = group_lines(&lines, Stratify::CursorPath);
494
495        let total_lines: usize = groups.iter().map(|g| g.len()).sum();
496        assert_eq!(groups.len(), 2);
497        assert_eq!(total_lines, 3);
498    }
499
500    #[test]
501    fn test_run_split_basic() {
502        let input = create_temp_jsonl(&[
503            r#"{"repository_url": "repo1", "id": 1}"#,
504            r#"{"repository_url": "repo1", "id": 2}"#,
505            r#"{"repository_url": "repo2", "id": 3}"#,
506            r#"{"repository_url": "repo2", "id": 4}"#,
507            r#"{"repository_url": "repo3", "id": 5}"#,
508            r#"{"repository_url": "repo3", "id": 6}"#,
509            r#"{"repository_url": "repo4", "id": 7}"#,
510            r#"{"repository_url": "repo4", "id": 8}"#,
511        ]);
512
513        let temp_dir = tempfile::tempdir().unwrap();
514        let train_path = temp_dir.path().join("train.jsonl");
515        let valid_path = temp_dir.path().join("valid.jsonl");
516
517        let args = SplitArgs {
518            seed: Some(42),
519            stratify: Stratify::Repo,
520        };
521        let inputs = vec![
522            input.path().to_path_buf(),
523            PathBuf::from(format!("{}=50%", train_path.display())),
524            PathBuf::from(format!("{}=rest", valid_path.display())),
525        ];
526
527        run_split(&args, &inputs).unwrap();
528
529        let train_content = std::fs::read_to_string(&train_path).unwrap();
530        let valid_content = std::fs::read_to_string(&valid_path).unwrap();
531
532        let train_lines: Vec<&str> = train_content.lines().collect();
533        let valid_lines: Vec<&str> = valid_content.lines().collect();
534
535        assert_eq!(train_lines.len() + valid_lines.len(), 8);
536
537        let get_repo = |line: &str| -> Option<String> {
538            let value: Value = serde_json::from_str(line).ok()?;
539            value
540                .get("repository_url")
541                .and_then(|v| v.as_str())
542                .map(|s| s.to_string())
543        };
544
545        let train_repos: std::collections::HashSet<_> =
546            train_lines.iter().filter_map(|l| get_repo(l)).collect();
547        let valid_repos: std::collections::HashSet<_> =
548            valid_lines.iter().filter_map(|l| get_repo(l)).collect();
549
550        assert!(
551            train_repos.is_disjoint(&valid_repos),
552            "train and valid should have non-overlapping repos"
553        );
554    }
555
556    #[test]
557    fn test_multiple_rest_fails() {
558        let specs = vec![
559            SplitSpec {
560                path: PathBuf::from("a"),
561                size: SplitSize::Rest,
562            },
563            SplitSpec {
564                path: PathBuf::from("b"),
565                size: SplitSize::Rest,
566            },
567        ];
568        assert!(compute_split_counts(&specs, 100).is_err());
569    }
570
571    #[test]
572    fn test_absolute_targets_lines_not_groups() {
573        // 5 repos × 3 lines each = 15 total lines.
574        // `train=6` should target ~6 lines (2 groups), NOT 6 groups (all 15 lines).
575        let input = create_temp_jsonl(&[
576            r#"{"repository_url": "r1", "id": 1}"#,
577            r#"{"repository_url": "r1", "id": 2}"#,
578            r#"{"repository_url": "r1", "id": 3}"#,
579            r#"{"repository_url": "r2", "id": 4}"#,
580            r#"{"repository_url": "r2", "id": 5}"#,
581            r#"{"repository_url": "r2", "id": 6}"#,
582            r#"{"repository_url": "r3", "id": 7}"#,
583            r#"{"repository_url": "r3", "id": 8}"#,
584            r#"{"repository_url": "r3", "id": 9}"#,
585            r#"{"repository_url": "r4", "id": 10}"#,
586            r#"{"repository_url": "r4", "id": 11}"#,
587            r#"{"repository_url": "r4", "id": 12}"#,
588            r#"{"repository_url": "r5", "id": 13}"#,
589            r#"{"repository_url": "r5", "id": 14}"#,
590            r#"{"repository_url": "r5", "id": 15}"#,
591        ]);
592
593        let temp_dir = tempfile::tempdir().unwrap();
594        let train_path = temp_dir.path().join("train.jsonl");
595        let valid_path = temp_dir.path().join("valid.jsonl");
596
597        let args = SplitArgs {
598            seed: Some(42),
599            stratify: Stratify::Repo,
600        };
601        let inputs = vec![
602            input.path().to_path_buf(),
603            PathBuf::from(format!("{}=6", train_path.display())),
604            PathBuf::from(format!("{}=rest", valid_path.display())),
605        ];
606
607        run_split(&args, &inputs).unwrap();
608
609        let train_content = std::fs::read_to_string(&train_path).unwrap();
610        let valid_content = std::fs::read_to_string(&valid_path).unwrap();
611
612        let train_lines: Vec<&str> = train_content.lines().collect();
613        let valid_lines: Vec<&str> = valid_content.lines().collect();
614
615        // With 3-line groups, train should get 2 groups (6 lines) to meet the
616        // target of 6, NOT 6 groups (which don't even exist). Valid gets the rest.
617        assert_eq!(train_lines.len(), 6);
618        assert_eq!(valid_lines.len(), 9);
619    }
620
621    #[test]
622    fn test_stratify_by_project() {
623        // 5 repos × 3 lines each = 15 total lines.
624        // `train=6` should target ~6 lines (2 groups), NOT 6 groups (all 15 lines).
625        let input = create_temp_jsonl(&[
626            r#"{"cursor_path": "project1/some/file.rs", "id": 1}"#,
627            r#"{"cursor_path": "project2/some/file.rs", "id": 2}"#,
628            r#"{"cursor_path": "project3/some/file.rs", "id": 3}"#,
629            r#"{"cursor_path": "project1/other/file.rs", "id": 4}"#,
630            r#"{"cursor_path": "project2/other/file.rs", "id": 5}"#,
631            r#"{"cursor_path": "project3/other/file.rs", "id": 6}"#,
632            r#"{"cursor_path": "project3/another/file.rs", "id": 7}"#,
633            r#"{"cursor_path": "project3/even/more.rs", "id": 8}"#,
634        ]);
635
636        let temp_dir = tempfile::tempdir().unwrap();
637        let train_path = temp_dir.path().join("train.jsonl");
638        let valid_path = temp_dir.path().join("valid.jsonl");
639
640        let args = SplitArgs {
641            seed: Some(1),
642            stratify: Stratify::Project,
643        };
644        let inputs = vec![
645            input.path().to_path_buf(),
646            PathBuf::from(format!("{}=4", train_path.display())),
647            PathBuf::from(format!("{}=rest", valid_path.display())),
648        ];
649
650        run_split(&args, &inputs).unwrap();
651
652        let train_content = std::fs::read_to_string(&train_path).unwrap();
653        let valid_content = std::fs::read_to_string(&valid_path).unwrap();
654
655        // Make sure project 1 and project 2 are in the train set, and project 3 is in the valid set.
656        let mut train_ids: Vec<u64> = train_content
657            .lines()
658            .map(|l| {
659                serde_json::from_str::<serde_json::Value>(l).unwrap()["id"]
660                    .as_u64()
661                    .unwrap()
662            })
663            .collect();
664        let mut valid_ids: Vec<u64> = valid_content
665            .lines()
666            .map(|l| {
667                serde_json::from_str::<serde_json::Value>(l).unwrap()["id"]
668                    .as_u64()
669                    .unwrap()
670            })
671            .collect();
672
673        train_ids.sort();
674        valid_ids.sort();
675
676        assert_eq!(train_ids, vec![1, 2, 4, 5]);
677        assert_eq!(valid_ids, vec![3, 6, 7, 8]);
678    }
679}
680