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
160
161
162
163
164
165
166
167
168

//! Build-side topic rebalancing: cap over-represented topics by dropping their
//! fringe (farthest-from-centroid) rows, against centroids clustered in-memory from
//! the curated train — which are then persisted as clusters.json (measure == cap).

use std::path::Path;

/// Indices to drop so no topic exceeds `cap` rows. For each topic in `0..k` whose
/// member count exceeds `cap`, drops the `(count - cap)` members with the lowest
/// `sims` (cosine of a row to its assigned centroid); within a topic, members rank
/// by `sims` descending, ties broken by index ascending (lowest index kept). Returns
/// drop indices sorted ascending. `cap >= 1` (the async caller clamps with `max(1, …)`).
pub fn topic_cap_drops(assign: &[usize], sims: &[f32], k: usize, cap: usize) -> Vec<usize> {
    debug_assert!(cap >= 1, "cap must be >= 1");
    let mut members: Vec<Vec<usize>> = vec![Vec::new(); k];
    for (i, &t) in assign.iter().enumerate() {
        if t < k {
            members[t].push(i);
        }
    }
    let mut drops: Vec<usize> = Vec::new();
    for m in members.iter_mut() {
        if m.len() <= cap {
            continue;
        }
        m.sort_by(|&a, &b| {
            sims[b].partial_cmp(&sims[a]).unwrap_or(std::cmp::Ordering::Equal).then(a.cmp(&b))
        });
        drops.extend_from_slice(&m[cap..]);
    }
    drops.sort_unstable();
    drops
}

/// The result of an inline rebalance: which train rows to drop, and the clustering
/// (train-only centroids, post-cap sizes) to persist as clusters.json.
pub struct RebalanceOutcome {
    pub drops: Vec<usize>,
    pub clusters: crate::cluster::ClusterResult,
}

/// Cluster the curated train answers in-memory, cap over-represented topics against
/// those centroids, and return the drop indices plus the clustering to persist.
/// Fail-soft: rebalance off / empty answers / empty base_url / embed error → `Ok(None)`.
pub async fn rebalance_inline(repo_root: &Path, train_answers: &[String]) -> std::io::Result<Option<RebalanceOutcome>> {
    let cfg = crate::config::load_config(&repo_root.join(crate::config::CONFIG_FILE));
    if !cfg.cluster.rebalance || train_answers.is_empty() || cfg.understand.embed.base_url.is_empty() {
        return Ok(None);
    }
    let store_dir = repo_root.join(&cfg.understand.embed.store);
    let vecs = match crate::embed::EndpointEmbedder::new(&cfg.understand.embed, cfg.network.proxy.as_deref()) {
        Ok(e) => match crate::vectors::get_or_embed(&e, &store_dir, &cfg.understand.embed.model, train_answers, cfg.understand.embed.batch_size).await {
            Ok(v) => v,
            Err(err) => { eprintln!("kibble: rebalance skipped (embed failed): {err}"); return Ok(None); }
        },
        Err(err) => { eprintln!("kibble: rebalance skipped (embed init failed): {err}"); return Ok(None); }
    };
    if vecs.is_empty() {
        return Ok(None);
    }
    let c = crate::cluster::cluster_from(train_answers, &vecs, cfg.cluster.k, &cfg.understand.embed.model, cfg.cluster.min_cluster_size);
    let assign = crate::cluster::assign_topics(&c.centroids, &vecs);
    let sims: Vec<f32> = assign.iter().enumerate()
        .map(|(i, &t)| crate::vectors::cosine(&vecs[i], &c.centroids[t]))
        .collect();
    let cap = ((cfg.cluster.max_topic_share * train_answers.len() as f64).floor() as usize).max(1);
    let drops = topic_cap_drops(&assign, &sims, c.centroids.len(), cap);

    // Recompute post-cap sizes + a surviving sample per topic; keep centroids/labels.
    let drop_set: std::collections::HashSet<usize> = drops.iter().copied().collect();
    let kk = c.centroids.len();
    let mut sizes = vec![0usize; kk];
    let mut first: Vec<Option<usize>> = vec![None; kk];
    for (i, &t) in assign.iter().enumerate() {
        if drop_set.contains(&i) { continue; }
        sizes[t] += 1;
        if first[t].is_none() { first[t] = Some(i); }
    }
    let samples: Vec<String> = (0..kk)
        .map(|t| first[t].map(|i| train_answers[i].chars().take(80).collect::<String>()).unwrap_or_default())
        .collect();
    let mut clusters = crate::cluster::ClusterResult {
        k: c.k, model: c.model, centroids: c.centroids, labels: c.labels, sizes, samples,
    };
    crate::cluster::apply_llm_labels(&cfg, &mut clusters).await;
    Ok(Some(RebalanceOutcome { drops, clusters }))
}

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

    #[test]
    fn drops_lowest_sim_over_cap() {
        // topic 0: rows 0,1,2 sims 0.9,0.5,0.7 ; topic 1: row 3 ; cap 2
        let assign = vec![0, 0, 0, 1];
        let sims = vec![0.9, 0.5, 0.7, 0.99];
        assert_eq!(topic_cap_drops(&assign, &sims, 2, 2), vec![1]);
    }

    #[test]
    fn multi_topic_only_over_cap_dropped() {
        // topic 0: 3 rows (over cap 2) ; topic 1: 2 rows (== cap, untouched)
        let assign = vec![0, 0, 0, 1, 1];
        let sims = vec![0.8, 0.6, 0.9, 0.7, 0.65];
        assert_eq!(topic_cap_drops(&assign, &sims, 2, 2), vec![1]);
    }

    #[test]
    fn tie_break_keeps_lower_index() {
        // topic 0: rows 0,1,2 all sim 0.5 ; cap 2 → keep indices 0,1, drop 2
        let assign = vec![0, 0, 0];
        let sims = vec![0.5, 0.5, 0.5];
        assert_eq!(topic_cap_drops(&assign, &sims, 1, 2), vec![2]);
    }

    #[test]
    fn no_drops_when_under_cap() {
        let assign = vec![0, 1, 0, 1];
        let sims = vec![0.5, 0.5, 0.5, 0.5];
        assert!(topic_cap_drops(&assign, &sims, 2, 10).is_empty());
    }

    #[test]
    fn inline_pipeline_caps_dominant_topic() {
        // Two clear groups: 5 near [1,0], 2 near [0,1]. cluster_from(k=2) recovers them;
        // capping the dominant group (5) at cap=3 drops its 2 lowest-similarity members.
        let answers: Vec<String> = (0..7).map(|i| format!("doc {i}")).collect();
        let vecs = vec![
            vec![1.0, 0.0], vec![0.98, 0.02], vec![0.95, 0.05], vec![0.9, 0.1], vec![0.85, 0.15],
            vec![0.0, 1.0], vec![0.05, 0.95],
        ];
        let c = crate::cluster::cluster_from(&answers, &vecs, 2, "m", 0);
        let assign = crate::cluster::assign_topics(&c.centroids, &vecs);
        let sims: Vec<f32> = assign.iter().enumerate()
            .map(|(i, &t)| crate::vectors::cosine(&vecs[i], &c.centroids[t])).collect();
        let cap = 3usize;
        let drops = topic_cap_drops(&assign, &sims, c.centroids.len(), cap);
        // dominant topic had 5 members → 2 dropped; small topic (2) untouched
        assert_eq!(drops.len(), 2);
        // surviving count per topic ≤ cap
        let dropset: std::collections::HashSet<usize> = drops.iter().copied().collect();
        let mut counts = vec![0usize; c.centroids.len()];
        for (i, &t) in assign.iter().enumerate() { if !dropset.contains(&i) { counts[t] += 1; } }
        assert!(counts.iter().all(|&n| n <= cap), "counts {counts:?} exceed cap {cap}");
    }

    #[tokio::test]
    async fn rebalance_inline_failsoft() {
        let dir = std::env::temp_dir().join(format!("kibble_rbi_{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&dir);
        std::fs::create_dir_all(&dir).unwrap();
        let ans = vec!["a".to_string(), "b".to_string()];

        // rebalance disabled (default) → None
        std::fs::write(dir.join(crate::config::CONFIG_FILE), "[understand.embed]\nbase_url=\"http://127.0.0.1:1/v1\"\n").unwrap();
        assert!(rebalance_inline(&dir, &ans).await.unwrap().is_none());

        // enabled but embed backend unreachable → None (fail-soft, no panic)
        std::fs::write(dir.join(crate::config::CONFIG_FILE), "[cluster]\nrebalance=true\n[understand.embed]\nbase_url=\"http://127.0.0.1:1/v1\"\nmodel=\"nomic\"\n").unwrap();
        assert!(rebalance_inline(&dir, &ans).await.unwrap().is_none());

        // enabled but empty base_url → None
        std::fs::write(dir.join(crate::config::CONFIG_FILE), "[cluster]\nrebalance=true\n[understand.embed]\nbase_url=\"\"\n").unwrap();
        assert!(rebalance_inline(&dir, &ans).await.unwrap().is_none());

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