backport small upstream improvements

src/corebpe.rs

@@ -8,11 +8,10 @@ use std::sync::Arc;
 
 pub type Rank = u32;
 
-fn _byte_pair_merge<T>(
-    piece: &[u8],
+fn _byte_pair_merge(
     ranks: &HashMap<Vec<u8>, Rank>,
-    f: impl Fn(std::ops::Range<usize>) -> T,
-) -> Vec<T> {
+    piece: &[u8],
+) -> Vec<(usize, Rank)> {
     // This is a vector of (start, rank).
     // The rank is of the byte pair starting at position start.
     // The rank of the last item in the vector is not a valid value.
@@ -87,25 +86,26 @@ fn _byte_pair_merge<T>(
             break;
         }
     }
-    let mut out: Vec<T> = Vec::with_capacity(parts.len() - 1);
-    for i in 0..parts.len() - 1 {
-        out.push(f(parts[i].0..parts[i + 1].0));
-    }
-    out
+
+    parts
 }
 
 pub fn byte_pair_encode(piece: &[u8], ranks: &HashMap<Vec<u8>, Rank>) -> Vec<Rank> {
     if piece.len() == 1 {
         return vec![ranks[piece]];
     }
-    _byte_pair_merge(piece, ranks, |p| ranks[&piece[p.start..p.end]])
+    _byte_pair_merge(ranks, piece)
+        .windows(2)
+        .map(|part| ranks[&piece[part[0].0..part[1].0]])
+        .collect()
 }
 
 pub fn byte_pair_split<'a>(piece: &'a [u8], ranks: &HashMap<Vec<u8>, Rank>) -> Vec<&'a [u8]> {
-    if piece.len() == 1 {
-        return vec![piece];
-    }
-    _byte_pair_merge(piece, ranks, |p| &piece[p.start..p.end])
+    assert!(piece.len() > 1);
+    _byte_pair_merge(ranks, piece)
+        .windows(2)
+        .map(|part| &piece[part[0].0..part[1].0])
+        .collect()
 }
 
 // Various performance notes:
@@ -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"]);
+    }
+
 }