Support None to reset pre_tokenizers and normalizers, and index seq…

…uences (#1590) * initial commit * support None * fix clippy * cleanup * clean? * propagate to pre_tokenizer * fix test * fix rust tests * fix node * propagate to decoder and post processor * fix calls * lint * fmt * node be happy I am fixing you * initial commit * support None * fix clippy * cleanup * clean? * propagate to pre_tokenizer * fix test * fix rust tests * fix node * propagate to decoder and post processor * fix calls * lint * fmt * node be happy I am fixing you * add a small test * styling * style merge * fix merge test * fmt * nits * update tset
huggingface · Aug 7, 2024 · bded212 · bded212
1 parent eea8e1a
commit bded212
Show file tree

Hide file tree

Showing 13 changed files with 134 additions and 67 deletions.
diff --git a/bindings/node/src/tokenizer.rs b/bindings/node/src/tokenizer.rs
@@ -208,7 +208,7 @@ impl Tokenizer {
       .tokenizer
       .write()
       .unwrap()
-      .with_pre_tokenizer((*pre_tokenizer).clone());
+      .with_pre_tokenizer(Some((*pre_tokenizer).clone()));
   }
 
   #[napi]
@@ -217,7 +217,7 @@ impl Tokenizer {
       .tokenizer
       .write()
       .unwrap()
-      .with_decoder((*decoder).clone());
+      .with_decoder(Some((*decoder).clone()));
   }
 
   #[napi]
@@ -231,7 +231,7 @@ impl Tokenizer {
       .tokenizer
       .write()
       .unwrap()
-      .with_post_processor((*post_processor).clone());
+      .with_post_processor(Some((*post_processor).clone()));
   }
 
   #[napi]
@@ -240,7 +240,7 @@ impl Tokenizer {
       .tokenizer
       .write()
       .unwrap()
-      .with_normalizer((*normalizer).clone());
+      .with_normalizer(Some((*normalizer).clone()));
   }
 
   #[napi]

diff --git a/bindings/python/src/normalizers.rs b/bindings/python/src/normalizers.rs
@@ -1,8 +1,6 @@
-use std::sync::{Arc, RwLock};
-
-use pyo3::exceptions;
-use pyo3::prelude::*;
 use pyo3::types::*;
+use pyo3::{exceptions, prelude::*};
+use std::sync::{Arc, RwLock};
 
 use crate::error::ToPyResult;
 use crate::utils::{PyNormalizedString, PyNormalizedStringRefMut, PyPattern};
@@ -354,6 +352,7 @@ impl PyNFKC {
 ///         A list of Normalizer to be run as a sequence
 #[pyclass(extends=PyNormalizer, module = "tokenizers.normalizers", name = "Sequence")]
 pub struct PySequence {}
+
 #[pymethods]
 impl PySequence {
     #[new]
@@ -380,6 +379,22 @@ impl PySequence {
     fn __len__(&self) -> usize {
         0
     }
+
+    fn __getitem__(self_: PyRef<'_, Self>, py: Python<'_>, index: usize) -> PyResult<Py<PyAny>> {
+        match &self_.as_ref().normalizer {
+            PyNormalizerTypeWrapper::Sequence(inner) => match inner.get(index) {
+                Some(item) => PyNormalizer::new(PyNormalizerTypeWrapper::Single(Arc::clone(item)))
+                    .get_as_subtype(py),
+                _ => Err(PyErr::new::<pyo3::exceptions::PyIndexError, _>(
+                    "Index not found",
+                )),
+            },
+            PyNormalizerTypeWrapper::Single(inner) => {
+                PyNormalizer::new(PyNormalizerTypeWrapper::Single(Arc::clone(inner)))
+                    .get_as_subtype(py)
+            }
+        }
+    }
 }
 
 /// Lowercase Normalizer

diff --git a/bindings/python/src/pre_tokenizers.rs b/bindings/python/src/pre_tokenizers.rs
@@ -463,6 +463,24 @@ impl PySequence {
     fn __getnewargs__<'p>(&self, py: Python<'p>) -> Bound<'p, PyTuple> {
         PyTuple::new_bound(py, [PyList::empty_bound(py)])
     }
+
+    fn __getitem__(self_: PyRef<'_, Self>, py: Python<'_>, index: usize) -> PyResult<Py<PyAny>> {
+        match &self_.as_ref().pretok {
+            PyPreTokenizerTypeWrapper::Sequence(inner) => match inner.get(index) {
+                Some(item) => {
+                    PyPreTokenizer::new(PyPreTokenizerTypeWrapper::Single(Arc::clone(item)))
+                        .get_as_subtype(py)
+                }
+                _ => Err(PyErr::new::<pyo3::exceptions::PyIndexError, _>(
+                    "Index not found",
+                )),
+            },
+            PyPreTokenizerTypeWrapper::Single(inner) => {
+                PyPreTokenizer::new(PyPreTokenizerTypeWrapper::Single(Arc::clone(inner)))
+                    .get_as_subtype(py)
+            }
+        }
+    }
 }
 
 pub(crate) fn from_string(string: String) -> Result<PrependScheme, PyErr> {

diff --git a/bindings/python/src/tokenizer.rs b/bindings/python/src/tokenizer.rs
@@ -1371,8 +1371,9 @@ impl PyTokenizer {
 
     /// Set the :class:`~tokenizers.normalizers.Normalizer`
     #[setter]
-    fn set_normalizer(&mut self, normalizer: PyRef<PyNormalizer>) {
-        self.tokenizer.with_normalizer(normalizer.clone());
+    fn set_normalizer(&mut self, normalizer: Option<PyRef<PyNormalizer>>) {
+        let normalizer_option = normalizer.map(|norm| norm.clone());
+        self.tokenizer.with_normalizer(normalizer_option);
     }
 
     /// The `optional` :class:`~tokenizers.pre_tokenizers.PreTokenizer` in use by the Tokenizer
@@ -1387,8 +1388,9 @@ impl PyTokenizer {
 
     /// Set the :class:`~tokenizers.normalizers.Normalizer`
     #[setter]
-    fn set_pre_tokenizer(&mut self, pretok: PyRef<PyPreTokenizer>) {
-        self.tokenizer.with_pre_tokenizer(pretok.clone());
+    fn set_pre_tokenizer(&mut self, pretok: Option<PyRef<PyPreTokenizer>>) {
+        self.tokenizer
+            .with_pre_tokenizer(pretok.map(|pre| pre.clone()));
     }
 
     /// The `optional` :class:`~tokenizers.processors.PostProcessor` in use by the Tokenizer
@@ -1403,8 +1405,9 @@ impl PyTokenizer {
 
     /// Set the :class:`~tokenizers.processors.PostProcessor`
     #[setter]
-    fn set_post_processor(&mut self, processor: PyRef<PyPostProcessor>) {
-        self.tokenizer.with_post_processor(processor.clone());
+    fn set_post_processor(&mut self, processor: Option<PyRef<PyPostProcessor>>) {
+        self.tokenizer
+            .with_post_processor(processor.map(|p| p.clone()));
     }
 
     /// The `optional` :class:`~tokenizers.decoders.Decoder` in use by the Tokenizer
@@ -1419,8 +1422,8 @@ impl PyTokenizer {
 
     /// Set the :class:`~tokenizers.decoders.Decoder`
     #[setter]
-    fn set_decoder(&mut self, decoder: PyRef<PyDecoder>) {
-        self.tokenizer.with_decoder(decoder.clone());
+    fn set_decoder(&mut self, decoder: Option<PyRef<PyDecoder>>) {
+        self.tokenizer.with_decoder(decoder.map(|d| d.clone()));
     }
 }
 
@@ -1436,10 +1439,12 @@ mod test {
     #[test]
     fn serialize() {
         let mut tokenizer = Tokenizer::new(PyModel::from(BPE::default()));
-        tokenizer.with_normalizer(PyNormalizer::new(PyNormalizerTypeWrapper::Sequence(vec![
-            Arc::new(RwLock::new(NFKC.into())),
-            Arc::new(RwLock::new(Lowercase.into())),
-        ])));
+        tokenizer.with_normalizer(Some(PyNormalizer::new(PyNormalizerTypeWrapper::Sequence(
+            vec![
+                Arc::new(RwLock::new(NFKC.into())),
+                Arc::new(RwLock::new(Lowercase.into())),
+            ],
+        ))));
 
         let tmp = NamedTempFile::new().unwrap().into_temp_path();
         tokenizer.save(&tmp, false).unwrap();
@@ -1450,10 +1455,12 @@ mod test {
     #[test]
     fn serde_pyo3() {
         let mut tokenizer = Tokenizer::new(PyModel::from(BPE::default()));
-        tokenizer.with_normalizer(PyNormalizer::new(PyNormalizerTypeWrapper::Sequence(vec![
-            Arc::new(RwLock::new(NFKC.into())),
-            Arc::new(RwLock::new(Lowercase.into())),
-        ])));
+        tokenizer.with_normalizer(Some(PyNormalizer::new(PyNormalizerTypeWrapper::Sequence(
+            vec![
+                Arc::new(RwLock::new(NFKC.into())),
+                Arc::new(RwLock::new(Lowercase.into())),
+            ],
+        ))));
 
         let output = crate::utils::serde_pyo3::to_string(&tokenizer).unwrap();
         assert_eq!(output, "Tokenizer(version=\"1.0\", truncation=None, padding=None, added_tokens=[], normalizer=Sequence(normalizers=[NFKC(), Lowercase()]), pre_tokenizer=None, post_processor=None, decoder=None, model=BPE(dropout=None, unk_token=None, continuing_subword_prefix=None, end_of_word_suffix=None, fuse_unk=False, byte_fallback=False, ignore_merges=False, vocab={}, merges=[]))");

diff --git a/bindings/python/tests/bindings/test_normalizers.py b/bindings/python/tests/bindings/test_normalizers.py
@@ -67,6 +67,14 @@ def test_can_make_sequences(self):
         output = normalizer.normalize_str("  HELLO  ")
         assert output == "hello"
 
+    def test_items(self):
+        normalizers = Sequence([BertNormalizer(True, True), Prepend()])
+        assert normalizers[1].__class__ == Prepend
+        normalizers[0].lowercase = False
+        assert not normalizers[0].lowercase
+        with pytest.raises(IndexError):
+            print(normalizers[2])
+
 
 class TestLowercase:
     def test_instantiate(self):

diff --git a/bindings/python/tests/bindings/test_pre_tokenizers.py b/bindings/python/tests/bindings/test_pre_tokenizers.py
@@ -169,6 +169,13 @@ def test_bert_like(self):
             ("?", (29, 30)),
         ]
 
+    def test_items(self):
+        pre_tokenizers = Sequence([Metaspace("a", "never", split=True), Punctuation()])
+        assert pre_tokenizers[1].__class__ == Punctuation
+        assert pre_tokenizers[0].__class__ == Metaspace
+        pre_tokenizers[0].split = False
+        assert not pre_tokenizers[0].split
+
 
 class TestDigits:
     def test_instantiate(self):

diff --git a/bindings/python/tests/bindings/test_tokenizer.py b/bindings/python/tests/bindings/test_tokenizer.py
@@ -6,10 +6,11 @@
 from tokenizers import AddedToken, Encoding, Tokenizer
 from tokenizers.implementations import BertWordPieceTokenizer
 from tokenizers.models import BPE, Model, Unigram
-from tokenizers.pre_tokenizers import ByteLevel
+from tokenizers.pre_tokenizers import ByteLevel, Metaspace
 from tokenizers.processors import RobertaProcessing, TemplateProcessing
 from tokenizers.normalizers import Strip, Lowercase, Sequence
 
+
 from ..utils import bert_files, data_dir, multiprocessing_with_parallelism, roberta_files
 
 
@@ -551,6 +552,16 @@ def test_decode_special(self):
         assert output == "name is john"
         assert tokenizer.get_added_tokens_decoder()[0] == AddedToken("my", special=True)
 
+    def test_setting_to_none(self):
+        tokenizer = Tokenizer(BPE())
+        tokenizer.normalizer = Strip()
+        tokenizer.normalizer = None
+        assert tokenizer.normalizer == None
+
+        tokenizer.pre_tokenizer = Metaspace()
+        tokenizer.pre_tokenizer = None
+        assert tokenizer.pre_tokenizer == None
+
 
 class TestTokenizerRepr:
     def test_repr(self):

diff --git a/tokenizers/benches/bert_benchmark.rs b/tokenizers/benches/bert_benchmark.rs
@@ -34,13 +34,13 @@ fn create_bert_tokenizer(wp: WordPiece) -> BertTokenizer {
     let sep_id = *wp.get_vocab().get("[SEP]").unwrap();
     let cls_id = *wp.get_vocab().get("[CLS]").unwrap();
     let mut tokenizer = TokenizerImpl::new(wp);
-    tokenizer.with_pre_tokenizer(BertPreTokenizer);
-    tokenizer.with_normalizer(BertNormalizer::default());
-    tokenizer.with_decoder(decoders::wordpiece::WordPiece::default());
-    tokenizer.with_post_processor(BertProcessing::new(
+    tokenizer.with_pre_tokenizer(Some(BertPreTokenizer));
+    tokenizer.with_normalizer(Some(BertNormalizer::default()));
+    tokenizer.with_decoder(Some(decoders::wordpiece::WordPiece::default()));
+    tokenizer.with_post_processor(Some(BertProcessing::new(
         ("[SEP]".to_string(), sep_id),
         ("[CLS]".to_string(), cls_id),
-    ));
+    )));
     tokenizer
 }
 
@@ -81,7 +81,7 @@ fn bench_train(c: &mut Criterion) {
         DecoderWrapper,
     >;
     let mut tokenizer = Tok::new(WordPiece::default());
-    tokenizer.with_pre_tokenizer(Whitespace {});
+    tokenizer.with_pre_tokenizer(Some(Whitespace {}));
     c.bench_function("WordPiece Train vocabulary (small)", |b| {
         b.iter_custom(|iters| {
             iter_bench_train(
@@ -94,7 +94,7 @@ fn bench_train(c: &mut Criterion) {
     });
 
     let mut tokenizer = Tok::new(WordPiece::default());
-    tokenizer.with_pre_tokenizer(Whitespace {});
+    tokenizer.with_pre_tokenizer(Some(Whitespace {}));
     c.bench_function("WordPiece Train vocabulary (big)", |b| {
         b.iter_custom(|iters| {
             iter_bench_train(

diff --git a/tokenizers/benches/bpe_benchmark.rs b/tokenizers/benches/bpe_benchmark.rs
@@ -22,8 +22,8 @@ static BATCH_SIZE: usize = 1_000;
 
 fn create_gpt2_tokenizer(bpe: BPE) -> Tokenizer {
     let mut tokenizer = Tokenizer::new(bpe);
-    tokenizer.with_pre_tokenizer(ByteLevel::default());
-    tokenizer.with_decoder(ByteLevel::default());
+    tokenizer.with_pre_tokenizer(Some(ByteLevel::default()));
+    tokenizer.with_decoder(Some(ByteLevel::default()));
     tokenizer.add_tokens(&[AddedToken::from("ing", false).single_word(false)]);
     tokenizer.add_special_tokens(&[AddedToken::from("[ENT]", true).single_word(true)]);
     tokenizer
@@ -74,7 +74,7 @@ fn bench_train(c: &mut Criterion) {
         .build()
         .into();
     let mut tokenizer = Tokenizer::new(BPE::default()).into_inner();
-    tokenizer.with_pre_tokenizer(Whitespace {});
+    tokenizer.with_pre_tokenizer(Some(Whitespace {}));
     c.bench_function("BPE Train vocabulary (small)", |b| {
         b.iter_custom(|iters| {
             iter_bench_train(
@@ -87,7 +87,7 @@ fn bench_train(c: &mut Criterion) {
     });
 
     let mut tokenizer = Tokenizer::new(BPE::default()).into_inner();
-    tokenizer.with_pre_tokenizer(Whitespace {});
+    tokenizer.with_pre_tokenizer(Some(Whitespace {}));
     c.bench_function("BPE Train vocabulary (big)", |b| {
         b.iter_custom(|iters| {
             iter_bench_train(

diff --git a/tokenizers/src/tokenizer/mod.rs b/tokenizers/src/tokenizer/mod.rs
@@ -550,19 +550,18 @@ where
     }
 
     /// Set the normalizer
-    pub fn with_normalizer(&mut self, normalizer: impl Into<N>) -> &mut Self {
-        self.normalizer = Some(normalizer.into());
+    pub fn with_normalizer(&mut self, normalizer: Option<impl Into<N>>) -> &mut Self {
+        self.normalizer = normalizer.map(|norm| norm.into());
         self
     }
-
     /// Get the normalizer
     pub fn get_normalizer(&self) -> Option<&N> {
         self.normalizer.as_ref()
     }
 
     /// Set the pre tokenizer
-    pub fn with_pre_tokenizer(&mut self, pre_tokenizer: impl Into<PT>) -> &mut Self {
-        self.pre_tokenizer = Some(pre_tokenizer.into());
+    pub fn with_pre_tokenizer(&mut self, pre_tokenizer: Option<impl Into<PT>>) -> &mut Self {
+        self.pre_tokenizer = pre_tokenizer.map(|tok| tok.into());
         self
     }
 
@@ -572,8 +571,8 @@ where
     }
 
     /// Set the post processor
-    pub fn with_post_processor(&mut self, post_processor: impl Into<PP>) -> &mut Self {
-        self.post_processor = Some(post_processor.into());
+    pub fn with_post_processor(&mut self, post_processor: Option<impl Into<PP>>) -> &mut Self {
+        self.post_processor = post_processor.map(|post_proc| post_proc.into());
         self
     }
 
@@ -583,8 +582,8 @@ where
     }
 
     /// Set the decoder
-    pub fn with_decoder(&mut self, decoder: impl Into<D>) -> &mut Self {
-        self.decoder = Some(decoder.into());
+    pub fn with_decoder(&mut self, decoder: Option<impl Into<D>>) -> &mut Self {
+        self.decoder = decoder.map(|dec| dec.into());
         self
     }
 

diff --git a/tokenizers/tests/common/mod.rs b/tokenizers/tests/common/mod.rs
@@ -23,9 +23,11 @@ pub fn get_byte_level_bpe() -> BPE {
 pub fn get_byte_level(add_prefix_space: bool, trim_offsets: bool) -> Tokenizer {
     let mut tokenizer = Tokenizer::new(get_byte_level_bpe());
     tokenizer
-        .with_pre_tokenizer(ByteLevel::default().add_prefix_space(add_prefix_space))
-        .with_decoder(ByteLevel::default())
-        .with_post_processor(ByteLevel::default().trim_offsets(trim_offsets));
+        .with_pre_tokenizer(Some(
+            ByteLevel::default().add_prefix_space(add_prefix_space),
+        ))
+        .with_decoder(Some(ByteLevel::default()))
+        .with_post_processor(Some(ByteLevel::default().trim_offsets(trim_offsets)));
 
     tokenizer
 }
@@ -43,13 +45,13 @@ pub fn get_bert() -> Tokenizer {
     let sep = tokenizer.get_model().token_to_id("[SEP]").unwrap();
     let cls = tokenizer.get_model().token_to_id("[CLS]").unwrap();
     tokenizer
-        .with_normalizer(BertNormalizer::default())
-        .with_pre_tokenizer(BertPreTokenizer)
-        .with_decoder(WordPieceDecoder::default())
-        .with_post_processor(BertProcessing::new(
+        .with_normalizer(Some(BertNormalizer::default()))
+        .with_pre_tokenizer(Some(BertPreTokenizer))
+        .with_decoder(Some(WordPieceDecoder::default()))
+        .with_post_processor(Some(BertProcessing::new(
             (String::from("[SEP]"), sep),
             (String::from("[CLS]"), cls),
-        ));
+        )));
 
     tokenizer
 }