cognis 0.2.0

LLM application framework built on cognis-core
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

use super::{merge_splits, TextSplitter};

/// Splits text on a single character separator.
pub struct CharacterTextSplitter {
    pub separator: String,
    pub chunk_size: usize,
    pub chunk_overlap: usize,
    pub keep_separator: bool,
}

impl Default for CharacterTextSplitter {
    fn default() -> Self {
        Self {
            separator: "\n\n".to_string(),
            chunk_size: 4000,
            chunk_overlap: 200,
            keep_separator: false,
        }
    }
}

impl CharacterTextSplitter {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn with_separator(mut self, sep: impl Into<String>) -> Self {
        self.separator = sep.into();
        self
    }

    pub fn with_chunk_size(mut self, size: usize) -> Self {
        self.chunk_size = size;
        self
    }

    pub fn with_chunk_overlap(mut self, overlap: usize) -> Self {
        self.chunk_overlap = overlap;
        self
    }
}

impl TextSplitter for CharacterTextSplitter {
    fn split_text(&self, text: &str) -> Vec<String> {
        let splits: Vec<&str> = if self.separator.is_empty() {
            text.chars().map(|_| "").collect() // fallback
        } else {
            text.split(&self.separator).collect()
        };

        let good_splits: Vec<&str> = splits
            .iter()
            .map(|s| s.trim())
            .filter(|s| !s.is_empty())
            .collect();

        merge_splits(
            &good_splits,
            &self.separator,
            self.chunk_size,
            self.chunk_overlap,
        )
    }

    fn chunk_size(&self) -> usize {
        self.chunk_size
    }

    fn chunk_overlap(&self) -> usize {
        self.chunk_overlap
    }
}