77 lines
2.2 KiB
Rust
77 lines
2.2 KiB
Rust
//! Naive but predictable text chunking for the knowledge base documents.
|
|
//!
|
|
//! Splits on paragraph boundaries first and packs paragraphs into
|
|
//! `chunk_size`-ish chunks with a small overlap, rather than cutting mid
|
|
//! sentence — the documents here are short policy pages, not novels, so
|
|
//! this stays simple on purpose.
|
|
|
|
pub fn chunk_text(text: &str, chunk_size: usize, overlap: usize) -> Vec<String> {
|
|
let paragraphs: Vec<&str> = text
|
|
.split("\n\n")
|
|
.map(str::trim)
|
|
.filter(|p| !p.is_empty())
|
|
.collect();
|
|
|
|
if paragraphs.is_empty() {
|
|
return Vec::new();
|
|
}
|
|
|
|
let mut chunks = Vec::new();
|
|
let mut current = String::new();
|
|
|
|
for paragraph in ¶graphs {
|
|
if !current.is_empty() && current.len() + paragraph.len() + 2 > chunk_size {
|
|
chunks.push(current.clone());
|
|
current = tail(¤t, overlap);
|
|
}
|
|
if !current.is_empty() {
|
|
current.push_str("\n\n");
|
|
}
|
|
current.push_str(paragraph);
|
|
}
|
|
if !current.trim().is_empty() {
|
|
chunks.push(current);
|
|
}
|
|
|
|
chunks
|
|
}
|
|
|
|
/// Returns the last `overlap` characters of `s`, snapped to a char boundary.
|
|
fn tail(s: &str, overlap: usize) -> String {
|
|
if s.len() <= overlap {
|
|
return s.to_string();
|
|
}
|
|
let start = s.len() - overlap;
|
|
let boundary = (start..s.len())
|
|
.find(|&i| s.is_char_boundary(i))
|
|
.unwrap_or(s.len());
|
|
s[boundary..].to_string()
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn splits_long_text_into_multiple_chunks() {
|
|
let text = "Paragraph one is here.\n\nParagraph two is here.\n\nParagraph three is here.\n\nParagraph four is here.";
|
|
let chunks = chunk_text(text, 50, 10);
|
|
assert!(chunks.len() > 1, "expected multiple chunks, got {chunks:?}");
|
|
for chunk in &chunks {
|
|
assert!(!chunk.is_empty());
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn keeps_short_text_as_a_single_chunk() {
|
|
let text = "Just one short paragraph.";
|
|
let chunks = chunk_text(text, 500, 50);
|
|
assert_eq!(chunks, vec!["Just one short paragraph.".to_string()]);
|
|
}
|
|
|
|
#[test]
|
|
fn empty_input_produces_no_chunks() {
|
|
assert!(chunk_text(" \n\n ", 100, 10).is_empty());
|
|
}
|
|
}
|