kibble 0.1.0

chew through any source into clean datasets — a fast ingestion, RAG & fine-tuning toolkit
Documentation
use crate::clean::clean_text;
use md5::{Digest as Md5Digest, Md5};
use sha2::Sha256;

pub fn topic_from_title(title: &str) -> String {
    // Use the stem; replace separators; collapse whitespace; truncate to 120.
    let stem = std::path::Path::new(title)
        .file_stem()
        .and_then(|s| s.to_str())
        .unwrap_or(title);
    let replaced = stem.replace(['_', '-'], " ");
    let collapsed = replaced.split_whitespace().collect::<Vec<_>>().join(" ");
    let base = if collapsed.chars().count() > 120 {
        let truncated: String = collapsed.chars().take(117).collect();
        format!("{truncated}...")
    } else {
        collapsed
    };
    if base.is_empty() {
        "this topic".to_string()
    } else {
        base
    }
}

fn chunk_chars(chars: Vec<char>, max_chars: usize) -> Vec<String> {
    if chars.len() <= max_chars {
        let t: String = chars.iter().collect();
        let t = t.trim().to_string();
        return if t.is_empty() { vec![] } else { vec![t] };
    }
    let overlap = max_chars / 8;
    let mut chunks = Vec::new();
    let mut start = 0usize;
    while start < chars.len() {
        let mut end = (start + max_chars).min(chars.len());
        if end < chars.len() {
            // Prefer a paragraph break after the halfway point.
            let half = start + max_chars / 2;
            let window: String = chars[start..end].iter().collect();
            if let Some(rel) = window.rfind("\n\n") {
                let break_at = start + window[..rel].chars().count();
                if break_at > half {
                    end = break_at;
                }
            }
        }
        let piece: String = chars[start..end].iter().collect();
        let piece = piece.trim().to_string();
        if !piece.is_empty() {
            chunks.push(piece);
        }
        if end >= chars.len() {
            break;
        }
        start = (end.saturating_sub(overlap)).max(start + 1);
    }
    chunks
}

pub fn chunk_text(text: &str, max_chars: usize) -> Vec<String> {
    let text = clean_text(text);
    chunk_chars(text.chars().collect(), max_chars)
}

pub fn chunk_text_raw(text: &str, max_chars: usize) -> Vec<String> {
    chunk_chars(text.chars().collect(), max_chars)
}

pub fn split_bucket(doc_id: &str) -> u8 {
    let digest = Sha256::digest(doc_id.as_bytes());
    // Replicate Python's int(hexdigest,16) % 100 by reducing mod 100 byte-by-byte.
    let mut acc: u32 = 0;
    for b in digest.iter() {
        acc = (acc * 256 + *b as u32) % 100;
    }
    acc as u8
}

pub fn pick_template<'a>(doc_id: &str, chunk_idx: usize, templates: &'a [&'a str]) -> &'a str {
    let key = format!("{doc_id}:{chunk_idx}");
    let digest = Md5::digest(key.as_bytes());
    let n = u128::from_be_bytes(digest.into());
    templates[(n % templates.len() as u128) as usize]
}

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

    #[test]
    fn topic_from_title_normalizes() {
        assert_eq!(topic_from_title("my_cool-post"), "my cool post");
        assert_eq!(topic_from_title(""), "this topic");
    }

    #[test]
    fn short_text_is_single_chunk() {
        assert_eq!(chunk_text("hello world", 3500), vec!["hello world".to_string()]);
        assert!(chunk_text("   ", 3500).is_empty());
    }

    #[test]
    fn long_text_splits_with_overlap() {
        let para = "x".repeat(2000);
        let text = format!("{para}\n\n{para}"); // ~4002 chars, one paragraph break
        let chunks = chunk_text(&text, 3500);
        assert!(chunks.len() >= 2);
    }

    #[test]
    fn split_bucket_is_deterministic_and_in_range() {
        let b = split_bucket("twitter:111");
        assert_eq!(b, split_bucket("twitter:111"));
        assert!(b < 100);
    }

    #[test]
    fn split_bucket_matches_python_full_sha256_mod_100() {
        // Expected values precomputed in Python:
        //   python3 -c "import hashlib; print(int(hashlib.sha256(b'twitter:111').hexdigest(),16)%100)"
        // twitter:111 -> 46 ; blog:hello -> 56 ; textfile:phrack_1 -> 66
        assert_eq!(split_bucket("twitter:111"), 46);
        assert_eq!(split_bucket("blog:hello"), 56);
        assert_eq!(split_bucket("textfile:phrack_1"), 66);
    }

    #[test]
    fn pick_template_is_deterministic() {
        let templates = ["a {topic}", "b {topic}", "c {topic}"];
        let first = pick_template("doc:1", 0, &templates);
        assert_eq!(first, pick_template("doc:1", 0, &templates));
    }
}