gigatoken 0.1.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (75) hide show
  1. checksums.yaml +7 -0
  2. data/Cargo.lock +3016 -0
  3. data/Cargo.toml +135 -0
  4. data/LICENSE +21 -0
  5. data/README.md +141 -0
  6. data/exe/gigatoken +9 -0
  7. data/ext/gigatoken/Cargo.toml +24 -0
  8. data/ext/gigatoken/extconf.rb +19 -0
  9. data/ext/gigatoken/src/error.rs +20 -0
  10. data/ext/gigatoken/src/gvl.rs +122 -0
  11. data/ext/gigatoken/src/lib.rs +50 -0
  12. data/ext/gigatoken/src/sentencepiece.rs +205 -0
  13. data/ext/gigatoken/src/sources.rs +207 -0
  14. data/ext/gigatoken/src/tokenizer.rs +571 -0
  15. data/lib/gigatoken/cli/bench.rb +64 -0
  16. data/lib/gigatoken/cli/support.rb +132 -0
  17. data/lib/gigatoken/cli/validate.rb +55 -0
  18. data/lib/gigatoken/cli.rb +18 -0
  19. data/lib/gigatoken/hub.rb +242 -0
  20. data/lib/gigatoken/packed_result.rb +48 -0
  21. data/lib/gigatoken/tokenizer.rb +117 -0
  22. data/lib/gigatoken/version.rb +5 -0
  23. data/lib/gigatoken.rb +29 -0
  24. data/rust-toolchain.toml +8 -0
  25. data/src/batch.rs +1808 -0
  26. data/src/bindings/bridge.rs +396 -0
  27. data/src/bindings/hub.rs +42 -0
  28. data/src/bindings/matcher.rs +114 -0
  29. data/src/bindings/mod.rs +14 -0
  30. data/src/bindings/padding.rs +177 -0
  31. data/src/bindings/pretokenize.rs +53 -0
  32. data/src/bindings/sources.rs +273 -0
  33. data/src/bindings/train.rs +125 -0
  34. data/src/bpe/mod.rs +1217 -0
  35. data/src/bpe/pretoken_cache.rs +495 -0
  36. data/src/bpe/sentencepiece.rs +1485 -0
  37. data/src/bpe/tiktoken.rs +2555 -0
  38. data/src/bpe_train.rs +351 -0
  39. data/src/input/decompress.rs +11 -0
  40. data/src/input/file_source.rs +514 -0
  41. data/src/input/jsonl.rs +94 -0
  42. data/src/input/mod.rs +333 -0
  43. data/src/input/parquet.rs +303 -0
  44. data/src/lib.rs +578 -0
  45. data/src/load_tokenizer/hf.rs +1036 -0
  46. data/src/load_tokenizer/hub.rs +344 -0
  47. data/src/load_tokenizer/mod.rs +3 -0
  48. data/src/load_tokenizer/tiktoken.rs +87 -0
  49. data/src/main.rs +95 -0
  50. data/src/pretokenize/fast/cl100k.rs +426 -0
  51. data/src/pretokenize/fast/cl100k_family.rs +891 -0
  52. data/src/pretokenize/fast/deepseek_v3.rs +605 -0
  53. data/src/pretokenize/fast/kimi.rs +281 -0
  54. data/src/pretokenize/fast/mask.rs +1486 -0
  55. data/src/pretokenize/fast/mod.rs +446 -0
  56. data/src/pretokenize/fast/nemotron.rs +138 -0
  57. data/src/pretokenize/fast/o200k.rs +347 -0
  58. data/src/pretokenize/fast/o200k_family.rs +1734 -0
  59. data/src/pretokenize/fast/olmo3.rs +505 -0
  60. data/src/pretokenize/fast/qwen2.rs +429 -0
  61. data/src/pretokenize/fast/qwen3_5.rs +541 -0
  62. data/src/pretokenize/fast/r50k.rs +1250 -0
  63. data/src/pretokenize/mod.rs +1079 -0
  64. data/src/pretokenize/options.rs +188 -0
  65. data/src/pretokenize/pretoken.rs +20 -0
  66. data/src/pretokenize/pretokenize_traits.rs +49 -0
  67. data/src/pretokenize/reference/avx512.rs +522 -0
  68. data/src/pretokenize/reference/combinator.rs +572 -0
  69. data/src/pretokenize/reference/mod.rs +28 -0
  70. data/src/pretokenize/reference/simd.rs +852 -0
  71. data/src/pretokenize/reference/state_machine.rs +365 -0
  72. data/src/pretokenize/unicode.rs +546 -0
  73. data/src/test_hub.rs +28 -0
  74. data/src/token.rs +42 -0
  75. metadata +161 -0
@@ -0,0 +1,132 @@
1
+ # frozen_string_literal: true
2
+
3
+ require "etc"
4
+
5
+ module Gigatoken
6
+ module CLI
7
+ # Helpers shared by the bench and validate commands: tokenizer loading,
8
+ # byte-size parsing, document splitting, and CPU identification.
9
+ module Support
10
+ SIZE_UNITS = {"" => 1, "k" => 10**3, "m" => 10**6, "g" => 10**9, "t" => 10**12}.freeze
11
+ private_constant :SIZE_UNITS
12
+
13
+ SIZE_PATTERN = /\A\s*(\d+(?:\.\d+)?)\s*([kmgt]?)i?b?\s*\z/i
14
+ private_constant :SIZE_PATTERN
15
+
16
+ class << self
17
+ # Load TOKENIZER: a tokenizer.json path/directory, a HuggingFace
18
+ # repo id, or a .tiktoken file — see Gigatoken::Tokenizer.load.
19
+ def load_tokenizer(spec)
20
+ Gigatoken::Tokenizer.load(spec)
21
+ end
22
+
23
+ # Parse a decimal byte size like "100MB", "2.5GB", or "1000000";
24
+ # "none"/"unlimited" means no limit.
25
+ def parse_size(text)
26
+ return nil if ["none", "unlimited"].include?(text.strip.downcase)
27
+
28
+ match = SIZE_PATTERN.match(text)
29
+ raise Gigatoken::Error, "cannot parse size #{text.inspect}; expected something like 100MB (or 'none')" unless match
30
+
31
+ (match[1].to_f * SIZE_UNITS.fetch(match[2].downcase)).to_i
32
+ end
33
+
34
+ # A Native::TextFileSource for FILES, split on `separator` when
35
+ # given.
36
+ def text_file_source(files, separator)
37
+ Gigatoken::Native::TextFileSource.new(files.map(&:to_s), separator: separator)
38
+ end
39
+
40
+ # Whole files as raw bytes, one document per file, or (with a
41
+ # separator) the separator-split pieces of each file in order, empty
42
+ # documents skipped.
43
+ def split_docs(files, separator)
44
+ raws = files.map { |file| File.binread(file.to_s) }
45
+ return raws if separator.nil?
46
+
47
+ sep = separator.b
48
+ raws.flat_map { |raw| raw.split(sep).reject(&:empty?) }
49
+ end
50
+
51
+ # The prefix of `docs` totalling at most `limit_bytes`, byte-
52
+ # truncating the final document to fill the budget. Unlike a
53
+ # text-comparison tool, gigatoken encodes raw bytes and does not
54
+ # require the cut to land on a UTF-8 character boundary.
55
+ def subset_docs(docs, limit_bytes)
56
+ return docs if limit_bytes.nil?
57
+
58
+ subset = []
59
+ used = 0
60
+ docs.each do |doc|
61
+ room = limit_bytes - used
62
+ if doc.bytesize <= room
63
+ subset << doc
64
+ used += doc.bytesize
65
+ else
66
+ subset << doc.byteslice(0, room) if room > 0
67
+ break
68
+ end
69
+ end
70
+ subset
71
+ end
72
+
73
+ # The benchmark machine's CPU as "name, N cores", plus ", M sockets"
74
+ # when there is more than one socket.
75
+ def cpu_info
76
+ name, cores, sockets =
77
+ case RbConfig::CONFIG["host_os"]
78
+ when /darwin/ then darwin_cpu_info
79
+ when /linux/ then linux_cpu_info
80
+ end
81
+ name ||= RbConfig::CONFIG["host_cpu"] || "unknown CPU"
82
+ cores ||= Etc.nprocessors
83
+ parts = [name, "#{cores} core#{"s" unless cores == 1}"]
84
+ parts << "#{sockets} sockets" if sockets && sockets > 1
85
+ parts.join(", ")
86
+ end
87
+
88
+ private
89
+
90
+ def darwin_cpu_info
91
+ name = sysctl("machdep.cpu.brand_string")
92
+ [name, sysctl_int("hw.physicalcpu"), sysctl_int("hw.packages")]
93
+ end
94
+
95
+ def sysctl(key)
96
+ output = IO.popen(["sysctl", "-n", key], err: File::NULL, &:read)
97
+ output.strip unless output.nil? || output.empty? || !$?.success?
98
+ rescue Errno::ENOENT
99
+ nil
100
+ end
101
+
102
+ def sysctl_int(key)
103
+ value = sysctl(key)
104
+ Integer(value) if value&.match?(/\A\d+\z/)
105
+ end
106
+
107
+ # Within each processor block "physical id" precedes "core id", so
108
+ # (socket, core) pairs count physical cores across sockets.
109
+ def linux_cpu_info
110
+ name = nil
111
+ physical_ids = Set.new
112
+ socket_core_ids = Set.new
113
+ physical_id = ""
114
+ File.foreach("/proc/cpuinfo") do |line|
115
+ key, _, value = line.partition(":")
116
+ key, value = key.strip, value.strip
117
+ case key
118
+ when "model name" then name ||= value
119
+ when "physical id"
120
+ physical_id = value
121
+ physical_ids << value
122
+ when "core id" then socket_core_ids << [physical_id, value]
123
+ end
124
+ end
125
+ [name, (socket_core_ids.size unless socket_core_ids.empty?), (physical_ids.size unless physical_ids.empty?)]
126
+ rescue Errno::ENOENT
127
+ [nil, nil, nil]
128
+ end
129
+ end
130
+ end
131
+ end
132
+ end
@@ -0,0 +1,55 @@
1
+ # frozen_string_literal: true
2
+
3
+ require_relative "support"
4
+
5
+ module Gigatoken
6
+ module CLI
7
+ # `gigatoken validate TOKENIZER FILES...` — a Ruby-vs-Ruby consistency
8
+ # check: encode FILES via `Tokenizer#encode_files` (native file loading
9
+ # and splitting) and via a Ruby-side split plus `Tokenizer#encode_batch`,
10
+ # and confirm the two paths agree per document. Cross-library validation
11
+ # against another tokenizer implementation is out of scope for v1
12
+ # (BRIEF §3.4).
13
+ class Validate < Dry::CLI::Command
14
+ desc "Check that encode_files agrees with a Ruby-side split plus encode_batch on FILES"
15
+
16
+ argument :tokenizer, required: true, desc: "tokenizer.json path or directory, HuggingFace repo id, or .tiktoken file"
17
+ argument :files, type: :array, required: true, desc: "UTF-8 text files to encode"
18
+
19
+ option :doc_separator, desc: 'document separator to split the files on, e.g. "<|endoftext|>"; whole files are single documents otherwise'
20
+
21
+ def call(tokenizer:, files:, doc_separator: nil, **)
22
+ gt_tokenizer = Support.load_tokenizer(tokenizer)
23
+
24
+ via_files = gt_tokenizer.encode_files(Support.text_file_source(files, doc_separator))
25
+ via_batch = gt_tokenizer.encode_batch(Support.split_docs(files, doc_separator))
26
+
27
+ if via_files.length != via_batch.length
28
+ err.puts "validation FAILED: encode_files produced #{via_files.length} documents, encode_batch produced #{via_batch.length}"
29
+ exit(1)
30
+ end
31
+
32
+ via_files.each_with_index do |doc, index|
33
+ next if doc == via_batch[index]
34
+
35
+ report_mismatch(index, doc, via_batch[index])
36
+ exit(1)
37
+ end
38
+
39
+ out.puts "validation OK: #{via_files.length} documents match"
40
+ rescue Gigatoken::Error => e
41
+ err.puts "error: #{e.message}"
42
+ exit(1)
43
+ end
44
+
45
+ private
46
+
47
+ def report_mismatch(index, from_files, from_batch)
48
+ at = from_files.zip(from_batch).index { |a, b| a != b } || [from_files.length, from_batch.length].min
49
+ err.puts "validation FAILED: document #{index}: first mismatch at token #{at} " \
50
+ "(encode_files #{from_files[at, 5]}... vs encode_batch #{from_batch[at, 5]}..., " \
51
+ "lengths #{from_files.length} vs #{from_batch.length})"
52
+ end
53
+ end
54
+ end
55
+ end
@@ -0,0 +1,18 @@
1
+ # frozen_string_literal: true
2
+
3
+ require "gigatoken"
4
+ require "dry/cli"
5
+
6
+ require_relative "cli/bench"
7
+ require_relative "cli/validate"
8
+
9
+ module Gigatoken
10
+ module CLI
11
+ module Commands
12
+ extend Dry::CLI::Registry
13
+
14
+ register "bench", Gigatoken::CLI::Bench
15
+ register "validate", Gigatoken::CLI::Validate
16
+ end
17
+ end
18
+ end
@@ -0,0 +1,242 @@
1
+ # frozen_string_literal: true
2
+
3
+ require "async"
4
+ require "async/http"
5
+ require "fileutils"
6
+ require "pathname"
7
+
8
+ module Gigatoken
9
+ # HuggingFace Hub file fetch, mirroring `huggingface_hub.hf_hub_download`:
10
+ # same endpoint and URL layout, same token discovery (HF_TOKEN env var,
11
+ # then the token file written by `hf auth login`), same cache directory
12
+ # resolution — without requiring huggingface_hub, tokenizers, or
13
+ # transformers. Files already present in the standard HF cache are served
14
+ # with a pure-filesystem lookup (no network); on a miss the file is
15
+ # downloaded straight into the shared cache, so later loads (ours and
16
+ # huggingface_hub's) are served from it.
17
+ #
18
+ # Network I/O runs on async-http, on the reactor: #hub_file wraps its
19
+ # fetch in `Sync`, so it composes whether the caller is already inside a
20
+ # reactor or is plain sync code.
21
+ class Hub
22
+ # Filename suffixes of local tokenizer files (tokenizer.json contents
23
+ # and raw sentencepiece models). A name ending in one of these is never
24
+ # treated as a Hub repo id, so a mistyped local path fails fast instead
25
+ # of hitting the network. Keep in sync with `src/load_tokenizer/hub.rs`'s
26
+ # TOKENIZER_FILE_SUFFIXES.
27
+ TOKENIZER_FILE_SUFFIXES = [".json", ".model"].freeze
28
+
29
+ DEFAULT_ENDPOINT = "https://huggingface.co"
30
+ private_constant :DEFAULT_ENDPOINT
31
+
32
+ MAX_REDIRECTS = 10
33
+ private_constant :MAX_REDIRECTS
34
+
35
+ class << self
36
+ # Whether `name` is shaped like a HuggingFace Hub repo id: `org/name`,
37
+ # or a bare legacy repo name like `gpt2`. At most one slash, and not
38
+ # something that is obviously a filesystem path to a local tokenizer
39
+ # file.
40
+ def looks_like_repo_id?(name)
41
+ parts = name.split("/", -1)
42
+ return false if parts.empty? || parts.size > 2
43
+
44
+ org, rest = parts
45
+ return false unless word_part?(org, first_alnum: true)
46
+ return false if rest && !word_part?(rest, first_alnum: false)
47
+
48
+ TOKENIZER_FILE_SUFFIXES.none? { |suffix| name.end_with?(suffix) }
49
+ end
50
+
51
+ # $HF_HOME, defaulting to $XDG_CACHE_HOME/huggingface then
52
+ # ~/.cache/huggingface — the root for both the hub cache and the token
53
+ # file.
54
+ def hf_home
55
+ Pathname.new(env("HF_HOME") || File.join(env("XDG_CACHE_HOME") || File.join(Dir.home, ".cache"), "huggingface"))
56
+ end
57
+
58
+ # The standard HuggingFace hub cache directory, resolved like
59
+ # huggingface_hub does it: HF_HUB_CACHE, then $HF_HOME/hub.
60
+ def hf_hub_cache_dir
61
+ Pathname.new(env("HF_HUB_CACHE") || hf_home.join("hub").to_s)
62
+ end
63
+
64
+ # The HuggingFace access token, discovered like huggingface_hub does
65
+ # it: the HF_TOKEN (or legacy HUGGING_FACE_HUB_TOKEN) environment
66
+ # variable, then the token file (HF_TOKEN_PATH, default $HF_HOME/token).
67
+ def hf_token
68
+ token = env("HF_TOKEN") || env("HUGGING_FACE_HUB_TOKEN")
69
+ return token.strip if token
70
+
71
+ token_path = env("HF_TOKEN_PATH") || hf_home.join("token").to_s
72
+ return nil unless File.file?(token_path)
73
+
74
+ token = File.read(token_path).strip
75
+ token unless token.empty?
76
+ end
77
+
78
+ # Path of `filename` in the local HF cache, or nil when not cached. A
79
+ # pure-filesystem lookup — no request is made. `revision` may be a
80
+ # commit hash (used directly as the snapshot name) or a branch/tag
81
+ # name (followed through the cached ref).
82
+ def cached_file(repo_id, filename, revision)
83
+ repo_dir = repo_cache_dir(repo_id)
84
+ commit = commit_hash?(revision) ? revision : cached_ref(repo_dir, revision)
85
+ return nil unless commit
86
+
87
+ path = repo_dir.join("snapshots", commit, filename)
88
+ path if path.file?
89
+ end
90
+
91
+ # The cache directory of a repo (`models--org--name`).
92
+ def repo_cache_dir(repo_id)
93
+ hf_hub_cache_dir.join("models--#{repo_id.gsub("/", "--")}")
94
+ end
95
+
96
+ # A full git commit hash: cache snapshot directories are named by
97
+ # these.
98
+ def commit_hash?(revision)
99
+ revision.match?(/\A[0-9a-f]{40}\z/)
100
+ end
101
+
102
+ private
103
+
104
+ def env(key)
105
+ value = ENV[key]
106
+ value unless value.nil? || value.empty?
107
+ end
108
+
109
+ def word_part?(part, first_alnum:)
110
+ return false if part.nil? || part.empty?
111
+
112
+ first_ok = first_alnum ? part[0].match?(/[A-Za-z0-9]/) : word_char?(part[0])
113
+ first_ok && part[1..].chars.all? { |c| word_char?(c) }
114
+ end
115
+
116
+ def word_char?(char)
117
+ char.match?(/[A-Za-z0-9_.-]/)
118
+ end
119
+
120
+ def cached_ref(repo_dir, revision)
121
+ ref_path = repo_dir.join("refs", revision)
122
+ File.read(ref_path).strip if ref_path.file?
123
+ end
124
+ end
125
+
126
+ # @parameter endpoint [String] the Hub endpoint to fetch from — override
127
+ # for pointing at a local server in tests (dependency injection, not a
128
+ # mock).
129
+ def initialize(endpoint: DEFAULT_ENDPOINT)
130
+ @endpoint = endpoint.chomp("/")
131
+ @internet = Async::HTTP::Internet.new
132
+ end
133
+
134
+ # Path of `filename` from Hub repo `repo_id` at `revision`, served from
135
+ # the standard HF cache, downloading into it first when absent.
136
+ def hub_file(repo_id, filename = "tokenizer.json", revision: "main")
137
+ self.class.cached_file(repo_id, filename, revision) ||
138
+ Sync { fetch(repo_id, filename, revision) }
139
+ end
140
+
141
+ private
142
+
143
+ # GET `endpoint/repo/resolve/revision/filename` and stream the body into
144
+ # the cache snapshot named by the `x-repo-commit` response header,
145
+ # recording the branch ref so later lookups (ours and
146
+ # huggingface_hub's) resolve it.
147
+ def fetch(repo_id, filename, revision)
148
+ url = "#{@endpoint}/#{repo_id}/resolve/#{revision}/#{filename}"
149
+ token = self.class.hf_token
150
+ response = @internet.get(url, auth_headers(token))
151
+ # Unlisted headers parse as a Header::Generic (an Array of values);
152
+ # x-repo-commit is always a single value, so flatten it to a String.
153
+ commit = response.headers["x-repo-commit"]&.to_s
154
+
155
+ # Redirects are followed by hand: resolve/ URLs answer with the
156
+ # x-repo-commit header and a redirect to a CDN for LFS files, and the
157
+ # Authorization header must not travel to the other host.
158
+ hops = 0
159
+ while (300...400).cover?(response.status)
160
+ ensure_ok!(url, response.status, !!token)
161
+ hops += 1
162
+ raise Error, "#{url}: too many redirects" if hops > MAX_REDIRECTS
163
+
164
+ location = response.headers["location"] ||
165
+ raise(Error, "#{url}: redirect with no Location header")
166
+ response.close
167
+ url = absolutize(location, url)
168
+ response = @internet.get(url, {"user-agent" => "gigatoken"})
169
+ end
170
+ ensure_ok!(url, response.status, !!token)
171
+
172
+ write_to_cache(repo_id, filename, revision, commit || revision, response)
173
+ end
174
+
175
+ def auth_headers(token)
176
+ headers = {"user-agent" => "gigatoken"}
177
+ headers["authorization"] = "Bearer #{token}" if token
178
+ headers
179
+ end
180
+
181
+ # Stream the response body to a sibling temp file, then rename into
182
+ # place: concurrent downloaders race benignly and readers never observe
183
+ # a partial file.
184
+ def write_to_cache(repo_id, filename, revision, commit, response)
185
+ repo_dir = self.class.repo_cache_dir(repo_id)
186
+ target = repo_dir.join("snapshots", commit, filename)
187
+ FileUtils.mkdir_p(target.dirname)
188
+
189
+ tmp = target.dirname.join(".#{target.basename}.#{Process.pid}.tmp")
190
+ begin
191
+ response.save(tmp.to_s)
192
+ rescue
193
+ FileUtils.rm_f(tmp)
194
+ raise
195
+ ensure
196
+ response.close
197
+ end
198
+ File.rename(tmp, target)
199
+
200
+ if !self.class.commit_hash?(revision) && revision != commit
201
+ refs_dir = repo_dir.join("refs")
202
+ FileUtils.mkdir_p(refs_dir)
203
+ File.write(refs_dir.join(revision), commit)
204
+ end
205
+
206
+ target
207
+ end
208
+
209
+ def ensure_ok!(url, status, had_token)
210
+ return if (200...400).cover?(status)
211
+
212
+ case status
213
+ when 404
214
+ raise Error, "#{url}: HTTP 404 — no such repo with that file, and no such local file either"
215
+ when 401, 403
216
+ token_note = had_token ? "the request used the discovered token" : "no token was found"
217
+ raise Error,
218
+ "#{url}: HTTP #{status} — the repo may be private or gated (#{token_note}; set HF_TOKEN or run " \
219
+ "`hf auth login`, and accept the repo's terms on huggingface.co if it is gated)"
220
+ else
221
+ raise Error, "#{url}: HTTP #{status}"
222
+ end
223
+ end
224
+
225
+ # A redirect Location resolved against the request URL: absolute URLs
226
+ # pass through, host-relative (`/x/y`) and path-relative ones join the
227
+ # base.
228
+ def absolutize(location, base)
229
+ return location if location.include?("://")
230
+
231
+ origin_end = base.index("://") ? base.index("://") + 3 : 0
232
+ origin_end = base.index("/", origin_end) || base.length
233
+
234
+ if location.start_with?("/")
235
+ "#{base[0...origin_end]}#{location}"
236
+ else
237
+ dir_end = base.rindex("/") || base.length
238
+ "#{base[0...[dir_end, origin_end].max]}/#{location}"
239
+ end
240
+ end
241
+ end
242
+ end
@@ -0,0 +1,48 @@
1
+ # frozen_string_literal: true
2
+
3
+ module Gigatoken
4
+ # A packed batch-encode result: one `IO::Buffer` of u32 token ids (native
5
+ # byte order) for every document, plus `lens`, each document's token count.
6
+ # Returned by `Tokenizer#encode_batch`/`#encode_files` with `packed: true`
7
+ # — the Ruby analog of the numpy/awkward-array path the native ext's
8
+ # header comment says was dropped for lack of a Ruby equivalent. Documents
9
+ # are unpacked from the buffer on demand via `#[]`/`#each`.
10
+ class PackedResult
11
+ include Enumerable
12
+
13
+ attr_reader :buffer, :lens
14
+
15
+ def initialize(buffer, lens)
16
+ @buffer = buffer
17
+ @lens = lens
18
+ @offsets = lens.inject([0]) { |offsets, len| offsets << offsets.last + len }
19
+ end
20
+
21
+ # Number of documents.
22
+ def size
23
+ lens.size
24
+ end
25
+
26
+ # Total token count across every document.
27
+ def token_count
28
+ lens.sum
29
+ end
30
+
31
+ # Array of token ids for document `i`, materialized on demand.
32
+ def [](i)
33
+ buffer.get_values(Array.new(lens[i], :u32), @offsets[i] * 4)
34
+ end
35
+
36
+ def each
37
+ return enum_for(:each) unless block_given?
38
+
39
+ lens.each_index { |i| yield self[i] }
40
+ end
41
+
42
+ # A ragged Array of Arrays, one per document — the same shape
43
+ # `encode_batch`/`encode_files` return with `packed: false`.
44
+ def to_a
45
+ each.to_a
46
+ end
47
+ end
48
+ end
@@ -0,0 +1,117 @@
1
+ # frozen_string_literal: true
2
+
3
+ require "json"
4
+
5
+ module Gigatoken
6
+ # A tokenizer: encode, batch encode, decode, and vocabulary introspection
7
+ # over a native `Gigatoken::Native::BPETokenizer` or
8
+ # `Gigatoken::Native::SentencePieceTokenizer`.
9
+ class Tokenizer
10
+ TIKTOKEN_ENDOFTEXT = "<|endoftext|>"
11
+ private_constant :TIKTOKEN_ENDOFTEXT
12
+
13
+ FILE_SOURCE_CLASSES = [Native::TextFileSource, Native::JsonlFileSource, Native::ParquetFileSource].freeze
14
+ private_constant :FILE_SOURCE_CLASSES
15
+
16
+ # Load from in-memory tokenizer.json contents (String or bytes). Backed
17
+ # by a BPETokenizer or a SentencePieceTokenizer, per the model's
18
+ # byte_fallback flag.
19
+ def self.from_json(data)
20
+ native = Native.load_hf_json(data)
21
+ new(native, special_tokens: special_tokens_from_json(data))
22
+ end
23
+
24
+ # Load from a tokenizer.json path, or a directory containing one.
25
+ def self.from_file(path)
26
+ path = File.join(path, "tokenizer.json") if File.directory?(path)
27
+ from_json(File.binread(path))
28
+ end
29
+
30
+ # Load from a .tiktoken mergeable-ranks file.
31
+ def self.from_tiktoken(path)
32
+ native = Native::BPETokenizer.from_tiktoken(path.to_s)
33
+ new(native, special_tokens: {TIKTOKEN_ENDOFTEXT => native.vocab_size - 1})
34
+ end
35
+
36
+ # Load tokenizer.json from HuggingFace Hub repo `repo_id` at `revision`
37
+ # (downloaded directly; huggingface_hub is not required).
38
+ def self.from_hub(repo_id, revision: "main", hub: Hub.new)
39
+ from_file(hub.hub_file(repo_id, "tokenizer.json", revision: revision))
40
+ end
41
+
42
+ # Load from any of the supported source shapes: an existing file or
43
+ # directory path (a tokenizer.json, or a directory containing one), a
44
+ # .tiktoken vocabulary file, or a HuggingFace Hub repo id like
45
+ # "openai-community/gpt2".
46
+ def self.load(source, revision: "main", hub: Hub.new)
47
+ source = source.to_s
48
+ return from_tiktoken(source) if source.end_with?(".tiktoken")
49
+ return from_file(source) if File.exist?(source)
50
+ return from_hub(source, revision: revision, hub: hub) if Hub.looks_like_repo_id?(source)
51
+
52
+ raise Error, "#{source.inspect}: no such file or directory, not a .tiktoken path, and doesn't look like a HuggingFace Hub repo id"
53
+ end
54
+
55
+ def self.special_tokens_from_json(data)
56
+ added = JSON.parse(data.dup.force_encoding(Encoding::UTF_8))["added_tokens"] || []
57
+ added.each_with_object({}) { |t, h| h[t["content"]] = t["id"] if t["special"] }
58
+ end
59
+ private_class_method :special_tokens_from_json
60
+
61
+ def initialize(native, special_tokens: {})
62
+ @native = native
63
+ @special_tokens = special_tokens
64
+ end
65
+
66
+ def encode(text)
67
+ @native.encode(text)
68
+ end
69
+
70
+ # Returns a ragged Array of Arrays of token ids, one row per document —
71
+ # or, with `packed: true`, a Gigatoken::PackedResult (one IO::Buffer of
72
+ # token ids plus per-document lengths), avoiding the per-token Ruby
73
+ # array materialization the ragged shape costs.
74
+ def encode_batch(texts, packed: false)
75
+ if packed
76
+ PackedResult.new(*@native.encode_batch_packed(texts))
77
+ else
78
+ @native.encode_batch(texts)
79
+ end
80
+ end
81
+
82
+ # Tokenize whole files in Rust: reads and encodes them in one fused pass
83
+ # without the documents ever becoming Ruby objects. `source` is a
84
+ # Native::{Text,Jsonl,Parquet}FileSource, a single path, or an array of
85
+ # paths; bare path(s) are wrapped in a TextFileSource (with `separator`,
86
+ # if given). Returns a ragged Array of Arrays of token ids, one row per
87
+ # document — or, with `packed: true`, a Gigatoken::PackedResult. `parallel:
88
+ # false` loads and encodes everything on the calling thread instead, with
89
+ # identical output, never touching the core worker pool.
90
+ def encode_files(source, separator: nil, parallel: true, packed: false)
91
+ source = Native::TextFileSource.new(Array(source).map(&:to_s), separator: separator) unless FILE_SOURCE_CLASSES.any? { |klass| source.is_a?(klass) }
92
+ if packed
93
+ PackedResult.new(*@native.encode_files_packed(source, parallel: parallel))
94
+ else
95
+ @native.encode_files(source, parallel: parallel)
96
+ end
97
+ end
98
+
99
+ def decode(ids)
100
+ @native.decode(ids)
101
+ end
102
+
103
+ def vocab_size
104
+ @native.vocab_size
105
+ end
106
+
107
+ def vocab
108
+ @native.vocab
109
+ end
110
+
111
+ def merges
112
+ @native.merges
113
+ end
114
+
115
+ attr_reader :special_tokens
116
+ end
117
+ end
@@ -0,0 +1,5 @@
1
+ # frozen_string_literal: true
2
+
3
+ module Gigatoken
4
+ VERSION = "0.1.0"
5
+ end
data/lib/gigatoken.rb ADDED
@@ -0,0 +1,29 @@
1
+ # frozen_string_literal: true
2
+
3
+ require_relative "gigatoken/version"
4
+
5
+ module Gigatoken
6
+ # Raised for tokenizer load and encode failures surfaced from the native
7
+ # extension — never a raw Rust panic across the Ruby boundary.
8
+ class Error < StandardError; end
9
+
10
+ NATIVE_EXTENSIONS = %w[.bundle .so .rb].freeze
11
+
12
+ # Precompiled native gems ship per-ABI subdirs (`gigatoken/4.0/...`),
13
+ # the source-gem `rake compile` build lands flat (`gigatoken/...`).
14
+ # Pick whichever exists for the current Ruby ABI, with the per-ABI path
15
+ # winning when both are present.
16
+ def self.locate_native(base, ruby_version: RUBY_VERSION)
17
+ abi = ruby_version[/\d+\.\d+/]
18
+ candidates = [File.join(base, abi, "gigatoken_rb"), File.join(base, "gigatoken_rb")]
19
+ candidates.find { |stem| NATIVE_EXTENSIONS.any? { |ext| File.exist?(stem + ext) } }
20
+ end
21
+ end
22
+
23
+ native = Gigatoken.locate_native(File.expand_path("gigatoken", __dir__))
24
+ raise LoadError, "could not locate gigatoken native extension" unless native
25
+ require native
26
+
27
+ require_relative "gigatoken/hub"
28
+ require_relative "gigatoken/packed_result"
29
+ require_relative "gigatoken/tokenizer"
@@ -0,0 +1,8 @@
1
+ [toolchain]
2
+ channel = "nightly"
3
+ # Cross triples for the precompiled-gem builds (ruby-gem.yml). rustup's
4
+ # toolchain-file auto-install is the only thing that provisions nightly
5
+ # inside the rb-sys-dock containers, so the targets must be declared here —
6
+ # the cross-gem action's `pre-script` input is dead plumbing (undeclared at
7
+ # v1.4.4, and no released rb-sys-dock reads INPUT_PRE_SCRIPT).
8
+ targets = ["aarch64-apple-darwin", "aarch64-unknown-linux-gnu", "x86_64-unknown-linux-gnu"]