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
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198

//! `kibble` auto-topics (Spec B) — assign each catalog document a hierarchical topic from the
//! train answer-cluster space. See docs/CATALOG.md. Opt-in via `[classify].enabled`.

use crate::catalog::CatalogEntry;
use crate::cluster::ClusterResult;
use crate::config::KibbleConfig;
use std::collections::HashMap;
use std::path::Path;

#[derive(Debug, serde::Serialize, PartialEq)]
pub struct TopicSummary {
    pub path: String,
    pub size: usize,
    pub samples: Vec<String>,
}

pub fn assign_doc_topics(
    assign: &[usize],
    doc_ids: &[String],
    labels: &[String],
    n_clusters: usize,
    entries: &mut [CatalogEntry],
) -> Vec<TopicSummary> {
    // Per-document cluster counts.
    let mut per_doc: HashMap<&str, HashMap<usize, usize>> = HashMap::new();
    for (i, &c) in assign.iter().enumerate() {
        *per_doc.entry(doc_ids[i].as_str()).or_default().entry(c).or_default() += 1;
    }
    // Majority cluster per document (ties → lowest cluster index; deterministic regardless of map order).
    let mut doc_topic: HashMap<&str, (usize, f32)> = HashMap::new();
    for (doc, counts) in &per_doc {
        let total: usize = counts.values().sum();
        let mut best_c = usize::MAX;
        let mut best_n = 0usize;
        for (&c, &n) in counts {
            if n > best_n || (n == best_n && c < best_c) {
                best_n = n;
                best_c = c;
            }
        }
        doc_topic.insert(doc, (best_c, best_n as f32 / total as f32));
    }
    // Write onto entries; accumulate the topics.json summary.
    let mut sizes = vec![0usize; n_clusters];
    let mut samples: Vec<Vec<String>> = vec![Vec::new(); n_clusters];
    for e in entries.iter_mut() {
        if let Some(&(c, conf)) = doc_topic.get(e.doc_id.as_str()) {
            e.auto_topic = Some(labels.get(c).cloned().unwrap_or_default());
            e.topic_confidence = Some(conf);
            if c < n_clusters {
                sizes[c] += 1;
                if samples[c].len() < 3 {
                    samples[c].push(e.title.clone());
                }
            }
        }
    }
    (0..n_clusters)
        .map(|c| TopicSummary {
            path: labels.get(c).cloned().unwrap_or_default(),
            size: sizes[c],
            samples: std::mem::take(&mut samples[c]),
        })
        .collect()
}

/// Opt-in pass: assign each catalog document a hierarchical auto-topic from the train
/// answer-cluster space, and write `<cat_dir>/topics.json`. Fail-soft: no embed backend /
/// embed error / empty centroids → warn and return Ok, leaving `entries` untouched.
pub async fn apply_auto_topics(
    cfg: &KibbleConfig,
    repo_root: &Path,
    cat_dir: &Path,
    curated: &crate::build::CuratedSplits,
    inline_clusters: Option<&ClusterResult>,
    entries: &mut [CatalogEntry],
) -> std::io::Result<()> {
    let embed = &cfg.understand.embed;
    if embed.base_url.is_empty() {
        eprintln!("kibble: auto-topics skipped (no embed backend — set [understand.embed].base_url)");
        return Ok(());
    }
    // 1. Collect (answer, doc_id) for every row across splits; record train indices.
    let mut answers: Vec<String> = Vec::new();
    let mut doc_ids: Vec<String> = Vec::new();
    let mut train_idx: Vec<usize> = Vec::new();
    for (si, split) in [&curated.train, &curated.valid, &curated.test].iter().enumerate() {
        for (r, id) in split.clean.iter().zip(split.clean_doc_ids.iter()) {
            if si == 0 { train_idx.push(answers.len()); }
            answers.push(crate::build::clean_text(crate::build::last_content(r, "assistant")));
            doc_ids.push(id.clone());
        }
        for (r, id) in split.raw.iter().zip(split.raw_doc_ids.iter()) {
            if si == 0 { train_idx.push(answers.len()); }
            answers.push(crate::build::last_content(r, "assistant").to_string());
            doc_ids.push(id.clone());
        }
    }
    if answers.is_empty() {
        eprintln!("kibble: auto-topics skipped (no rows)");
        return Ok(());
    }
    // 2. Embed once (cached).
    let store = repo_root.join(&embed.store);
    let vecs = match crate::embed::EndpointEmbedder::new(embed, cfg.network.proxy.as_deref()) {
        Ok(e) => match crate::vectors::get_or_embed(&e, &store, &embed.model, &answers, embed.batch_size).await {
            Ok(v) => v,
            Err(err) => { eprintln!("kibble: auto-topics skipped (embed failed): {err}"); return Ok(()); }
        },
        Err(err) => { eprintln!("kibble: auto-topics skipped (embed init failed): {err}"); return Ok(()); }
    };
    // 3. Centroids: reuse rebalance's in-memory clustering, else cluster fresh from train.
    let cr: ClusterResult = if let Some(c) = inline_clusters {
        c.clone()
    } else {
        let train_answers: Vec<String> = train_idx.iter().map(|&i| answers[i].clone()).collect();
        let train_vecs: Vec<Vec<f32>> = train_idx.iter().map(|&i| vecs[i].clone()).collect();
        if train_answers.is_empty() {
            eprintln!("kibble: auto-topics skipped (no train rows)");
            return Ok(());
        }
        crate::cluster::cluster_from(&train_answers, &train_vecs, cfg.cluster.k, &embed.model, cfg.cluster.min_cluster_size)
    };
    if cr.centroids.is_empty() {
        eprintln!("kibble: auto-topics skipped (no centroids)");
        return Ok(());
    }
    // 4. Hierarchical LLM names — required. Without a working [ask] LLM (or on an
    //    unparseable/wrong-count reply) we skip rather than assign flat term labels as
    //    `auto_topic` (the field's contract is a hierarchical path, or absent).
    let labels = match crate::cluster::hier_labels(cfg, &cr.labels, &cr.samples).await {
        Some(names) => names,
        None => {
            eprintln!("kibble: auto-topics skipped (no [ask] LLM or unparseable hierarchical names)");
            return Ok(());
        }
    };
    // 5. Assign every row to the reused centroids.
    let assign = crate::cluster::assign_topics(&cr.centroids, &vecs);
    // 6-7. Majority per doc → set entries; write topics.json.
    let topics = assign_doc_topics(&assign, &doc_ids, &labels, cr.centroids.len(), entries);
    std::fs::write(cat_dir.join("topics.json"), serde_json::to_string_pretty(&topics)?)?;
    Ok(())
}

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

    fn entry(doc_id: &str, title: &str) -> CatalogEntry {
        CatalogEntry {
            doc_id: doc_id.into(), source: "s".into(), title: title.into(), role: "style".into(),
            topics: vec![], lora_bucket: "style".into(), rag: false, is_reference: false, chars: 1,
            auto_topic: None, topic_confidence: None,
        }
    }

    #[test]
    fn majority_and_confidence() {
        // doc A: rows in clusters [0,0,1] -> majority 0, conf 2/3. doc B: [1] -> 1, conf 1.
        let assign = vec![0, 0, 1, 1];
        let doc_ids = vec!["A".into(), "A".into(), "A".into(), "B".into()];
        let labels = vec!["T0".into(), "T1".into()];
        let mut es = vec![entry("A", "Atitle"), entry("B", "Btitle")];
        let topics = assign_doc_topics(&assign, &doc_ids, &labels, 2, &mut es);
        assert_eq!(es[0].auto_topic.as_deref(), Some("T0"));
        assert_eq!(es[0].topic_confidence, Some(2.0 / 3.0));
        assert_eq!(es[1].auto_topic.as_deref(), Some("T1"));
        assert_eq!(es[1].topic_confidence, Some(1.0));
        // topics.json summary: T0 has 1 doc (A), T1 has 1 doc (B)
        assert_eq!(topics[0], TopicSummary { path: "T0".into(), size: 1, samples: vec!["Atitle".into()] });
        assert_eq!(topics[1].size, 1);
    }

    #[test]
    fn tie_breaks_to_lowest_cluster_index() {
        // doc A: one row cluster 1, one row cluster 0 -> tie -> lowest index 0.
        let assign = vec![1, 0];
        let doc_ids = vec!["A".into(), "A".into()];
        let labels = vec!["T0".into(), "T1".into()];
        let mut es = vec![entry("A", "t")];
        assign_doc_topics(&assign, &doc_ids, &labels, 2, &mut es);
        assert_eq!(es[0].auto_topic.as_deref(), Some("T0"));
        assert_eq!(es[0].topic_confidence, Some(0.5));
    }

    #[test]
    fn zero_row_document_stays_none() {
        // entry "Z" has no rows in assign -> stays None.
        let assign = vec![0];
        let doc_ids = vec!["A".into()];
        let labels = vec!["T0".into()];
        let mut es = vec![entry("A", "a"), entry("Z", "z")];
        assign_doc_topics(&assign, &doc_ids, &labels, 1, &mut es);
        assert!(es[1].auto_topic.is_none() && es[1].topic_confidence.is_none());
    }
}