tokenizers 0.2.2 → 0.3.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 197131371ec438d82623bc0aacb8fe82ba255904e847eeb9259358f38a7063f0
-  data.tar.gz: 42ef490120e56fbb79d847ec1eb2b0a6b0ca7aa8f2ad90c09d2053d167491350
+  metadata.gz: 4ff4d1ad7b56010f603ead7a4794c003c5294e50f1b33de62c8089ddf150d5ad
+  data.tar.gz: 295aaabb720971f2ddcc832ab0d5deedf1e0ed8dab03aca96ac1d396b5723de7
 SHA512:
-  metadata.gz: 0a21b4811cc9e31565209eb514e55d6b22302c350371a76205aeb3b67cf94ea6dabf85074cebd48c65f9eca56e8e750b83a1df841807e53afb1275961bca50ce
-  data.tar.gz: 222bb9d759e3a2cc00ad7a4950c821fdbad1bbf6d4413f237bcf9cdc0698c2011022890b3f306be6df3d70b05abd446ad43066851ffa6c27387ddf3191f7557d
+  metadata.gz: e14207004cddeef40590229ea2c8a9bf54e5c5b75cdbcdd32cd6f23c24feb8544fcabe86fa9bced32cb41f2581ee0df4d36ed2b6a58ef2fc668aa33c270659df
+  data.tar.gz: b2bb202c8c37bdd0d14ca64be147e99b224128c1461f56761f1d58d9326b40768e7b903bbbb4c2a0363bd4b1c9ef5a66be53210ad801d5e45c7d86dd0945bd82

data/CHANGELOG.md

@@ -1,3 +1,19 @@
+## 0.3.0 (2022-02-07)
+
+- Added support for training tokenizers
+- Added more methods to `Tokenizer`
+- Added `encode_batch` method to `Encoding`
+- Added `pair` argument to `encode` method
+- Changed `encode` method to include special tokens by default
+- Changed how offsets are calculated for strings with multibyte characters
+
+## 0.2.3 (2022-01-22)
+
+- Added `add_special_tokens` option to `encode` method
+- Added warning about `encode` method including special tokens by default in 0.3.0
+- Added more methods to `Encoding`
+- Fixed error with precompiled gem on Mac ARM
+
 ## 0.2.2 (2022-01-15)
 
 - Added precompiled gem for Linux ARM

data/Cargo.lock

@@ -50,9 +50,9 @@ checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a"
 
 [[package]]
 name = "cc"
-version = "1.0.78"
+version = "1.0.79"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a20104e2335ce8a659d6dd92a51a767a0c062599c73b343fd152cb401e828c3d"
+checksum = "50d30906286121d95be3d479533b458f87493b30a4b5f79a607db8f5d11aa91f"
 
 [[package]]
 name = "cexpr"
@@ -138,9 +138,9 @@ dependencies = [
 
 [[package]]
 name = "darling"
-version = "0.14.2"
+version = "0.14.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b0dd3cd20dc6b5a876612a6e5accfe7f3dd883db6d07acfbf14c128f61550dfa"
+checksum = "c0808e1bd8671fb44a113a14e13497557533369847788fa2ae912b6ebfce9fa8"
 dependencies = [
  "darling_core",
  "darling_macro",
@@ -148,9 +148,9 @@ dependencies = [
 
 [[package]]
 name = "darling_core"
-version = "0.14.2"
+version = "0.14.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a784d2ccaf7c98501746bf0be29b2022ba41fd62a2e622af997a03e9f972859f"
+checksum = "001d80444f28e193f30c2f293455da62dcf9a6b29918a4253152ae2b1de592cb"
 dependencies = [
  "fnv",
  "ident_case",
@@ -162,9 +162,9 @@ dependencies = [
 
 [[package]]
 name = "darling_macro"
-version = "0.14.2"
+version = "0.14.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7618812407e9402654622dd402b0a89dff9ba93badd6540781526117b92aab7e"
+checksum = "b36230598a2d5de7ec1c6f51f72d8a99a9208daff41de2084d06e3fd3ea56685"
 dependencies = [
  "darling_core",
  "quote",
@@ -202,31 +202,11 @@ dependencies = [
  "syn",
 ]
 
-[[package]]
-name = "dirs"
-version = "3.0.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "30baa043103c9d0c2a57cf537cc2f35623889dc0d405e6c3cccfadbc81c71309"
-dependencies = [
- "dirs-sys",
-]
-
-[[package]]
-name = "dirs-sys"
-version = "0.3.7"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "1b1d1d91c932ef41c0f2663aa8b0ca0342d444d842c06914aa0a7e352d0bada6"
-dependencies = [
- "libc",
- "redox_users",
- "winapi",
-]
-
 [[package]]
 name = "either"
-version = "1.8.0"
+version = "1.8.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "90e5c1c8368803113bf0c9584fc495a58b86dc8a29edbf8fe877d21d9507e797"
+checksum = "7fcaabb2fef8c910e7f4c7ce9f67a1283a1715879a7c230ca9d6d1ae31f16d91"
 
 [[package]]
 name = "encode_unicode"
@@ -372,9 +352,8 @@ checksum = "58093314a45e00c77d5c508f76e77c3396afbbc0d01506e7fae47b018bac2b1d"
 
 [[package]]
 name = "magnus"
-version = "0.4.4"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "fc87660cd7daa49fddbfd524c836de54d5c927d520cd163f43700c5087c57d6c"
+version = "0.5.0"
+source = "git+https://github.com/matsadler/magnus#eda735faa7e03da2443eaf2c4058a184917d6b87"
 dependencies = [
  "magnus-macros",
  "rb-sys",
@@ -384,8 +363,7 @@ dependencies = [
 [[package]]
 name = "magnus-macros"
 version = "0.3.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "206cb23bfeea05180c97522ef6a3e52a4eb17b0ed2f30ee3ca9c4f994d2378ae"
+source = "git+https://github.com/matsadler/magnus#eda735faa7e03da2443eaf2c4058a184917d6b87"
 dependencies = [
  "proc-macro2",
  "quote",
@@ -415,9 +393,9 @@ checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a"
 
 [[package]]
 name = "nom"
-version = "7.1.2"
+version = "7.1.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e5507769c4919c998e69e49c839d9dc6e693ede4cc4290d6ad8b41d4f09c548c"
+checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a"
 dependencies = [
  "memchr",
  "minimal-lexical",
@@ -493,9 +471,9 @@ checksum = "5b40af805b3121feab8a3c29f04d8ad262fa8e0561883e7653e024ae4479e6de"
 
 [[package]]
 name = "proc-macro2"
-version = "1.0.49"
+version = "1.0.51"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "57a8eca9f9c4ffde41714334dee777596264c7825420f521abc92b5b5deb63a5"
+checksum = "5d727cae5b39d21da60fa540906919ad737832fe0b1c165da3a34d6548c849d6"
 dependencies = [
  "unicode-ident",
 ]
@@ -562,9 +540,9 @@ dependencies = [
 
 [[package]]
 name = "rayon-core"
-version = "1.10.1"
+version = "1.10.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "cac410af5d00ab6884528b4ab69d1e8e146e8d471201800fa1b4524126de6ad3"
+checksum = "356a0625f1954f730c0201cdab48611198dc6ce21f4acff55089b5a78e6e835b"
 dependencies = [
  "crossbeam-channel",
  "crossbeam-deque",
@@ -574,18 +552,18 @@ dependencies = [
 
 [[package]]
 name = "rb-sys"
-version = "0.9.56"
+version = "0.9.64"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ef82428221475c6f9e7893fe30b88d45ac86bdb12e58e7c92055ba4bceb78a69"
+checksum = "cc8945662df8083245deda89e236647173cc7ad750f481ddcd7bbfd3afe3fa5e"
 dependencies = [
  "rb-sys-build",
 ]
 
 [[package]]
 name = "rb-sys-build"
-version = "0.9.56"
+version = "0.9.64"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "950bfc239d2e7704576abe4d37b008876bbfd70a99196a188c5caeae2ba7344a"
+checksum = "ae8c3cdf9edc3908ee1555b7a1bca58ee1b499439b32cd1c1ec3e66736a8df48"
 dependencies = [
  "bindgen",
  "regex",
@@ -594,29 +572,9 @@ dependencies = [
 
 [[package]]
 name = "rb-sys-env"
-version = "0.1.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "74c38752410925faeb82c400c06ba2fd9ee6aa8f719dd33994c9e53f5242d25f"
-
-[[package]]
-name = "redox_syscall"
-version = "0.2.16"
+version = "0.1.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "fb5a58c1855b4b6819d59012155603f0b22ad30cad752600aadfcb695265519a"
-dependencies = [
- "bitflags",
-]
-
-[[package]]
-name = "redox_users"
-version = "0.4.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b033d837a7cf162d7993aded9304e30a83213c648b6e389db233191f891e5c2b"
-dependencies = [
- "getrandom",
- "redox_syscall",
- "thiserror",
-]
+checksum = "a35802679f07360454b418a5d1735c89716bde01d35b1560fc953c1415a0b3bb"
 
 [[package]]
 name = "regex"
@@ -675,9 +633,9 @@ dependencies = [
 
 [[package]]
 name = "serde_json"
-version = "1.0.91"
+version = "1.0.92"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "877c235533714907a8c2464236f5c4b2a17262ef1bd71f38f35ea592c8da6883"
+checksum = "7434af0dc1cbd59268aa98b4c22c131c0584d2232f6fb166efb993e2832e896a"
 dependencies = [
  "itoa",
  "ryu",
@@ -753,20 +711,21 @@ dependencies = [
 
 [[package]]
 name = "tokenizers"
-version = "0.2.2"
+version = "0.2.3"
 dependencies = [
  "magnus",
+ "onig",
+ "serde",
  "tokenizers 0.13.2",
 ]
 
 [[package]]
 name = "tokenizers"
 version = "0.13.2"
-source = "git+https://github.com/huggingface/tokenizers#fe4ae7dc38be11a5c93ae703816c869f993c21ab"
+source = "git+https://github.com/huggingface/tokenizers#fa66caf0abff16bae2213658ffa3e969c5445750"
 dependencies = [
  "aho-corasick",
  "derive_builder",
- "dirs",
  "esaxx-rs",
  "getrandom",
  "indicatif",
@@ -807,9 +766,9 @@ dependencies = [
 
 [[package]]
 name = "unicode-segmentation"
-version = "1.10.0"
+version = "1.10.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "0fdbf052a0783de01e944a6ce7a8cb939e295b1e7be835a1112c3b9a7f047a5a"
+checksum = "1dd624098567895118886609431a7c3b8f516e41d30e0643f03d94592a147e36"
 
 [[package]]
 name = "unicode-width"

data/README.md

@@ -40,6 +40,10 @@ Load a tokenizer from files
 tokenizer = Tokenizers::CharBPETokenizer.new("vocab.json", "merges.txt")
 ```
 
+## Training
+
+Check out the [Quicktour](https://huggingface.co/docs/tokenizers/quicktour) and equivalent [Ruby code](https://github.com/ankane/tokenizers-ruby/blob/master/test/quicktour_test.rb#L8)
+
 ## History
 
 View the [changelog](https://github.com/ankane/tokenizers-ruby/blob/master/CHANGELOG.md)

data/ext/tokenizers/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "tokenizers"
-version = "0.2.2"
+version = "0.2.3"
 license = "Apache-2.0"
 authors = ["Andrew Kane <andrew@ankane.org>"]
 edition = "2021"
@@ -10,7 +10,9 @@ publish = false
 crate-type = ["cdylib"]
 
 [dependencies]
-magnus = "0.4"
+magnus = { git = "https://github.com/matsadler/magnus" }
+onig = { version = "6.0", default-features = false }
+serde = { version = "1.0", features = ["rc", "derive"] }
 
 [dependencies.tokenizers]
 version = "0.13.2" # also update in from_pretrained.rb

data/ext/tokenizers/src/decoders.rs

@@ -1,14 +1,283 @@
+use std::sync::{Arc, RwLock};
+
+use magnus::typed_data::DataTypeBuilder;
+use magnus::{
+    function, memoize, method, Class, DataType, DataTypeFunctions, Module, Object, RClass, RModule,
+    TypedData,
+};
+use serde::{Deserialize, Serialize};
 use tk::decoders::bpe::BPEDecoder;
+use tk::decoders::byte_level::ByteLevel;
+use tk::decoders::ctc::CTC;
+use tk::decoders::metaspace::Metaspace;
+use tk::decoders::wordpiece::WordPiece;
+use tk::decoders::DecoderWrapper;
+use tk::Decoder;
+
+use super::RbResult;
+
+#[derive(DataTypeFunctions, Clone, Deserialize, Serialize)]
+pub struct RbDecoder {
+    #[serde(flatten)]
+    pub(crate) decoder: RbDecoderWrapper,
+}
+
+impl Decoder for RbDecoder {
+    fn decode_chain(&self, tokens: Vec<String>) -> tk::Result<Vec<String>> {
+        self.decoder.decode_chain(tokens)
+    }
+}
+
+macro_rules! getter {
+    ($self: ident, $variant: ident, $($name: tt)+) => {{
+        let decoder = &$self.decoder;
+        let RbDecoderWrapper::Wrapped(ref wrap) = decoder;
+        if let DecoderWrapper::$variant(ref dec) = *wrap.read().unwrap() {
+            dec.$($name)+
+        } else {
+            unreachable!()
+        }
+    }};
+}
 
-#[magnus::wrap(class = "Tokenizers::BPEDecoder")]
-pub struct RbBPEDecoder {
-    pub decoder: BPEDecoder,
+macro_rules! setter {
+    ($self: ident, $variant: ident, $name: ident, $value: expr) => {{
+        let decoder = &$self.decoder;
+        let RbDecoderWrapper::Wrapped(ref wrap) = decoder;
+        if let DecoderWrapper::$variant(ref mut dec) = *wrap.write().unwrap() {
+            dec.$name = $value;
+        }
+    }};
+    ($self: ident, $variant: ident, @$name: ident, $value: expr) => {{
+        let decoder = &$self.decoder;
+        let RbDecoderWrapper::Wrapped(ref wrap) = decoder;
+        if let DecoderWrapper::$variant(ref mut dec) = *wrap.write().unwrap() {
+            dec.$name($value);
+        }
+    }};
 }
+impl RbDecoder {
+    pub fn bpe_suffix(&self) -> String {
+        getter!(self, BPE, suffix.clone())
+    }
+
+    pub fn bpe_set_suffix(&self, suffix: String) {
+        setter!(self, BPE, suffix, suffix);
+    }
+
+    pub fn ctc_cleanup(&self) -> bool {
+        getter!(self, CTC, cleanup)
+    }
+
+    pub fn ctc_set_cleanup(&self, cleanup: bool) {
+        setter!(self, CTC, cleanup, cleanup);
+    }
+
+    pub fn ctc_pad_token(&self) -> String {
+        getter!(self, CTC, pad_token.clone())
+    }
+
+    pub fn ctc_set_pad_token(&self, pad_token: String) {
+        setter!(self, CTC, pad_token, pad_token);
+    }
+
+    pub fn ctc_word_delimiter_token(&self) -> String {
+        getter!(self, CTC, word_delimiter_token.clone())
+    }
+
+    pub fn ctc_set_word_delimiter_token(&self, word_delimiter_token: String) {
+        setter!(self, CTC, word_delimiter_token, word_delimiter_token);
+    }
+
+    pub fn metaspace_replacement(&self) -> char {
+        getter!(self, Metaspace, get_replacement().clone())
+    }
+
+    pub fn metaspace_set_replacement(&self, replacement: char) {
+        setter!(self, Metaspace, @set_replacement, replacement);
+    }
+
+    pub fn metaspace_add_prefix_space(&self) -> bool {
+        getter!(self, Metaspace, add_prefix_space)
+    }
+
+    pub fn metaspace_set_add_prefix_space(&self, add_prefix_space: bool) {
+        setter!(self, Metaspace, add_prefix_space, add_prefix_space);
+    }
+
+    pub fn word_piece_cleanup(&self) -> bool {
+        getter!(self, WordPiece, cleanup)
+    }
+
+    pub fn word_piece_set_cleanup(&self, cleanup: bool) {
+        setter!(self, WordPiece, cleanup, cleanup);
+    }
+
+    pub fn word_piece_prefix(&self) -> String {
+        getter!(self, WordPiece, prefix.clone())
+    }
+
+    pub fn word_piece_set_prefix(&self, prefix: String) {
+        setter!(self, WordPiece, prefix, prefix);
+    }
+}
+
+pub struct RbBPEDecoder {}
 
 impl RbBPEDecoder {
-    pub fn new() -> Self {
-        RbBPEDecoder {
-            decoder: BPEDecoder::default(),
+    pub fn new(suffix: String) -> RbDecoder {
+        BPEDecoder::new(suffix).into()
+    }
+}
+
+pub struct RbByteLevelDecoder {}
+
+impl RbByteLevelDecoder {
+    pub fn new() -> RbDecoder {
+        ByteLevel::default().into()
+    }
+}
+
+pub struct RbCTC {}
+
+impl RbCTC {
+    pub fn new(pad_token: String, word_delimiter_token: String, cleanup: bool) -> RbDecoder {
+        CTC::new(pad_token, word_delimiter_token, cleanup).into()
+    }
+}
+
+pub struct RbMetaspaceDecoder {}
+
+impl RbMetaspaceDecoder {
+    pub fn new(replacement: char, add_prefix_space: bool) -> RbDecoder {
+        Metaspace::new(replacement, add_prefix_space).into()
+    }
+}
+
+pub struct RbWordPieceDecoder {}
+
+impl RbWordPieceDecoder {
+    pub fn new(prefix: String, cleanup: bool) -> RbDecoder {
+        WordPiece::new(prefix, cleanup).into()
+    }
+}
+
+#[derive(Clone, Deserialize, Serialize)]
+#[serde(untagged)]
+pub(crate) enum RbDecoderWrapper {
+    // Custom(Arc<RwLock<CustomDecoder>>),
+    Wrapped(Arc<RwLock<DecoderWrapper>>),
+}
+
+impl<I> From<I> for RbDecoderWrapper
+where
+    I: Into<DecoderWrapper>,
+{
+    fn from(norm: I) -> Self {
+        RbDecoderWrapper::Wrapped(Arc::new(RwLock::new(norm.into())))
+    }
+}
+
+impl<I> From<I> for RbDecoder
+where
+    I: Into<DecoderWrapper>,
+{
+    fn from(dec: I) -> Self {
+        RbDecoder {
+            decoder: dec.into().into(),
+        }
+    }
+}
+
+impl Decoder for RbDecoderWrapper {
+    fn decode_chain(&self, tokens: Vec<String>) -> tk::Result<Vec<String>> {
+        match self {
+            RbDecoderWrapper::Wrapped(inner) => inner.read().unwrap().decode_chain(tokens),
+            // RbDecoderWrapper::Custom(inner) => inner.read().unwrap().decode_chain(tokens),
+        }
+    }
+}
+
+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 data_type() -> &'static DataType {
+        memoize!(DataType: DataTypeBuilder::<RbDecoder>::new("Tokenizers::Decoders::Decoder").build())
+    }
+
+    fn class_for(value: &Self) -> RClass {
+        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
+                }),
+                _ => todo!(),
+            },
         }
     }
 }
+
+pub fn decoders(module: &RModule) -> RbResult<()> {
+    let decoder = module.define_class("Decoder", Default::default())?;
+
+    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("ByteLevel", decoder)?;
+    class.define_singleton_method("new", function!(RbByteLevelDecoder::new, 0))?;
+
+    let class = module.define_class("CTC", decoder)?;
+    class.define_singleton_method("_new", function!(RbCTC::new, 3))?;
+    class.define_method("cleanup", method!(RbDecoder::ctc_cleanup, 0))?;
+    class.define_method("cleanup=", method!(RbDecoder::ctc_set_cleanup, 1))?;
+    class.define_method("pad_token", method!(RbDecoder::ctc_pad_token, 0))?;
+    class.define_method("pad_token=", method!(RbDecoder::ctc_set_pad_token, 1))?;
+    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("Metaspace", decoder)?;
+    class.define_singleton_method("_new", function!(RbMetaspaceDecoder::new, 2))?;
+    class.define_method("add_prefix_space", method!(RbDecoder::metaspace_add_prefix_space, 0))?;
+    class.define_method("add_prefix_space=", method!(RbDecoder::metaspace_set_add_prefix_space, 1))?;
+    class.define_method("replacement", method!(RbDecoder::metaspace_replacement, 0))?;
+    class.define_method("replacement=", method!(RbDecoder::metaspace_set_replacement, 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))?;
+    class.define_method("cleanup=", method!(RbDecoder::word_piece_set_cleanup, 1))?;
+    class.define_method("prefix", method!(RbDecoder::word_piece_prefix, 0))?;
+    class.define_method("prefix=", method!(RbDecoder::word_piece_set_prefix, 1))?;
+
+    Ok(())
+}

data/ext/tokenizers/src/encoding.rs

@@ -1,16 +1,91 @@
-use tk::Encoding;
+use magnus::RArray;
+use tk::{Encoding, Offsets};
 
 #[magnus::wrap(class = "Tokenizers::Encoding")]
+#[repr(transparent)]
 pub struct RbEncoding {
     pub encoding: Encoding,
 }
 
+impl From<Encoding> for RbEncoding {
+    fn from(v: Encoding) -> Self {
+        Self { encoding: v }
+    }
+}
+
 impl RbEncoding {
+    pub fn n_sequences(&self) -> usize {
+        self.encoding.n_sequences()
+    }
+
     pub fn ids(&self) -> Vec<u32> {
-        self.encoding.get_ids().into()
+        self.encoding.get_ids().to_vec()
     }
 
     pub fn tokens(&self) -> Vec<String> {
-        self.encoding.get_tokens().into()
+        self.encoding.get_tokens().to_vec()
+    }
+
+    pub fn word_ids(&self) -> Vec<Option<u32>> {
+        self.encoding.get_word_ids().to_vec()
+    }
+
+    pub fn sequence_ids(&self) -> Vec<Option<usize>> {
+        self.encoding.get_sequence_ids()
+    }
+
+    pub fn type_ids(&self) -> Vec<u32> {
+        self.encoding.get_type_ids().to_vec()
+    }
+
+    pub fn offsets(&self) -> Vec<(usize, usize)> {
+        self.encoding.get_offsets().to_vec()
+    }
+
+    pub fn special_tokens_mask(&self) -> Vec<u32> {
+        self.encoding.get_special_tokens_mask().to_vec()
+    }
+
+    pub fn attention_mask(&self) -> Vec<u32> {
+        self.encoding.get_attention_mask().to_vec()
+    }
+
+    pub fn overflowing(&self) -> RArray {
+        self.encoding
+            .get_overflowing()
+            .clone()
+            .into_iter()
+            .map(Into::<RbEncoding>::into)
+            .collect()
+    }
+
+    pub fn word_to_tokens(&self, word_index: u32, sequence_index: usize) -> Option<(usize, usize)> {
+        self.encoding.word_to_tokens(word_index, sequence_index)
+    }
+
+    pub fn word_to_chars(&self, word_index: u32, sequence_index: usize) -> Option<Offsets> {
+        self.encoding.word_to_chars(word_index, sequence_index)
+    }
+
+    pub fn token_to_sequence(&self, token_index: usize) -> Option<usize> {
+        self.encoding.token_to_sequence(token_index)
+    }
+
+    pub fn token_to_chars(&self, token_index: usize) -> Option<Offsets> {
+        let (_, offsets) = self.encoding.token_to_chars(token_index)?;
+        Some(offsets)
+    }
+
+    pub fn token_to_word(&self, token_index: usize) -> Option<u32> {
+        let (_, word_idx) = self.encoding.token_to_word(token_index)?;
+        Some(word_idx)
+    }
+
+    pub fn char_to_token(&self, char_pos: usize, sequence_index: usize) -> Option<usize> {
+        self.encoding.char_to_token(char_pos, sequence_index)
+    }
+
+    pub fn char_to_word(&self, char_pos: usize, sequence_index: usize) -> Option<u32> {
+        self.encoding.char_to_word(char_pos, sequence_index)
     }
 }