Avoid calling byte_pair_encode for existing tokens

This was byte_pair_encode can be optimized further, assuming we'll always have at least 2 tokens

(cherry-picked from openai/tiktoken@b4c687e)

Author: Lőrinc

src/corebpe.rs

@@ -209,11 +209,10 @@ impl CoreBPE {
         let mut ret = vec![];
         for mat in regex.find_iter(text) {
             let piece = mat.unwrap().as_str().as_bytes();
-            if let Some(token) = self.encoder.get(piece) {
-                ret.push(*token);
-                continue;
+            match self.encoder.get(piece) {
+                Some(token) => ret.push(*token),
+                None => ret.extend(&byte_pair_encode(piece, &self.encoder)),
             }
-            ret.extend(&byte_pair_encode(piece, &self.encoder));
         }
         ret
     }
@@ -525,7 +524,10 @@ impl CoreBPE {
                     unstable_bytes.extend_from_slice(&bytes[e.valid_up_to()..]);
 
                     tokens.truncate(tokens.len() - last_piece_token_len);
-                    tokens.extend(byte_pair_encode(&unstable_bytes, &self.encoder));
+                    match self.encoder.get(&unstable_bytes) {
+                        Some(token) => tokens.push(*token),
+                        None => tokens.extend(&byte_pair_encode(&unstable_bytes, &self.encoder)),
+                    }
                 }
                 tokens
             }
@@ -590,15 +592,26 @@ impl CoreBPE {
 mod tests {
     use rustc_hash::FxHashMap as HashMap;
 
-    use crate::corebpe::byte_pair_split;
+    use crate::corebpe::{byte_pair_split, Rank};
 
-    #[test]
-    fn very_simple_test() {
-        let mut ranks = HashMap::default();
-        ranks.insert(b"ab".to_vec(), 1);
-        ranks.insert(b"cd".to_vec(), 2);
+    fn setup_ranks() -> HashMap<Vec<u8>, Rank> {
+        HashMap::from_iter([
+            (b"ab".to_vec(), 0),
+            (b"cd".to_vec(), 1),
+        ])
+    }
 
+    #[test]
+    fn test_simple_characters() {
+        let ranks = setup_ranks();
         let res = byte_pair_split(b"abcd", &ranks);
         assert_eq!(res, vec![b"ab", b"cd"]);
     }
+    #[test]
+    fn test_repeated_characters() {
+        let ranks = setup_ranks();
+        let res = byte_pair_split(b"abab", &ranks);
+        assert_eq!(res, vec![b"ab", b"ab"]);
+    }
+
 }