tokenize_text/
tokenize_text.rs1use std::collections::BTreeMap;
10
11use tpt_tokenizer_core::{BpeTokenizer, Tokenizer, WordPieceTokenizer};
12
13fn main() {
14 let mut bpe_vocab = BTreeMap::new();
16 for (i, t) in ["l", "o", "w", "lo", "low", "e", "r", "er", "<unk>"]
17 .into_iter()
18 .enumerate()
19 {
20 bpe_vocab.insert(t.to_string(), i as u32);
21 }
22 let bpe_merges = vec![
23 ("l".to_string(), "o".to_string()),
24 ("lo".to_string(), "w".to_string()),
25 ("e".to_string(), "r".to_string()),
26 ];
27 let bpe = BpeTokenizer::from_vocab_merges(bpe_vocab, bpe_merges);
28
29 let text = "low lower";
30 let bpe_ids = bpe.encode(text).unwrap();
31 println!("BPE encode({text:?}) = {bpe_ids:?}");
32 println!(
33 "BPE decode({bpe_ids:?}) = {:?}",
34 bpe.decode(&bpe_ids).unwrap()
35 );
36
37 let mut wp_vocab = BTreeMap::new();
39 for (i, t) in ["[UNK]", "un", "##aff", "##able", "play", "##ing"]
40 .into_iter()
41 .enumerate()
42 {
43 wp_vocab.insert(t.to_string(), i as u32);
44 }
45 let wp = WordPieceTokenizer::from_vocab(wp_vocab, "[UNK]").unwrap();
46
47 let text2 = "unaffable playing";
48 let wp_ids = wp.encode(text2).unwrap();
49 println!("WP encode({text2:?}) = {wp_ids:?}");
50 println!(
51 "WP decode({wp_ids:?}) = {:?}",
52 wp.decode(&wp_ids).unwrap()
53 );
54}