Repository: zurawiki/tiktoken-rs (language: rust)

Issue:
cl100k usage
Greetings.

How can we use cl100k with this library? I can see that the file for it is in the source directory, but there is no function to load it. Would it be the same as loading the others? If so I can do it myself, I'm just new to this topic and I don't know if there's something more complex that would have to be done.

Thanks.


Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index 4e7336a..a4c3e34 100644
--- a/README.md
+++ b/README.md
@@ -27,12 +27,26 @@ let tokens = bpe.encode_with_special_tokens("This is an example");
 println!("Token count: {}", tokens.len());
 ```
 
-See the examples in the repo for usecases.
+`tiktoken` supports three encodings used by OpenAI models:
+
+| Encoding name           | OpenAI models                                       |
+|-------------------------|-----------------------------------------------------|
+| `cl100k_base`           | ChatGPT models, `text-embedding-ada-002`            |
+| `p50k_base`             | Code models, `text-davinci-002`, `text-davinci-003` |
+| `r50k_base` (or `gpt2`) | GPT-3 models like `davinci`                         |
+
+
+See the [examples](./examples/) in the repo for use cases. For more context on the different tokenizers, see the [OpenAI Cookbook](https://github.com/openai/openai-cookbook/blob/66b988407d8d13cad5060a881dc8c892141f2d5c/examples/How_to_count_tokens_with_tiktoken.ipynb)
+
 
 ## Encountered any bugs?
 
 If you encounter any bugs or have any suggestions for improvements, please open an issue on the repository.
 
+## Acknowledgements
+
+- Thanks for @spolu for the original code, and `.tiktoken` files.
+
 ## License
 
 This project is licensed under the [MIT License](./LICENSE).
diff --git a/src/tiktoken.rs b/src/tiktoken.rs
index ad29ce6..fad3bfe 100644
--- a/src/tiktoken.rs
+++ b/src/tiktoken.rs
@@ -10,41 +10,38 @@ use rustc_hash::FxHashMap as HashMap;
 use std::collections::HashSet;
 use std::sync::Arc;
 
+/// Initializes and returns a new instance of the r50k_base tokenizer (also known as `gpt2`)
+/// Use for GPT-3 models like `davinci`
 pub fn r50k_base() -> Result<CoreBPE> {
-    let r50k_base = include_str!("r50k_base.tiktoken");
-
-    let mut encoder = HashMap::default();
-    for line in r50k_base.lines() {
-        let mut parts = line.split(' ');
-        let token = &general_purpose::STANDARD.decode(parts.next().unwrap())?;
-        let rank: usize = parts.next().unwrap().parse().unwrap();
-        encoder.insert(token.clone(), rank);
-    }
-
-    let mut special_tokens = HashMap::default();
-    special_tokens.insert(String::from("
"), 50256);
-
-    CoreBPE::new(
-        encoder,
-        special_tokens,
-        "'s|'t|'re|'ve|'m|'ll|'d| ?\\p{L}+| ?\\p{N}+| ?[^\\s\\p{L}\\p{N}]+|\\s+(?!\\S)|\\s+",
-    )
+    create_bpe(include_str!("r50k_base.tiktoken"))
 }
 
+/// Initializes and returns a new instance of the p50k_base tokenizer.
+/// Use for Code models, `text-davinci-002`, `text-davinci-003`
 pub fn p50k_base() -> Result<CoreBPE> {
-    let p50k_base = include_str!("p50k_base.tiktoken");
+    create_bpe(include_str!("p50k_base.tiktoken"))
+}
+
+/// Initializes and returns a new instance of the cl100k_base tokenizer.
+/// Use for ChatGPT models, `text-embedding-ada-002`
+pub fn cl100k_base() -> Result<CoreBPE> {
+    create_bpe(include_str!("cl100k_base.tiktoken"))
+}
 
+fn create_bpe(input_file: &'static str) -> Result<CoreBPE> {
     let mut encoder = HashMap::default();
-    for line in p50k_base.lines() {
+    let mut max_rank = 0usize;
+    for line in input_file.lines() {
         let mut parts = line.split(' ');
         let raw = parts.next().unwrap();
         let token = &general_purpose::STANDARD.decode(raw)?;
         let rank: usize = parts.next().unwrap().parse().unwrap();
         encoder.insert(token.clone(), rank);
+        max_rank = max_rank.max(rank);
     }
 
     let mut special_tokens = HashMap::default();
-    special_tokens.insert(String::from("
"), 50256);
+    special_tokens.insert(String::from("
"), max_rank + 1);
 
     CoreBPE::new(
         encoder,
@@ -53,6 +50,10 @@ pub fn p50k_base() -> Result<CoreBPE> {
     )
 }
 
+/// Returns a singleton instance of the r50k_base tokenizer. (also known as `gpt2`)
+/// Use for GPT-3 models like `davinci`
+///
+/// This function will only initialize the tokenizer once, and then return a reference the tokenizer
 pub fn r50k_base_singleton() -> Arc<Mutex<CoreBPE>> {
     lazy_static! {
         static ref R50K_BASE: Arc<Mutex<CoreBPE>> = Arc::new(Mutex::new(r50k_base().unwrap()));
@@ -60,6 +61,10 @@ pub fn r50k_base_singleton() -> Arc<Mutex<CoreBPE>> {
     R50K_BASE.clone()
 }
 
+/// Returns a singleton instance of the p50k_base tokenizer.
+/// Use for Code models, `text-davinci-002`, `text-davinci-003`
+///
+/// This function will only initialize the tokenizer once, and then return a reference the tokenizer.
 pub fn p50k_base_singleton() -> Arc<Mutex<CoreBPE>> {
     lazy_static! {
         static ref P50K_BASE: Arc<Mutex<CoreBPE>> = Arc::new(Mutex::new(p50k_base().unwrap()));
@@ -67,6 +72,17 @@ pub fn p50k_base_singleton() -> Arc<Mutex<CoreBPE>> {
     P50K_BASE.clone()
 }
 
+/// Returns a singleton instance of the cl100k_base tokenizer.
+/// Use for ChatGPT models, `text-embedding-ada-002`
+///
+/// This function will only initialize the tokenizer once, and then return a reference the tokenizer
+pub fn cl100k_base_singleton() -> Arc<Mutex<CoreBPE>> {
+    lazy_static! {
+        static ref CL100K_BASE: Arc<Mutex<CoreBPE>> = Arc::new(Mutex::new(cl100k_base().unwrap()));
+    }
+    CL100K_BASE.clone()
+}
+
 fn _byte_pair_merge(piece: &[u8], ranks: &HashMap<Vec<u8>, usize>) -> Vec<std::ops::Range<usize>> {
     let mut parts: Vec<_> = (0..piece.len()).map(|i| i..i + 1).collect();
 


Test patch:
diff --git a/tests/tiktoken.rs b/tests/tiktoken.rs
index d57f98f..26047ed 100644
--- a/tests/tiktoken.rs
+++ b/tests/tiktoken.rs
@@ -12,30 +12,70 @@ fn very_simple_test() {
     assert_eq!(res, vec![b"ab", b"cd"]);
 }
 
+fn test_roundtrip(bpe: &CoreBPE, text: &str) {
+    let tokens = bpe.encode_with_special_tokens(text);
+    let decoded = bpe.decode(tokens).unwrap();
+    assert_eq!(decoded, text);
+}
+
+fn test_decode(bpe: &CoreBPE, text: &str, exected_tokens: Vec<usize>) {
+    let tokens = bpe.encode_with_special_tokens(text);
+    assert_eq!(tokens, exected_tokens,);
+}
 #[test]
 fn p50k_base_test() {
     let bpe = p50k_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();
-    assert_eq!(decoded, "This is a test         with a lot of spaces");
-    assert_eq!(
-        tokens,
-        vec![1212, 318, 257, 1332, 50263, 351, 257, 1256, 286, 9029]
+    test_roundtrip(&bpe, "This is a test         with a lot of spaces");
+    test_decode(
+        &bpe,
+        "This is a test         with a lot of spaces",
+        vec![1212, 318, 257, 1332, 50263, 351, 257, 1256, 286, 9029],
+    );
+    test_decode(
+        &bpe,
+        "This is a test         with a lot of spaces
",
+        vec![
+            1212, 318, 257, 1332, 50263, 351, 257, 1256, 286, 9029, 50281,
+        ],
     );
 }
 
 #[test]
 fn r50k_base_test() {
     let bpe = r50k_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();
-    assert_eq!(decoded, "This is a test         with a lot of spaces");
-    println!("{tokens:?}");
-    assert_eq!(
-        tokens,
+    test_roundtrip(&bpe, "This is a test         with a lot of spaces");
+    test_decode(
+        &bpe,
+        "This is a test         with a lot of spaces",
+        vec![
+            1212, 318, 257, 1332, 220, 220, 220, 220, 220, 220, 220, 220, 351, 257, 1256, 286, 9029,
+        ],
+    );
+    test_decode(
+        &bpe,
+        "This is a test         with a lot of spaces
",
+        vec![
+            1212, 318, 257, 1332, 220, 220, 220, 220, 220, 220, 220, 220, 351, 257, 1256, 286,
+            9029, 50256,
+        ],
+    );
+}
+
+#[test]
+fn cl100k_base_test() {
+    let bpe = cl100k_base().unwrap();
+    test_roundtrip(&bpe, "This is a test         with a lot of spaces");
+    test_decode(
+        &bpe,
+        "This is a test         with a lot of spaces",
+        vec![2028, 374, 264, 1296, 260, 449, 264, 2763, 315, 12908],
+    );
+    test_decode(
+        &bpe,
+        "This is a test         with a lot of spaces
",
         vec![
-            1212, 318, 257, 1332, 220, 220, 220, 220, 220, 220, 220, 220, 351, 257, 1256, 286, 9029
-        ]
+            2028, 374, 264, 1296, 260, 449, 264, 2763, 315, 12908, 100256,
+        ],
     );
 }