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.
- checksums.yaml +7 -0
- data/Cargo.lock +3016 -0
- data/Cargo.toml +135 -0
- data/LICENSE +21 -0
- data/README.md +141 -0
- data/exe/gigatoken +9 -0
- data/ext/gigatoken/Cargo.toml +24 -0
- data/ext/gigatoken/extconf.rb +19 -0
- data/ext/gigatoken/src/error.rs +20 -0
- data/ext/gigatoken/src/gvl.rs +122 -0
- data/ext/gigatoken/src/lib.rs +50 -0
- data/ext/gigatoken/src/sentencepiece.rs +205 -0
- data/ext/gigatoken/src/sources.rs +207 -0
- data/ext/gigatoken/src/tokenizer.rs +571 -0
- data/lib/gigatoken/cli/bench.rb +64 -0
- data/lib/gigatoken/cli/support.rb +132 -0
- data/lib/gigatoken/cli/validate.rb +55 -0
- data/lib/gigatoken/cli.rb +18 -0
- data/lib/gigatoken/hub.rb +242 -0
- data/lib/gigatoken/packed_result.rb +48 -0
- data/lib/gigatoken/tokenizer.rb +117 -0
- data/lib/gigatoken/version.rb +5 -0
- data/lib/gigatoken.rb +29 -0
- data/rust-toolchain.toml +8 -0
- data/src/batch.rs +1808 -0
- data/src/bindings/bridge.rs +396 -0
- data/src/bindings/hub.rs +42 -0
- data/src/bindings/matcher.rs +114 -0
- data/src/bindings/mod.rs +14 -0
- data/src/bindings/padding.rs +177 -0
- data/src/bindings/pretokenize.rs +53 -0
- data/src/bindings/sources.rs +273 -0
- data/src/bindings/train.rs +125 -0
- data/src/bpe/mod.rs +1217 -0
- data/src/bpe/pretoken_cache.rs +495 -0
- data/src/bpe/sentencepiece.rs +1485 -0
- data/src/bpe/tiktoken.rs +2555 -0
- data/src/bpe_train.rs +351 -0
- data/src/input/decompress.rs +11 -0
- data/src/input/file_source.rs +514 -0
- data/src/input/jsonl.rs +94 -0
- data/src/input/mod.rs +333 -0
- data/src/input/parquet.rs +303 -0
- data/src/lib.rs +578 -0
- data/src/load_tokenizer/hf.rs +1036 -0
- data/src/load_tokenizer/hub.rs +344 -0
- data/src/load_tokenizer/mod.rs +3 -0
- data/src/load_tokenizer/tiktoken.rs +87 -0
- data/src/main.rs +95 -0
- data/src/pretokenize/fast/cl100k.rs +426 -0
- data/src/pretokenize/fast/cl100k_family.rs +891 -0
- data/src/pretokenize/fast/deepseek_v3.rs +605 -0
- data/src/pretokenize/fast/kimi.rs +281 -0
- data/src/pretokenize/fast/mask.rs +1486 -0
- data/src/pretokenize/fast/mod.rs +446 -0
- data/src/pretokenize/fast/nemotron.rs +138 -0
- data/src/pretokenize/fast/o200k.rs +347 -0
- data/src/pretokenize/fast/o200k_family.rs +1734 -0
- data/src/pretokenize/fast/olmo3.rs +505 -0
- data/src/pretokenize/fast/qwen2.rs +429 -0
- data/src/pretokenize/fast/qwen3_5.rs +541 -0
- data/src/pretokenize/fast/r50k.rs +1250 -0
- data/src/pretokenize/mod.rs +1079 -0
- data/src/pretokenize/options.rs +188 -0
- data/src/pretokenize/pretoken.rs +20 -0
- data/src/pretokenize/pretokenize_traits.rs +49 -0
- data/src/pretokenize/reference/avx512.rs +522 -0
- data/src/pretokenize/reference/combinator.rs +572 -0
- data/src/pretokenize/reference/mod.rs +28 -0
- data/src/pretokenize/reference/simd.rs +852 -0
- data/src/pretokenize/reference/state_machine.rs +365 -0
- data/src/pretokenize/unicode.rs +546 -0
- data/src/test_hub.rs +28 -0
- data/src/token.rs +42 -0
- 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
|
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"
|
data/rust-toolchain.toml
ADDED
|
@@ -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"]
|