Fix split on special tokens & bump version

bindings/python/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "tokenizers-python"
-version = "0.1.0"
+version = "0.1.1"
 authors = ["Anthony MOI <[email protected]>"]
 edition = "2018"
 

bindings/python/setup.py

@@ -3,7 +3,7 @@
 
 setup(
     name="tokenizers",
-    version="0.1.0",
+    version="0.1.1",
     description="Fast and Customizable Tokenizers",
     long_description=open("README.md", "r", encoding="utf-8").read(),
     long_description_content_type="text/markdown",

bindings/python/tokenizers/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.1.0"
+__version__ = "0.1.1"
 
 from .tokenizers import Tokenizer, Encoding
 from .tokenizers import decoders

tokenizers/Cargo.toml

@@ -2,7 +2,7 @@
 authors = ["Anthony MOI <[email protected]>"]
 edition = "2018"
 name = "tokenizers"
-version = "0.6.0"
+version = "0.6.1"
 homepage = "https://github.com/huggingface/tokenizers"
 repository = "https://github.com/huggingface/tokenizers"
 documentation = "https://docs.rs/tokenizers/"

tokenizers/src/tokenizer/mod.rs

@@ -565,6 +565,8 @@ impl Tokenizer {
             }
         }
 
+        self.refresh_added_tokens();
+
         added
     }
 
@@ -591,11 +593,27 @@ impl Tokenizer {
                 .or_insert_with(|| token.clone());
         }
 
+        self.refresh_added_tokens();
+
+        // Return the number of added tokens
+        tokens.len() - ignored
+    }
+
+    fn refresh_added_tokens(&mut self) {
         // We rebuild the regex here everytime on purpose, because the added tokens may
         // have changed
+        let special_tokens = self
+            .special_tokens
+            .keys()
+            .map(|t| AddedToken {
+                content: t.to_owned(),
+                single_word: true,
+            })
+            .collect::<Vec<_>>();
         let added_tokens = self
             .added_tokens
             .keys()
+            .chain(special_tokens.iter())
             .map(|token| {
                 if token.single_word {
                     let first_b = token
@@ -635,9 +653,6 @@ impl Tokenizer {
             self.split_re =
                 Some(regex::Regex::new(&format!(r"({})", added_tokens.join("|"))).unwrap());
         }
-
-        // Return the number of added tokens
-        tokens.len() - ignored
     }
 
     /// Split the given sentence on multiple parts, finding the added tokens and their id in the process
@@ -677,10 +692,13 @@ impl Tokenizer {
                     .into_iter()
                     .map(|(start, end)| unsafe {
                         let s = sentence.get_unchecked(start..end).to_owned();
-                        let id = self.added_tokens.get(&AddedToken {
-                            content: s.clone(),
-                            ..Default::default()
-                        });
+                        let mut id = self.special_tokens.get(&s);
+                        if id.is_none() {
+                            id = self.added_tokens.get(&AddedToken {
+                                content: s.clone(),
+                                ..Default::default()
+                            });
+                        }
                         (s, id.copied())
                     })
                     .collect()