Merge pull request #42 from github/aneubeck/count_till_limit

fix count_till_limit function

Author: aneubeck

crates/bpe-openai/src/lib.rs

@@ -92,12 +92,11 @@ impl Tokenizer {
     /// Otherwise, it returns none. This function can be faster than [`Self::count`]` when the
     /// token limit is much smaller than the provided text. Applies pre-tokenization before counting.
     pub fn count_till_limit(&self, text: &str, token_limit: usize) -> Option<usize> {
-        self.split(text)
-            .try_fold(token_limit, |token_limit, piece| {
-                self.bpe
-                    .count_till_limit(piece.as_bytes(), token_limit)
-                    .map(|piece_count| token_limit - piece_count)
-            })
+        self.split(text).try_fold(0, |consumed, piece| {
+            self.bpe
+                .count_till_limit(piece.as_bytes(), token_limit - consumed)
+                .map(|piece_count| consumed + piece_count)
+        })
     }
 
     /// Returns the tokens for the encoding of the given text. Applies pre-tokenization before
@@ -231,4 +230,12 @@ mod tests {
             }
         }
     }
+
+    #[test]
+    fn test_count_till_limit() {
+        assert_eq!(cl100k_base().count_till_limit("abc", 3), Some(1));
+        assert_eq!(cl100k_base().count_till_limit("abcabc", 3), Some(2));
+        assert_eq!(cl100k_base().count_till_limit("abcabcabc", 3), Some(3));
+        assert_eq!(cl100k_base().count_till_limit("abcabcabcabc", 3), None);
+    }
 }