From 97fc7de6754c8b5c383d7d5e3f9e26b3bbfb08fb Mon Sep 17 00:00:00 2001 From: Anthony Rubick <68485672+AnthonyMichaelTDM@users.noreply.github.com> Date: Sun, 9 Jul 2023 20:07:09 -0700 Subject: fix: minor code cleanup and pedantic lints (#134) --- src/utils/tiktoken.rs | 63 +++++++++++++++++++++++++-------------------------- 1 file changed, 31 insertions(+), 32 deletions(-) (limited to 'src/utils/tiktoken.rs') diff --git a/src/utils/tiktoken.rs b/src/utils/tiktoken.rs index eb8819d..940d6f9 100644 --- a/src/utils/tiktoken.rs +++ b/src/utils/tiktoken.rs @@ -1,6 +1,6 @@ //! Use tiktoken for count tokens //! -//! Copy from https://github.com/dust-tt/dust/tree/main/core/src/providers/tiktoken +//! Copy from [https://github.com/dust-tt/dust/tree/main/core/src/providers/tiktoken](https://github.com/dust-tt/dust/tree/main/core/src/providers/tiktoken) #![allow(unused)] @@ -26,7 +26,7 @@ pub fn text_to_tokens(text: &str) -> Vec { } /// Convert tokens to plan text -pub fn tokens_to_text(tokens: Vec) -> Result { +pub fn tokens_to_text(tokens: &[usize]) -> Result { cl100k_base_singleton().lock().decode(tokens) } @@ -43,11 +43,11 @@ pub fn cl100k_base() -> Result { } let mut special_tokens = HashMap::default(); - special_tokens.insert(String::from("<|endoftext|>"), 100257); - special_tokens.insert(String::from("<|fim_prefix|>"), 100258); - special_tokens.insert(String::from("<|fim_middle|>"), 100259); - special_tokens.insert(String::from("<|fim_suffix|>"), 100260); - special_tokens.insert(String::from("<|endofprompt|>"), 100276); + special_tokens.insert(String::from("<|endoftext|>"), 100_257); + special_tokens.insert(String::from("<|fim_prefix|>"), 100_258); + special_tokens.insert(String::from("<|fim_middle|>"), 100_259); + special_tokens.insert(String::from("<|fim_suffix|>"), 100_260); + special_tokens.insert(String::from("<|endofprompt|>"), 100_276); CoreBPE::new( encoder, @@ -64,6 +64,7 @@ pub fn cl100k_base_singleton() -> Arc> { } fn _byte_pair_merge(piece: &[u8], ranks: &HashMap, usize>) -> Vec> { + #[allow(clippy::range_plus_one)] let mut parts: Vec<_> = (0..piece.len()).map(|i| i..i + 1).collect(); // If you have n parts and m merges, this does O(mn) work @@ -156,11 +157,11 @@ pub struct CoreBPE { } impl CoreBPE { - fn _get_regex(&self) -> &Regex { + const fn _get_regex(&self) -> &Regex { &self.regex } - fn _get_special_regex(&self) -> &Regex { + const fn _get_special_regex(&self) -> &Regex { &self.special_regex } @@ -262,15 +263,12 @@ impl CoreBPE { // Here is a quick and dirty fix: { let token_is_all_space = |token| { - self.decoder - .get(token) - .map(|token_bytes| { - token_bytes - .iter() - .rev() - .all(|&b| [b' ', b'\n', b'\t'].contains(&b)) - }) - .unwrap_or(false) + self.decoder.get(token).map_or(false, |token_bytes| { + token_bytes + .iter() + .rev() + .all(|&b| [b' ', b'\n', b'\t'].contains(&b)) + }) }; if last_piece_token_len > 0 && token_is_all_space(&tokens[tokens.len() - last_piece_token_len]) @@ -342,6 +340,7 @@ impl CoreBPE { && self.sorted_token_bytes[point].starts_with(suffix) { let possibility = [prefix, self.sorted_token_bytes[point].as_slice()].concat(); + #[allow(clippy::option_if_let_else)] let encoded = match std::str::from_utf8(&possibility) { // Morally, this is byte_pair_encode(&possibility, &self.encoder) // But we might have introduced a regex split which would prevent merges. @@ -386,7 +385,7 @@ impl CoreBPE { if unstable_bytes.len() > 1 { let last_decoded = bstr::decode_last_utf8(unstable_bytes.as_slice()); if unstable_bytes.len() - last_decoded.1 > 0 - && last_decoded.0.map_or(false, |c| c.is_whitespace()) + && last_decoded.0.map_or(false, char::is_whitespace) { let mut reencoded = byte_pair_encode( &unstable_bytes[..unstable_bytes.len() - last_decoded.1], @@ -413,11 +412,11 @@ impl CoreBPE { let regex = Regex::new(pattern)?; let special_regex = { - let _parts = special_tokens_encoder + let parts = special_tokens_encoder .keys() .map(|s| fancy_regex::escape(s)) .collect::>(); - Regex::new(&_parts.join("|"))? + Regex::new(&parts.join("|"))? }; let decoder: HashMap> = @@ -434,7 +433,7 @@ impl CoreBPE { let mut sorted_token_bytes: Vec> = encoder.keys().cloned().collect(); sorted_token_bytes.sort(); - Ok(CoreBPE { + Ok(Self { encoder, special_tokens_encoder, decoder, @@ -453,15 +452,15 @@ impl CoreBPE { self._encode_ordinary_native(text) } - pub fn encode(&self, text: &str, allowed_special: HashSet<&str>) -> Vec { - self._encode_native(text, &allowed_special).0 + pub fn encode(&self, text: &str, allowed_special: &HashSet<&str>) -> Vec { + self._encode_native(text, allowed_special).0 } pub fn encode_with_special_tokens(&self, text: &str) -> Vec { let allowed_special = self .special_tokens_encoder .keys() - .map(|s| s.as_str()) + .map(std::string::String::as_str) .collect(); self._encode_native(text, &allowed_special).0 } @@ -495,9 +494,9 @@ impl CoreBPE { fn encode_with_unstable( &self, text: &str, - allowed_special: HashSet<&str>, + allowed_special: &HashSet<&str>, ) -> (Vec, HashSet>) { - self._encode_unstable_native(text, &allowed_special) + self._encode_unstable_native(text, allowed_special) } #[allow(dead_code)] @@ -525,12 +524,12 @@ impl CoreBPE { // Decoding // ==================== - pub fn decode_bytes(&self, tokens: Vec) -> Vec { - self._decode_native(&tokens) + pub fn decode_bytes(&self, tokens: &[usize]) -> Vec { + self._decode_native(tokens) } - pub fn decode(&self, tokens: Vec) -> Result { - match String::from_utf8(self._decode_native(&tokens)) { + pub fn decode(&self, tokens: &[usize]) -> Result { + match String::from_utf8(self._decode_native(tokens)) { Ok(text) => Ok(text), Err(e) => Err(anyhow!("Unable to decode into a valid UTF-8 string: {}", e)), } @@ -576,7 +575,7 @@ mod tests { fn cl100k_base_test() { let bpe = cl100k_base().unwrap(); let tokens = bpe.encode_with_special_tokens("This is a test with a lot of spaces"); - let decoded = bpe.decode(tokens.clone()).unwrap(); + let decoded = bpe.decode(&tokens).unwrap(); assert_eq!(decoded, "This is a test with a lot of spaces"); assert_eq!( tokens, -- cgit v1.2.3