Add other token sets as well

crates/bpe-openai/README.md

@@ -7,6 +7,8 @@ For convencience it re-exports the `bpe` crate so that depending on this crate i
 
 Supported token sets:
 
+- r50k
+- p50k
 - cl100k
 - o200k
 

crates/bpe-openai/build.rs

@@ -7,6 +7,24 @@ use serde::Serialize;
 use tiktoken_rs::CoreBPE;
 
 fn main() {
+    serialize_tokens(
+        "r50k",
+        &tiktoken_rs::r50k_base().expect("tiktoken initialization must not fail!"),
+        50256,
+        1,
+    );
+    serialize_tokens(
+        "p50k",
+        &tiktoken_rs::p50k_base().expect("tiktoken initialization must not fail!"),
+        50280,
+        1,
+    );
+    serialize_tokens(
+        "cl100k",
+        &tiktoken_rs::cl100k_base().expect("tiktoken initialization must not fail!"),
+        100256,
+        17846336922010275747,
+    );
     serialize_tokens(
         "cl100k",
         &tiktoken_rs::cl100k_base().expect("tiktoken initialization must not fail!"),

crates/bpe-openai/src/lib.rs

@@ -2,6 +2,16 @@ use std::sync::LazyLock;
 
 use bpe::byte_pair_encoding::BytePairEncoding;
 
+static BPE_R50K: LazyLock<BytePairEncoding> = LazyLock::new(|| {
+    let bytes = include_bytes!(concat!(env!("OUT_DIR"), "/bpe_r50k.dict"));
+    rmp_serde::from_slice(bytes).expect("")
+});
+
+static BPE_P50K: LazyLock<BytePairEncoding> = LazyLock::new(|| {
+    let bytes = include_bytes!(concat!(env!("OUT_DIR"), "/bpe_p50k.dict"));
+    rmp_serde::from_slice(bytes).expect("")
+});
+
 static BPE_CL100K: LazyLock<BytePairEncoding> = LazyLock::new(|| {
     let bytes = include_bytes!(concat!(env!("OUT_DIR"), "/bpe_cl100k.dict"));
     rmp_serde::from_slice(bytes).expect("")
@@ -14,6 +24,14 @@ static BPE_O200K: LazyLock<BytePairEncoding> = LazyLock::new(|| {
 
 pub use bpe::*;
 
+pub fn r50k() -> &'static BytePairEncoding {
+    &BPE_R50K
+}
+
+pub fn p50k() -> &'static BytePairEncoding {
+    &BPE_P50K
+}
+
 pub fn cl100k() -> &'static BytePairEncoding {
     &BPE_CL100K
 }
@@ -26,6 +44,16 @@ pub fn o200k() -> &'static BytePairEncoding {
 mod tests {
     use super::*;
 
+    #[test]
+    fn can_load_r50k() {
+        r50k().count("".as_bytes());
+    }
+
+    #[test]
+    fn can_load_p50k() {
+        p50k().count("".as_bytes());
+    }
+
     #[test]
     fn can_load_cl100k() {
         cl100k().count("".as_bytes());