add multithreading encoding (batching)

bindings/python/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "pyrwkv_tokenizer"
-version = "0.8.5"
+version = "0.9.0"
 edition = "2021"
 authors = ["Cahya Wirawan <[email protected]>"]
 description = "A fast RWKV Tokenizer"
@@ -22,4 +22,5 @@ crate-type = ["cdylib"]
 
 [dependencies]
 pyo3 = "0.21.2"
-rwkv-tokenizer = "0.8.5"
+rwkv-tokenizer = { path = "../../rwkv-tokenizer" }
+rayon = "1.10.0"

bindings/python/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "maturin"
 
 [project]
 name = "pyrwkv_tokenizer"
-version = "0.8.5"
+version = "0.9.0"
 requires-python = ">=3.8"
 description = "RWKV Tokenizer"
 readme = "README.md"

bindings/python/pyrwkv_tokenizer/__init__.py

@@ -7,7 +7,7 @@
 from .pyrwkv_tokenizer import __doc__  # noqa: F401
 
 __all__ = __all__ + ["RWKVTokenizer"]
-__version__ = "0.8.5"
+__version__ = "0.9.0"
 
 
 class RWKVTokenizer:
@@ -22,6 +22,16 @@ def encode(self, text: str):
         tokens_ids = self.tokenizer.encode(text)
         return tokens_ids
 
+    def encode_batch(self, text_list: [str]):
+        tokens_ids = self.tokenizer.encode_batch(text_list)
+        return tokens_ids
+
     def decode(self, tokens_ids):
         text = self.tokenizer.decode(tokens_ids)
         return text
+
+    def vocab_size(self):
+        return self.tokenizer.vocab_size()
+
+    def get_vocab(self):
+        return self.tokenizer.get_vocab()

bindings/python/src/lib.rs

@@ -22,6 +22,10 @@ impl WorldTokenizer {
         self.tokenizer.encode(word)
     }
 
+    pub(crate) fn encode_batch(&self, word_list: Vec<String>) -> Vec<Vec<u16>> {
+        self.tokenizer.encode_batch(word_list)
+    }
+
     pub(crate) fn decode(&self, vec: Vec<u16>) -> String {
         return self.tokenizer.decode(vec);
     }
@@ -30,7 +34,7 @@ impl WorldTokenizer {
         return self.tokenizer.vocab_size();
     }
 
-    pub(crate) fn get_vocab(&self) -> HashMap<&Vec<u8>, usize> {
+    pub(crate) fn get_vocab(&self) -> HashMap<String, usize> {
         return self.tokenizer.get_vocab();
     }
 }

rwkv-tokenizer/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "rwkv-tokenizer"
-version = "0.8.5"
+version = "0.9.0"
 edition = "2021"
 authors = ["Cahya Wirawan <[email protected]>"]
 description = "A fast RWKV Tokenizer"
@@ -11,3 +11,4 @@ exclude = []
 [dependencies]
 regex = "1.10.4"
 unescape = "0.1.0"
+rayon = "1.10.0"

rwkv-tokenizer/src/lib.rs

@@ -7,6 +7,8 @@ use std::path::{Path};
 use regex::Regex;
 use trie::Trie;
 use unescape::unescape;
+use rayon::prelude::*;
+
 
 #[derive(Debug)]
 pub struct WorldTokenizer {
@@ -58,6 +60,10 @@ impl WorldTokenizer {
         self.trie.tokenize(word)
     }
 
+    pub fn encode_batch(&self, word_list: Vec<String>) -> Vec<Vec<u16>> {
+        word_list.par_iter().map(|word| self.trie.tokenize(word)).collect()
+    }
+
     pub fn decode(&self, vec: Vec<u16>) -> String {
         let mut result: Vec<u8> = Vec::new();
         for index in vec.iter() {
@@ -71,10 +77,11 @@ impl WorldTokenizer {
         self.tokens.len()
     }
 
-    pub fn get_vocab(&self) -> HashMap<&Vec<u8>, usize> {
-        let mut vocabularies: HashMap<&Vec<u8>, usize> = HashMap::new();
+    pub fn get_vocab(&self) -> HashMap<String, usize> {
+        let mut vocabularies: HashMap<String, usize> = HashMap::new();
         for (index, value) in self.tokens.iter().enumerate() {
-            vocabularies.insert(value, index);
+            let text: String = String::from_utf8((*value).to_owned()).unwrap_or_else(|_e| "Binary string (TODO)".to_string());
+            vocabularies.insert(text, index);
         }
         vocabularies
     }
@@ -838,4 +845,13 @@ Nórdicg: Ljœr ye caudran créneþ ý jor cẃran."#;
         let text = tokenizer.decode(token_ids);
         assert_eq!(text, LONG_UTF8_TEXT);
     }
+
+    #[test]
+    fn test_get_vocab() {
+        let tokenizer = WorldTokenizer::new(None).unwrap();
+        let vocab = tokenizer.get_vocab();
+        // The vocab size should be 65529, but currently, the binary keys/strings are not included,
+        // therefore it is only 65044. It will be added later.
+        assert_eq!(vocab.len(), 65044);
+    }
 }