use std::collections::HashMap;
use anyhow::{Context as _, Result, anyhow, ensure};
use unicode_categories::UnicodeCategories as _;
use unicode_normalization::{UnicodeNormalization as _, char::is_combining_mark};
use super::gguf::Gguf;
const MAX_WORD_CHARACTERS: usize = 100;
pub(super) struct WordPieceTokenizer {
tokens: HashMap<String, u32>,
unknown: u32,
cls: u32,
separator: u32,
max_piece_bytes: usize,
}
impl WordPieceTokenizer {
pub(super) fn from_bert(gguf: &Gguf) -> Result<Self> {
ensure!(
gguf.string("tokenizer.ggml.model")? == "bert",
"expected BERT tokenizer model"
);
let vocabulary = gguf.strings("tokenizer.ggml.tokens")?;
let mut tokens = HashMap::with_capacity(vocabulary.len());
let mut max_piece_bytes = 0;
for (index, text) in vocabulary.iter().enumerate() {
let token = u32::try_from(index).context("token index exceeds u32")?;
ensure!(
tokens.insert(text.clone(), token).is_none(),
"duplicate tokenizer token {text:?}"
);
max_piece_bytes = max_piece_bytes.max(text.len());
}
let unknown = special_token(&tokens, "[UNK]")?;
let cls = special_token(&tokens, "[CLS]")?;
let separator = special_token(&tokens, "[SEP]")?;
ensure!(
gguf.u32("tokenizer.ggml.unknown_token_id")? == unknown,
"unknown token ID differs"
);
ensure!(
gguf.u32("tokenizer.ggml.seperator_token_id")? == separator,
"separator token ID differs"
);
Ok(Self {
tokens,
unknown,
cls,
separator,
max_piece_bytes,
})
}
pub(super) fn encode(&self, text: &str, max_tokens: usize) -> Result<Vec<u32>> {
ensure!(max_tokens >= 2, "BERT context must fit CLS and SEP tokens");
let mut output = Vec::with_capacity(max_tokens.min(64));
output.push(self.cls);
'words: for word in normalized_words(text) {
for token in self.encode_word(&word) {
if output.len() + 1 >= max_tokens {
break 'words;
}
output.push(token);
}
}
output.push(self.separator);
Ok(output)
}
fn encode_word(&self, word: &str) -> Vec<u32> {
if word.chars().count() > MAX_WORD_CHARACTERS {
return vec![self.unknown];
}
let text = format!("▁{word}");
let boundaries = text
.char_indices()
.map(|(offset, _)| offset)
.chain(std::iter::once(text.len()))
.collect::<Vec<_>>();
let mut output = Vec::new();
let mut start = 0;
while start < text.len() {
let maximum = start.saturating_add(self.max_piece_bytes).min(text.len());
let candidate_count = boundaries.partition_point(|end| *end <= maximum);
let matched = boundaries[..candidate_count]
.iter()
.copied()
.rev()
.take_while(|end| *end > start)
.find_map(|end| {
self.tokens
.get(&text[start..end])
.copied()
.map(|token| (end, token))
});
let Some((end, token)) = matched else {
return vec![self.unknown];
};
output.push(token);
start = end;
}
output
}
}
fn special_token(tokens: &HashMap<String, u32>, text: &str) -> Result<u32> {
tokens
.get(text)
.copied()
.ok_or_else(|| anyhow!("missing tokenizer token {text:?}"))
}
fn normalized_words(text: &str) -> Vec<String> {
let mut words = Vec::new();
let mut word = String::new();
for character in text.nfd() {
if is_combining_mark(character) || character == '\0' || character == '\u{fffd}' {
continue;
}
for character in character.to_lowercase() {
if character.is_whitespace() || character.is_separator() {
push_word(&mut words, &mut word);
} else if !(character.is_other_control() || character.is_other_format()) {
if separates_word(character) {
push_word(&mut words, &mut word);
words.push(character.to_string());
} else {
word.push(character);
}
}
}
}
push_word(&mut words, &mut word);
words
}
fn push_word(words: &mut Vec<String>, word: &mut String) {
if !word.is_empty() {
words.push(std::mem::take(word));
}
}
fn separates_word(character: char) -> bool {
character.is_ascii_punctuation() || character.is_punctuation() || is_cjk(character)
}
fn is_cjk(character: char) -> bool {
matches!(
u32::from(character),
0x3400..=0x4dbf
| 0x4e00..=0x9fff
| 0xf900..=0xfaff
| 0x20000..=0x2a6df
| 0x2a700..=0x2b73f
| 0x2b740..=0x2b81f
| 0x2b920..=0x2ceaf
| 0x2f800..=0x2fa1f
)
}