tokenizer-ruby 0.1.1 → 0.2.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 1eea09dc0997579eab072e4fc3705f8e5e5161110126f80478922ce65732d6a5
-  data.tar.gz: 4e6534b235e6477742a5170fc872cca3b59bbb44f32960fcb656a20d70d00c9f
+  metadata.gz: 0bd2566191cac59c7f91316990a1d2bf872cebd006c429de2bc34b5f217e37f0
+  data.tar.gz: aadf725da91f4976a0b393f3021288587e3da099cc9fae3c2a7f49f855db8a05
 SHA512:
-  metadata.gz: 4b7f5258322ef3d76d3a16a889b0e233f0078ed4970fff3b06d63f9d47fe3aeb39690fe370a37d0e52faf08bf808919d2bfe415a22bafa5bc31f38a6bcf70642
-  data.tar.gz: d2702b426ec82453372e9452a103921f025d68d44296e466c57e5b86d041a6c5c6b6558087a779767ac02080089a849cfee5262f863b0777331a42e51bb6b15a
+  metadata.gz: 19d5ca2f2cde59867f53f151ab819e73a18e3efc2cccd4c407ae71162b9689282219f5de51b9ffc017fb5dfdb9dc119acaaeb48f6714252341aafd1b509030d5
+  data.tar.gz: 7c8d23d28db0f55b9e9befab6fa55e111f308f87d56ffb72ce7fcb7192513b9b8f582bfcf8287fb3fb39a7c66abca70bb52e240f60aea07c95c18fbf77b687e4

data/CLAUDE.md

@@ -174,7 +174,7 @@ Follow the same pattern as zvec-ruby:
 
 ## Publishing
 
-- RubyGems.org: `GEM_HOST_API_KEY=rubygems_5d46e91ceb51fb455e98a7f491a2321bb6879f9be35d6842 gem push tokenizer-ruby-*.gem`
+- RubyGems.org: `GEM_HOST_API_KEY=rubygems_b63436a0904abc9d8c595793bcd53901397650205d0daaf0 gem push tokenizer-ruby-*.gem`
 - gem.coop: `GEM_HOST_API_KEY=hjncPswY8PbGDfLPw4RMj928 gem push tokenizer-ruby-*.gem --host https://beta.gem.coop/@johannesdwicahyo`
 
 ## Notes from zvec-ruby Experience

data/ext/tokenizer_ruby/src/lib.rs

@@ -2,22 +2,22 @@ use magnus::{
     define_module, function, method, prelude::*, Error, RHash, Ruby,
     RArray,
 };
-use std::cell::RefCell;
+use std::sync::Mutex;
 use tokenizers::Tokenizer;
 
 #[magnus::wrap(class = "TokenizerRuby::InternalTokenizer", free_immediately)]
-struct RubyTokenizer(RefCell<Tokenizer>);
+struct RubyTokenizer(Mutex<Tokenizer>);
 
 fn from_pretrained(ruby: &Ruby, identifier: String) -> Result<RubyTokenizer, Error> {
     let tokenizer = Tokenizer::from_pretrained(&identifier, None)
         .map_err(|e| Error::new(ruby.exception_runtime_error(), format!("{}", e)))?;
-    Ok(RubyTokenizer(RefCell::new(tokenizer)))
+    Ok(RubyTokenizer(Mutex::new(tokenizer)))
 }
 
 fn from_file(ruby: &Ruby, path: String) -> Result<RubyTokenizer, Error> {
     let tokenizer = Tokenizer::from_file(&path)
         .map_err(|e| Error::new(ruby.exception_runtime_error(), format!("{}", e)))?;
-    Ok(RubyTokenizer(RefCell::new(tokenizer)))
+    Ok(RubyTokenizer(Mutex::new(tokenizer)))
 }
 
 fn encoding_to_hash(ruby: &Ruby, encoding: &tokenizers::Encoding) -> Result<RHash, Error> {
@@ -38,27 +38,51 @@ fn encoding_to_hash(ruby: &Ruby, encoding: &tokenizers::Encoding) -> Result<RHas
     let mask: Vec<i64> = encoding.get_attention_mask().iter().map(|&m| m as i64).collect();
     hash.aset(ruby.sym_new("attention_mask"), RArray::from_vec(mask))?;
 
+    // type_ids
+    let type_ids: Vec<i64> = encoding.get_type_ids().iter().map(|&t| t as i64).collect();
+    hash.aset(ruby.sym_new("type_ids"), RArray::from_vec(type_ids))?;
+
+    // special_tokens_mask
+    let special_mask: Vec<i64> = encoding.get_special_tokens_mask().iter().map(|&m| m as i64).collect();
+    hash.aset(ruby.sym_new("special_tokens_mask"), RArray::from_vec(special_mask))?;
+
+    // word_ids
+    let word_ids_array = RArray::new();
+    for word_id in encoding.get_word_ids() {
+        match word_id {
+            Some(id) => word_ids_array.push(*id as i64)?,
+            None => word_ids_array.push(ruby.qnil())?,
+        }
+    }
+    hash.aset(ruby.sym_new("word_ids"), word_ids_array)?;
+
     Ok(hash)
 }
 
 impl RubyTokenizer {
-    fn encode(ruby: &Ruby, rb_self: &Self, text: String) -> Result<RHash, Error> {
-        let encoding = rb_self.0.borrow()
-            .encode(text.as_str(), false)
+    fn encode(ruby: &Ruby, rb_self: &Self, text: String, add_special_tokens: bool) -> Result<RHash, Error> {
+        let guard = rb_self.0.lock()
+            .map_err(|e| Error::new(ruby.exception_runtime_error(), format!("lock poisoned: {}", e)))?;
+        let encoding = guard
+            .encode(text.as_str(), add_special_tokens)
             .map_err(|e| Error::new(ruby.exception_runtime_error(), format!("{}", e)))?;
         encoding_to_hash(ruby, &encoding)
     }
 
-    fn decode(ruby: &Ruby, rb_self: &Self, ids: Vec<u32>) -> Result<String, Error> {
-        rb_self.0.borrow()
-            .decode(&ids, true)
+    fn decode(ruby: &Ruby, rb_self: &Self, ids: Vec<u32>, skip_special_tokens: bool) -> Result<String, Error> {
+        let guard = rb_self.0.lock()
+            .map_err(|e| Error::new(ruby.exception_runtime_error(), format!("lock poisoned: {}", e)))?;
+        guard
+            .decode(&ids, skip_special_tokens)
             .map_err(|e| Error::new(ruby.exception_runtime_error(), format!("{}", e)))
     }
 
-    fn encode_batch(ruby: &Ruby, rb_self: &Self, texts: Vec<String>) -> Result<RArray, Error> {
+    fn encode_batch(ruby: &Ruby, rb_self: &Self, texts: Vec<String>, add_special_tokens: bool) -> Result<RArray, Error> {
         let inputs: Vec<&str> = texts.iter().map(|s| s.as_str()).collect();
-        let encodings = rb_self.0.borrow()
-            .encode_batch(inputs, false)
+        let guard = rb_self.0.lock()
+            .map_err(|e| Error::new(ruby.exception_runtime_error(), format!("lock poisoned: {}", e)))?;
+        let encodings = guard
+            .encode_batch(inputs, add_special_tokens)
             .map_err(|e| Error::new(ruby.exception_runtime_error(), format!("{}", e)))?;
         let result = RArray::new();
         for encoding in &encodings {
@@ -67,40 +91,68 @@ impl RubyTokenizer {
         Ok(result)
     }
 
-    fn decode_batch(ruby: &Ruby, rb_self: &Self, ids_array: Vec<Vec<u32>>) -> Result<Vec<String>, Error> {
+    fn decode_batch(ruby: &Ruby, rb_self: &Self, ids_array: Vec<Vec<u32>>, skip_special_tokens: bool) -> Result<Vec<String>, Error> {
         let refs: Vec<&[u32]> = ids_array.iter().map(|v| v.as_slice()).collect();
-        rb_self.0.borrow()
-            .decode_batch(&refs, true)
+        let guard = rb_self.0.lock()
+            .map_err(|e| Error::new(ruby.exception_runtime_error(), format!("lock poisoned: {}", e)))?;
+        guard
+            .decode_batch(&refs, skip_special_tokens)
             .map_err(|e| Error::new(ruby.exception_runtime_error(), format!("{}", e)))
     }
 
-    fn vocab_size(&self) -> usize {
-        self.0.borrow().get_vocab_size(true)
+    fn vocab_size(ruby: &Ruby, rb_self: &Self) -> Result<usize, Error> {
+        let guard = rb_self.0.lock()
+            .map_err(|e| Error::new(ruby.exception_runtime_error(), format!("lock poisoned: {}", e)))?;
+        Ok(guard.get_vocab_size(true))
     }
 
-    fn token_to_id(&self, token: String) -> Option<u32> {
-        self.0.borrow().token_to_id(&token)
+    fn token_to_id(ruby: &Ruby, rb_self: &Self, token: String) -> Result<Option<u32>, Error> {
+        let guard = rb_self.0.lock()
+            .map_err(|e| Error::new(ruby.exception_runtime_error(), format!("lock poisoned: {}", e)))?;
+        Ok(guard.token_to_id(&token))
     }
 
-    fn id_to_token(&self, id: u32) -> Option<String> {
-        self.0.borrow().id_to_token(id)
+    fn id_to_token(ruby: &Ruby, rb_self: &Self, id: u32) -> Result<Option<String>, Error> {
+        let guard = rb_self.0.lock()
+            .map_err(|e| Error::new(ruby.exception_runtime_error(), format!("lock poisoned: {}", e)))?;
+        Ok(guard.id_to_token(id))
     }
 
-    fn enable_truncation(&self, max_length: usize) {
+    fn enable_truncation(ruby: &Ruby, rb_self: &Self, max_length: usize) -> Result<(), Error> {
         let params = tokenizers::TruncationParams {
             max_length,
             ..Default::default()
         };
-        let _ = self.0.borrow_mut().with_truncation(Some(params));
+        let mut guard = rb_self.0.lock()
+            .map_err(|e| Error::new(ruby.exception_runtime_error(), format!("lock poisoned: {}", e)))?;
+        let _ = guard.with_truncation(Some(params));
+        Ok(())
     }
 
-    fn enable_padding(&self, length: usize, pad_token: String) {
+    fn disable_truncation(ruby: &Ruby, rb_self: &Self) -> Result<(), Error> {
+        let mut guard = rb_self.0.lock()
+            .map_err(|e| Error::new(ruby.exception_runtime_error(), format!("lock poisoned: {}", e)))?;
+        let _ = guard.with_truncation(None);
+        Ok(())
+    }
+
+    fn enable_padding(ruby: &Ruby, rb_self: &Self, length: usize, pad_token: String) -> Result<(), Error> {
         let params = tokenizers::PaddingParams {
             strategy: tokenizers::PaddingStrategy::Fixed(length),
             pad_token,
             ..Default::default()
         };
-        self.0.borrow_mut().with_padding(Some(params));
+        let mut guard = rb_self.0.lock()
+            .map_err(|e| Error::new(ruby.exception_runtime_error(), format!("lock poisoned: {}", e)))?;
+        guard.with_padding(Some(params));
+        Ok(())
+    }
+
+    fn disable_padding(ruby: &Ruby, rb_self: &Self) -> Result<(), Error> {
+        let mut guard = rb_self.0.lock()
+            .map_err(|e| Error::new(ruby.exception_runtime_error(), format!("lock poisoned: {}", e)))?;
+        guard.with_padding(None);
+        Ok(())
     }
 }
 
@@ -111,15 +163,17 @@ fn init(ruby: &Ruby) -> Result<(), Error> {
     let class = module.define_class("InternalTokenizer", ruby.class_object())?;
     class.define_singleton_method("from_pretrained", function!(from_pretrained, 1))?;
     class.define_singleton_method("from_file", function!(from_file, 1))?;
-    class.define_method("_encode", method!(RubyTokenizer::encode, 1))?;
-    class.define_method("_decode", method!(RubyTokenizer::decode, 1))?;
-    class.define_method("_encode_batch", method!(RubyTokenizer::encode_batch, 1))?;
-    class.define_method("_decode_batch", method!(RubyTokenizer::decode_batch, 1))?;
+    class.define_method("_encode", method!(RubyTokenizer::encode, 2))?;
+    class.define_method("_decode", method!(RubyTokenizer::decode, 2))?;
+    class.define_method("_encode_batch", method!(RubyTokenizer::encode_batch, 2))?;
+    class.define_method("_decode_batch", method!(RubyTokenizer::decode_batch, 2))?;
     class.define_method("vocab_size", method!(RubyTokenizer::vocab_size, 0))?;
     class.define_method("token_to_id", method!(RubyTokenizer::token_to_id, 1))?;
     class.define_method("id_to_token", method!(RubyTokenizer::id_to_token, 1))?;
     class.define_method("_enable_truncation", method!(RubyTokenizer::enable_truncation, 1))?;
+    class.define_method("_disable_truncation", method!(RubyTokenizer::disable_truncation, 0))?;
     class.define_method("_enable_padding", method!(RubyTokenizer::enable_padding, 2))?;
+    class.define_method("_disable_padding", method!(RubyTokenizer::disable_padding, 0))?;
 
     Ok(())
 }

data/lib/tokenizer_ruby/tokenizer.rb

@@ -18,46 +18,52 @@ module TokenizerRuby
       new(InternalTokenizer.from_file(path))
     end
 
-    def encode(text)
+    def encode(text, add_special_tokens: false)
       raise TokenizerRuby::Error, "encode expects a String, got #{text.class}" unless text.is_a?(String)
 
       begin
-        result = @inner._encode(text)
+        result = @inner._encode(text, add_special_tokens)
       rescue => e
-        raise TokenizerRuby::Error, "failed to encode text: #{e.message}"
+        raise TokenizerRuby::TokenizationError, "failed to encode text: #{e.message}"
       end
       Encoding.new(
         ids: result[:ids],
         tokens: result[:tokens],
         offsets: result[:offsets],
-        attention_mask: result[:attention_mask]
+        attention_mask: result[:attention_mask],
+        type_ids: result[:type_ids],
+        special_tokens_mask: result[:special_tokens_mask],
+        word_ids: result[:word_ids]
       )
     end
 
-    def decode(ids)
+    def decode(ids, skip_special_tokens: true)
       raise TokenizerRuby::Error, "decode expects an Array, got #{ids.class}" unless ids.is_a?(Array)
 
       begin
-        @inner._decode(ids)
+        @inner._decode(ids, skip_special_tokens)
       rescue => e
-        raise TokenizerRuby::Error, "failed to decode ids: #{e.message}"
+        raise TokenizerRuby::TokenizationError, "failed to decode ids: #{e.message}"
       end
     end
 
-    def encode_batch(texts)
-      results = @inner._encode_batch(texts)
+    def encode_batch(texts, add_special_tokens: false)
+      results = @inner._encode_batch(texts, add_special_tokens)
       results.map do |result|
         Encoding.new(
           ids: result[:ids],
           tokens: result[:tokens],
           offsets: result[:offsets],
-          attention_mask: result[:attention_mask]
+          attention_mask: result[:attention_mask],
+          type_ids: result[:type_ids],
+          special_tokens_mask: result[:special_tokens_mask],
+          word_ids: result[:word_ids]
         )
       end
     end
 
-    def decode_batch(ids_array)
-      @inner._decode_batch(ids_array)
+    def decode_batch(ids_array, skip_special_tokens: true)
+      @inner._decode_batch(ids_array, skip_special_tokens)
     end
 
     def vocab_size
@@ -77,7 +83,7 @@ module TokenizerRuby
     end
 
     def truncate(text, max_tokens:)
-      raise TokenizerRuby::Error, "max_tokens must be positive, got #{max_tokens}" unless max_tokens > 0
+      raise TokenizerRuby::ConfigurationError, "max_tokens must be positive, got #{max_tokens}" unless max_tokens > 0
 
       encoding = encode(text)
       return text if encoding.length <= max_tokens
@@ -90,8 +96,16 @@ module TokenizerRuby
       @inner._enable_truncation(max_length)
     end
 
+    def disable_truncation
+      @inner._disable_truncation
+    end
+
     def enable_padding(length:, pad_token: "[PAD]")
       @inner._enable_padding(length, pad_token)
     end
+
+    def disable_padding
+      @inner._disable_padding
+    end
   end
 end

data/lib/tokenizer_ruby/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module TokenizerRuby
-  VERSION = "0.1.1"
+  VERSION = "0.2.0"
 end

data/lib/tokenizer_ruby.rb

@@ -13,4 +13,7 @@ end
 
 module TokenizerRuby
   class Error < StandardError; end
+  class TokenizationError < Error; end
+  class FileNotFoundError < Error; end
+  class ConfigurationError < Error; end
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: tokenizer-ruby
 version: !ruby/object:Gem::Version
-  version: 0.1.1
+  version: 0.2.0
 platform: ruby
 authors:
 - Johannes Dwi Cahyo