kibble 0.1.0

chew through any source into clean datasets — a fast ingestion, RAG & fine-tuning toolkit
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159

//! `kibble tune` — generate the Unsloth sprint curriculum + config.yaml from the built
//! dataset, mapping phases to sources via [tune].phase.

use std::io;
use std::path::Path;

/// Group the dataset rows into phases by source. Returns `(phase_name, rows)` for each
/// non-empty phase, in `[[tune.phase]]` order; a phase whose sources contribute no rows is
/// skipped. `rows[i]` corresponds to `sources[i]` (parallel).
pub fn plan_phases(rows: &[String], sources: &[String], phases: &[crate::config::TunePhase]) -> Vec<(String, Vec<String>)> {
    let mut by_source: std::collections::BTreeMap<&str, Vec<&str>> = std::collections::BTreeMap::new();
    for (r, s) in rows.iter().zip(sources.iter()) {
        by_source.entry(s.as_str()).or_default().push(r.as_str());
    }
    let mut out = Vec::new();
    for p in phases {
        let mut lines: Vec<String> = Vec::new();
        for src in &p.sources {
            if let Some(rs) = by_source.get(src.as_str()) {
                lines.extend(rs.iter().map(|s| s.to_string()));
            }
        }
        if !lines.is_empty() {
            out.push((p.name.clone(), lines));
        }
    }
    out
}

/// Read the built dataset + source manifest, write the sprint curriculum + config.yaml.
pub fn run_tune(repo_root: &Path) -> io::Result<()> {
    let cfg = crate::config::load_config(&repo_root.join(crate::config::CONFIG_FILE));
    let tune = cfg.tune.ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput,
        "no [tune] table in kibble.toml — see docs/TUNE.md"))?;
    let ds = repo_root.join(&cfg.paths.dataset_dir);
    let rows: Vec<String> = std::fs::read_to_string(ds.join("train.jsonl"))
        .map_err(|_| io::Error::new(io::ErrorKind::NotFound, "missing train.jsonl — run `kibble build` first"))?
        .lines().map(|l| l.to_string()).collect();
    let sources: Vec<String> = std::fs::read_to_string(ds.join("train.sources.jsonl"))
        .map_err(|_| io::Error::new(io::ErrorKind::NotFound, "missing train.sources.jsonl — run `kibble build` first"))?
        .lines().map(|l| l.to_string()).collect();
    if rows.len() != sources.len() {
        return Err(io::Error::new(io::ErrorKind::InvalidData,
            format!("train.jsonl ({}) / train.sources.jsonl ({}) out of sync — re-run `kibble build`", rows.len(), sources.len())));
    }
    for src in tune.phase.iter().flat_map(|p| &p.sources) {
        if !sources.iter().any(|s| s == src) {
            eprintln!("kibble: tune — source '{src}' has no rows in the dataset");
        }
    }

    let emitted = plan_phases(&rows, &sources, &tune.phase);
    if emitted.is_empty() {
        return Err(io::Error::new(io::ErrorKind::InvalidInput,
            "no rows matched any [tune.phase].sources — check your source names"));
    }

    let out_dir = repo_root.join(&tune.out);
    for (name, lines) in &emitted {
        let dir = out_dir.join("sprint").join(name);
        std::fs::create_dir_all(&dir)?;
        std::fs::write(dir.join("train.jsonl"), lines.join("\n") + "\n")?;
        println!("  {name}: {} rows", lines.len());
    }
    // phase0-smoke: first smoke_rows of the first non-empty phase.
    let smoke: Vec<&String> = emitted[0].1.iter().take(tune.smoke_rows).collect();
    let smoke_dir = out_dir.join("sprint").join("phase0-smoke");
    std::fs::create_dir_all(&smoke_dir)?;
    std::fs::write(smoke_dir.join("train.jsonl"),
        smoke.iter().map(|s| s.as_str()).collect::<Vec<_>>().join("\n") + "\n")?;

    // config.yaml (pretty JSON == valid YAML for the trainer's yaml.safe_load).
    let by_name: std::collections::HashMap<&str, &crate::config::TunePhase> =
        tune.phase.iter().map(|p| (p.name.as_str(), p)).collect();
    let phase_json: Vec<serde_json::Value> = emitted.iter().map(|(name, _)| {
        let p = by_name[name.as_str()];
        serde_json::json!({
            "name": name,
            "train_path": format!("sprint/{name}/train.jsonl"),
            "output_dir": name,
            "max_steps": p.max_steps,
            "resume_from": p.resume_from,
        })
    }).collect();
    let config = serde_json::json!({
        "model_name": tune.model,
        "trust_remote_code": true,
        "load_in_4bit": true,
        "r": tune.r,
        "lora_alpha": tune.lora_alpha,
        "lora_dropout": tune.lora_dropout,
        "sprint": {
            "max_seq_length": tune.max_seq_length,
            "per_device_train_batch_size": 1,
            "gradient_accumulation_steps": 8,
            "learning_rate": tune.learning_rate,
            "phases": phase_json,
        }
    });
    std::fs::write(out_dir.join("config.yaml"), serde_json::to_string_pretty(&config)?)?;
    println!("Wrote {} -> sprint/ + config.yaml. Next: kibble pack -> upload -> run the trainer.", out_dir.display());
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::config::TunePhase;

    fn phase(name: &str, sources: &[&str]) -> TunePhase {
        TunePhase { name: name.into(), sources: sources.iter().map(|s| s.to_string()).collect(), max_steps: 100, resume_from: None }
    }

    #[test]
    fn plan_phases_groups_by_source_in_order() {
        let rows: Vec<String> = vec!["R0".into(), "R1".into(), "R2".into(), "R3".into()];
        let sources = vec!["a".to_string(), "a".into(), "b".into(), "b".into()];
        let phases = vec![phase("p1", &["a"]), phase("p2", &["b"]), phase("p3", &["missing"])];
        let out = plan_phases(&rows, &sources, &phases);
        assert_eq!(out.len(), 2, "empty phase (missing source) is skipped");
        assert_eq!(out[0], ("p1".to_string(), vec!["R0".to_string(), "R1".to_string()]));
        assert_eq!(out[1], ("p2".to_string(), vec!["R2".to_string(), "R3".to_string()]));
    }

    #[tokio::test]
    async fn run_tune_writes_sprint_and_config() {
        let dir = std::env::temp_dir().join(format!("kibble_tune_{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&dir);
        let ds = dir.join("data/datasets/unsloth");
        std::fs::create_dir_all(&ds).unwrap();
        std::fs::write(ds.join("train.jsonl"),
            "{\"messages\":[{\"role\":\"user\",\"content\":\"q0\"}]}\n{\"messages\":[{\"role\":\"user\",\"content\":\"q1\"}]}\n{\"messages\":[{\"role\":\"user\",\"content\":\"q2\"}]}\n").unwrap();
        std::fs::write(ds.join("train.sources.jsonl"), "a\na\nb\n").unwrap();
        std::fs::write(dir.join(crate::config::CONFIG_FILE), concat!(
            "[tune]\nmodel = \"m\"\nsmoke_rows = 1\nout = \"data/tune\"\n",
            "[[tune.phase]]\nname = \"phase1\"\nsources = [\"a\"]\nmax_steps = 300\n",
            "[[tune.phase]]\nname = \"phase2\"\nsources = [\"b\"]\nmax_steps = 400\nresume_from = \"phase1\"\n",
        )).unwrap();

        run_tune(&dir).unwrap();
        let p1 = std::fs::read_to_string(dir.join("data/tune/sprint/phase1/train.jsonl")).unwrap();
        assert_eq!(p1.lines().count(), 2, "phase1 = the two 'a' rows");
        assert!(p1.contains("q0") && p1.contains("q1") && !p1.contains("q2"));
        let p2 = std::fs::read_to_string(dir.join("data/tune/sprint/phase2/train.jsonl")).unwrap();
        assert_eq!(p2.lines().count(), 1);
        assert!(p2.contains("q2"));
        let smoke = std::fs::read_to_string(dir.join("data/tune/sprint/phase0-smoke/train.jsonl")).unwrap();
        assert_eq!(smoke.lines().count(), 1, "smoke = first smoke_rows of first non-empty phase");
        let cfg: serde_json::Value = serde_json::from_str(
            &std::fs::read_to_string(dir.join("data/tune/config.yaml")).unwrap()).unwrap(); // valid JSON == valid YAML
        assert_eq!(cfg["model_name"], "m");
        let phases = cfg["sprint"]["phases"].as_array().unwrap();
        assert_eq!(phases.len(), 2);
        assert_eq!(phases[0]["train_path"], "sprint/phase1/train.jsonl");
        assert_eq!(phases[1]["resume_from"], "phase1");

        std::fs::remove_dir_all(&dir).ok();
    }
}