Tweak API

Author: Hendrik van Antwerpen

crates/bpe/bindings/python/pyproject.toml

@@ -19,3 +19,9 @@ dynamic = ["version"]
 
 [tool.maturin]
 features = ["pyo3/extension-module"]
+
+[dependency-groups]
+dev = [
+    "maturin>=1.8.2",
+    "pip>=25.0.1",
+]

crates/bpe/bindings/python/src/lib.rs

@@ -2,59 +2,66 @@ use std::borrow::Cow;
 
 use pyo3::prelude::*;
 
-#[pyclass]
-struct BytePairEncoding(Cow<'static, ::bpe::byte_pair_encoding::BytePairEncoding>);
+#[pymodule]
+mod bpe {
+    use super::*;
 
-#[pyclass]
-struct Tokenizer(Cow<'static, ::bpe_openai::Tokenizer>);
+    #[pyclass]
+    struct BytePairEncoding(&'static ::bpe::byte_pair_encoding::BytePairEncoding);
 
-#[pymethods]
-impl BytePairEncoding {
-    fn count(&self, input: &[u8]) -> usize {
-        self.0.count(input)
-    }
+    #[pymethods]
+    impl BytePairEncoding {
+        fn count(&self, input: &[u8]) -> usize {
+            self.0.count(input)
+        }
 
-    fn encode_via_backtracking(&self, input: &[u8]) -> Vec<u32> {
-        self.0.encode_via_backtracking(input)
-    }
+        fn encode_via_backtracking(&self, input: &[u8]) -> Vec<u32> {
+            self.0.encode_via_backtracking(input)
+        }
 
-    fn decode_tokens(&self, tokens: Vec<u32>) -> Vec<u8> {
-        self.0.decode_tokens(&tokens)
+        fn decode_tokens(&self, tokens: Vec<u32>) -> Vec<u8> {
+            self.0.decode_tokens(&tokens)
+        }
     }
-}
 
-#[pymethods]
-impl Tokenizer {
-    fn count(&self, input: &str) -> usize {
-        self.0.count(&input)
-    }
+    #[pymodule]
+    mod openai {
+        use super::*;
 
-    fn count_till_limit(&self, input: Cow<str>, limit: usize) -> Option<usize> {
-        self.0.count_till_limit(&input, limit)
-    }
+        #[pyclass]
+        struct Tokenizer(&'static ::bpe_openai::Tokenizer);
 
-    fn encode(&self, input: Cow<str>) -> Vec<u32> {
-        self.0.encode(&input)
-    }
+        #[pymethods]
+        impl Tokenizer {
+            fn count(&self, input: &str) -> usize {
+                self.0.count(&input)
+            }
 
-    fn decode(&self, tokens: Vec<u32>) -> Option<String> {
-        self.0.decode(&tokens)
-    }
-}
+            fn count_till_limit(&self, input: Cow<str>, limit: usize) -> Option<usize> {
+                self.0.count_till_limit(&input, limit)
+            }
 
-#[pyfunction]
-fn cl100k_base() -> PyResult<Tokenizer> {
-    Ok(Tokenizer(Cow::Borrowed(::bpe_openai::cl100k_base())))
-}
+            fn encode(&self, input: Cow<str>) -> Vec<u32> {
+                self.0.encode(&input)
+            }
 
-#[pyfunction]
-fn o200k_base() -> PyResult<Tokenizer> {
-    Ok(Tokenizer(Cow::Borrowed(::bpe_openai::o200k_base())))
-}
+            fn decode(&self, tokens: Vec<u32>) -> Option<String> {
+                self.0.decode(&tokens)
+            }
 
-#[pymodule]
-fn bpe_openai(m: &Bound<'_, PyModule>) -> PyResult<()> {
-    m.add_function(wrap_pyfunction!(cl100k_base, m)?)?;
-    m.add_function(wrap_pyfunction!(o200k_base, m)?)?;
-    Ok(())
+            fn bpe(&self) -> BytePairEncoding {
+                BytePairEncoding(&self.0.bpe)
+            }
+        }
+
+        #[pyfunction]
+        fn cl100k_base() -> PyResult<Tokenizer> {
+            Ok(Tokenizer(::bpe_openai::cl100k_base()))
+        }
+
+        #[pyfunction]
+        fn o200k_base() -> PyResult<Tokenizer> {
+            Ok(Tokenizer(::bpe_openai::o200k_base()))
+        }
+    }
 }

crates/bpe/bindings/python/test.py

@@ -2,11 +2,17 @@
 
 import bpe
 
-tok = bpe.cl100k_base()
+tok = bpe.openai.cl100k_base()
+
+## Use tokenizer
 
 enc = tok.encode("Hello, world!")
 print(enc)
 cnt = tok.count("Hello, world!")
 print(cnt)
 dec = tok.decode(enc)
 print(dec)
+
+## Use underlying BPE instance
+
+bpe = tok.bpe()