tokenizers 0.5.3 → 0.5.4

data/ext/tokenizers/src/pre_tokenizers.rs

@@ -1,8 +1,8 @@
 use std::sync::{Arc, RwLock};
 
 use magnus::{
-    data_type_builder, exception, function, method, value::Lazy, Class, DataType, DataTypeFunctions, Error, Module, Object,
-    RArray, RClass, RModule, Ruby, TryConvert, TypedData,
+    data_type_builder, exception, function, method, value::Lazy, Class, DataType,
+    DataTypeFunctions, Error, Module, Object, RArray, RClass, RModule, Ruby, TryConvert, TypedData,
 };
 
 use serde::ser::SerializeStruct;
@@ -22,7 +22,7 @@ use tk::tokenizer::Offsets;
 use tk::{PreTokenizedString, PreTokenizer};
 
 use super::utils::*;
-use super::{PRE_TOKENIZERS, RbError, RbResult};
+use super::{RbError, RbResult, PRE_TOKENIZERS};
 
 #[derive(DataTypeFunctions, Clone, Serialize, Deserialize)]
 pub struct RbPreTokenizer {
@@ -34,7 +34,9 @@ impl RbPreTokenizer {
     fn pre_tokenize_str(&self, s: String) -> RbResult<Vec<(String, Offsets)>> {
         let mut pretokenized = tk::tokenizer::PreTokenizedString::from(s);
 
-        self.pretok.pre_tokenize(&mut pretokenized).map_err(RbError::from)?;
+        self.pretok
+            .pre_tokenize(&mut pretokenized)
+            .map_err(RbError::from)?;
 
         Ok(pretokenized
             .get_splits(tk::OffsetReferential::Original, tk::OffsetType::Char)
@@ -195,11 +197,7 @@ impl RbDigits {
 pub struct RbMetaspace {}
 
 impl RbMetaspace {
-    fn new(
-        replacement: char,
-        prepend_scheme: String,
-        split: bool,
-    ) -> RbResult<RbPreTokenizer> {
+    fn new(replacement: char, prepend_scheme: String, split: bool) -> RbResult<RbPreTokenizer> {
         let prepend_scheme = from_string(prepend_scheme)?;
         Ok(Metaspace::new(replacement, prepend_scheme, split).into())
     }
@@ -216,8 +214,14 @@ impl RbPunctuation {
 pub struct RbSplit {}
 
 impl RbSplit {
-    pub fn new(pattern: RbPattern, behavior: RbSplitDelimiterBehavior, invert: bool) -> RbResult<RbPreTokenizer> {
-        Split::new(pattern, behavior.into(), invert).map(|v| v.into()).map_err(RbError::from)
+    pub fn new(
+        pattern: RbPattern,
+        behavior: RbSplitDelimiterBehavior,
+        invert: bool,
+    ) -> RbResult<RbPreTokenizer> {
+        Split::new(pattern, behavior.into(), invert)
+            .map(|v| v.into())
+            .map_err(RbError::from)
     }
 }
 
@@ -258,16 +262,18 @@ pub struct RbSequence {}
 impl RbSequence {
     fn new(pre_tokenizers: RArray) -> RbResult<RbPreTokenizer> {
         let mut sequence = Vec::with_capacity(pre_tokenizers.len());
-        for n in pre_tokenizers.into_iter() {
+        for n in pre_tokenizers {
             let pretokenizer: &RbPreTokenizer = TryConvert::try_convert(n)?;
             match &pretokenizer.pretok {
                 RbPreTokenizerTypeWrapper::Sequence(inner) => {
-                    sequence.extend(inner.iter().cloned())
+                    sequence.extend(inner.iter().cloned());
                 }
                 RbPreTokenizerTypeWrapper::Single(inner) => sequence.push(inner.clone()),
             }
         }
-        Ok(RbPreTokenizer::new(RbPreTokenizerTypeWrapper::Sequence(sequence)))
+        Ok(RbPreTokenizer::new(RbPreTokenizerTypeWrapper::Sequence(
+            sequence,
+        )))
     }
 }
 
@@ -277,10 +283,13 @@ pub(crate) fn from_string(string: String) -> RbResult<PrependScheme> {
         "never" => PrependScheme::Never,
         "always" => PrependScheme::Always,
         _ => {
-            return Err(Error::new(exception::arg_error(), format!(
-                "{} is an unknown variant, should be one of ['first', 'never', 'always']",
-                string
-            )));
+            return Err(Error::new(
+                exception::arg_error(),
+                format!(
+                    "{} is an unknown variant, should be one of ['first', 'never', 'always']",
+                    string
+                ),
+            ));
         }
     };
     Ok(scheme)
@@ -381,7 +390,10 @@ impl PreTokenizer for RbPreTokenizerWrapper {
 unsafe impl TypedData for RbPreTokenizer {
     fn class(ruby: &Ruby) -> RClass {
         static CLASS: Lazy<RClass> = Lazy::new(|ruby| {
-            let class: RClass = ruby.get_inner(&PRE_TOKENIZERS).const_get("PreTokenizer").unwrap();
+            let class: RClass = ruby
+                .get_inner(&PRE_TOKENIZERS)
+                .const_get("PreTokenizer")
+                .unwrap();
             class.undef_default_alloc_func();
             class
         });
@@ -389,28 +401,41 @@ unsafe impl TypedData for RbPreTokenizer {
     }
 
     fn data_type() -> &'static DataType {
-        static DATA_TYPE: DataType = data_type_builder!(RbPreTokenizer, "Tokenizers::PreTokenizers::PreTokenizer").build();
+        static DATA_TYPE: DataType =
+            data_type_builder!(RbPreTokenizer, "Tokenizers::PreTokenizers::PreTokenizer").build();
         &DATA_TYPE
     }
 
     fn class_for(ruby: &Ruby, value: &Self) -> RClass {
         static SEQUENCE: Lazy<RClass> = Lazy::new(|ruby| {
-            let class: RClass = ruby.get_inner(&PRE_TOKENIZERS).const_get("Sequence").unwrap();
+            let class: RClass = ruby
+                .get_inner(&PRE_TOKENIZERS)
+                .const_get("Sequence")
+                .unwrap();
             class.undef_default_alloc_func();
             class
         });
         static BERT_PRE_TOKENIZER: Lazy<RClass> = Lazy::new(|ruby| {
-            let class: RClass = ruby.get_inner(&PRE_TOKENIZERS).const_get("BertPreTokenizer").unwrap();
+            let class: RClass = ruby
+                .get_inner(&PRE_TOKENIZERS)
+                .const_get("BertPreTokenizer")
+                .unwrap();
             class.undef_default_alloc_func();
             class
         });
         static BYTE_LEVEL: Lazy<RClass> = Lazy::new(|ruby| {
-            let class: RClass = ruby.get_inner(&PRE_TOKENIZERS).const_get("ByteLevel").unwrap();
+            let class: RClass = ruby
+                .get_inner(&PRE_TOKENIZERS)
+                .const_get("ByteLevel")
+                .unwrap();
             class.undef_default_alloc_func();
             class
         });
         static CHAR_DELIMITER_SPLIT: Lazy<RClass> = Lazy::new(|ruby| {
-            let class: RClass = ruby.get_inner(&PRE_TOKENIZERS).const_get("CharDelimiterSplit").unwrap();
+            let class: RClass = ruby
+                .get_inner(&PRE_TOKENIZERS)
+                .const_get("CharDelimiterSplit")
+                .unwrap();
             class.undef_default_alloc_func();
             class
         });
@@ -420,12 +445,18 @@ unsafe impl TypedData for RbPreTokenizer {
             class
         });
         static METASPACE: Lazy<RClass> = Lazy::new(|ruby| {
-            let class: RClass = ruby.get_inner(&PRE_TOKENIZERS).const_get("Metaspace").unwrap();
+            let class: RClass = ruby
+                .get_inner(&PRE_TOKENIZERS)
+                .const_get("Metaspace")
+                .unwrap();
             class.undef_default_alloc_func();
             class
         });
         static PUNCTUATION: Lazy<RClass> = Lazy::new(|ruby| {
-            let class: RClass = ruby.get_inner(&PRE_TOKENIZERS).const_get("Punctuation").unwrap();
+            let class: RClass = ruby
+                .get_inner(&PRE_TOKENIZERS)
+                .const_get("Punctuation")
+                .unwrap();
             class.undef_default_alloc_func();
             class
         });
@@ -435,17 +466,26 @@ unsafe impl TypedData for RbPreTokenizer {
             class
         });
         static UNICODE_SCRIPTS: Lazy<RClass> = Lazy::new(|ruby| {
-            let class: RClass = ruby.get_inner(&PRE_TOKENIZERS).const_get("UnicodeScripts").unwrap();
+            let class: RClass = ruby
+                .get_inner(&PRE_TOKENIZERS)
+                .const_get("UnicodeScripts")
+                .unwrap();
             class.undef_default_alloc_func();
             class
         });
         static WHITESPACE: Lazy<RClass> = Lazy::new(|ruby| {
-            let class: RClass = ruby.get_inner(&PRE_TOKENIZERS).const_get("Whitespace").unwrap();
+            let class: RClass = ruby
+                .get_inner(&PRE_TOKENIZERS)
+                .const_get("Whitespace")
+                .unwrap();
             class.undef_default_alloc_func();
             class
         });
         static WHITESPACE_SPLIT: Lazy<RClass> = Lazy::new(|ruby| {
-            let class: RClass = ruby.get_inner(&PRE_TOKENIZERS).const_get("WhitespaceSplit").unwrap();
+            let class: RClass = ruby
+                .get_inner(&PRE_TOKENIZERS)
+                .const_get("WhitespaceSplit")
+                .unwrap();
             class.undef_default_alloc_func();
             class
         });
@@ -472,7 +512,10 @@ unsafe impl TypedData for RbPreTokenizer {
 
 pub fn init_pre_tokenizers(ruby: &Ruby, module: &RModule) -> RbResult<()> {
     let pre_tokenizer = module.define_class("PreTokenizer", ruby.class_object())?;
-    pre_tokenizer.define_method("pre_tokenize_str", method!(RbPreTokenizer::pre_tokenize_str, 1))?;
+    pre_tokenizer.define_method(
+        "pre_tokenize_str",
+        method!(RbPreTokenizer::pre_tokenize_str, 1),
+    )?;
 
     let class = module.define_class("Sequence", pre_tokenizer)?;
     class.define_singleton_method("new", function!(RbSequence::new, 1))?;
@@ -483,27 +526,63 @@ pub fn init_pre_tokenizers(ruby: &Ruby, module: &RModule) -> RbResult<()> {
     let class = module.define_class("ByteLevel", pre_tokenizer)?;
     class.define_singleton_method("_new", function!(RbByteLevel::new, 2))?;
     class.define_singleton_method("alphabet", function!(RbByteLevel::alphabet, 0))?;
-    class.define_method("add_prefix_space", method!(RbPreTokenizer::byte_level_add_prefix_space, 0))?;
-    class.define_method("add_prefix_space=", method!(RbPreTokenizer::byte_level_set_add_prefix_space, 1))?;
-    class.define_method("use_regex", method!(RbPreTokenizer::byte_level_use_regex, 0))?;
-    class.define_method("use_regex=", method!(RbPreTokenizer::byte_level_set_use_regex, 1))?;
+    class.define_method(
+        "add_prefix_space",
+        method!(RbPreTokenizer::byte_level_add_prefix_space, 0),
+    )?;
+    class.define_method(
+        "add_prefix_space=",
+        method!(RbPreTokenizer::byte_level_set_add_prefix_space, 1),
+    )?;
+    class.define_method(
+        "use_regex",
+        method!(RbPreTokenizer::byte_level_use_regex, 0),
+    )?;
+    class.define_method(
+        "use_regex=",
+        method!(RbPreTokenizer::byte_level_set_use_regex, 1),
+    )?;
 
     let class = module.define_class("CharDelimiterSplit", pre_tokenizer)?;
     class.define_singleton_method("new", function!(RbCharDelimiterSplit::new, 1))?;
-    class.define_method("delimiter", method!(RbPreTokenizer::char_delimiter_split_delimiter, 0))?;
-    class.define_method("delimiter=", method!(RbPreTokenizer::char_delimiter_split_set_delimiter, 1))?;
+    class.define_method(
+        "delimiter",
+        method!(RbPreTokenizer::char_delimiter_split_delimiter, 0),
+    )?;
+    class.define_method(
+        "delimiter=",
+        method!(RbPreTokenizer::char_delimiter_split_set_delimiter, 1),
+    )?;
 
     let class = module.define_class("Digits", pre_tokenizer)?;
     class.define_singleton_method("_new", function!(RbDigits::new, 1))?;
-    class.define_method("individual_digits", method!(RbPreTokenizer::digits_individual_digits, 0))?;
-    class.define_method("individual_digits=", method!(RbPreTokenizer::digits_set_individual_digits, 1))?;
+    class.define_method(
+        "individual_digits",
+        method!(RbPreTokenizer::digits_individual_digits, 0),
+    )?;
+    class.define_method(
+        "individual_digits=",
+        method!(RbPreTokenizer::digits_set_individual_digits, 1),
+    )?;
 
     let class = module.define_class("Metaspace", pre_tokenizer)?;
     class.define_singleton_method("_new", function!(RbMetaspace::new, 3))?;
-    class.define_method("prepend_scheme", method!(RbPreTokenizer::metaspace_prepend_scheme, 0))?;
-    class.define_method("prepend_scheme=", method!(RbPreTokenizer::metaspace_set_prepend_scheme, 1))?;
-    class.define_method("replacement", method!(RbPreTokenizer::metaspace_replacement, 0))?;
-    class.define_method("replacement=", method!(RbPreTokenizer::metaspace_set_replacement, 1))?;
+    class.define_method(
+        "prepend_scheme",
+        method!(RbPreTokenizer::metaspace_prepend_scheme, 0),
+    )?;
+    class.define_method(
+        "prepend_scheme=",
+        method!(RbPreTokenizer::metaspace_set_prepend_scheme, 1),
+    )?;
+    class.define_method(
+        "replacement",
+        method!(RbPreTokenizer::metaspace_replacement, 0),
+    )?;
+    class.define_method(
+        "replacement=",
+        method!(RbPreTokenizer::metaspace_set_replacement, 1),
+    )?;
     class.define_method("split", method!(RbPreTokenizer::metaspace_split, 0))?;
     class.define_method("split=", method!(RbPreTokenizer::metaspace_set_split, 1))?;
 

data/ext/tokenizers/src/processors.rs

@@ -1,8 +1,8 @@
 use std::sync::Arc;
 
 use magnus::{
-    data_type_builder, function, value::Lazy, Class, DataType, DataTypeFunctions, Module, Object, RClass, RModule,
-    Ruby, TryConvert, TypedData, Value,
+    data_type_builder, function, value::Lazy, Class, DataType, DataTypeFunctions, Module, Object,
+    RClass, RModule, Ruby, TryConvert, TypedData, Value,
 };
 use serde::{Deserialize, Serialize};
 use tk::processors::bert::BertProcessing;
@@ -12,7 +12,7 @@ use tk::processors::template::{SpecialToken, Template};
 use tk::processors::PostProcessorWrapper;
 use tk::{Encoding, PostProcessor};
 
-use super::{PROCESSORS, RbResult};
+use super::{RbResult, PROCESSORS};
 
 #[derive(DataTypeFunctions, Clone, Deserialize, Serialize)]
 pub struct RbPostProcessor {
@@ -106,7 +106,6 @@ impl RbByteLevel {
         }
         RbPostProcessor::new(Arc::new(byte_level.into()))
     }
-
 }
 
 pub struct RbRobertaProcessing {}
@@ -117,7 +116,7 @@ impl RbRobertaProcessing {
         cls: (String, u32),
         trim_offsets: bool,
         add_prefix_space: bool,
-    ) ->  RbPostProcessor {
+    ) -> RbPostProcessor {
         let proc = RobertaProcessing::new(sep, cls)
             .trim_offsets(trim_offsets)
             .add_prefix_space(add_prefix_space);
@@ -153,7 +152,10 @@ impl RbTemplateProcessing {
 unsafe impl TypedData for RbPostProcessor {
     fn class(ruby: &Ruby) -> RClass {
         static CLASS: Lazy<RClass> = Lazy::new(|ruby| {
-            let class: RClass = ruby.get_inner(&PROCESSORS).const_get("PostProcessor").unwrap();
+            let class: RClass = ruby
+                .get_inner(&PROCESSORS)
+                .const_get("PostProcessor")
+                .unwrap();
             class.undef_default_alloc_func();
             class
         });
@@ -161,13 +163,17 @@ unsafe impl TypedData for RbPostProcessor {
     }
 
     fn data_type() -> &'static DataType {
-        static DATA_TYPE: DataType = data_type_builder!(RbPostProcessor, "Tokenizers::Processors::PostProcessor").build();
+        static DATA_TYPE: DataType =
+            data_type_builder!(RbPostProcessor, "Tokenizers::Processors::PostProcessor").build();
         &DATA_TYPE
     }
 
     fn class_for(ruby: &Ruby, value: &Self) -> RClass {
         static BERT_PROCESSING: Lazy<RClass> = Lazy::new(|ruby| {
-            let class: RClass = ruby.get_inner(&PROCESSORS).const_get("BertProcessing").unwrap();
+            let class: RClass = ruby
+                .get_inner(&PROCESSORS)
+                .const_get("BertProcessing")
+                .unwrap();
             class.undef_default_alloc_func();
             class
         });
@@ -177,12 +183,18 @@ unsafe impl TypedData for RbPostProcessor {
             class
         });
         static ROBERTA_PROCESSING: Lazy<RClass> = Lazy::new(|ruby| {
-            let class: RClass = ruby.get_inner(&PROCESSORS).const_get("RobertaProcessing").unwrap();
+            let class: RClass = ruby
+                .get_inner(&PROCESSORS)
+                .const_get("RobertaProcessing")
+                .unwrap();
             class.undef_default_alloc_func();
             class
         });
         static TEMPLATE_PROCESSING: Lazy<RClass> = Lazy::new(|ruby| {
-            let class: RClass = ruby.get_inner(&PROCESSORS).const_get("TemplateProcessing").unwrap();
+            let class: RClass = ruby
+                .get_inner(&PROCESSORS)
+                .const_get("TemplateProcessing")
+                .unwrap();
             class.undef_default_alloc_func();
             class
         });

data/ext/tokenizers/src/tokenizer.rs

@@ -6,8 +6,8 @@ use std::str::FromStr;
 use magnus::prelude::*;
 use magnus::{exception, Error, RArray, RHash, RString, Symbol, TryConvert, Value};
 use tk::tokenizer::{
-    Model, PaddingDirection, PaddingParams, PaddingStrategy,
-    TruncationDirection, TruncationParams, TruncationStrategy, TokenizerImpl
+    Model, PaddingDirection, PaddingParams, PaddingStrategy, TokenizerImpl, TruncationDirection,
+    TruncationParams, TruncationStrategy,
 };
 use tk::AddedToken;
 
@@ -284,7 +284,10 @@ impl RbTokenizer {
     }
 
     pub fn to_str(&self, pretty: bool) -> RbResult<String> {
-        self.tokenizer.borrow().to_string(pretty).map_err(RbError::from)
+        self.tokenizer
+            .borrow()
+            .to_string(pretty)
+            .map_err(RbError::from)
     }
 
     pub fn add_special_tokens(&self, tokens: Vec<String>) -> usize {
@@ -383,7 +386,11 @@ impl RbTokenizer {
             .map_err(RbError::from)
     }
 
-    pub fn decode_batch(&self, sequences: Vec<Vec<u32>>, skip_special_tokens: bool) -> RbResult<Vec<String>> {
+    pub fn decode_batch(
+        &self,
+        sequences: Vec<Vec<u32>>,
+        skip_special_tokens: bool,
+    ) -> RbResult<Vec<String>> {
         let slices = sequences.iter().map(|v| &v[..]).collect::<Vec<&[u32]>>();
         self.tokenizer
             .borrow()
@@ -455,7 +462,12 @@ impl RbTokenizer {
             params.direction = match dir_str.as_str() {
                 "left" => PaddingDirection::Left,
                 "right" => PaddingDirection::Right,
-                _ => return Err(Error::new(exception::arg_error(), "The direction value must be 'left' or 'right'")),
+                _ => {
+                    return Err(Error::new(
+                        exception::arg_error(),
+                        "The direction value must be 'left' or 'right'",
+                    ))
+                }
             }
         }
 
@@ -501,24 +513,27 @@ impl RbTokenizer {
     }
 
     pub fn padding(&self) -> RbResult<Option<RHash>> {
-        self.tokenizer.borrow().get_padding().map_or(Ok(None), |params| {
-            let ret_hash = RHash::new();
-
-            ret_hash.aset(
-                "length",
-                match params.strategy {
-                    tk::PaddingStrategy::BatchLongest => None,
-                    tk::PaddingStrategy::Fixed(size) => Some(size),
-                },
-            )?;
-            ret_hash.aset("pad_to_multiple_of", params.pad_to_multiple_of)?;
-            ret_hash.aset("pad_id", params.pad_id)?;
-            ret_hash.aset("pad_token", &*params.pad_token)?;
-            ret_hash.aset("pad_type_id", params.pad_type_id)?;
-            ret_hash.aset("direction", params.direction.as_ref())?;
-
-            Ok(Some(ret_hash))
-        })
+        self.tokenizer
+            .borrow()
+            .get_padding()
+            .map_or(Ok(None), |params| {
+                let ret_hash = RHash::new();
+
+                ret_hash.aset(
+                    "length",
+                    match params.strategy {
+                        tk::PaddingStrategy::BatchLongest => None,
+                        tk::PaddingStrategy::Fixed(size) => Some(size),
+                    },
+                )?;
+                ret_hash.aset("pad_to_multiple_of", params.pad_to_multiple_of)?;
+                ret_hash.aset("pad_id", params.pad_id)?;
+                ret_hash.aset("pad_token", &*params.pad_token)?;
+                ret_hash.aset("pad_type_id", params.pad_type_id)?;
+                ret_hash.aset("direction", params.direction.as_ref())?;
+
+                Ok(Some(ret_hash))
+            })
     }
 
     pub fn enable_truncation(&self, max_length: usize, kwargs: RHash) -> RbResult<()> {
@@ -539,7 +554,10 @@ impl RbTokenizer {
                 "longest_first" => TruncationStrategy::LongestFirst,
                 "only_first" => TruncationStrategy::OnlyFirst,
                 "only_second" => TruncationStrategy::OnlySecond,
-                _ => return Err(Error::new(exception::arg_error(), "The strategy value must be 'longest_first', 'only_first', or 'only_second'")),
+                _ => return Err(Error::new(
+                    exception::arg_error(),
+                    "The strategy value must be 'longest_first', 'only_first', or 'only_second'",
+                )),
             }
         }
 
@@ -549,7 +567,12 @@ impl RbTokenizer {
             params.direction = match dir_str.as_str() {
                 "left" => TruncationDirection::Left,
                 "right" => TruncationDirection::Right,
-                _ => return Err(Error::new(exception::arg_error(), "The direction value must be 'left' or 'right'")),
+                _ => {
+                    return Err(Error::new(
+                        exception::arg_error(),
+                        "The direction value must be 'left' or 'right'",
+                    ))
+                }
             }
         }
 
@@ -559,7 +582,10 @@ impl RbTokenizer {
         }
 
         if let Err(error_message) = self.tokenizer.borrow_mut().with_truncation(Some(params)) {
-            return Err(Error::new(exception::arg_error(), error_message.to_string()));
+            return Err(Error::new(
+                exception::arg_error(),
+                error_message.to_string(),
+            ));
         }
 
         Ok(())
@@ -573,16 +599,19 @@ impl RbTokenizer {
     }
 
     pub fn truncation(&self) -> RbResult<Option<RHash>> {
-        self.tokenizer.borrow().get_truncation().map_or(Ok(None), |params| {
-            let ret_hash = RHash::new();
+        self.tokenizer
+            .borrow()
+            .get_truncation()
+            .map_or(Ok(None), |params| {
+                let ret_hash = RHash::new();
 
-            ret_hash.aset("max_length", params.max_length)?;
-            ret_hash.aset("stride", params.stride)?;
-            ret_hash.aset("strategy", params.strategy.as_ref())?;
-            ret_hash.aset("direction", params.direction.as_ref())?;
+                ret_hash.aset("max_length", params.max_length)?;
+                ret_hash.aset("stride", params.stride)?;
+                ret_hash.aset("strategy", params.strategy.as_ref())?;
+                ret_hash.aset("direction", params.direction.as_ref())?;
 
-            Ok(Some(ret_hash))
-        })
+                Ok(Some(ret_hash))
+            })
     }
 
     pub fn num_special_tokens_to_add(&self, is_pair: bool) -> usize {