🚨 Support updating template processors (#1652)

* current updates

* simplify

* set_item works, but `tokenizer._tokenizer.post_processor[1].single = ["$0", "</s>"]` does not !

* fix: `normalizers` deserialization and other refactoring

* fix: `pre_tokenizer` deserialization

* feat: add `__len__` implementation for `normalizer::PySequence`

* feat: add `__setitem__` impl for `normalizers::PySequence`

* feat: add `__setitem__` impl to `pre_tokenizer::PySequence`

* feat: add `__setitem__` impl to `post_processor::PySequence`

* test: add normalizer sequence setter check

* refactor: allow unused `processors::setter` macro

* test: add `__setitem__` test for processors & pretok

* refactor: `unwrap` -> `PyException::new_err()?`

* refactor: fmt

* refactor: remove unnecessary `pub`

* feat(bindings): add missing getters & setters for pretoks

* feat(bindings): add missing getters & setters for processors

* refactor(bindings): rewrite RwLock poison error msg

* refactor: remove debug print

* feat(bindings): add description as to why custom deser is needed

* feat: make post proc sequence elements mutable

* fix(binding): serialization

---------

Co-authored-by: Luc Georges <[email protected]>

Author: ArthurZucker

.github/workflows/python.yml

@@ -55,7 +55,6 @@ jobs:
     steps:
       - name: Checkout repository
         uses: actions/checkout@v4
-        
 
       - name: Install Rust
         uses: actions-rs/toolchain@v1

bindings/python/Cargo.toml

@@ -14,7 +14,7 @@ serde = { version = "1.0", features = ["rc", "derive"] }
 serde_json = "1.0"
 libc = "0.2"
 env_logger = "0.11"
-pyo3 = { version = "0.23", features = ["abi3", "abi3-py39"] }
+pyo3 = { version = "0.23", features = ["abi3", "abi3-py39", "py-clone"] }
 numpy = "0.23"
 ndarray = "0.16"
 itertools = "0.12"

bindings/python/pyproject.toml

@@ -2,8 +2,8 @@
 name = 'tokenizers'
 requires-python = '>=3.9'
 authors = [
-    {name = 'Nicolas Patry', email = '[email protected]'},
-    {name = 'Anthony Moi', email = '[email protected]'}
+  { name = 'Nicolas Patry', email = '[email protected]' },
+  { name = 'Anthony Moi', email = '[email protected]' },
 ]
 classifiers = [
   "Development Status :: 5 - Production/Stable",
@@ -21,12 +21,7 @@ classifiers = [
   "Topic :: Scientific/Engineering :: Artificial Intelligence",
 ]
 keywords = ["NLP", "tokenizer", "BPE", "transformer", "deep learning"]
-dynamic = [
-    'description',
-    'license',
-    'readme',
-    'version',
-]
+dynamic = ['description', 'license', 'readme', 'version']
 dependencies = ["huggingface_hub>=0.16.4,<1.0"]
 
 [project.urls]
@@ -58,16 +53,16 @@ target-version = ['py35']
 line-length = 119
 target-version = "py311"
 lint.ignore = [
-    # a == None in tests vs is None.
-    "E711",
-    # a == False in tests vs is False.
-    "E712",
-    # try.. import except.. pattern without using the lib.
-    "F401",
-    # Raw type equality is required in asserts
-    "E721",
-    # Import order
-    "E402",
-    # Fixtures unused import
-    "F811",
+  # a == None in tests vs is None.
+  "E711",
+  # a == False in tests vs is False.
+  "E712",
+  # try.. import except.. pattern without using the lib.
+  "F401",
+  # Raw type equality is required in asserts
+  "E721",
+  # Import order
+  "E402",
+  # Fixtures unused import
+  "F811",
 ]

bindings/python/src/normalizers.rs

@@ -1,3 +1,4 @@
+use pyo3::exceptions::PyException;
 use pyo3::types::*;
 use pyo3::{exceptions, prelude::*};
 use std::sync::{Arc, RwLock};
@@ -41,7 +42,7 @@ impl PyNormalizedStringMut<'_> {
 /// This class is not supposed to be instantiated directly. Instead, any implementation of a
 /// Normalizer will return an instance of this class when instantiated.
 #[pyclass(dict, module = "tokenizers.normalizers", name = "Normalizer", subclass)]
-#[derive(Clone, Serialize, Deserialize)]
+#[derive(Clone, Debug, Serialize, Deserialize)]
 #[serde(transparent)]
 pub struct PyNormalizer {
     pub(crate) normalizer: PyNormalizerTypeWrapper,
@@ -58,7 +59,11 @@ impl PyNormalizer {
                 .into_pyobject(py)?
                 .into_any()
                 .into(),
-            PyNormalizerTypeWrapper::Single(ref inner) => match &*inner.as_ref().read().unwrap() {
+            PyNormalizerTypeWrapper::Single(ref inner) => match &*inner
+                .as_ref()
+                .read()
+                .map_err(|_| PyException::new_err("RwLock synchronisation primitive is poisoned, cannot get subtype of PyNormalizer"))?
+            {
                 PyNormalizerWrapper::Custom(_) => {
                     Py::new(py, base)?.into_pyobject(py)?.into_any().into()
                 }
@@ -218,7 +223,9 @@ macro_rules! getter {
     ($self: ident, $variant: ident, $name: ident) => {{
         let super_ = $self.as_ref();
         if let PyNormalizerTypeWrapper::Single(ref norm) = super_.normalizer {
-            let wrapper = norm.read().unwrap();
+            let wrapper = norm.read().expect(
+                "RwLock synchronisation primitive is poisoned, cannot get subtype of PyNormalizer",
+            );
             if let PyNormalizerWrapper::Wrapped(NormalizerWrapper::$variant(o)) = (&*wrapper) {
                 o.$name.clone()
             } else {
@@ -234,7 +241,9 @@ macro_rules! setter {
     ($self: ident, $variant: ident, $name: ident, $value: expr) => {{
         let super_ = $self.as_ref();
         if let PyNormalizerTypeWrapper::Single(ref norm) = super_.normalizer {
-            let mut wrapper = norm.write().unwrap();
+            let mut wrapper = norm.write().expect(
+                "RwLock synchronisation primitive is poisoned, cannot get subtype of PyNormalizer",
+            );
             if let PyNormalizerWrapper::Wrapped(NormalizerWrapper::$variant(ref mut o)) = *wrapper {
                 o.$name = $value;
             }
@@ -410,25 +419,55 @@ impl PySequence {
         PyTuple::new(py, [PyList::empty(py)])
     }
 
-    fn __len__(&self) -> usize {
-        0
+    fn __len__(self_: PyRef<'_, Self>) -> usize {
+        match &self_.as_ref().normalizer {
+            PyNormalizerTypeWrapper::Sequence(inner) => inner.len(),
+            PyNormalizerTypeWrapper::Single(_) => 1,
+        }
     }
 
     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)))
+                Some(item) => PyNormalizer::new(PyNormalizerTypeWrapper::Single(item.clone()))
                     .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)
+                PyNormalizer::new(PyNormalizerTypeWrapper::Single(inner.clone())).get_as_subtype(py)
             }
         }
     }
+
+    fn __setitem__(self_: PyRef<'_, Self>, index: usize, value: Bound<'_, PyAny>) -> PyResult<()> {
+        let norm: PyNormalizer = value.extract()?;
+        let PyNormalizerTypeWrapper::Single(norm) = norm.normalizer else {
+            return Err(PyException::new_err("normalizer should not be a sequence"));
+        };
+        match &self_.as_ref().normalizer {
+            PyNormalizerTypeWrapper::Sequence(inner) => match inner.get(index) {
+                Some(item) => {
+                    *item
+                        .write()
+                        .map_err(|_| PyException::new_err("RwLock synchronisation primitive is poisoned, cannot get subtype of PyNormalizer"))? = norm
+                        .read()
+                        .map_err(|_| PyException::new_err("RwLock synchronisation primitive is poisoned, cannot get subtype of PyNormalizer"))?
+                        .clone();
+                }
+                _ => {
+                    return Err(PyErr::new::<pyo3::exceptions::PyIndexError, _>(
+                        "Index not found",
+                    ))
+                }
+            },
+            PyNormalizerTypeWrapper::Single(_) => {
+                return Err(PyException::new_err("normalizer is not a sequence"))
+            }
+        };
+        Ok(())
+    }
 }
 
 /// Lowercase Normalizer
@@ -570,9 +609,31 @@ impl PyReplace {
             ToPyResult(Replace::new(pattern, content)).into_py()?.into(),
         ))
     }
+
+    #[getter]
+    fn get_pattern(_self: PyRef<Self>) -> PyResult<()> {
+        Err(PyException::new_err("Cannot get pattern"))
+    }
+
+    #[setter]
+    fn set_pattern(_self: PyRef<Self>, _pattern: PyPattern) -> PyResult<()> {
+        Err(PyException::new_err(
+            "Cannot set pattern, please instantiate a new replace pattern instead",
+        ))
+    }
+
+    #[getter]
+    fn get_content(self_: PyRef<Self>) -> String {
+        getter!(self_, Replace, content)
+    }
+
+    #[setter]
+    fn set_content(self_: PyRef<Self>, content: String) {
+        setter!(self_, Replace, content, content)
+    }
 }
 
-#[derive(Debug)]
+#[derive(Clone, Debug)]
 pub(crate) struct CustomNormalizer {
     inner: PyObject,
 }
@@ -615,7 +676,7 @@ impl<'de> Deserialize<'de> for CustomNormalizer {
     }
 }
 
-#[derive(Debug, Deserialize)]
+#[derive(Clone, Debug, Deserialize)]
 #[serde(untagged)]
 pub(crate) enum PyNormalizerWrapper {
     Custom(CustomNormalizer),
@@ -634,13 +695,27 @@ impl Serialize for PyNormalizerWrapper {
     }
 }
 
-#[derive(Debug, Clone, Deserialize)]
-#[serde(untagged)]
+#[derive(Debug, Clone)]
 pub(crate) enum PyNormalizerTypeWrapper {
     Sequence(Vec<Arc<RwLock<PyNormalizerWrapper>>>),
     Single(Arc<RwLock<PyNormalizerWrapper>>),
 }
 
+/// XXX: we need to manually implement deserialize here because of the structure of the
+/// PyNormalizerTypeWrapper enum. Given the underlying PyNormalizerWrapper can contain a Sequence,
+/// default deserialization will give us a PyNormalizerTypeWrapper::Single(Sequence) when we'd like
+/// it to be PyNormalizerTypeWrapper::Sequence(// ...).
+impl<'de> Deserialize<'de> for PyNormalizerTypeWrapper {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        let wrapper = NormalizerWrapper::deserialize(deserializer)?;
+        let py_wrapper: PyNormalizerWrapper = wrapper.into();
+        Ok(py_wrapper.into())
+    }
+}
+
 impl Serialize for PyNormalizerTypeWrapper {
     fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
     where
@@ -672,7 +747,17 @@ where
     I: Into<PyNormalizerWrapper>,
 {
     fn from(norm: I) -> Self {
-        PyNormalizerTypeWrapper::Single(Arc::new(RwLock::new(norm.into())))
+        let norm = norm.into();
+        match norm {
+            PyNormalizerWrapper::Wrapped(NormalizerWrapper::Sequence(seq)) => {
+                PyNormalizerTypeWrapper::Sequence(
+                    seq.into_iter()
+                        .map(|e| Arc::new(RwLock::new(PyNormalizerWrapper::Wrapped(e.clone()))))
+                        .collect(),
+                )
+            }
+            _ => PyNormalizerTypeWrapper::Single(Arc::new(RwLock::new(norm))),
+        }
     }
 }
 
@@ -690,10 +775,15 @@ where
 impl Normalizer for PyNormalizerTypeWrapper {
     fn normalize(&self, normalized: &mut NormalizedString) -> tk::Result<()> {
         match self {
-            PyNormalizerTypeWrapper::Single(inner) => inner.read().unwrap().normalize(normalized),
-            PyNormalizerTypeWrapper::Sequence(inner) => inner
-                .iter()
-                .try_for_each(|n| n.read().unwrap().normalize(normalized)),
+            PyNormalizerTypeWrapper::Single(inner) => inner
+                .read()
+                .map_err(|_| PyException::new_err("RwLock synchronisation primitive is poisoned, cannot get subtype of PyNormalizer"))?
+                .normalize(normalized),
+            PyNormalizerTypeWrapper::Sequence(inner) => inner.iter().try_for_each(|n| {
+                n.read()
+                    .map_err(|_| PyException::new_err("RwLock synchronisation primitive is poisoned, cannot get subtype of PyNormalizer"))?
+                    .normalize(normalized)
+            }),
         }
     }
 }
@@ -793,18 +883,14 @@ mod test {
         let normalizer: PyNormalizer = serde_json::from_str(&sequence_string).unwrap();
 
         match normalizer.normalizer {
-            PyNormalizerTypeWrapper::Single(inner) => match &*inner.as_ref().read().unwrap() {
-                PyNormalizerWrapper::Wrapped(NormalizerWrapper::Sequence(sequence)) => {
-                    let normalizers = sequence.get_normalizers();
-                    assert_eq!(normalizers.len(), 1);
-                    match normalizers[0] {
-                        NormalizerWrapper::NFKC(_) => {}
-                        _ => panic!("Expected NFKC"),
-                    }
-                }
-                _ => panic!("Expected sequence"),
-            },
-            _ => panic!("Expected single"),
+            PyNormalizerTypeWrapper::Sequence(inner) => {
+                assert_eq!(inner.len(), 1);
+                match *inner[0].as_ref().read().unwrap() {
+                    PyNormalizerWrapper::Wrapped(NormalizerWrapper::NFKC(_)) => {}
+                    _ => panic!("Expected NFKC"),
+                };
+            }
+            _ => panic!("Expected sequence"),
         };
     }
 }