Skip to main content

tokenize_text/
tokenize_text.rs

1//! Tokenizes a small built-in sample with both BPE and WordPiece.
2//!
3//! Run with:
4//!
5//! ```sh
6//! cargo run -p tpt-tokenizer-core --example tokenize_text
7//! ```
8
9use std::collections::BTreeMap;
10
11use tpt_tokenizer_core::{BpeTokenizer, Tokenizer, WordPieceTokenizer};
12
13fn main() {
14    // --- BPE (GPT-2 style) ---
15    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    // --- WordPiece (BERT style) ---
38    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}