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

use std::fs;
use std::path::Path;

#[derive(Debug, Clone)]
pub struct Document {
    pub doc_id: String,
    pub text: String,
    pub source: String,
    pub title: String,
    pub is_reference: bool,
}

/// Map a raw-doc subdir name to (canonical source, is_reference).
fn source_for_dir(dir_name: &str) -> (String, bool) {
    match dir_name {
        "textfiles" => ("textfile".to_string(), true),
        other => (other.to_string(), false),
    }
}

pub fn load_raw_docs(raw_root: &Path) -> std::io::Result<Vec<Document>> {
    let mut docs = Vec::new();
    if !raw_root.is_dir() {
        return Ok(docs);
    }
    for sub in fs::read_dir(raw_root)? {
        let sub = sub?.path();
        if !sub.is_dir() {
            continue;
        }
        let dir_name = sub.file_name().and_then(|n| n.to_str()).unwrap_or("").to_string();
        let (source, is_reference) = source_for_dir(&dir_name);
        for entry in fs::read_dir(&sub)? {
            let path = entry?.path();
            if path.extension().and_then(|e| e.to_str()) != Some("txt") {
                continue;
            }
            let stem = path.file_stem().and_then(|s| s.to_str()).unwrap_or("").to_string();
            let text = match fs::read_to_string(&path) {
                Ok(t) => t,
                Err(e) => {
                    eprintln!("kibble: skipping unreadable {}: {e}", path.display());
                    continue;
                }
            };
            if text.trim().is_empty() {
                continue;
            }
            docs.push(Document {
                doc_id: format!("{source}:{stem}"),
                text,
                source: source.clone(),
                title: stem,
                is_reference,
            });
        }
    }
    Ok(docs)
}

#[cfg(test)]
mod tests {
    use super::load_raw_docs;
    use std::fs;

    #[test]
    fn loads_docs_with_source_and_reference_flags() {
        let dir = std::env::temp_dir().join(format!("kibble_corpus_test_{}", std::process::id()));
        let _ = fs::remove_dir_all(&dir);
        let raw = dir.join("raw");
        fs::create_dir_all(raw.join("twitter")).unwrap();
        fs::create_dir_all(raw.join("textfiles")).unwrap();
        fs::write(raw.join("twitter").join("tweet_111.txt"), "hello world").unwrap();
        fs::write(raw.join("twitter").join("manifest.json"), "{}").unwrap();
        fs::write(raw.join("textfiles").join("phrack_1.txt"), "old school text").unwrap();

        let mut docs = load_raw_docs(&raw).unwrap();
        docs.sort_by(|a, b| a.doc_id.cmp(&b.doc_id));
        assert_eq!(docs.len(), 2); // manifest.json excluded
        let tw = docs.iter().find(|d| d.source == "twitter").unwrap();
        assert_eq!(tw.doc_id, "twitter:tweet_111");
        assert!(!tw.is_reference);
        let tf = docs.iter().find(|d| d.source == "textfile").unwrap();
        assert!(tf.is_reference);
        assert_eq!(tf.title, "phrack_1");
    }
}