text_splitter/splitter/
text.rs

1/*!
2# [`TextSplitter`]
3Semantic splitting of text documents.
4*/
5
6use std::ops::Range;
7
8use itertools::Itertools;
9use memchr::memchr2_iter;
10
11use crate::{
12    splitter::{SemanticLevel, Splitter},
13    ChunkConfig, ChunkSizer, TryChunkSizer,
14};
15
16use super::{fallback::GRAPHEME_SEGMENTER, ChunkCharIndex};
17
18/// Default plain-text splitter. Recursively splits chunks into the largest
19/// semantic units that fit within the chunk size. Also will attempt to merge
20/// neighboring chunks if they can fit within the given chunk size.
21#[derive(Debug)]
22pub struct TextSplitter<Sizer>
23where
24    Sizer: TryChunkSizer,
25{
26    /// Method of determining chunk sizes.
27    chunk_config: ChunkConfig<Sizer>,
28}
29
30impl<Sizer> TextSplitter<Sizer>
31where
32    Sizer: TryChunkSizer,
33{
34    /// Creates a new [`TextSplitter`].
35    ///
36    /// ```
37    /// use text_splitter::TextSplitter;
38    ///
39    /// // By default, the chunk sizer is based on characters.
40    /// let splitter = TextSplitter::new(512);
41    /// ```
42    #[must_use]
43    pub fn new(chunk_config: impl Into<ChunkConfig<Sizer>>) -> Self {
44        Self {
45            chunk_config: chunk_config.into(),
46        }
47    }
48}
49
50impl<Sizer: ChunkSizer> TextSplitter<Sizer> {
51    /// Generate a list of chunks from a given text. Each chunk will be up to the `chunk_capacity`.
52    ///
53    /// ## Method
54    ///
55    /// To preserve as much semantic meaning within a chunk as possible, each chunk is composed of the largest semantic units that can fit in the next given chunk. For each splitter type, there is a defined set of semantic levels. Here is an example of the steps used:
56    ///
57    /// 1. Split the text by a increasing semantic levels.
58    /// 2. Check the first item for each level and select the highest level whose first item still fits within the chunk size.
59    /// 3. Merge as many of these neighboring sections of this level or above into a chunk to maximize chunk length.
60    ///    Boundaries of higher semantic levels are always included when merging, so that the chunk doesn't inadvertantly cross semantic boundaries.
61    ///
62    /// The boundaries used to split the text if using the `chunks` method, in ascending order:
63    ///
64    /// 1. Characters
65    /// 2. [Unicode Grapheme Cluster Boundaries](https://www.unicode.org/reports/tr29/#Grapheme_Cluster_Boundaries)
66    /// 3. [Unicode Word Boundaries](https://www.unicode.org/reports/tr29/#Word_Boundaries)
67    /// 4. [Unicode Sentence Boundaries](https://www.unicode.org/reports/tr29/#Sentence_Boundaries)
68    /// 5. Ascending sequence length of newlines. (Newline is `\r\n`, `\n`, or `\r`)
69    ///    Each unique length of consecutive newline sequences is treated as its own semantic level. So a sequence of 2 newlines is a higher level than a sequence of 1 newline, and so on.
70    ///
71    /// Splitting doesn't occur below the character level, otherwise you could get partial bytes of a char, which may not be a valid unicode str.
72    ///
73    /// ```
74    /// use text_splitter::TextSplitter;
75    ///
76    /// let splitter = TextSplitter::new(10);
77    /// let text = "Some text\n\nfrom a\ndocument";
78    /// let chunks = splitter.chunks(text).collect::<Vec<_>>();
79    ///
80    /// assert_eq!(vec!["Some text", "from a", "document"], chunks);
81    /// ```
82    pub fn chunks<'splitter, 'text: 'splitter>(
83        &'splitter self,
84        text: &'text str,
85    ) -> impl Iterator<Item = &'text str> + 'splitter {
86        Splitter::<_>::try_chunks(self, text).map(super::unwrap_infallible)
87    }
88
89    /// Returns an iterator over chunks of the text and their byte offsets.
90    /// Each chunk will be up to the `chunk_capacity`.
91    ///
92    /// See [`TextSplitter::chunks`] for more information.
93    ///
94    /// ```
95    /// use text_splitter::{Characters, TextSplitter};
96    ///
97    /// let splitter = TextSplitter::new(10);
98    /// let text = "Some text\n\nfrom a\ndocument";
99    /// let chunks = splitter.chunk_indices(text).collect::<Vec<_>>();
100    ///
101    /// assert_eq!(vec![(0, "Some text"), (11, "from a"), (18, "document")], chunks);
102    /// ```
103    pub fn chunk_indices<'splitter, 'text: 'splitter>(
104        &'splitter self,
105        text: &'text str,
106    ) -> impl Iterator<Item = (usize, &'text str)> + 'splitter {
107        Splitter::<_>::try_chunk_indices(self, text).map(super::unwrap_infallible)
108    }
109
110    /// Returns an iterator over chunks of the text with their byte and character offsets.
111    /// Each chunk will be up to the `chunk_capacity`.
112    ///
113    /// See [`TextSplitter::chunks`] for more information.
114    ///
115    /// This will be more expensive than just byte offsets, and for most usage in Rust, just
116    /// having byte offsets is sufficient. But when interfacing with other languages or systems
117    /// that require character offsets, this will track the character offsets for you,
118    /// accounting for any trimming that may have occurred.
119    ///
120    /// ```
121    /// use text_splitter::{Characters, ChunkCharIndex, TextSplitter};
122    ///
123    /// let splitter = TextSplitter::new(10);
124    /// let text = "Some text\n\nfrom a\ndocument";
125    /// let chunks = splitter.chunk_char_indices(text).collect::<Vec<_>>();
126    ///
127    /// assert_eq!(vec![ChunkCharIndex { chunk: "Some text", byte_offset: 0, char_offset: 0 }, ChunkCharIndex { chunk: "from a", byte_offset: 11, char_offset: 11 }, ChunkCharIndex { chunk: "document", byte_offset: 18, char_offset: 18 }], chunks);
128    pub fn chunk_char_indices<'splitter, 'text: 'splitter>(
129        &'splitter self,
130        text: &'text str,
131    ) -> impl Iterator<Item = ChunkCharIndex<'text>> + 'splitter {
132        Splitter::<_>::try_chunk_char_indices(self, text).map(super::unwrap_infallible)
133    }
134}
135
136impl<Sizer: TryChunkSizer> TextSplitter<Sizer> {
137    /// Fallible version of [`Self::chunks`] for a [`TryChunkSizer`].
138    ///
139    /// Yields the first sizing error, then ends without calling the sizer again.
140    /// Collect into `Result<Vec<_>, _>` to stop at that error and discard partial
141    /// output. See [`TryChunkSizer`] for an example.
142    pub fn try_chunks<'splitter, 'text: 'splitter>(
143        &'splitter self,
144        text: &'text str,
145    ) -> impl Iterator<Item = Result<&'text str, Sizer::Error>> + 'splitter {
146        Splitter::<_>::try_chunks(self, text)
147    }
148
149    /// Fallible version of [`Self::chunk_indices`], including byte offsets.
150    ///
151    /// Yields the first sizing error, then ends without calling the sizer again.
152    pub fn try_chunk_indices<'splitter, 'text: 'splitter>(
153        &'splitter self,
154        text: &'text str,
155    ) -> impl Iterator<Item = Result<(usize, &'text str), Sizer::Error>> + 'splitter {
156        Splitter::<_>::try_chunk_indices(self, text)
157    }
158
159    /// Fallible version of [`Self::chunk_char_indices`], including byte and
160    /// character offsets.
161    ///
162    /// Yields the first sizing error, then ends without calling the sizer again.
163    pub fn try_chunk_char_indices<'splitter, 'text: 'splitter>(
164        &'splitter self,
165        text: &'text str,
166    ) -> impl Iterator<Item = Result<ChunkCharIndex<'text>, Sizer::Error>> + 'splitter {
167        Splitter::<_>::try_chunk_char_indices(self, text)
168    }
169}
170
171impl<Sizer> Splitter<Sizer> for TextSplitter<Sizer>
172where
173    Sizer: TryChunkSizer,
174{
175    type Level = LineBreaks;
176
177    fn chunk_config(&self) -> &ChunkConfig<Sizer> {
178        &self.chunk_config
179    }
180
181    fn parse(&self, text: &str) -> Vec<(Self::Level, Range<usize>)> {
182        memchr2_iter(b'\n', b'\r', text.as_bytes())
183            .map(|i| i..i + 1)
184            .coalesce(|a, b| {
185                if a.end == b.start {
186                    Ok(a.start..b.end)
187                } else {
188                    Err((a, b))
189                }
190            })
191            .map(|range| {
192                let level = GRAPHEME_SEGMENTER
193                    .segment_str(text.get(range.start..range.end).unwrap())
194                    .tuple_windows::<(usize, usize)>()
195                    .count();
196                (
197                    match level {
198                        0 => unreachable!("regex should always match at least one newline"),
199                        n => LineBreaks(n),
200                    },
201                    range,
202                )
203            })
204            .collect()
205    }
206}
207
208/// Different semantic levels that text can be split by.
209/// Each level provides a method of splitting text into chunks of a given level
210/// as well as a fallback in case a given fallback is too large.
211///
212/// Split by given number of linebreaks, either `\n`, `\r`, or `\r\n`.
213#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd)]
214pub struct LineBreaks(usize);
215
216impl SemanticLevel for LineBreaks {}
217
218#[cfg(test)]
219mod tests {
220    use std::cmp::min;
221
222    use fake::{Fake, Faker};
223
224    use crate::{splitter::SemanticSplitRanges, ChunkCharIndex};
225
226    use super::*;
227
228    #[test]
229    fn returns_one_chunk_if_text_is_shorter_than_max_chunk_size() {
230        let text = Faker.fake::<String>();
231        let chunks = TextSplitter::new(ChunkConfig::new(text.chars().count()).with_trim(false))
232            .chunks(&text)
233            .collect::<Vec<_>>();
234
235        assert_eq!(vec![&text], chunks);
236    }
237
238    #[test]
239    fn returns_two_chunks_if_text_is_longer_than_max_chunk_size() {
240        let text1 = Faker.fake::<String>();
241        let text2 = Faker.fake::<String>();
242        let text = format!("{text1}{text2}");
243        // Round up to one above half so it goes to 2 chunks
244        let max_chunk_size = text.chars().count() / 2 + 1;
245        let chunks = TextSplitter::new(ChunkConfig::new(max_chunk_size).with_trim(false))
246            .chunks(&text)
247            .collect::<Vec<_>>();
248
249        assert!(chunks.iter().all(|c| c.chars().count() <= max_chunk_size));
250
251        // Check that beginning of first chunk and text 1 matches
252        let len = min(text1.len(), chunks[0].len());
253        assert_eq!(text1[..len], chunks[0][..len]);
254        // Check that end of second chunk and text 2 matches
255        let len = min(text2.len(), chunks[1].len());
256        assert_eq!(
257            text2[(text2.len() - len)..],
258            chunks[1][chunks[1].len() - len..]
259        );
260
261        assert_eq!(chunks.join(""), text);
262    }
263
264    #[test]
265    fn empty_string() {
266        let text = "";
267        let chunks = TextSplitter::new(ChunkConfig::new(100).with_trim(false))
268            .chunks(text)
269            .collect::<Vec<_>>();
270
271        assert!(chunks.is_empty());
272    }
273
274    #[test]
275    fn can_handle_unicode_characters() {
276        let text = "éé"; // Char that is more than one byte
277        let chunks = TextSplitter::new(ChunkConfig::new(1).with_trim(false))
278            .chunks(text)
279            .collect::<Vec<_>>();
280        assert_eq!(vec!["é", "é"], chunks);
281    }
282
283    // Just for testing
284    struct Str;
285
286    impl ChunkSizer for Str {
287        fn size(&self, chunk: &str) -> usize {
288            chunk.len()
289        }
290    }
291
292    #[test]
293    fn custom_len_function() {
294        let text = "éé"; // Char that is two bytes each
295        let chunks = TextSplitter::new(ChunkConfig::new(2).with_sizer(Str).with_trim(false))
296            .chunks(text)
297            .collect::<Vec<_>>();
298
299        assert_eq!(vec!["é", "é"], chunks);
300    }
301
302    #[test]
303    fn handles_char_bigger_than_len() {
304        let text = "éé"; // Char that is two bytes each
305        let chunks = TextSplitter::new(ChunkConfig::new(1).with_sizer(Str).with_trim(false))
306            .chunks(text)
307            .collect::<Vec<_>>();
308
309        // We can only go so small
310        assert_eq!(vec!["é", "é"], chunks);
311    }
312
313    #[test]
314    fn chunk_by_graphemes() {
315        let text = "a̐éö̲\r\n";
316        let chunks = TextSplitter::new(ChunkConfig::new(3).with_trim(false))
317            .chunks(text)
318            .collect::<Vec<_>>();
319
320        // \r\n is grouped together not separated
321        assert_eq!(vec!["a̐é", "ö̲", "\r\n"], chunks);
322    }
323
324    #[test]
325    fn trim_char_indices() {
326        let text = " a b ";
327        let chunks = TextSplitter::new(1).chunk_indices(text).collect::<Vec<_>>();
328
329        assert_eq!(vec![(1, "a"), (3, "b")], chunks);
330    }
331
332    #[test]
333    fn chunk_char_indices() {
334        let text = " a b ";
335        let chunks = TextSplitter::new(1)
336            .chunk_char_indices(text)
337            .collect::<Vec<_>>();
338
339        assert_eq!(
340            vec![
341                ChunkCharIndex {
342                    chunk: "a",
343                    byte_offset: 1,
344                    char_offset: 1
345                },
346                ChunkCharIndex {
347                    chunk: "b",
348                    byte_offset: 3,
349                    char_offset: 3,
350                },
351            ],
352            chunks
353        );
354    }
355
356    #[test]
357    fn graphemes_fallback_to_chars() {
358        let text = "a̐éö̲\r\n";
359        let chunks = TextSplitter::new(ChunkConfig::new(1).with_trim(false))
360            .chunks(text)
361            .collect::<Vec<_>>();
362        assert_eq!(
363            vec!["a", "\u{310}", "é", "ö", "\u{332}", "\r", "\n"],
364            chunks
365        );
366    }
367
368    #[test]
369    fn trim_grapheme_indices() {
370        let text = "\r\na̐éö̲\r\n";
371        let chunks = TextSplitter::new(3).chunk_indices(text).collect::<Vec<_>>();
372
373        assert_eq!(vec![(2, "a̐é"), (7, "ö̲")], chunks);
374    }
375
376    #[test]
377    fn grapheme_char_indices() {
378        let text = "\r\na̐éö̲\r\n";
379        let chunks = TextSplitter::new(3)
380            .chunk_char_indices(text)
381            .collect::<Vec<_>>();
382
383        assert_eq!(
384            vec![
385                ChunkCharIndex {
386                    chunk: "a̐é",
387                    byte_offset: 2,
388                    char_offset: 2
389                },
390                ChunkCharIndex {
391                    chunk: "ö̲",
392                    byte_offset: 7,
393                    char_offset: 5
394                }
395            ],
396            chunks
397        );
398    }
399
400    #[test]
401    fn chunk_by_words() {
402        let text = "The quick (\"brown\") fox can't jump 32.3 feet, right?";
403        let chunks = TextSplitter::new(ChunkConfig::new(10).with_trim(false))
404            .chunks(text)
405            .collect::<Vec<_>>();
406
407        assert_eq!(
408            vec![
409                "The quick ",
410                "(\"brown\") ",
411                "fox can't ",
412                "jump 32.3 ",
413                "feet, ",
414                "right?"
415            ],
416            chunks
417        );
418    }
419
420    #[test]
421    fn words_fallback_to_graphemes() {
422        let text = "Thé quick\r\n";
423        let chunks = TextSplitter::new(ChunkConfig::new(2).with_trim(false))
424            .chunks(text)
425            .collect::<Vec<_>>();
426        assert_eq!(vec!["Th", "é ", "qu", "ic", "k", "\r\n"], chunks);
427    }
428
429    #[test]
430    fn trim_word_indices() {
431        let text = "Some text from a document";
432        let chunks = TextSplitter::new(10)
433            .chunk_indices(text)
434            .collect::<Vec<_>>();
435        assert_eq!(
436            vec![(0, "Some text"), (10, "from a"), (17, "document")],
437            chunks
438        );
439    }
440
441    #[test]
442    fn chunk_by_sentences() {
443        let text = "Mr. Fox jumped. [...] The dog was too lazy.";
444        let chunks = TextSplitter::new(ChunkConfig::new(21).with_trim(false))
445            .chunks(text)
446            .collect::<Vec<_>>();
447        assert_eq!(
448            vec!["Mr. Fox jumped. ", "[...] ", "The dog was too lazy."],
449            chunks
450        );
451    }
452
453    #[test]
454    fn sentences_falls_back_to_words() {
455        let text = "Mr. Fox jumped. [...] The dog was too lazy.";
456        let chunks = TextSplitter::new(ChunkConfig::new(16).with_trim(false))
457            .chunks(text)
458            .collect::<Vec<_>>();
459        assert_eq!(
460            vec!["Mr. Fox jumped. ", "[...] ", "The dog was too ", "lazy."],
461            chunks
462        );
463    }
464
465    #[test]
466    fn trim_sentence_indices() {
467        let text = "Some text. From a document.";
468        let chunks = TextSplitter::new(10)
469            .chunk_indices(text)
470            .collect::<Vec<_>>();
471        assert_eq!(
472            vec![(0, "Some text."), (11, "From a"), (18, "document.")],
473            chunks
474        );
475    }
476
477    #[test]
478    fn trim_paragraph_indices() {
479        let text = "Some text\n\nfrom a\ndocument";
480        let chunks = TextSplitter::new(10)
481            .chunk_indices(text)
482            .collect::<Vec<_>>();
483        assert_eq!(
484            vec![(0, "Some text"), (11, "from a"), (18, "document")],
485            chunks
486        );
487    }
488
489    #[test]
490    fn correctly_determines_newlines() {
491        let text = "\r\n\r\ntext\n\n\ntext2";
492        let splitter = TextSplitter::new(10);
493        let linebreaks = SemanticSplitRanges::new(splitter.parse(text));
494        assert_eq!(
495            vec![(LineBreaks(2), 0..4), (LineBreaks(3), 8..11)],
496            linebreaks.ranges
497        );
498    }
499}
text_splitter/splitter/text.rs

text_splitter/splitter/
text.rs
