langchain_rust/text_splitter/
markdown_splitter.rs1use async_trait::async_trait;
2use text_splitter::ChunkConfig;
3use tiktoken_rs::tokenizer::Tokenizer;
4
5use super::{SplitterOptions, TextSplitter, TextSplitterError};
6
7pub struct MarkdownSplitter {
8 splitter_options: SplitterOptions,
9}
10
11impl Default for MarkdownSplitter {
12 fn default() -> Self {
13 MarkdownSplitter::new(SplitterOptions::default())
14 }
15}
16
17impl MarkdownSplitter {
18 pub fn new(options: SplitterOptions) -> MarkdownSplitter {
19 MarkdownSplitter {
20 splitter_options: options,
21 }
22 }
23
24 #[deprecated = "Use `SplitterOptions::get_tokenizer_from_str` instead"]
25 pub fn get_tokenizer_from_str(&self, s: &str) -> Option<Tokenizer> {
26 match s.to_lowercase().as_str() {
27 "cl100k_base" => Some(Tokenizer::Cl100kBase),
28 "p50k_base" => Some(Tokenizer::P50kBase),
29 "r50k_base" => Some(Tokenizer::R50kBase),
30 "p50k_edit" => Some(Tokenizer::P50kEdit),
31 "gpt2" => Some(Tokenizer::Gpt2),
32 _ => None,
33 }
34 }
35}
36
37#[async_trait]
38impl TextSplitter for MarkdownSplitter {
39 async fn split_text(&self, text: &str) -> Result<Vec<String>, TextSplitterError> {
40 let chunk_config = ChunkConfig::try_from(&self.splitter_options)?;
41 Ok(text_splitter::MarkdownSplitter::new(chunk_config)
42 .chunks(text)
43 .map(|x| x.to_string())
44 .collect())
45 }
46}