tokenizers 0.3.3 → 0.4.0

data/ext/tokenizers/src/processors.rs

@@ -1,9 +1,8 @@
 use std::sync::Arc;
 
-use magnus::typed_data::DataTypeBuilder;
 use magnus::{
-    function, memoize, Class, DataType, DataTypeFunctions, Module, Object, RClass, RModule,
-    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;
@@ -13,7 +12,7 @@ use tk::processors::template::{SpecialToken, Template};
 use tk::processors::PostProcessorWrapper;
 use tk::{Encoding, PostProcessor};
 
-use super::RbResult;
+use super::{PROCESSORS, RbResult};
 
 #[derive(DataTypeFunctions, Clone, Deserialize, Serialize)]
 pub struct RbPostProcessor {
@@ -53,9 +52,9 @@ impl From<RbSpecialToken> for SpecialToken {
 
 impl TryConvert for RbSpecialToken {
     fn try_convert(ob: Value) -> RbResult<Self> {
-        if let Ok(v) = ob.try_convert::<(String, u32)>() {
+        if let Ok(v) = <(String, u32)>::try_convert(ob) {
             Ok(Self(v.into()))
-        } else if let Ok(v) = ob.try_convert::<(u32, String)>() {
+        } else if let Ok(v) = <(u32, String)>::try_convert(ob) {
             Ok(Self(v.into()))
         } else {
             todo!()
@@ -74,11 +73,11 @@ impl From<RbTemplate> for Template {
 
 impl TryConvert for RbTemplate {
     fn try_convert(ob: Value) -> RbResult<Self> {
-        if let Ok(s) = ob.try_convert::<String>() {
+        if let Ok(s) = String::try_convert(ob) {
             Ok(Self(
                 s.try_into().unwrap(), //.map_err(RbError::from)?,
             ))
-        } else if let Ok(s) = ob.try_convert::<Vec<String>>() {
+        } else if let Ok(s) = <Vec<String>>::try_convert(ob) {
             Ok(Self(
                 s.try_into().unwrap(), //.map_err(RbError::from)?,
             ))
@@ -152,47 +151,53 @@ impl RbTemplateProcessing {
 }
 
 unsafe impl TypedData for RbPostProcessor {
-    fn class() -> RClass {
-        *memoize!(RClass: {
-          let class: RClass = crate::processors().const_get("PostProcessor").unwrap();
-          class.undef_alloc_func();
-          class
-        })
+    fn class(ruby: &Ruby) -> RClass {
+        static CLASS: Lazy<RClass> = Lazy::new(|ruby| {
+            let class: RClass = ruby.get_inner(&PROCESSORS).const_get("PostProcessor").unwrap();
+            class.undef_default_alloc_func();
+            class
+        });
+        ruby.get_inner(&CLASS)
     }
 
     fn data_type() -> &'static DataType {
-        memoize!(DataType: DataTypeBuilder::<RbPostProcessor>::new("Tokenizers::Processors::PostProcessor").build())
+        static DATA_TYPE: DataType = data_type_builder!(RbPostProcessor, "Tokenizers::Processors::PostProcessor").build();
+        &DATA_TYPE
     }
 
-    fn class_for(value: &Self) -> RClass {
+    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();
+            class.undef_default_alloc_func();
+            class
+        });
+        static BYTE_LEVEL: Lazy<RClass> = Lazy::new(|ruby| {
+            let class: RClass = ruby.get_inner(&PROCESSORS).const_get("ByteLevel").unwrap();
+            class.undef_default_alloc_func();
+            class
+        });
+        static ROBERTA_PROCESSING: Lazy<RClass> = Lazy::new(|ruby| {
+            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();
+            class.undef_default_alloc_func();
+            class
+        });
         match *value.processor {
-            PostProcessorWrapper::Bert(_) => *memoize!(RClass: {
-                let class: RClass = crate::processors().const_get("BertProcessing").unwrap();
-                class.undef_alloc_func();
-                class
-            }),
-            PostProcessorWrapper::ByteLevel(_) => *memoize!(RClass: {
-                let class: RClass = crate::processors().const_get("ByteLevel").unwrap();
-                class.undef_alloc_func();
-                class
-            }),
-            PostProcessorWrapper::Roberta(_) => *memoize!(RClass: {
-                let class: RClass = crate::processors().const_get("RobertaProcessing").unwrap();
-                class.undef_alloc_func();
-                class
-            }),
-            PostProcessorWrapper::Template(_) => *memoize!(RClass: {
-                let class: RClass = crate::processors().const_get("TemplateProcessing").unwrap();
-                class.undef_alloc_func();
-                class
-            }),
+            PostProcessorWrapper::Bert(_) => ruby.get_inner(&BERT_PROCESSING),
+            PostProcessorWrapper::ByteLevel(_) => ruby.get_inner(&BYTE_LEVEL),
+            PostProcessorWrapper::Roberta(_) => ruby.get_inner(&ROBERTA_PROCESSING),
+            PostProcessorWrapper::Template(_) => ruby.get_inner(&TEMPLATE_PROCESSING),
             _ => todo!(),
         }
     }
 }
 
-pub fn processors(module: &RModule) -> RbResult<()> {
-    let post_processor = module.define_class("PostProcessor", Default::default())?;
+pub fn init_processors(ruby: &Ruby, module: &RModule) -> RbResult<()> {
+    let post_processor = module.define_class("PostProcessor", ruby.class_object())?;
 
     let class = module.define_class("BertProcessing", post_processor)?;
     class.define_singleton_method("new", function!(RbBertProcessing::new, 2))?;

data/ext/tokenizers/src/tokenizer.rs

@@ -2,6 +2,7 @@ use std::cell::RefCell;
 use std::collections::HashMap;
 use std::path::PathBuf;
 
+use magnus::prelude::*;
 use magnus::{exception, Error, RArray, RHash, Symbol, TryConvert, Value};
 use tk::tokenizer::{
     Model, PaddingDirection, PaddingParams, PaddingStrategy,
@@ -78,7 +79,7 @@ struct TextInputSequence<'s>(tk::InputSequence<'s>);
 
 impl<'s> TryConvert for TextInputSequence<'s> {
     fn try_convert(ob: Value) -> RbResult<Self> {
-        Ok(Self(ob.try_convert::<String>()?.into()))
+        Ok(Self(String::try_convert(ob)?.into()))
     }
 }
 
@@ -92,7 +93,7 @@ struct RbArrayStr(Vec<String>);
 
 impl TryConvert for RbArrayStr {
     fn try_convert(ob: Value) -> RbResult<Self> {
-        let seq = ob.try_convert::<Vec<String>>()?;
+        let seq = <Vec<String>>::try_convert(ob)?;
         Ok(Self(seq))
     }
 }
@@ -107,7 +108,7 @@ struct PreTokenizedInputSequence<'s>(tk::InputSequence<'s>);
 
 impl<'s> TryConvert for PreTokenizedInputSequence<'s> {
     fn try_convert(ob: Value) -> RbResult<Self> {
-        if let Ok(seq) = ob.try_convert::<RbArrayStr>() {
+        if let Ok(seq) = RbArrayStr::try_convert(ob) {
             return Ok(Self(seq.into()));
         }
         todo!()
@@ -124,14 +125,14 @@ struct TextEncodeInput<'s>(tk::EncodeInput<'s>);
 
 impl<'s> TryConvert for TextEncodeInput<'s> {
     fn try_convert(ob: Value) -> RbResult<Self> {
-        if let Ok(i) = ob.try_convert::<TextInputSequence>() {
+        if let Ok(i) = TextInputSequence::try_convert(ob) {
             return Ok(Self(i.into()));
         }
-        if let Ok((i1, i2)) = ob.try_convert::<(TextInputSequence, TextInputSequence)>() {
+        if let Ok((i1, i2)) = <(TextInputSequence, TextInputSequence)>::try_convert(ob) {
             return Ok(Self((i1, i2).into()));
         }
         // TODO check if this branch is needed
-        if let Ok(arr) = ob.try_convert::<RArray>() {
+        if let Ok(arr) = RArray::try_convert(ob) {
             if arr.len() == 2 {
                 let first = arr.entry::<TextInputSequence>(0).unwrap();
                 let second = arr.entry::<TextInputSequence>(1).unwrap();
@@ -155,16 +156,16 @@ struct PreTokenizedEncodeInput<'s>(tk::EncodeInput<'s>);
 
 impl<'s> TryConvert for PreTokenizedEncodeInput<'s> {
     fn try_convert(ob: Value) -> RbResult<Self> {
-        if let Ok(i) = ob.try_convert::<PreTokenizedInputSequence>() {
+        if let Ok(i) = PreTokenizedInputSequence::try_convert(ob) {
             return Ok(Self(i.into()));
         }
         if let Ok((i1, i2)) =
-            ob.try_convert::<(PreTokenizedInputSequence, PreTokenizedInputSequence)>()
+            <(PreTokenizedInputSequence, PreTokenizedInputSequence)>::try_convert(ob)
         {
             return Ok(Self((i1, i2).into()));
         }
         // TODO check if this branch is needed
-        if let Ok(arr) = ob.try_convert::<RArray>() {
+        if let Ok(arr) = RArray::try_convert(ob) {
             if arr.len() == 2 {
                 let first = arr.entry::<PreTokenizedInputSequence>(0).unwrap();
                 let second = arr.entry::<PreTokenizedInputSequence>(1).unwrap();
@@ -251,16 +252,16 @@ impl RbTokenizer {
         add_special_tokens: bool,
     ) -> RbResult<RbEncoding> {
         let sequence: tk::InputSequence = if is_pretokenized {
-            sequence.try_convert::<PreTokenizedInputSequence>()?.into()
+            PreTokenizedInputSequence::try_convert(sequence)?.into()
         } else {
-            sequence.try_convert::<TextInputSequence>()?.into()
+            TextInputSequence::try_convert(sequence)?.into()
         };
         let input = match pair {
             Some(pair) => {
                 let pair: tk::InputSequence = if is_pretokenized {
-                    pair.try_convert::<PreTokenizedInputSequence>()?.into()
+                    PreTokenizedInputSequence::try_convert(pair)?.into()
                 } else {
-                    pair.try_convert::<TextInputSequence>()?.into()
+                    TextInputSequence::try_convert(pair)?.into()
                 };
                 tk::EncodeInput::Dual(sequence, pair)
             }
@@ -284,9 +285,9 @@ impl RbTokenizer {
             .each()
             .map(|o| {
                 let input: tk::EncodeInput = if is_pretokenized {
-                    o?.try_convert::<PreTokenizedEncodeInput>()?.into()
+                    PreTokenizedEncodeInput::try_convert(o?)?.into()
                 } else {
-                    o?.try_convert::<TextEncodeInput>()?.into()
+                    TextEncodeInput::try_convert(o?)?.into()
                 };
                 Ok(input)
             })
@@ -306,14 +307,15 @@ impl RbTokenizer {
     pub fn decode(&self, ids: Vec<u32>, skip_special_tokens: bool) -> RbResult<String> {
         self.tokenizer
             .borrow()
-            .decode(ids, skip_special_tokens)
+            .decode(&ids, skip_special_tokens)
             .map_err(RbError::from)
     }
 
     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()
-            .decode_batch(sequences, skip_special_tokens)
+            .decode_batch(&slices, skip_special_tokens)
             .map_err(RbError::from)
     }
 
@@ -353,7 +355,7 @@ impl RbTokenizer {
 
         let value: Value = kwargs.delete(Symbol::new("direction"))?;
         if !value.is_nil() {
-            let dir_str: String = value.try_convert()?;
+            let dir_str = String::try_convert(value)?;
             params.direction = match dir_str.as_str() {
                 "left" => PaddingDirection::Left,
                 "right" => PaddingDirection::Right,
@@ -363,29 +365,29 @@ impl RbTokenizer {
 
         let value: Value = kwargs.delete(Symbol::new("pad_to_multiple_of"))?;
         if !value.is_nil() {
-            params.pad_to_multiple_of = value.try_convert()?;
+            params.pad_to_multiple_of = TryConvert::try_convert(value)?;
         }
 
         let value: Value = kwargs.delete(Symbol::new("pad_id"))?;
         if !value.is_nil() {
-            params.pad_id = value.try_convert()?;
+            params.pad_id = TryConvert::try_convert(value)?;
         }
 
         let value: Value = kwargs.delete(Symbol::new("pad_type_id"))?;
         if !value.is_nil() {
-            params.pad_type_id = value.try_convert()?;
+            params.pad_type_id = TryConvert::try_convert(value)?;
         }
 
         let value: Value = kwargs.delete(Symbol::new("pad_token"))?;
         if !value.is_nil() {
-            params.pad_token = value.try_convert()?;
+            params.pad_token = TryConvert::try_convert(value)?;
         }
 
         let value: Value = kwargs.delete(Symbol::new("length"))?;
         if value.is_nil() {
             params.strategy = PaddingStrategy::BatchLongest;
         } else {
-            params.strategy = PaddingStrategy::Fixed(value.try_convert()?);
+            params.strategy = PaddingStrategy::Fixed(TryConvert::try_convert(value)?);
         }
 
         if !kwargs.is_empty() {
@@ -431,12 +433,12 @@ impl RbTokenizer {
 
         let value: Value = kwargs.delete(Symbol::new("stride"))?;
         if !value.is_nil() {
-            params.stride = value.try_convert()?;
+            params.stride = TryConvert::try_convert(value)?;
         }
 
         let value: Value = kwargs.delete(Symbol::new("strategy"))?;
         if !value.is_nil() {
-            let strategy_str: String = value.try_convert()?;
+            let strategy_str = String::try_convert(value)?;
             params.strategy = match strategy_str.as_str() {
                 "longest_first" => TruncationStrategy::LongestFirst,
                 "only_first" => TruncationStrategy::OnlyFirst,
@@ -447,7 +449,7 @@ impl RbTokenizer {
 
         let value: Value = kwargs.delete(Symbol::new("direction"))?;
         if !value.is_nil() {
-            let dir_str: String = value.try_convert()?;
+            let dir_str = String::try_convert(value)?;
             params.direction = match dir_str.as_str() {
                 "left" => TruncationDirection::Left,
                 "right" => TruncationDirection::Right,
@@ -460,13 +462,18 @@ impl RbTokenizer {
             return Err(Error::new(exception::arg_error(), "unknown keyword"));
         }
 
-        self.tokenizer.borrow_mut().with_truncation(Some(params));
+        if let Err(error_message) = self.tokenizer.borrow_mut().with_truncation(Some(params)) {
+            return Err(Error::new(exception::arg_error(), error_message.to_string()));
+        }
 
         Ok(())
     }
 
     pub fn no_truncation(&self) {
-        self.tokenizer.borrow_mut().with_truncation(None);
+        self.tokenizer
+            .borrow_mut()
+            .with_truncation(None)
+            .expect("Failed to set truncation to `None`! This should never happen");
     }
 
     pub fn truncation(&self) -> RbResult<Option<RHash>> {