tokenizers 0.3.2 → 0.4.0

data/ext/tokenizers/src/decoders.rs

@@ -1,20 +1,25 @@
 use std::sync::{Arc, RwLock};
 
-use magnus::typed_data::DataTypeBuilder;
+use magnus::value::Lazy;
 use magnus::{
-    function, memoize, method, Class, DataType, DataTypeFunctions, Module, Object, RClass, RModule,
-    TypedData,
+    data_type_builder, function, method, Class, DataType, DataTypeFunctions, Module, Object, RClass, RModule,
+    Ruby, TypedData,
 };
 use serde::{Deserialize, Serialize};
 use tk::decoders::bpe::BPEDecoder;
+use tk::decoders::byte_fallback::ByteFallback;
 use tk::decoders::byte_level::ByteLevel;
 use tk::decoders::ctc::CTC;
+use tk::decoders::fuse::Fuse;
 use tk::decoders::metaspace::Metaspace;
+use tk::decoders::strip::Strip;
 use tk::decoders::wordpiece::WordPiece;
 use tk::decoders::DecoderWrapper;
 use tk::Decoder;
+use tk::normalizers::replace::Replace;
 
-use super::RbResult;
+use super::utils::*;
+use super::{DECODERS, RbError, RbResult};
 
 #[derive(DataTypeFunctions, Clone, Deserialize, Serialize)]
 pub struct RbDecoder {
@@ -89,6 +94,30 @@ impl RbDecoder {
         setter!(self, CTC, word_delimiter_token, word_delimiter_token);
     }
 
+    fn strip_content(&self) -> char {
+        getter!(self, Strip, content)
+    }
+
+    fn strip_set_content(&self, content: char) {
+        setter!(self, Strip, content, content)
+    }
+
+    fn strip_start(&self) -> usize {
+        getter!(self, Strip, start)
+    }
+
+    fn strip_set_start(&self, start: usize) {
+        setter!(self, Strip, start, start)
+    }
+
+    fn strip_stop(&self) -> usize {
+        getter!(self, Strip, stop)
+    }
+
+    fn strip_set_stop(&self, stop: usize) {
+        setter!(self, Strip, stop, stop)
+    }
+
     pub fn metaspace_replacement(&self) -> char {
         getter!(self, Metaspace, get_replacement().clone())
     }
@@ -130,6 +159,14 @@ impl RbBPEDecoder {
     }
 }
 
+pub struct RbByteFallbackDecoder {}
+
+impl RbByteFallbackDecoder {
+    pub fn new() -> RbDecoder {
+        ByteFallback::default().into()
+    }
+}
+
 pub struct RbByteLevelDecoder {}
 
 impl RbByteLevelDecoder {
@@ -146,6 +183,14 @@ impl RbCTC {
     }
 }
 
+pub struct RbFuse {}
+
+impl RbFuse {
+    pub fn new() -> RbDecoder {
+        Fuse::default().into()
+    }
+}
+
 pub struct RbMetaspaceDecoder {}
 
 impl RbMetaspaceDecoder {
@@ -154,6 +199,22 @@ impl RbMetaspaceDecoder {
     }
 }
 
+pub struct RbReplaceDecoder {}
+
+impl RbReplaceDecoder {
+    pub fn new(pattern: RbPattern, content: String) -> RbResult<RbDecoder> {
+        Replace::new(pattern, content).map(|v| v.into()).map_err(RbError::from)
+    }
+}
+
+pub struct RbStripDecoder {}
+
+impl RbStripDecoder {
+    pub fn new(content: char, start: usize, stop: usize) -> RbDecoder {
+        Strip::new(content, start, stop).into()
+    }
+}
+
 pub struct RbWordPieceDecoder {}
 
 impl RbWordPieceDecoder {
@@ -199,60 +260,94 @@ impl Decoder for RbDecoderWrapper {
 }
 
 unsafe impl TypedData for RbDecoder {
-    fn class() -> RClass {
-        *memoize!(RClass: {
-          let class: RClass = crate::decoders().const_get("Decoder").unwrap();
-          class.undef_alloc_func();
-          class
-        })
+    fn class(ruby: &Ruby) -> RClass {
+        static CLASS: Lazy<RClass> = Lazy::new(|ruby| {
+            let class: RClass = ruby.get_inner(&DECODERS).const_get("Decoder").unwrap();
+            class.undef_default_alloc_func();
+            class
+        });
+        ruby.get_inner(&CLASS)
     }
 
     fn data_type() -> &'static DataType {
-        memoize!(DataType: DataTypeBuilder::<RbDecoder>::new("Tokenizers::Decoders::Decoder").build())
+        static DATA_TYPE: DataType = data_type_builder!(RbDecoder, "Tokenizers::Decoders::Decoder").build();
+        &DATA_TYPE
     }
 
-    fn class_for(value: &Self) -> RClass {
+    fn class_for(ruby: &Ruby, value: &Self) -> RClass {
+        static BPE_DECODER: Lazy<RClass> = Lazy::new(|ruby| {
+            let class: RClass = ruby.get_inner(&DECODERS).const_get("BPEDecoder").unwrap();
+            class.undef_default_alloc_func();
+            class
+        });
+        static BYTE_FALLBACK: Lazy<RClass> = Lazy::new(|ruby| {
+            let class: RClass = ruby.get_inner(&DECODERS).const_get("ByteFallback").unwrap();
+            class.undef_default_alloc_func();
+            class
+        });
+        static BYTE_LEVEL: Lazy<RClass> = Lazy::new(|ruby| {
+            let class: RClass = ruby.get_inner(&DECODERS).const_get("ByteLevel").unwrap();
+            class.undef_default_alloc_func();
+            class
+        });
+        static CTC: Lazy<RClass> = Lazy::new(|ruby| {
+            let class: RClass = ruby.get_inner(&DECODERS).const_get("CTC").unwrap();
+            class.undef_default_alloc_func();
+            class
+        });
+        static FUSE: Lazy<RClass> = Lazy::new(|ruby| {
+            let class: RClass = ruby.get_inner(&DECODERS).const_get("Fuse").unwrap();
+            class.undef_default_alloc_func();
+            class
+        });
+        static METASPACE: Lazy<RClass> = Lazy::new(|ruby| {
+            let class: RClass = ruby.get_inner(&DECODERS).const_get("Metaspace").unwrap();
+            class.undef_default_alloc_func();
+            class
+        });
+        static REPLACE: Lazy<RClass> = Lazy::new(|ruby| {
+            let class: RClass = ruby.get_inner(&DECODERS).const_get("Replace").unwrap();
+            class.undef_default_alloc_func();
+            class
+        });
+        static STRIP: Lazy<RClass> = Lazy::new(|ruby| {
+            let class: RClass = ruby.get_inner(&DECODERS).const_get("Strip").unwrap();
+            class.undef_default_alloc_func();
+            class
+        });
+        static WORD_PIECE: Lazy<RClass> = Lazy::new(|ruby| {
+            let class: RClass = ruby.get_inner(&DECODERS).const_get("WordPiece").unwrap();
+            class.undef_default_alloc_func();
+            class
+        });
         match &value.decoder {
             RbDecoderWrapper::Wrapped(inner) => match *inner.read().unwrap() {
-                DecoderWrapper::BPE(_) => *memoize!(RClass: {
-                    let class: RClass = crate::decoders().const_get("BPEDecoder").unwrap();
-                    class.undef_alloc_func();
-                    class
-                }),
-                DecoderWrapper::ByteLevel(_) => *memoize!(RClass: {
-                    let class: RClass = crate::decoders().const_get("ByteLevel").unwrap();
-                    class.undef_alloc_func();
-                    class
-                }),
-                DecoderWrapper::CTC(_) => *memoize!(RClass: {
-                    let class: RClass = crate::decoders().const_get("CTC").unwrap();
-                    class.undef_alloc_func();
-                    class
-                }),
-                DecoderWrapper::Metaspace(_) => *memoize!(RClass: {
-                    let class: RClass = crate::decoders().const_get("Metaspace").unwrap();
-                    class.undef_alloc_func();
-                    class
-                }),
-                DecoderWrapper::WordPiece(_) => *memoize!(RClass: {
-                    let class: RClass = crate::decoders().const_get("WordPiece").unwrap();
-                    class.undef_alloc_func();
-                    class
-                }),
+                DecoderWrapper::BPE(_) => ruby.get_inner(&BPE_DECODER),
+                DecoderWrapper::ByteFallback(_) => ruby.get_inner(&BYTE_FALLBACK),
+                DecoderWrapper::ByteLevel(_) => ruby.get_inner(&BYTE_LEVEL),
+                DecoderWrapper::CTC(_) => ruby.get_inner(&CTC),
+                DecoderWrapper::Fuse(_) => ruby.get_inner(&FUSE),
+                DecoderWrapper::Metaspace(_) => ruby.get_inner(&METASPACE),
+                DecoderWrapper::Replace(_) => ruby.get_inner(&REPLACE),
+                DecoderWrapper::Strip(_) => ruby.get_inner(&STRIP),
+                DecoderWrapper::WordPiece(_) => ruby.get_inner(&WORD_PIECE),
                 _ => todo!(),
             },
         }
     }
 }
 
-pub fn decoders(module: &RModule) -> RbResult<()> {
-    let decoder = module.define_class("Decoder", Default::default())?;
+pub fn init_decoders(ruby: &Ruby, module: &RModule) -> RbResult<()> {
+    let decoder = module.define_class("Decoder", ruby.class_object())?;
 
     let class = module.define_class("BPEDecoder", decoder)?;
     class.define_singleton_method("_new", function!(RbBPEDecoder::new, 1))?;
     class.define_method("suffix", method!(RbDecoder::bpe_suffix, 0))?;
     class.define_method("suffix=", method!(RbDecoder::bpe_set_suffix, 1))?;
 
+    let class = module.define_class("ByteFallback", decoder)?;
+    class.define_singleton_method("new", function!(RbByteFallbackDecoder::new, 0))?;
+
     let class = module.define_class("ByteLevel", decoder)?;
     class.define_singleton_method("new", function!(RbByteLevelDecoder::new, 0))?;
 
@@ -265,6 +360,9 @@ pub fn decoders(module: &RModule) -> RbResult<()> {
     class.define_method("word_delimiter_token", method!(RbDecoder::ctc_word_delimiter_token, 0))?;
     class.define_method("word_delimiter_token=", method!(RbDecoder::ctc_set_word_delimiter_token, 1))?;
 
+    let class = module.define_class("Fuse", decoder)?;
+    class.define_singleton_method("new", function!(RbFuse::new, 0))?;
+
     let class = module.define_class("Metaspace", decoder)?;
     class.define_singleton_method("_new", function!(RbMetaspaceDecoder::new, 2))?;
     class.define_method("add_prefix_space", method!(RbDecoder::metaspace_add_prefix_space, 0))?;
@@ -272,6 +370,18 @@ pub fn decoders(module: &RModule) -> RbResult<()> {
     class.define_method("replacement", method!(RbDecoder::metaspace_replacement, 0))?;
     class.define_method("replacement=", method!(RbDecoder::metaspace_set_replacement, 1))?;
 
+    let class = module.define_class("Replace", decoder)?;
+    class.define_singleton_method("new", function!(RbReplaceDecoder::new, 2))?;
+
+    let class = module.define_class("Strip", decoder)?;
+    class.define_singleton_method("_new", function!(RbStripDecoder::new, 3))?;
+    class.define_method("content", method!(RbDecoder::strip_content, 0))?;
+    class.define_method("content=", method!(RbDecoder::strip_set_content, 1))?;
+    class.define_method("start", method!(RbDecoder::strip_start, 0))?;
+    class.define_method("start=", method!(RbDecoder::strip_set_start, 1))?;
+    class.define_method("stop", method!(RbDecoder::strip_stop, 0))?;
+    class.define_method("stop=", method!(RbDecoder::strip_set_stop, 1))?;
+
     let class = module.define_class("WordPiece", decoder)?;
     class.define_singleton_method("_new", function!(RbWordPieceDecoder::new, 2))?;
     class.define_method("cleanup", method!(RbDecoder::word_piece_cleanup, 0))?;

data/ext/tokenizers/src/error.rs

@@ -1,6 +1,6 @@
-use magnus::{memoize, Error, ExceptionClass, Module};
+use magnus::{prelude::*, value::Lazy, Error, ExceptionClass, Ruby};
 
-use super::module;
+use super::TOKENIZERS;
 
 pub struct RbError {}
 
@@ -11,6 +11,8 @@ impl RbError {
     }
 }
 
+static ERROR: Lazy<ExceptionClass> = Lazy::new(|ruby| ruby.get_inner(&TOKENIZERS).const_get("Error").unwrap());
+
 fn error() -> ExceptionClass {
-    *memoize!(ExceptionClass: module().const_get("Error").unwrap())
+    Ruby::get().unwrap().get_inner(&ERROR)
 }

data/ext/tokenizers/src/lib.rs

@@ -1,3 +1,5 @@
+#![allow(clippy::new_ret_no_self)]
+
 extern crate tokenizers as tk;
 
 mod decoders;
@@ -16,43 +18,29 @@ use error::RbError;
 use tokenizer::RbTokenizer;
 use utils::RbRegex;
 
-use magnus::{define_module, function, memoize, method, prelude::*, Error, RModule};
+use magnus::{function, method, prelude::*, value::Lazy, Error, RModule, Ruby};
 
 type RbResult<T> = Result<T, Error>;
 
-fn module() -> RModule {
-    *memoize!(RModule: define_module("Tokenizers").unwrap())
-}
+static TOKENIZERS: Lazy<RModule> = Lazy::new(|ruby| ruby.class_object().const_get("Tokenizers").unwrap());
 
-fn decoders() -> RModule {
-    *memoize!(RModule: module().const_get("Decoders").unwrap())
-}
+static DECODERS: Lazy<RModule> = Lazy::new(|ruby| ruby.get_inner(&TOKENIZERS).const_get("Decoders").unwrap());
 
-fn models() -> RModule {
-    *memoize!(RModule: module().const_get("Models").unwrap())
-}
+static MODELS: Lazy<RModule> = Lazy::new(|ruby| ruby.get_inner(&TOKENIZERS).const_get("Models").unwrap());
 
-fn normalizers() -> RModule {
-    *memoize!(RModule: module().const_get("Normalizers").unwrap())
-}
+static NORMALIZERS: Lazy<RModule> = Lazy::new(|ruby| ruby.get_inner(&TOKENIZERS).const_get("Normalizers").unwrap());
 
-fn pre_tokenizers() -> RModule {
-    *memoize!(RModule: module().const_get("PreTokenizers").unwrap())
-}
+static PRE_TOKENIZERS: Lazy<RModule> = Lazy::new(|ruby| ruby.get_inner(&TOKENIZERS).const_get("PreTokenizers").unwrap());
 
-fn processors() -> RModule {
-    *memoize!(RModule: module().const_get("Processors").unwrap())
-}
+static PROCESSORS: Lazy<RModule> = Lazy::new(|ruby| ruby.get_inner(&TOKENIZERS).const_get("Processors").unwrap());
 
-fn trainers() -> RModule {
-    *memoize!(RModule: module().const_get("Trainers").unwrap())
-}
+static TRAINERS: Lazy<RModule> = Lazy::new(|ruby| ruby.get_inner(&TOKENIZERS).const_get("Trainers").unwrap());
 
 #[magnus::init]
-fn init() -> RbResult<()> {
-    let module = module();
+fn init(ruby: &Ruby) -> RbResult<()> {
+    let module = ruby.get_inner(&TOKENIZERS);
 
-    let class = module.define_class("Tokenizer", Default::default())?;
+    let class = module.define_class("Tokenizer", ruby.class_object())?;
     class.define_singleton_method("new", function!(RbTokenizer::from_model, 1))?;
     class.define_singleton_method("from_file", function!(RbTokenizer::from_file, 1))?;
     class.define_method(
@@ -86,7 +74,7 @@ fn init() -> RbResult<()> {
     class.define_method("_vocab_size", method!(RbTokenizer::vocab_size, 1))?;
     class.define_method("_to_s", method!(RbTokenizer::to_str, 1))?;
 
-    let class = module.define_class("Encoding", Default::default())?;
+    let class = module.define_class("Encoding", ruby.class_object())?;
     class.define_method("n_sequences", method!(RbEncoding::n_sequences, 0))?;
     class.define_method("ids", method!(RbEncoding::ids, 0))?;
     class.define_method("tokens", method!(RbEncoding::tokens, 0))?;
@@ -111,7 +99,7 @@ fn init() -> RbResult<()> {
     class.define_method("_char_to_token", method!(RbEncoding::char_to_token, 2))?;
     class.define_method("_char_to_word", method!(RbEncoding::char_to_word, 2))?;
 
-    let class = module.define_class("Regex", Default::default())?;
+    let class = module.define_class("Regex", ruby.class_object())?;
     class.define_singleton_method("new", function!(RbRegex::new, 1))?;
 
     let models = module.define_module("Models")?;
@@ -121,12 +109,12 @@ fn init() -> RbResult<()> {
     let normalizers = module.define_module("Normalizers")?;
     let trainers = module.define_module("Trainers")?;
 
-    models::models(&models)?;
-    pre_tokenizers::pre_tokenizers(&pre_tokenizers)?;
-    decoders::decoders(&decoders)?;
-    processors::processors(&processors)?;
-    normalizers::normalizers(&normalizers)?;
-    trainers::trainers(&trainers)?;
+    models::init_models(ruby, &models)?;
+    pre_tokenizers::init_pre_tokenizers(ruby, &pre_tokenizers)?;
+    decoders::init_decoders(ruby, &decoders)?;
+    processors::init_processors(ruby, &processors)?;
+    normalizers::init_normalizers(ruby, &normalizers)?;
+    trainers::init_trainers(ruby, &trainers)?;
 
     Ok(())
 }

data/ext/tokenizers/src/models.rs

@@ -3,10 +3,10 @@ use std::path::{Path, PathBuf};
 use std::sync::{Arc, RwLock};
 
 use crate::trainers::RbTrainer;
-use magnus::typed_data::DataTypeBuilder;
+use magnus::prelude::*;
 use magnus::{
-    exception, function, memoize, method, Class, DataType, DataTypeFunctions, Error, Module, Object,
-    RClass, RHash, RModule, Symbol, TypedData, Value,
+    data_type_builder, exception, function, method, value::Lazy, Class, DataType, DataTypeFunctions, Error, Module, Object,
+    RClass, RHash, RModule, Ruby, Symbol, TryConvert, TypedData, Value,
 };
 use serde::{Deserialize, Serialize};
 use tk::models::bpe::{BpeBuilder, Merges, Vocab, BPE};
@@ -16,7 +16,7 @@ use tk::models::wordlevel::WordLevel;
 use tk::models::wordpiece::{WordPiece, WordPieceBuilder};
 use tk::{Model, Token};
 
-use super::{RbError, RbResult};
+use super::{MODELS, RbError, RbResult};
 
 #[derive(DataTypeFunctions, Clone, Serialize, Deserialize)]
 pub struct RbModel {
@@ -73,32 +73,37 @@ impl RbBPE {
     fn with_builder(mut builder: BpeBuilder, kwargs: RHash) -> RbResult<RbModel> {
         let value: Value = kwargs.delete(Symbol::new("cache_capacity"))?;
         if !value.is_nil() {
-            builder = builder.cache_capacity(value.try_convert()?);
+            builder = builder.cache_capacity(TryConvert::try_convert(value)?);
         }
 
         let value: Value = kwargs.delete(Symbol::new("dropout"))?;
         if !value.is_nil() {
-            builder = builder.dropout(value.try_convert()?);
+            builder = builder.dropout(TryConvert::try_convert(value)?);
         }
 
         let value: Value = kwargs.delete(Symbol::new("unk_token"))?;
         if !value.is_nil() {
-            builder = builder.unk_token(value.try_convert()?);
+            builder = builder.unk_token(TryConvert::try_convert(value)?);
         }
 
         let value: Value = kwargs.delete(Symbol::new("continuing_subword_prefix"))?;
         if !value.is_nil() {
-            builder = builder.continuing_subword_prefix(value.try_convert()?);
+            builder = builder.continuing_subword_prefix(TryConvert::try_convert(value)?);
         }
 
         let value: Value = kwargs.delete(Symbol::new("end_of_word_suffix"))?;
         if !value.is_nil() {
-            builder = builder.end_of_word_suffix(value.try_convert()?);
+            builder = builder.end_of_word_suffix(TryConvert::try_convert(value)?);
         }
 
         let value: Value = kwargs.delete(Symbol::new("fuse_unk"))?;
         if !value.is_nil() {
-            builder = builder.fuse_unk(value.try_convert()?);
+            builder = builder.fuse_unk(TryConvert::try_convert(value)?);
+        }
+
+        let value: Value = kwargs.delete(Symbol::new("byte_fallback"))?;
+        if !value.is_nil() {
+            builder = builder.byte_fallback(TryConvert::try_convert(value)?);
         }
 
         if !kwargs.is_empty() {
@@ -169,6 +174,14 @@ impl RbModel {
         setter!(self, BPE, fuse_unk, fuse_unk);
     }
 
+    pub fn bpe_byte_fallback(&self) -> bool {
+        getter!(self, BPE, byte_fallback)
+    }
+
+    pub fn bpe_set_byte_fallback(&self, byte_fallback: bool) {
+        setter!(self, BPE, byte_fallback, byte_fallback);
+    }
+
     pub fn bpe_continuing_subword_prefix(&self) -> Option<String> {
         getter!(self, BPE, continuing_subword_prefix.clone())
     }
@@ -221,13 +234,13 @@ impl RbModel {
 pub struct RbUnigram {}
 
 impl RbUnigram {
-    fn new(vocab: Option<Vec<(String, f64)>>, unk_id: Option<usize>) -> RbResult<RbModel> {
-        match (vocab, unk_id) {
-            (Some(vocab), unk_id) => {
-                let model = Unigram::from(vocab, unk_id).map_err(RbError::from)?;
+    fn new(vocab: Option<Vec<(String, f64)>>, unk_id: Option<usize>, byte_fallback: Option<bool>) -> RbResult<RbModel> {
+        match (vocab, unk_id, byte_fallback) {
+            (Some(vocab), unk_id, byte_fallback) => {
+                let model = Unigram::from(vocab, unk_id, byte_fallback.unwrap_or(false)).map_err(RbError::from)?;
                 Ok(model.into())
             }
-            (None, None) => Ok(Unigram::default().into()),
+            (None, None, _) => Ok(Unigram::default().into()),
             _ => Err(Error::new(exception::arg_error(), "`vocab` and `unk_id` must be both specified")),
         }
     }
@@ -264,17 +277,17 @@ impl RbWordPiece {
     fn with_builder(mut builder: WordPieceBuilder, kwargs: RHash) -> RbResult<RbModel> {
         let value: Value = kwargs.delete(Symbol::new("unk_token"))?;
         if !value.is_nil() {
-            builder = builder.unk_token(value.try_convert()?);
+            builder = builder.unk_token(TryConvert::try_convert(value)?);
         }
 
         let value: Value = kwargs.delete(Symbol::new("max_input_chars_per_word"))?;
         if !value.is_nil() {
-            builder = builder.max_input_chars_per_word(value.try_convert()?);
+            builder = builder.max_input_chars_per_word(TryConvert::try_convert(value)?);
         }
 
         let value: Value = kwargs.delete(Symbol::new("continuing_subword_prefix"))?;
         if !value.is_nil() {
-            builder = builder.continuing_subword_prefix(value.try_convert()?);
+            builder = builder.continuing_subword_prefix(TryConvert::try_convert(value)?);
         }
 
         if !kwargs.is_empty() {
@@ -301,46 +314,52 @@ impl RbWordPiece {
 }
 
 unsafe impl TypedData for RbModel {
-    fn class() -> RClass {
-        *memoize!(RClass: {
-          let class: RClass = crate::models().const_get("Model").unwrap();
-          class.undef_alloc_func();
+    fn class(ruby: &Ruby) -> RClass {
+        static CLASS: Lazy<RClass> = Lazy::new(|ruby| {
+          let class: RClass = ruby.get_inner(&MODELS).const_get("Model").unwrap();
+          class.undef_default_alloc_func();
           class
-        })
+        });
+        ruby.get_inner(&CLASS)
     }
 
     fn data_type() -> &'static DataType {
-        memoize!(DataType: DataTypeBuilder::<RbModel>::new("Tokenizers::Models::Model").build())
-    }
-
-    fn class_for(value: &Self) -> RClass {
+        static DATA_TYPE: DataType = data_type_builder!(RbModel, "Tokenizers::Models::Model").build();
+        &DATA_TYPE
+    }
+
+    fn class_for(ruby: &Ruby, value: &Self) -> RClass {
+        static BPE: Lazy<RClass> = Lazy::new(|ruby| {
+            let class: RClass = ruby.get_inner(&MODELS).const_get("BPE").unwrap();
+            class.undef_default_alloc_func();
+            class
+        });
+        static UNIGRAM: Lazy<RClass> = Lazy::new(|ruby| {
+            let class: RClass = ruby.get_inner(&MODELS).const_get("Unigram").unwrap();
+            class.undef_default_alloc_func();
+            class
+        });
+        static WORD_LEVEL: Lazy<RClass> = Lazy::new(|ruby| {
+            let class: RClass = ruby.get_inner(&MODELS).const_get("WordLevel").unwrap();
+            class.undef_default_alloc_func();
+            class
+        });
+        static WORD_PIECE: Lazy<RClass> = Lazy::new(|ruby| {
+            let class: RClass = ruby.get_inner(&MODELS).const_get("WordPiece").unwrap();
+            class.undef_default_alloc_func();
+            class
+        });
         match *value.model.read().unwrap() {
-            ModelWrapper::BPE(_) => *memoize!(RClass: {
-                let class: RClass = crate::models().const_get("BPE").unwrap();
-                class.undef_alloc_func();
-                class
-            }),
-            ModelWrapper::Unigram(_) => *memoize!(RClass: {
-                let class: RClass = crate::models().const_get("Unigram").unwrap();
-                class.undef_alloc_func();
-                class
-            }),
-            ModelWrapper::WordLevel(_) => *memoize!(RClass: {
-                let class: RClass = crate::models().const_get("WordLevel").unwrap();
-                class.undef_alloc_func();
-                class
-            }),
-            ModelWrapper::WordPiece(_) => *memoize!(RClass: {
-                let class: RClass = crate::models().const_get("WordPiece").unwrap();
-                class.undef_alloc_func();
-                class
-            }),
+            ModelWrapper::BPE(_) => ruby.get_inner(&BPE),
+            ModelWrapper::Unigram(_) => ruby.get_inner(&UNIGRAM),
+            ModelWrapper::WordLevel(_) => ruby.get_inner(&WORD_LEVEL),
+            ModelWrapper::WordPiece(_) => ruby.get_inner(&WORD_PIECE),
         }
     }
 }
 
-pub fn models(module: &RModule) -> RbResult<()> {
-    let model = module.define_class("Model", Default::default())?;
+pub fn init_models(ruby: &Ruby, module: &RModule) -> RbResult<()> {
+    let model = module.define_class("Model", ruby.class_object())?;
 
     let class = module.define_class("BPE", model)?;
     class.define_singleton_method("_new", function!(RbBPE::new, 3))?;
@@ -355,9 +374,11 @@ pub fn models(module: &RModule) -> RbResult<()> {
     class.define_method("end_of_word_suffix=", method!(RbModel::bpe_set_end_of_word_suffix, 1))?;
     class.define_method("fuse_unk", method!(RbModel::bpe_fuse_unk, 0))?;
     class.define_method("fuse_unk=", method!(RbModel::bpe_set_fuse_unk, 1))?;
+    class.define_method("byte_fallback", method!(RbModel::bpe_byte_fallback, 0))?;
+    class.define_method("byte_fallback=", method!(RbModel::bpe_set_byte_fallback, 1))?;
 
     let class = module.define_class("Unigram", model)?;
-    class.define_singleton_method("_new", function!(RbUnigram::new, 2))?;
+    class.define_singleton_method("_new", function!(RbUnigram::new, 3))?;
 
     let class = module.define_class("WordLevel", model)?;
     class.define_singleton_method("_new", function!(RbWordLevel::new, 2))?;