code-splitter 0.1.5

Split code into semantic chunks using tree-sitter
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

use code_splitter::{CharCounter, Sizer, Splitter, WordCounter};
use std::fs;

const TEST_FILE: &str = "tests/testdata/markdown.md";

fn read_test_file() -> Vec<u8> {
    fs::read(TEST_FILE).expect("Failed to read test file")
}

fn split_and_show<T>(sizer: T, max_size: usize)
where
    T: Sizer,
{
    let code = read_test_file();
    let lang = tree_sitter_md::language();

    let splitter = Splitter::new(lang, sizer)
        .expect("Failed to create markdown splitter")
        .with_max_size(max_size);
    let chunks = splitter
        .split(&code)
        .expect("Failed to split markdown code");

    for (i, chunk) in chunks.iter().enumerate() {
        println!("{i} {chunk}\n");
    }
}

#[test]
fn split_by_chars() {
    split_and_show(CharCounter, 200);
}

#[cfg(feature = "tokenizers")]
#[test]
fn split_by_tokens_huggingface() {
    use tokenizers::tokenizer::Tokenizer;
    let tokenizer = Tokenizer::from_pretrained("bert-base-cased", None)
        .expect("Failed to load HuggingFace tokenizer");
    split_and_show(tokenizer, 100);
}

#[cfg(feature = "tiktoken-rs")]
#[test]
fn split_by_tokens_tiktoken() {
    use tiktoken_rs::cl100k_base;
    let bpe = cl100k_base().expect("Failed to load tiktoken tokenizer");
    split_and_show(bpe, 100);
}

#[test]
fn split_by_words() {
    split_and_show(WordCounter, 50);
}