Update tiktoken-rs to v0.5.6

connorjacobsen · Nov 13, 2023 · 74f4bf5 · 74f4bf5
1 parent 9e5113a
commit 74f4bf5
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 4 deletions.
diff --git a/native/tiktoken/Cargo.lock b/native/tiktoken/Cargo.lock
diff --git a/native/tiktoken/Cargo.toml b/native/tiktoken/Cargo.toml
@@ -11,4 +11,4 @@ crate-type = ["cdylib"]
 
 [dependencies]
 rustler = "0.30.0"
-tiktoken-rs = "0.2.2"
+tiktoken-rs = "0.5.6"
diff --git a/native/tiktoken/src/lib.rs b/native/tiktoken/src/lib.rs
@@ -3,7 +3,13 @@ use std::vec::Vec;
 
 #[rustler::nif]
 fn encoding_for_model(model: &str) -> Option<&str> {
-    tiktoken_rs::encoding_for_model(model)
+    match tiktoken_rs::tokenizer::get_tokenizer(model) {
+        Some(tiktoken_rs::tokenizer::Tokenizer::Cl100kBase) => Some("cl100k_base"),
+        Some(tiktoken_rs::tokenizer::Tokenizer::P50kBase) => Some("p50k_base"),
+        Some(tiktoken_rs::tokenizer::Tokenizer::R50kBase) => Some("r50k_base"),
+        Some(tiktoken_rs::tokenizer::Tokenizer::P50kEdit) => Some("p50k_edit"),
+        _ => None,
+    }
 }
 
 // p50k