tiktoken_ruby 0.0.2 → 0.0.4

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 417d96c1d3bb9bac547fb3cc632e061894e13e7d7673a74e08ddc33d33bc3cbb
-  data.tar.gz: b27d96b688bc4fa5f31cdecfaa629306e262bd34960a25bda93e7b8364e69c57
+  metadata.gz: 1c461629e4b1174e825e05f4644e7bcdaac90849047997c6724358435c70f0bb
+  data.tar.gz: 244217e1b298a998c468723d2082bcf10bb46e6a80168d808a56ee4b618bbf6f
 SHA512:
-  metadata.gz: 03cdce8758407d412fe29e580ed5caed455d8153923c982b60ed0440f3537f73034c121fa758e01b86354368443ff2e1f2ed4b255833ccf385d99ff79f337058
-  data.tar.gz: 55b2e0db635609d0f756abef80e761b828eedc345b442b620c086f3e07072895b094a115993d4083633f34511e5a74ff38a2f0187a6308b28abd6f9d5483c016
+  metadata.gz: e2920fba51441a33440435c310e2389723ed5e63dd150e0303743e09ee2ef5f68499e14f9df6254daa7a355568fbdf54e91a9ef3eebba30182f757e7f190ba11
+  data.tar.gz: b0da8ae135800a6405e546a03f1c17319978096716872a3d4b2fbe7fe812f2061f73d02eb8540f866d41a009d21c075273e096c3073be6ae1b5562f799f806c1

data/Gemfile

@@ -13,4 +13,5 @@ gem "rb_sys"
 gem "rspec", "~> 3.0"
 
 gem "standard", "~> 1.3"
-gem 'pry', '~> 0.14.2'
+
+gem "yard-doctest", "~> 0.1.17"

data/Gemfile.lock

@@ -1,23 +1,19 @@
 PATH
   remote: .
   specs:
-    tiktoken_ruby (0.0.2)
+    tiktoken_ruby (0.0.4)
 
 GEM
   remote: https://rubygems.org/
   specs:
     ast (2.4.2)
-    coderay (1.1.3)
     diff-lcs (1.5.0)
     json (2.6.3)
     language_server-protocol (3.17.0.3)
-    method_source (1.0.0)
+    minitest (5.18.0)
     parallel (1.22.1)
     parser (3.2.1.1)
       ast (~> 2.4.1)
-    pry (0.14.2)
-      coderay (~> 1.1)
-      method_source (~> 1.0)
     rainbow (3.1.1)
     rake (13.0.6)
     rake-compiler (1.2.1)
@@ -59,18 +55,26 @@ GEM
       rubocop (= 1.48.1)
       rubocop-performance (= 1.16.0)
     unicode-display_width (2.4.2)
+    webrick (1.7.0)
+    yard (0.9.28)
+      webrick (~> 1.7.0)
+    yard-doctest (0.1.17)
+      minitest
+      yard
 
 PLATFORMS
   arm64-darwin-22
+  x86_64-darwin-22
+  x86_64-linux
 
 DEPENDENCIES
-  pry (~> 0.14.2)
   rake (~> 13.0)
   rake-compiler
   rb_sys
   rspec (~> 3.0)
   standard (~> 1.3)
   tiktoken_ruby!
+  yard-doctest (~> 0.1.17)
 
 BUNDLED WITH
    2.4.6

data/README.md

@@ -15,11 +15,22 @@ If bundler is not being used to manage dependencies, install the gem by executin
     $ gem install tiktoken_ruby
 
 ## Usage
+Usage should be very similar to the python library. Here's a simple example
 
+Encode and decode text
 ```ruby
-encoding = Tiktoken::Encoding.r50k_base
-tokens = encoding.encode("Hello world!")
-puts encoding.decode(tokens)
+require 'tiktoken_ruby'
+
+enc = Tiktoken.get_encoding("cl100k_base")
+enc.decode(enc.encode("hello world")) #=> "hello world"
+```
+
+Encoders can also be retrieved by model name
+```ruby
+require 'tiktoken_ruby'
+
+enc = Tiktoken.encoding_for_model("gpt-4")
+enc.encode("hello world").length #=> 2
 ```
 
 ## Development
@@ -32,6 +43,17 @@ To install this gem onto your local machine, run `bundle exec rake install`. To
 
 Bug reports and pull requests are welcome on GitHub at https://github.com/iapark/tiktoken_ruby.
 
+To get started with development:
+
+```sh
+git clone https://github.com/IAPark/tiktoken_ruby.git
+cd tiktoken_ruby
+bundle install
+bundle exec rake compile
+bundle exec rake spec
+```
+
+
 ## License
 
 The gem is available as open source under the terms of the [MIT License](https://opensource.org/licenses/MIT).

data/Rakefile

@@ -14,7 +14,6 @@ end
 
 RSpec::Core::RakeTask.new(:spec)
 
-
 task :native, [:platform] do |_t, platform:|
   sh "bundle", "exec", "rb-sys-dock", "--platform", platform, "--build"
 end

data/doctest_helper.rb

@@ -0,0 +1 @@
+require "lib/tiktoken_ruby"

data/ext/tiktoken_ruby/src/lib.rs

@@ -35,7 +35,7 @@ fn init() -> Result<(), Error> {
     let factory_module = module.define_module("BpeFactory")?;
     factory_module.define_singleton_method("r50k_base", function!(r50k_base, 0))?;
     factory_module.define_singleton_method("p50k_base", function!(p50k_base, 0))?;
-    factory_module.define_singleton_method("p50k_base", function!(p50k_edit, 0))?;
+    factory_module.define_singleton_method("p50k_edit", function!(p50k_edit, 0))?;
     factory_module.define_singleton_method("cl100k_base", function!(cl100k_base, 0))?;
 
 

data/lib/tiktoken_ruby/encoding.rb

@@ -1,23 +1,52 @@
 # frozen_string_literal: true
 
-class Tiktoken::Encoding 
-    def self.method_missing(method)
-        Tiktoken::Encoding.new(Tiktoken::BpeFactory.send(method))
-    end
-
-    def initialize(ext_base_bpe)
-        @ext_base_bpe = ext_base_bpe
-    end
-
-    def encode_ordinary(text)
-        @ext_base_bpe.encode_ordinary(text)
-    end
-
-    def encode(text, allowed_special: [])
-        @ext_base_bpe.encode(text, allowed_special)
-    end
-
-    def decode(tokens)
-        @ext_base_bpe.decode(tokens)
-    end
+class Tiktoken::Encoding
+  attr_reader :name
+
+  # This returns a new Tiktoken::Encoding instance for the requested encoding
+  # @param encoding [Symbol] The name of the encoding to load
+  # @return [Tiktoken::Encoding] The encoding instance
+  def self.for_name(encoding)
+    Tiktoken::Encoding.new(Tiktoken::BpeFactory.send(encoding.to_sym), encoding.to_sym)
+  end
+
+  # This returns a Tiktoken::Encoding instance for the requested encoding
+  # It will reuse an existing encoding if it's already been loaded
+  # @param encoding [Symbol] The name of the encoding to load
+  # @return [Tiktoken::Encoding] The encoding instance
+  def self.for_name_cached(encoding)
+    @encodings ||= {}
+    @encodings[encoding.to_sym] ||= Tiktoken::Encoding.for_name(encoding)
+  end
+
+  # Encodes the text as a list of integer tokens. This encoding will encode special non text tokens
+  # basically it's unescaped
+  # @param text [String] The text to encode
+  # @return [Array<Integer>] The encoded tokens
+  def encode_ordinary(text)
+    @ext_base_bpe.encode_ordinary(text)
+  end
+
+  # Encodes the text as a list of integer tokens. This encoding will treat special non text tokens
+  # as text unless they're in the allowed_special array. It's basically like the text was escaped
+  # @param text [String] The text to encode
+  # @param allowed_special [Array<String>] An array of special tokens to allow
+  # @return [Array<Integer>] The encoded tokens
+  def encode(text, allowed_special: [])
+    @ext_base_bpe.encode(text, allowed_special)
+  end
+
+  # Decodes the tokens back into text
+  # @param tokens [Array<Integer>] The tokens to decode
+  # @return [String] The decoded text
+  def decode(tokens)
+    @ext_base_bpe.decode(tokens)
+  end
+
+  private
+
+  def initialize(ext_base_bpe, name)
+    @ext_base_bpe = ext_base_bpe
+    @name = name
+  end
 end

data/lib/tiktoken_ruby/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Tiktoken
-  VERSION = "0.0.2"
+  VERSION = "0.0.4"
 end

data/lib/tiktoken_ruby.rb

@@ -1,7 +1,7 @@
 # frozen_string_literal: true
 
 require_relative "tiktoken_ruby/version"
-require_relative "tiktoken_ruby/encoding.rb"
+require_relative "tiktoken_ruby/encoding"
 
 begin
   RUBY_VERSION =~ /(\d+\.\d+)/
@@ -11,5 +11,104 @@ rescue LoadError
 end
 
 module Tiktoken
-  class Error < StandardError; end
+  class << self
+    # Returns an encoding by name. If the encoding is not already loaded it will be loaded, but otherwise
+    # it will reuse the instance of that type that was previous loaded
+    # @param name [Symbol|String] The name of the encoding to load
+    # @return [Tiktoken::Encoding] The encoding instance
+    # @example Encode and decode text
+    #   enc = Tiktoken.get_encoding("cl100k_base")
+    #   enc.decode(enc.encode("hello world")) #=> "hello world"
+    def get_encoding(name)
+      name = name.to_sym
+      return nil unless SUPPORTED_ENCODINGS.include?(name)
+
+      Tiktoken::Encoding.for_name_cached(name)
+    end
+
+    # Gets the encoding for an OpenAI model
+    # @param model_name [Symbol|String] The name of the model to get the encoding for
+    # @return [Tiktoken::Encoding] The encoding instance
+    # @example Count tokens for text
+    #   enc = Tiktoken.encoding_for_model("gpt-4")
+    #   enc.encode("hello world").length #=> 2
+    def encoding_for_model(model_name)
+      PREFIX_MODELS.each do |prefix|
+        if model_name.to_s.start_with?("#{prefix}-")
+          model_name = prefix
+          break
+        end
+      end
+
+      encoding_name = MODEL_TO_ENCODING_NAME[model_name.to_sym]
+      return nil unless encoding_name
+
+      get_encoding(encoding_name)
+    end
+
+    # Lists all the encodings that are supported
+    # @return [Array<Symbol>] The list of supported encodings
+    def list_encoding_names
+      SUPPORTED_ENCODINGS
+    end
+
+    # Lists all the models that are supported
+    # @return [Array<Symbol>] The list of supported models
+    def list_model_names
+      MODEL_TO_ENCODING_NAME.keys
+    end
+
+    private
+
+    SUPPORTED_ENCODINGS = [
+      :r50k_base,
+      :p50k_base,
+      :p50k_edit,
+      :cl100k_base
+    ]
+
+    # taken from the python library here https://github.com/openai/tiktoken/blob/main/tiktoken/model.py#L13-L53
+    # that is also MIT licensed but by OpenAI
+    MODEL_TO_ENCODING_NAME = {
+      "gpt-4": "cl100k_base",
+      "gpt-3.5-turbo": "cl100k_base",
+      # text
+      "text-davinci-003": "p50k_base",
+      "text-davinci-002": "p50k_base",
+      "text-davinci-001": "r50k_base",
+      "text-curie-001": "r50k_base",
+      "text-babbage-001": "r50k_base",
+      "text-ada-001": "r50k_base",
+      davinci: "r50k_base",
+      curie: "r50k_base",
+      babbage: "r50k_base",
+      ada: "r50k_base",
+      # code
+      "code-davinci-002": "p50k_base",
+      "code-davinci-001": "p50k_base",
+      "code-cushman-002": "p50k_base",
+      "code-cushman-001": "p50k_base",
+      "davinci-codex": "p50k_base",
+      "cushman-codex": "p50k_base",
+      # edit
+      "text-davinci-edit-001": "p50k_edit",
+      "code-davinci-edit-001": "p50k_edit",
+      # embeddings
+      "text-embedding-ada-002": "cl100k_base",
+      # old embeddings
+      "text-similarity-davinci-001": "r50k_base",
+      "text-similarity-curie-001": "r50k_base",
+      "text-similarity-babbage-001": "r50k_base",
+      "text-similarity-ada-001": "r50k_base",
+      "text-search-davinci-doc-001": "r50k_base",
+      "text-search-curie-doc-001": "r50k_base",
+      "text-search-babbage-doc-001": "r50k_base",
+      "text-search-ada-doc-001": "r50k_base",
+      "code-search-babbage-code-001": "r50k_base",
+      "code-search-ada-code-001": "r50k_base"
+    }
+
+    # these are models that have a versioned models that are otherwise identical
+    PREFIX_MODELS = ["gpt-4", "gpt-3.5-turbo"]
+  end
 end

data/tiktoken_ruby.gemspec

@@ -9,18 +9,21 @@ Gem::Specification.new do |spec|
   spec.email = ["isaac.a.park@gmail.com"]
 
   spec.summary = "Ruby wrapper for Tiktoken"
-  spec.description = "Unofficial Ruby wrapper for Tiktoken by way of the unofficial rust bindings"
+  spec.description = "An unofficial Ruby wrapper for Tiktoken, " \
+    "a BPE tokenizer written by and used by OpenAI. It can be used to " \
+    "count the number of tokens in text before sending it to OpenAI APIs."
+
   spec.homepage = "https://github.com/IAPark/tiktoken_ruby"
   spec.license = "MIT"
-  spec.required_ruby_version = ">= 2.6.0"
+  spec.required_ruby_version = ">= 2.7.0"
   spec.required_rubygems_version = ">= 3.1.0"
   spec.platform = Gem::Platform::RUBY
 
-  #spec.metadata["allowed_push_host"] = "TODO: Set to your gem server 'https://example.com'"
-
   spec.metadata["homepage_uri"] = spec.homepage
   spec.metadata["source_code_uri"] = "https://github.com/IAPark/tiktoken_ruby"
-  #spec.metadata["changelog_uri"] = "TODO: Put your gem's CHANGELOG.md URL here."
+  spec.metadata["documentation_uri"] = "https://rubydoc.info/github/IAPark/tiktoken_ruby/main"
+
+  # spec.metadata["changelog_uri"] = "TODO: Put your gem's CHANGELOG.md URL here."
 
   # Specify which files should be added to the gem when it is released.
   # The `git ls-files -z` loads the files in the RubyGem that have been added into git.
@@ -32,11 +35,7 @@ Gem::Specification.new do |spec|
   spec.bindir = "exe"
   spec.executables = spec.files.grep(%r{\Aexe/}) { |f| File.basename(f) }
   spec.require_paths = ["lib"]
-  spec.extensions = ["ext/tiktoken_ruby/Cargo.toml"]
-
-  # Uncomment to register a new dependency of your gem
-  # spec.add_dependency "example-gem", "~> 1.0"
-  # spec.add_dependency "rb_sys", "~> 0.9"
+  spec.extensions = ["ext/tiktoken_ruby/extconf.rb"]
 
   # For more information and examples about making a new gem, check out our
   # guide at: https://bundler.io/guides/creating_gem.html

metadata

@@ -1,21 +1,23 @@
 --- !ruby/object:Gem::Specification
 name: tiktoken_ruby
 version: !ruby/object:Gem::Version
-  version: 0.0.2
+  version: 0.0.4
 platform: ruby
 authors:
 - IAPark
 autorequire:
 bindir: exe
 cert_chain: []
-date: 2023-03-19 00:00:00.000000000 Z
+date: 2023-03-28 00:00:00.000000000 Z
 dependencies: []
-description: Unofficial Ruby wrapper for Tiktoken by way of the unofficial rust bindings
+description: An unofficial Ruby wrapper for Tiktoken, a BPE tokenizer written by and
+  used by OpenAI. It can be used to count the number of tokens in text before sending
+  it to OpenAI APIs.
 email:
 - isaac.a.park@gmail.com
 executables: []
 extensions:
-- ext/tiktoken_ruby/Cargo.toml
+- ext/tiktoken_ruby/extconf.rb
 extra_rdoc_files: []
 files:
 - ".rspec"
@@ -27,6 +29,7 @@ files:
 - LICENSE.txt
 - README.md
 - Rakefile
+- doctest_helper.rb
 - ext/tiktoken_ruby/Cargo.toml
 - ext/tiktoken_ruby/extconf.rb
 - ext/tiktoken_ruby/src/core_bpe_wrapper.rs
@@ -42,6 +45,7 @@ licenses:
 metadata:
   homepage_uri: https://github.com/IAPark/tiktoken_ruby
   source_code_uri: https://github.com/IAPark/tiktoken_ruby
+  documentation_uri: https://rubydoc.info/github/IAPark/tiktoken_ruby/main
 post_install_message:
 rdoc_options: []
 require_paths:
@@ -50,7 +54,7 @@ required_ruby_version: !ruby/object:Gem::Requirement
   requirements:
   - - ">="
     - !ruby/object:Gem::Version
-      version: 2.6.0
+      version: 2.7.0
 required_rubygems_version: !ruby/object:Gem::Requirement
   requirements:
   - - ">="