fix: minor code cleanup and pedantic lints (#134)

1 files changed, 31 insertions, 32 deletions

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<usize> {
 }
 
 /// Convert tokens to plan text
-pub fn tokens_to_text(tokens: Vec<usize>) -> Result<String> {
+pub fn tokens_to_text(tokens: &[usize]) -> Result<String> {
     cl100k_base_singleton().lock().decode(tokens)
 }
 
@@ -43,11 +43,11 @@ pub fn cl100k_base() -> Result<CoreBPE> {
     }
 
     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<Mutex<CoreBPE>> {
 }
 
 fn _byte_pair_merge(piece: &[u8], ranks: &HashMap<Vec<u8>, usize>) -> Vec<std::ops::Range<usize>> {
+    #[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::<Vec<_>>();
-            Regex::new(&_parts.join("|"))?
+            Regex::new(&parts.join("|"))?
         };
 
         let decoder: HashMap<usize, Vec<u8>> =
@@ -434,7 +433,7 @@ impl CoreBPE {
         let mut sorted_token_bytes: Vec<Vec<u8>> = 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<usize> {
-        self._encode_native(text, &allowed_special).0
+    pub fn encode(&self, text: &str, allowed_special: &HashSet<&str>) -> Vec<usize> {
+        self._encode_native(text, allowed_special).0
     }
 
     pub fn encode_with_special_tokens(&self, text: &str) -> Vec<usize> {
         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<usize>, HashSet<Vec<usize>>) {
-        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<usize>) -> Vec<u8> {
-        self._decode_native(&tokens)
+    pub fn decode_bytes(&self, tokens: &[usize]) -> Vec<u8> {
+        self._decode_native(tokens)
     }
 
-    pub fn decode(&self, tokens: Vec<usize>) -> Result<String> {
-        match String::from_utf8(self._decode_native(&tokens)) {
+    pub fn decode(&self, tokens: &[usize]) -> Result<String> {
+        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,