tokenizers 0.2.1 → 0.2.3

data/Cargo.toml

@@ -3,8 +3,3 @@ members = ["ext/tokenizers"]
 
 [profile.release]
 strip = true
-
-[patch.crates-io]
-number_prefix = { git = "https://github.com/ankane/rust-number-prefix", branch = "license-file" }
-rb-sys-env = { git = "https://github.com/oxidize-rb/rb-sys" }
-tokenizers = { git = "https://github.com/huggingface/tokenizers" }

data/README.md

@@ -24,8 +24,8 @@ Encode
 
 ```ruby
 encoded = tokenizer.encode("I can feel the magic, can you?")
-encoded.ids
 encoded.tokens
+encoded.ids
 ```
 
 Decode

data/ext/tokenizers/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "tokenizers"
-version = "0.2.1"
+version = "0.2.3"
 license = "Apache-2.0"
 authors = ["Andrew Kane <andrew@ankane.org>"]
 edition = "2021"
@@ -13,6 +13,7 @@ crate-type = ["cdylib"]
 magnus = "0.4"
 
 [dependencies.tokenizers]
-version = "0.13.2"
+version = "0.13.2" # also update in from_pretrained.rb
+git = "https://github.com/huggingface/tokenizers"
 default-features = false
-features = ["progressbar", "http", "onig", "esaxx_fast"]
+features = ["progressbar", "onig", "esaxx_fast"]

data/ext/tokenizers/src/encoding.rs

@@ -1,16 +1,90 @@
-use tk::Encoding;
+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) -> Vec<Self> {
+        self.encoding
+            .get_overflowing()
+            .clone()
+            .into_iter()
+            .map(|e| e.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)
     }
 }

data/ext/tokenizers/src/lib.rs

@@ -27,10 +27,7 @@ fn module() -> RModule {
 #[magnus::init]
 fn init() -> RbResult<()> {
     let module = module();
-    module.define_singleton_method(
-        "_from_pretrained",
-        function!(RbTokenizer::from_pretrained, 3),
-    )?;
+    module.define_singleton_method("from_file", function!(RbTokenizer::from_file, 1))?;
 
     let class = module.define_class("BPE", Default::default())?;
     class.define_singleton_method("new", function!(RbBPE::new, 2))?;
@@ -41,15 +38,39 @@ fn init() -> RbResult<()> {
         "add_special_tokens",
         method!(RbTokenizer::add_special_tokens, 1),
     )?;
-    class.define_method("encode", method!(RbTokenizer::encode, 1))?;
+    class.define_method("add_tokens", method!(RbTokenizer::add_tokens, 1))?;
+    class.define_method("_encode", method!(RbTokenizer::encode, 2))?;
     class.define_method("decode", method!(RbTokenizer::decode, 1))?;
     class.define_method("decoder=", method!(RbTokenizer::set_decoder, 1))?;
     class.define_method("pre_tokenizer=", method!(RbTokenizer::set_pre_tokenizer, 1))?;
     class.define_method("normalizer=", method!(RbTokenizer::set_normalizer, 1))?;
+    class.define_method("token_to_id", method!(RbTokenizer::token_to_id, 1))?;
+    class.define_method("id_to_token", method!(RbTokenizer::id_to_token, 1))?;
 
     let class = module.define_class("Encoding", Default::default())?;
+    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))?;
+    class.define_method("word_ids", method!(RbEncoding::word_ids, 0))?;
+    class.define_method("sequence_ids", method!(RbEncoding::sequence_ids, 0))?;
+    class.define_method("type_ids", method!(RbEncoding::type_ids, 0))?;
+    class.define_method("offsets", method!(RbEncoding::offsets, 0))?;
+    class.define_method(
+        "special_tokens_mask",
+        method!(RbEncoding::special_tokens_mask, 0),
+    )?;
+    class.define_method("attention_mask", method!(RbEncoding::attention_mask, 0))?;
+    class.define_method("overflowing", method!(RbEncoding::overflowing, 0))?;
+    class.define_method("_word_to_tokens", method!(RbEncoding::word_to_tokens, 2))?;
+    class.define_method("_word_to_chars", method!(RbEncoding::word_to_chars, 2))?;
+    class.define_method(
+        "token_to_sequence",
+        method!(RbEncoding::token_to_sequence, 1),
+    )?;
+    class.define_method("token_to_chars", method!(RbEncoding::token_to_chars, 1))?;
+    class.define_method("token_to_word", method!(RbEncoding::token_to_word, 1))?;
+    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("BPEDecoder", Default::default())?;
     class.define_singleton_method("new", function!(RbBPEDecoder::new, 0))?;

data/ext/tokenizers/src/tokenizer.rs

@@ -1,5 +1,5 @@
-use magnus::Module;
 use std::cell::RefCell;
+use std::path::PathBuf;
 use tk::tokenizer::Tokenizer;
 use tk::AddedToken;
 
@@ -8,7 +8,7 @@ use super::encoding::RbEncoding;
 use super::models::RbBPE;
 use super::normalizers::RbBertNormalizer;
 use super::pre_tokenizers::RbBertPreTokenizer;
-use super::{module, RbError, RbResult};
+use super::{RbError, RbResult};
 
 #[magnus::wrap(class = "Tokenizers::Tokenizer")]
 pub struct RbTokenizer {
@@ -22,22 +22,8 @@ impl RbTokenizer {
         }
     }
 
-    pub fn from_pretrained(
-        identifier: String,
-        revision: String,
-        auth_token: Option<String>,
-    ) -> RbResult<Self> {
-        let version = module().const_get("VERSION").unwrap();
-        let params = tk::FromPretrainedParameters {
-            revision,
-            auth_token,
-            user_agent: [("bindings", "Ruby".to_string()), ("version", version)]
-                .iter()
-                .map(|(k, v)| (k.to_string(), v.to_string()))
-                .collect(),
-        };
-
-        Tokenizer::from_pretrained(identifier, Some(params))
+    pub fn from_file(path: PathBuf) -> RbResult<Self> {
+        Tokenizer::from_file(path)
             .map(|v| RbTokenizer {
                 tokenizer: RefCell::new(v),
             })
@@ -50,10 +36,16 @@ impl RbTokenizer {
         // TODO return self
     }
 
-    pub fn encode(&self, text: String) -> RbResult<RbEncoding> {
+    pub fn add_tokens(&self, tokens: Vec<String>) {
+        let tokens: Vec<AddedToken> = tokens.iter().map(|t| AddedToken::from(t, true)).collect();
+        self.tokenizer.borrow_mut().add_tokens(&tokens);
+        // TODO return self
+    }
+
+    pub fn encode(&self, sequence: String, add_special_tokens: bool) -> RbResult<RbEncoding> {
         self.tokenizer
             .borrow()
-            .encode(text, false)
+            .encode(sequence, add_special_tokens)
             .map(|v| RbEncoding { encoding: v })
             .map_err(RbError::from)
     }
@@ -82,4 +74,12 @@ impl RbTokenizer {
             .borrow_mut()
             .with_normalizer(normalizer.normalizer);
     }
+
+    pub fn token_to_id(&self, token: String) -> Option<u32> {
+        self.tokenizer.borrow().token_to_id(&token)
+    }
+
+    pub fn id_to_token(&self, id: u32) -> Option<String> {
+        self.tokenizer.borrow().id_to_token(id)
+    }
 }

data/lib/tokenizers/char_bpe_tokenizer.rb

@@ -8,8 +8,8 @@ module Tokenizers
       @tokenizer.decoder = BPEDecoder.new
     end
 
-    def encode(text)
-      @tokenizer.encode(text)
+    def encode(text, **options)
+      @tokenizer.encode(text, **options)
     end
 
     def decode(ids)

data/lib/tokenizers/encoding.rb

@@ -0,0 +1,19 @@
+module Tokenizers
+  class Encoding
+    def word_to_tokens(word_index, sequence_index = 0)
+      _word_to_tokens(word_index, sequence_index)
+    end
+
+    def word_to_chars(word_index, sequence_index = 0)
+      _word_to_chars(word_index, sequence_index)
+    end
+
+    def char_to_token(char_pos, sequence_index = 0)
+      _char_to_token(char_pos, sequence_index)
+    end
+
+    def char_to_word(char_pos, sequence_index = 0)
+      _char_to_word(word_index, sequence_index)
+    end
+  end
+end

data/lib/tokenizers/from_pretrained.rb

@@ -0,0 +1,119 @@
+module Tokenizers
+  module FromPretrained
+    # for user agent
+    TOKENIZERS_VERSION = "0.13.2"
+
+    # use Ruby for downloads
+    # this avoids the need to vendor OpenSSL on Linux
+    # and reduces the extension size by about half
+    def from_pretrained(identifier, revision: "main", auth_token: nil)
+      require "cgi"
+      require "digest"
+      require "fileutils"
+      require "json"
+      require "open-uri"
+
+      cache_dir = ensure_cache_dir
+
+      # string options are headers
+      options = {
+        open_timeout: 3,
+        read_timeout: 30,
+        "User-Agent" => "tokenizers/#{TOKENIZERS_VERSION}; bindings/Ruby; version/#{VERSION}"
+      }
+      if auth_token
+        options["Authorization"] = "Bearer #{auth_token}"
+      end
+
+      url = "https://huggingface.co/%s/resolve/%s/tokenizer.json" % [identifier, revision].map { |v| CGI.escape(v) }
+
+      path =
+        begin
+          cached_path(cache_dir, url, options)
+        rescue OpenURI::HTTPError
+          raise Error, "Model \"#{identifier}\" on the Hub doesn't have a tokenizer"
+        end
+
+      from_file(path)
+    end
+
+    private
+
+    # use same storage format as Rust version
+    # https://github.com/epwalsh/rust-cached-path
+    def cached_path(cache_dir, url, options)
+      fsum = Digest::SHA256.hexdigest(url)
+      meta_paths = Dir[File.join(cache_dir, "#{fsum}.*.meta")]
+      meta = meta_paths.map { |f| JSON.load_file(f) }.max_by { |m| m["creation_time"] }
+      etag = meta["etag"] if meta
+
+      if etag
+        esum = Digest::SHA256.hexdigest(etag)
+        resource_path = File.join(cache_dir, "#{fsum}.#{esum}")
+        options["If-None-Match"] = etag if File.exist?(resource_path)
+      end
+
+      options[:content_length_proc] = -> (_) { puts "Downloading..." }
+
+      tempfile =
+        begin
+          URI.open(url, options)
+        rescue OpenURI::HTTPError => e
+          if e.message == "304 Not Modified"
+            return resource_path
+          else
+            raise e
+          end
+        end
+
+      etag = tempfile.meta["etag"]
+      esum = Digest::SHA256.hexdigest(etag)
+      resource_path = File.join(cache_dir, "#{fsum}.#{esum}")
+      meta_path = "#{resource_path}.meta"
+
+      meta = {
+        resource: url,
+        resource_path: resource_path,
+        meta_path: meta_path,
+        etag: etag,
+        expires: nil,
+        creation_time: Time.now.to_f
+      }
+
+      File.write("#{resource_path}.lock", "")
+      File.open(resource_path, "wb") { |f| IO.copy_stream(tempfile, f) }
+      File.write(meta_path, JSON.generate(meta))
+
+      resource_path
+    end
+
+    def cache_dir
+      if ENV["TOKENIZERS_CACHE"]
+        ENV["TOKENIZERS_CACHE"]
+      else
+        # use same directory as Rust version
+        # https://docs.rs/dirs/latest/dirs/fn.cache_dir.html
+        dir =
+          if Gem.win_platform?
+            ENV.fetch("LOCALAPPDATA")
+          elsif mac?
+            File.join(ENV.fetch("HOME"), "Library", "Caches")
+          else
+            ENV["XDG_CACHE_HOME"] || File.join(ENV.fetch("HOME"), ".cache")
+          end
+
+        File.join(dir, "huggingface", "tokenizers")
+      end
+    end
+
+    def ensure_cache_dir
+      dir = cache_dir
+      FileUtils.mkdir_p(dir)
+      dir
+    end
+
+    def mac?
+      RbConfig::CONFIG["host_os"] =~ /darwin/i
+    end
+  end
+end

data/lib/tokenizers/tokenizer.rb

@@ -0,0 +1,12 @@
+module Tokenizers
+  class Tokenizer
+    # TODO change add_special_tokens default to true in 0.3.0
+    def encode(sequence, add_special_tokens: nil)
+      if add_special_tokens.nil?
+        warn "[tokenizers] add_special_tokens will default to true in 0.3.0. Pass add_special_tokens: true/false to silence this warning."
+        add_special_tokens = false
+      end
+      _encode(sequence, add_special_tokens)
+    end
+  end
+end

data/lib/tokenizers/version.rb

@@ -1,3 +1,3 @@
 module Tokenizers
-  VERSION = "0.2.1"
+  VERSION = "0.2.3"
 end

data/lib/tokenizers.rb

@@ -1,18 +1,19 @@
 # ext
 begin
-  require "tokenizers/#{RUBY_VERSION.to_f}/tokenizers"
+  require_relative "tokenizers/#{RUBY_VERSION.to_f}/tokenizers"
 rescue LoadError
-  require "tokenizers/tokenizers"
+  require_relative "tokenizers/tokenizers"
 end
 
 # modules
-require "tokenizers/char_bpe_tokenizer"
-require "tokenizers/version"
+require_relative "tokenizers/char_bpe_tokenizer"
+require_relative "tokenizers/encoding"
+require_relative "tokenizers/from_pretrained"
+require_relative "tokenizers/tokenizer"
+require_relative "tokenizers/version"
 
 module Tokenizers
   class Error < StandardError; end
 
-  def self.from_pretrained(identifier, revision: "main", auth_token: nil)
-    _from_pretrained(identifier, revision, auth_token)
-  end
+  extend FromPretrained
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: tokenizers
 version: !ruby/object:Gem::Version
-  version: 0.2.1
+  version: 0.2.3
 platform: ruby
 authors:
 - Andrew Kane
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2023-01-12 00:00:00.000000000 Z
+date: 2023-01-22 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: rb_sys
@@ -48,6 +48,9 @@ files:
 - ext/tokenizers/src/tokenizer.rs
 - lib/tokenizers.rb
 - lib/tokenizers/char_bpe_tokenizer.rb
+- lib/tokenizers/encoding.rb
+- lib/tokenizers/from_pretrained.rb
+- lib/tokenizers/tokenizer.rb
 - lib/tokenizers/version.rb
 homepage: https://github.com/ankane/tokenizers-ruby
 licenses: