From 4fb2cb0b2b27cadebef129da1b52111758b929de Mon Sep 17 00:00:00 2001 From: Arthur Date: Tue, 14 Nov 2023 16:11:06 +0100 Subject: [PATCH] allow option string --- bindings/python/src/pre_tokenizers.rs | 18 +++++------------- 1 file changed, 5 insertions(+), 13 deletions(-) diff --git a/bindings/python/src/pre_tokenizers.rs b/bindings/python/src/pre_tokenizers.rs index e5dad363c..53dc30788 100644 --- a/bindings/python/src/pre_tokenizers.rs +++ b/bindings/python/src/pre_tokenizers.rs @@ -523,24 +523,19 @@ impl PyMetaspace { } #[new] - #[pyo3(signature = (replacement = PyChar('▁'), add_prefix_space = true, **_kwargs), text_signature = "(self, replacement=\"_\", add_prefix_space=True)")] + #[pyo3(signature = (replacement = PyChar('▁'), add_prefix_space = true, prepend_scheme=None, **_kwargs), text_signature = "(self, replacement=\"_\", add_prefix_space=True)")] fn new( replacement: PyChar, add_prefix_space: bool, + prepend_scheme: Option, _kwargs: Option<&PyDict>, ) -> PyResult<(Self, PyPreTokenizer)> { // Create a new Metaspace instance let mut new_instance: Metaspace = Metaspace::new(replacement.0, add_prefix_space); - if let Some(prepend_scheme) = _kwargs.and_then(|kwargs| { - kwargs - .get_item("prepend_scheme") - .and_then(|prepend_scheme| prepend_scheme.extract::().ok()) - }) { - // Unwrap the Option before calling to_lowercase - let prepend_scheme_value: String = prepend_scheme.to_lowercase(); - - let prepend_scheme_enum = match prepend_scheme_value.as_str() { + // If a prepend scheme is provided, set it + if let Some(prepend_scheme) = prepend_scheme { + let prepend_scheme_enum = match prepend_scheme.as_str() { "first" => PrependScheme::First, "never" => PrependScheme::Never, "always" => PrependScheme::Always, @@ -552,11 +547,8 @@ impl PyMetaspace { ))); } }; - - // Set the prepend_scheme in the new instance new_instance.set_prepend_scheme(prepend_scheme_enum); } - Ok((PyMetaspace {}, new_instance.into())) } }