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
data/Cargo.toml
ADDED
|
@@ -0,0 +1,135 @@
|
|
|
1
|
+
[workspace]
|
|
2
|
+
members = [".", "ext/gigatoken"]
|
|
3
|
+
|
|
4
|
+
[package]
|
|
5
|
+
name = "gigatoken"
|
|
6
|
+
version = "0.9.0"
|
|
7
|
+
edition = "2024"
|
|
8
|
+
license = "MIT"
|
|
9
|
+
|
|
10
|
+
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
|
|
11
|
+
[lib]
|
|
12
|
+
name = "gigatoken_rs"
|
|
13
|
+
crate-type = ["cdylib", "lib"]
|
|
14
|
+
|
|
15
|
+
[dependencies]
|
|
16
|
+
dashmap = "6.1"
|
|
17
|
+
indicatif = "0.18"
|
|
18
|
+
itertools = "0.15"
|
|
19
|
+
priority-queue = "2.0"
|
|
20
|
+
# For now uses abi3-py310, should use version-specific ABI in the future
|
|
21
|
+
pyo3 = { version = "0.29", features = ["abi3-py310", "eyre"], optional = true }
|
|
22
|
+
numpy = { version = "0.29", optional = true }
|
|
23
|
+
rayon = "1.11"
|
|
24
|
+
rustc-hash = "2.1"
|
|
25
|
+
memmap2 = { version = "0.9", features = ["stable_deref_trait"] }
|
|
26
|
+
base64 = "0.22"
|
|
27
|
+
icu = { version = "2.2", features = ["compiled_data", "datagen"] }
|
|
28
|
+
eyre = "0.6.12"
|
|
29
|
+
sonic-rs = "0.5.8"
|
|
30
|
+
memchr = "2.8.0"
|
|
31
|
+
aho-corasick = "1.1"
|
|
32
|
+
serde = { version = "1.0.228", features = ["derive"] }
|
|
33
|
+
flate2 = "1.1"
|
|
34
|
+
zstd = "0.13"
|
|
35
|
+
spm_precompiled = "0.1.4"
|
|
36
|
+
simdutf = "0.7.0"
|
|
37
|
+
winnow = { version = "1", features = ["simd"] }
|
|
38
|
+
parquet = { version = "59.1.0", default-features = false, features = [
|
|
39
|
+
"arrow",
|
|
40
|
+
"snap",
|
|
41
|
+
"zstd",
|
|
42
|
+
"flate2",
|
|
43
|
+
"lz4",
|
|
44
|
+
"brotli",
|
|
45
|
+
"flate2-rust_backend",
|
|
46
|
+
] }
|
|
47
|
+
arrow-array = "59"
|
|
48
|
+
arrow-schema = "59"
|
|
49
|
+
ureq = { version = "3.3.0", default-features = false, features = ["rustls"] }
|
|
50
|
+
|
|
51
|
+
# madvise/prctl on Linux (pretoken cache table, batch gather, bench THP
|
|
52
|
+
# setup); already in the tree transitively via memmap2/pyo3.
|
|
53
|
+
[target.'cfg(target_os = "linux")'.dependencies]
|
|
54
|
+
libc = "0.2"
|
|
55
|
+
|
|
56
|
+
[features]
|
|
57
|
+
default = ["python"]
|
|
58
|
+
python = ["dep:pyo3", "dep:numpy"]
|
|
59
|
+
|
|
60
|
+
[profile.release]
|
|
61
|
+
lto = "fat"
|
|
62
|
+
|
|
63
|
+
# Profiling parity: benches keep the release codegen (fat LTO, same inlining)
|
|
64
|
+
# but carry full debug info so profilers can resolve inlined frames and source
|
|
65
|
+
# lines. Debug info does not change generated code; parity is verified by
|
|
66
|
+
# comparing bench throughput against a stripped release build (within noise).
|
|
67
|
+
[profile.bench]
|
|
68
|
+
inherits = "release"
|
|
69
|
+
debug = true
|
|
70
|
+
strip = "none"
|
|
71
|
+
split-debuginfo = "packed"
|
|
72
|
+
|
|
73
|
+
[profile.profiling]
|
|
74
|
+
inherits = "release"
|
|
75
|
+
debug = 2
|
|
76
|
+
lto = false
|
|
77
|
+
codegen-units = 1
|
|
78
|
+
strip = "none"
|
|
79
|
+
split-debuginfo = "packed"
|
|
80
|
+
# Frame pointers make Instruments' stack unwinding reliable. Mangling is left at
|
|
81
|
+
# the toolchain default (v0); `cargo instruments` demangles v0 symbols itself.
|
|
82
|
+
rustflags = ["-C", "force-frame-pointers=yes"]
|
|
83
|
+
|
|
84
|
+
[profile.profiling.package."*"]
|
|
85
|
+
debug = 2
|
|
86
|
+
strip = "none"
|
|
87
|
+
|
|
88
|
+
[dev-dependencies]
|
|
89
|
+
criterion = "0.8"
|
|
90
|
+
fancy-regex = "0.18.0"
|
|
91
|
+
rand = "0.10"
|
|
92
|
+
simdutf = "0.7.0"
|
|
93
|
+
tempfile = "3.27.0"
|
|
94
|
+
|
|
95
|
+
|
|
96
|
+
[[bench]]
|
|
97
|
+
name = "pretokenize"
|
|
98
|
+
harness = false
|
|
99
|
+
|
|
100
|
+
[[bench]]
|
|
101
|
+
name = "encode"
|
|
102
|
+
harness = false
|
|
103
|
+
test = false
|
|
104
|
+
|
|
105
|
+
[[bench]]
|
|
106
|
+
name = "encode_st"
|
|
107
|
+
harness = false
|
|
108
|
+
test = false
|
|
109
|
+
|
|
110
|
+
[[bench]]
|
|
111
|
+
name = "encode_st_sp"
|
|
112
|
+
harness = false
|
|
113
|
+
test = false
|
|
114
|
+
|
|
115
|
+
[[bench]]
|
|
116
|
+
name = "encode_doc"
|
|
117
|
+
harness = false
|
|
118
|
+
test = false
|
|
119
|
+
|
|
120
|
+
[[bench]]
|
|
121
|
+
name = "pretokenize_profile"
|
|
122
|
+
harness = false
|
|
123
|
+
test = false
|
|
124
|
+
|
|
125
|
+
[[bench]]
|
|
126
|
+
name = "pretokenize_profile_all"
|
|
127
|
+
harness = false
|
|
128
|
+
test = false
|
|
129
|
+
|
|
130
|
+
[[bench]]
|
|
131
|
+
name = "simdutf_transcode"
|
|
132
|
+
harness = false
|
|
133
|
+
|
|
134
|
+
[lints.rust]
|
|
135
|
+
unused = "allow"
|
data/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Marcel Rød
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
data/README.md
ADDED
|
@@ -0,0 +1,141 @@
|
|
|
1
|
+
# gigatoken-rb
|
|
2
|
+
|
|
3
|
+
**12 GB/s / 2.8 billion tokens per second in Ruby.**
|
|
4
|
+
|
|
5
|
+
Ruby bindings for [marcelroed/gigatoken](https://github.com/marcelroed/gigatoken), the fastest open-source BPE tokenizer around — running **1.6x faster than upstream's own Python package**, on the same Rust engine.
|
|
6
|
+
|
|
7
|
+
| | Corpus | MB/s (median) | Gtok/s (median) |
|
|
8
|
+
|---|---|---|---|
|
|
9
|
+
| **gigatoken** (this gem, Ruby) | 11.9 GB | **12,278** | **2.78** |
|
|
10
|
+
| gigatoken (Python wheel, upstream) | 11.9 GB | 7,400 | 1.68 |
|
|
11
|
+
| tiktoken (Python) | 1.35 GB | 69.7 | 0.0158 |
|
|
12
|
+
| tiktoken_ruby | 1.35 GB | 30.7 | 0.0070 |
|
|
13
|
+
| tokenizers gem (ankane) | 1.35 GB | 10.0 | 0.0023 |
|
|
14
|
+
| tokenizers (Python, Hugging Face) | 1.35 GB | 5.6 | 0.0013 |
|
|
15
|
+
|
|
16
|
+
Mac Studio M4 Max, OpenWebText, GPT-2 tokenizer; every library produces the same tokenization, gigatoken just does it faster. **340x faster** than the fastest existing Ruby gem (tiktoken_ruby) and **1,050x faster** than the tokenizers gem. Full methodology, exact counts, and the caveats that matter: [docs/rb/benchmarks.md](docs/rb/benchmarks.md).
|
|
17
|
+
|
|
18
|
+
## Install
|
|
19
|
+
|
|
20
|
+
```bash
|
|
21
|
+
gem install gigatoken
|
|
22
|
+
```
|
|
23
|
+
|
|
24
|
+
Precompiled native gems ship for Apple Silicon macOS (`arm64-darwin`) and x86_64/aarch64 Linux — on those platforms RubyGems grabs the binary automatically, no Rust toolchain, no compile wait. In a Bundler project it's one command:
|
|
25
|
+
|
|
26
|
+
```bash
|
|
27
|
+
bundle add gigatoken
|
|
28
|
+
```
|
|
29
|
+
|
|
30
|
+
(or drop `gem "gigatoken"` into the Gemfile yourself).
|
|
31
|
+
|
|
32
|
+
Anywhere else (or with `--platform ruby` to opt out of the binary), the extension builds from source. That needs a Rust toolchain: `rust-toolchain.toml` pins the nightly, and `rustup` fetches it automatically on first build.
|
|
33
|
+
|
|
34
|
+
## Use
|
|
35
|
+
|
|
36
|
+
```ruby
|
|
37
|
+
require "gigatoken"
|
|
38
|
+
|
|
39
|
+
tok = Gigatoken::Tokenizer.load("openai-community/gpt2")
|
|
40
|
+
|
|
41
|
+
tok.encode("Hello, world!") # => [15496, 11, 995, 0]
|
|
42
|
+
tok.decode([15496, 11, 995, 0]) # => "Hello, world!"
|
|
43
|
+
tok.encode_batch(["Hello!", "Another"]) # => [[15496, 0], [6610]]
|
|
44
|
+
|
|
45
|
+
tok.vocab_size # => 50257
|
|
46
|
+
tok.special_tokens # => {"<|endoftext|>" => 50256}
|
|
47
|
+
```
|
|
48
|
+
|
|
49
|
+
`load` takes a `tokenizer.json` path, a directory holding one, a HuggingFace Hub repo id, or a `.tiktoken` mergeable-ranks file, and dispatches on shape. Hub downloads run over socketry's `async-http` — no Python anywhere. Know what you have? Skip the dispatch:
|
|
50
|
+
|
|
51
|
+
```ruby
|
|
52
|
+
Gigatoken::Tokenizer.from_file("tokenizer.json")
|
|
53
|
+
Gigatoken::Tokenizer.from_hub("openai-community/gpt2", revision: "main")
|
|
54
|
+
Gigatoken::Tokenizer.from_tiktoken("vocab.tiktoken")
|
|
55
|
+
Gigatoken::Tokenizer.from_json(File.binread("tokenizer.json"))
|
|
56
|
+
```
|
|
57
|
+
|
|
58
|
+
SentencePiece-BPE models (Llama, Gemma, Mistral — any `tokenizer.json` with `byte_fallback: true`) load through the same entry points and pick the right backend automatically. One difference: the SentencePiece core decodes text, so it validates input and raises `Gigatoken::Error` on invalid UTF-8 instead of guessing.
|
|
59
|
+
|
|
60
|
+
### Tokenize files without leaving Rust
|
|
61
|
+
|
|
62
|
+
`encode_files` reads and tokenizes files entirely on the native side — document contents never materialize as Ruby objects. `.gz` and `.zst` decompress transparently.
|
|
63
|
+
|
|
64
|
+
```ruby
|
|
65
|
+
tok.encode_files("owt_train.txt", separator: "<|endoftext|>")
|
|
66
|
+
|
|
67
|
+
jsonl = Gigatoken::Native::JsonlFileSource.new(["docs.jsonl"], field: "text")
|
|
68
|
+
parquet = Gigatoken::Native::ParquetFileSource.new(["docs.parquet"], column: "text")
|
|
69
|
+
tok.encode_files(jsonl)
|
|
70
|
+
```
|
|
71
|
+
|
|
72
|
+
### Packed results
|
|
73
|
+
|
|
74
|
+
Pass `packed: true` to `encode_batch` or `encode_files` and results land in a single `IO::Buffer` of u32 token ids instead of a ragged Array of Arrays — no per-token Ruby allocation, the fastest way out of the engine:
|
|
75
|
+
|
|
76
|
+
```ruby
|
|
77
|
+
packed = tok.encode_files("owt_train.txt", packed: true, separator: "<|endoftext|>")
|
|
78
|
+
|
|
79
|
+
packed.buffer # => one IO::Buffer, every document's ids back to back
|
|
80
|
+
packed.lens # => [12, 8, 41, ...] tokens per document
|
|
81
|
+
packed.token_count # => total tokens
|
|
82
|
+
packed[3] # => document 3's ids as an Array, on demand
|
|
83
|
+
```
|
|
84
|
+
|
|
85
|
+
### Async
|
|
86
|
+
|
|
87
|
+
`encode_batch` and `encode_files` release the GVL for the whole encode; the parallelism runs on the engine's rayon pool, not Ruby threads. Under `Async`, give the fiber scheduler a worker pool (`ASYNC_SCHEDULER_WORKER_POOL=true`) and the calling fiber yields to the reactor too. Design notes: [docs/rb/async.md](docs/rb/async.md).
|
|
88
|
+
|
|
89
|
+
## CLI
|
|
90
|
+
|
|
91
|
+
```bash
|
|
92
|
+
gigatoken bench openai-community/gpt2 owt_train.txt --doc-separator "<|endoftext|>"
|
|
93
|
+
gigatoken validate openai-community/gpt2 owt_train.txt --doc-separator "<|endoftext|>"
|
|
94
|
+
```
|
|
95
|
+
|
|
96
|
+
`bench` reports MB/s and Mtok/s (`--packed` for the fused packed path, `--no-parallel` for the serial core). `validate` confirms native split-and-encode agrees with a Ruby-side split through `encode_batch`.
|
|
97
|
+
|
|
98
|
+
## Development
|
|
99
|
+
|
|
100
|
+
```bash
|
|
101
|
+
bundle install
|
|
102
|
+
bundle exec rake compile # native extension (Rust nightly, via rust-toolchain.toml)
|
|
103
|
+
bundle exec rspec
|
|
104
|
+
bundle exec standardrb
|
|
105
|
+
```
|
|
106
|
+
|
|
107
|
+
The Ruby layer is fiber-first throughout — no `Thread`, no `Mutex`; all parallelism lives in the core's rayon pool. CI runs ubuntu + macos × Ruby 3.3/3.4/4.0, and `release.yml` cross-builds the precompiled native gems (arm64-darwin, x86_64-linux, aarch64-linux).
|
|
108
|
+
|
|
109
|
+
## Fork status
|
|
110
|
+
|
|
111
|
+
This fork exists because I need fast tokenization in Ruby. The Rust core is changed as little as possible from upstream. Most of the python shell has been removed from this fork, but you can still find it [upstream](https://github.com/marcelroed/gigatoken).
|
|
112
|
+
|
|
113
|
+
Not ported/no current plans:
|
|
114
|
+
- the HF/tiktoken Python compat shims
|
|
115
|
+
- padded-batch matrices
|
|
116
|
+
- and BPE training
|
|
117
|
+
|
|
118
|
+
SentencePiece works but — matching upstream — is less optimized than the BPE path.
|
|
119
|
+
|
|
120
|
+
## Citation
|
|
121
|
+
|
|
122
|
+
The engine is Marcel Rød's gigatoken. If it shows up in your research, cite that:
|
|
123
|
+
|
|
124
|
+
```bibtex
|
|
125
|
+
@software{roed2026gigatoken,
|
|
126
|
+
author = {Marcel R{\o}d},
|
|
127
|
+
title = {{G}igatoken: SIMD and Cache Hierarchies for 1000x Faster Byte-Pair Encoding Tokenization on Modern CPUs},
|
|
128
|
+
url = {https://github.com/marcelroed/gigatoken},
|
|
129
|
+
year = {2026},
|
|
130
|
+
}
|
|
131
|
+
```
|
|
132
|
+
|
|
133
|
+
---
|
|
134
|
+
|
|
135
|
+
<details>
|
|
136
|
+
<summary>AI Use Disclosure</summary>
|
|
137
|
+
|
|
138
|
+
The Rust engine is upstream's — see <a href="https://github.com/marcelroed/gigatoken#readme">upstream's AI-use disclosure</a> for how that was built (majority hand-crafted, AI-assisted toward the end).
|
|
139
|
+
|
|
140
|
+
The Ruby port in this fork is 100% AI generated using Fable 5 and Sonnet 5 via [space-architect](https://github.com/jetpks/space-architect) over ~24 hours.
|
|
141
|
+
</details>
|
data/exe/gigatoken
ADDED
|
@@ -0,0 +1,24 @@
|
|
|
1
|
+
[package]
|
|
2
|
+
name = "gigatoken-rb"
|
|
3
|
+
version = "0.1.0"
|
|
4
|
+
edition = "2021"
|
|
5
|
+
|
|
6
|
+
[lib]
|
|
7
|
+
crate-type = ["cdylib"]
|
|
8
|
+
|
|
9
|
+
[dependencies]
|
|
10
|
+
# `rb-sys`: lets us reach `magnus::rb_sys::{AsRawValue, protect}` for the raw
|
|
11
|
+
# rb-sys VALUE conversions and locktmp calls this crate needs directly (see
|
|
12
|
+
# tokenizer.rs) that magnus's safe API doesn't expose.
|
|
13
|
+
magnus = { version = "0.8", features = ["rb-sys"] }
|
|
14
|
+
gigatoken = { path = "../..", default-features = false }
|
|
15
|
+
rayon = "1.11"
|
|
16
|
+
mimalloc = { version = "0.1", default-features = false }
|
|
17
|
+
# Same version requirement and feature set magnus-0.8.2 already puts on its
|
|
18
|
+
# own rb-sys dependency (see magnus-0.8.2/Cargo.toml), so Cargo resolves both
|
|
19
|
+
# to the one rb-sys already in Cargo.lock rather than a second copy.
|
|
20
|
+
rb-sys = { version = ">=0.9.113", default-features = false, features = [
|
|
21
|
+
"bindgen-rbimpls",
|
|
22
|
+
"bindgen-deprecated-types",
|
|
23
|
+
"stable-api",
|
|
24
|
+
] }
|
|
@@ -0,0 +1,19 @@
|
|
|
1
|
+
require "mkmf"
|
|
2
|
+
require "rb_sys/mkmf"
|
|
3
|
+
|
|
4
|
+
# Hide install_name_tool from rb_sys's fixup_libnames (rb_sys-0.9.128's
|
|
5
|
+
# mkmf.rb): on some hosts `install_name_tool -id "" $(DLLIB)` corrupts the
|
|
6
|
+
# compiled bundle's LC_ID_DYLIB load command (dyld then refuses to load it:
|
|
7
|
+
# "load command #N string extends beyond end of load command"), reproduced
|
|
8
|
+
# independent of this crate on a trivial unrelated cdylib. Skipping the
|
|
9
|
+
# rewrite leaves the bundle's original (absolute) install name in place,
|
|
10
|
+
# which `Kernel#require`'s dlopen does not care about.
|
|
11
|
+
def find_executable(bin, path = nil)
|
|
12
|
+
return nil if bin == "install_name_tool"
|
|
13
|
+
super
|
|
14
|
+
end
|
|
15
|
+
|
|
16
|
+
# Target basename must match the compiled artifact: the crate's [package]
|
|
17
|
+
# name is "gigatoken-rb" (the workspace root package already owns
|
|
18
|
+
# "gigatoken"), so Cargo's default lib name is "gigatoken_rb".
|
|
19
|
+
create_rust_makefile("gigatoken/gigatoken_rb")
|
|
@@ -0,0 +1,20 @@
|
|
|
1
|
+
//! `Gigatoken::Error` is defined in Ruby (`lib/gigatoken.rb`), loaded before
|
|
2
|
+
//! this extension is required — looked up fresh at each raise site rather
|
|
3
|
+
//! than cached, so no magnus `Value` needs a GC-registered static home.
|
|
4
|
+
|
|
5
|
+
use magnus::{exception::ExceptionClass, prelude::*, Error, RModule, Ruby};
|
|
6
|
+
|
|
7
|
+
fn error_class(ruby: &Ruby) -> Result<ExceptionClass, Error> {
|
|
8
|
+
ruby.class_object()
|
|
9
|
+
.const_get::<_, RModule>("Gigatoken")?
|
|
10
|
+
.const_get("Error")
|
|
11
|
+
}
|
|
12
|
+
|
|
13
|
+
/// Raise `Gigatoken::Error` with `message` (core load/encode failures surface
|
|
14
|
+
/// through this — never a Rust panic across the Ruby boundary).
|
|
15
|
+
pub fn raise(ruby: &Ruby, message: impl Into<String>) -> Error {
|
|
16
|
+
match error_class(ruby) {
|
|
17
|
+
Ok(class) => Error::new(class, message.into()),
|
|
18
|
+
Err(e) => e,
|
|
19
|
+
}
|
|
20
|
+
}
|
|
@@ -0,0 +1,122 @@
|
|
|
1
|
+
//! Release Ruby's GVL (global VM lock) for the duration of a CPU-bound Rust
|
|
2
|
+
//! closure, so other Ruby threads/fibers keep running while the core worker
|
|
3
|
+
//! pool chews through a batch encode.
|
|
4
|
+
//!
|
|
5
|
+
//! Uses `rb_sys::rb_nogvl` (this crate's direct `rb-sys` dependency, the same
|
|
6
|
+
//! version magnus already resolves — see `ext/gigatoken/Cargo.toml`) with the
|
|
7
|
+
//! `RB_NOGVL_OFFLOAD_SAFE` flag, rather than plain `rb_thread_call_without_gvl`.
|
|
8
|
+
//! Under a `Fiber.scheduler` built with a worker pool (`Async::Scheduler.new`
|
|
9
|
+
//! with `ASYNC_SCHEDULER_WORKER_POOL=true` or an explicit `worker_pool:`,
|
|
10
|
+
//! `~/architect/src/github.com/socketry/async/lib/async/scheduler.rb:31-101`),
|
|
11
|
+
//! `OFFLOAD_SAFE` tells Ruby it's safe to hand `func`/`data1` to the
|
|
12
|
+
//! scheduler's `blocking_operation_wait` hook instead of just blocking this OS
|
|
13
|
+
//! thread: the scheduler runs the closure on an `IO::Event::WorkerPool`
|
|
14
|
+
//! background thread while transferring the calling *fiber* (not just other
|
|
15
|
+
//! Threads) back to the reactor, so other fibers on this thread keep running
|
|
16
|
+
//! for the encode's duration
|
|
17
|
+
//! (`~/architect/src/github.com/socketry/io-event/ext/io/event/worker_pool.c:309-316`).
|
|
18
|
+
//! Without such a scheduler (or with one lacking a worker pool), `rb_nogvl`
|
|
19
|
+
//! degrades to exactly today's behavior: release the GVL, block this thread.
|
|
20
|
+
//! See `docs/rb/async-design.md` and `docs/rb/async.md` for the full design
|
|
21
|
+
//! and gotchas (worker pool is opt-in, defaults to one background worker).
|
|
22
|
+
//!
|
|
23
|
+
//! `func`/`data1` may now run on a different OS thread than the caller (the
|
|
24
|
+
//! scheduler's worker pool), where `rb_thread_call_without_gvl` always ran
|
|
25
|
+
//! them on the calling thread itself. `without_gvl` therefore requires
|
|
26
|
+
//! `F: Send` (the closure crosses onto the worker thread) and `R: Send` (its
|
|
27
|
+
//! result crosses back) — the compiler enforces the cross-thread contract
|
|
28
|
+
//! instead of leaving it to call-site audit. Every current `without_gvl`
|
|
29
|
+
//! call site (`tokenizer.rs`'s `encode_batch`/`encode_files`) already shares
|
|
30
|
+
//! its `&Tokenizer`/`&WorkerPool` across `rayon` worker threads inside
|
|
31
|
+
//! `encode_docs_ragged`. The BPE batch paths' input marshal (`tokenizer.rs`'s
|
|
32
|
+
//! `marshal_inputs`) borrows heap `RString` buffers zero-copy where it can
|
|
33
|
+
//! (locked via `rb_str_locktmp`, or unlocked when the string is frozen) and
|
|
34
|
+
//! falls back to an owned `Vec<u8>` copy otherwise — either way, only raw
|
|
35
|
+
//! byte slices are ever captured by the closure passed here, never a Ruby
|
|
36
|
+
//! `VALUE` or thread-local state, so running one OS thread over instead of
|
|
37
|
+
//! another changes nothing about its safety, and the types involved satisfy
|
|
38
|
+
//! `Send`/`Sync` on their own merits.
|
|
39
|
+
|
|
40
|
+
use std::any::Any;
|
|
41
|
+
use std::ffi::c_void;
|
|
42
|
+
use std::os::raw::c_int;
|
|
43
|
+
use std::panic::{self, AssertUnwindSafe};
|
|
44
|
+
|
|
45
|
+
use rb_sys::rb_nogvl;
|
|
46
|
+
|
|
47
|
+
/// `RB_NOGVL_OFFLOAD_SAFE` (`ruby/thread.h:84` in a current Ruby checkout;
|
|
48
|
+
/// introduced by Ruby's `Fiber::Scheduler#blocking_operation_wait` support,
|
|
49
|
+
/// first released in Ruby 3.4.0). Defined locally rather than taken from
|
|
50
|
+
/// `rb_sys::` bindings: `rb-sys` bindgens its constants from the *building*
|
|
51
|
+
/// Ruby's own headers, and this gem's floor (`gigatoken.gemspec`, `>=
|
|
52
|
+
/// 3.3.0`) may build against 3.3 headers that don't declare this macro at
|
|
53
|
+
/// all. Verified against Ruby 3.3.0's `thread.c` (`rb_nogvl`, ~line 1508):
|
|
54
|
+
/// it only tests `RB_NOGVL_UBF_ASYNC_SAFE`/`RB_NOGVL_INTR_FAIL` against
|
|
55
|
+
/// `flags`, so an unrecognized bit here is silently ignored — passing it on
|
|
56
|
+
/// 3.3 degrades to plain blocking `rb_nogvl`, identical to the old
|
|
57
|
+
/// `rb_thread_call_without_gvl`. No `RUBY_VERSION` gate is needed.
|
|
58
|
+
const RB_NOGVL_OFFLOAD_SAFE: c_int = 0x4;
|
|
59
|
+
|
|
60
|
+
/// The result of running `f` inside `call_without_gvl`: either its return
|
|
61
|
+
/// value, or a caught panic payload to re-raise once we're back on ordinary
|
|
62
|
+
/// (non-`extern "C"`) Rust stack frames. Unwinding a panic directly across
|
|
63
|
+
/// the `extern "C"` trampoline `rb_nogvl` calls into is undefined behavior;
|
|
64
|
+
/// catching it here and resuming it from `without_gvl` below turns that into
|
|
65
|
+
/// an ordinary Rust panic, which magnus's own `method!`/`function!` call
|
|
66
|
+
/// trampolines already wrap in `catch_unwind` and convert into a fatal Ruby
|
|
67
|
+
/// exception (`magnus::error::Error::from_panic`) — the same outcome any
|
|
68
|
+
/// other panicking native method already gets, just carried safely across
|
|
69
|
+
/// the extra C boundary this one call adds.
|
|
70
|
+
enum Outcome<R> {
|
|
71
|
+
Value(R),
|
|
72
|
+
Panic(Box<dyn Any + Send + 'static>),
|
|
73
|
+
}
|
|
74
|
+
|
|
75
|
+
unsafe extern "C" fn call_without_gvl<F, R>(arg: *mut c_void) -> *mut c_void
|
|
76
|
+
where
|
|
77
|
+
F: FnOnce() -> R + Send,
|
|
78
|
+
R: Send,
|
|
79
|
+
{
|
|
80
|
+
// SAFETY: `arg` is the `*mut Option<F>` handed to `rb_nogvl` below, valid
|
|
81
|
+
// for the duration of that (synchronous) call, and this is the only
|
|
82
|
+
// place it's dereferenced.
|
|
83
|
+
let closure = unsafe { (*(arg as *mut Option<F>)).take() }
|
|
84
|
+
.expect("without_gvl callback invoked more than once");
|
|
85
|
+
let outcome = match panic::catch_unwind(AssertUnwindSafe(closure)) {
|
|
86
|
+
Ok(value) => Outcome::Value(value),
|
|
87
|
+
Err(payload) => Outcome::Panic(payload),
|
|
88
|
+
};
|
|
89
|
+
Box::into_raw(Box::new(outcome)) as *mut c_void
|
|
90
|
+
}
|
|
91
|
+
|
|
92
|
+
/// Run `f` with the GVL released: other Ruby threads may run while `f`
|
|
93
|
+
/// executes, and — under a fiber scheduler with a worker pool — the calling
|
|
94
|
+
/// fiber yields to the reactor while `f` runs on a background thread. `f`
|
|
95
|
+
/// must not touch any Ruby object (`VALUE`) — only plain Rust data — per the
|
|
96
|
+
/// Ruby C API's contract for this call. A panic inside `f` is caught and
|
|
97
|
+
/// re-raised here rather than left to unwind across the C trampoline.
|
|
98
|
+
pub fn without_gvl<F, R>(f: F) -> R
|
|
99
|
+
where
|
|
100
|
+
F: FnOnce() -> R + Send,
|
|
101
|
+
R: Send,
|
|
102
|
+
{
|
|
103
|
+
let mut slot = Some(f);
|
|
104
|
+
let arg = &mut slot as *mut Option<F> as *mut c_void;
|
|
105
|
+
let result = unsafe {
|
|
106
|
+
rb_nogvl(
|
|
107
|
+
Some(call_without_gvl::<F, R>),
|
|
108
|
+
arg,
|
|
109
|
+
None,
|
|
110
|
+
std::ptr::null_mut(),
|
|
111
|
+
RB_NOGVL_OFFLOAD_SAFE,
|
|
112
|
+
)
|
|
113
|
+
};
|
|
114
|
+
// SAFETY: `result` is the `Box::into_raw(Box::new(Outcome<R>))` pointer
|
|
115
|
+
// produced by the callback above, which always runs exactly once before
|
|
116
|
+
// `rb_nogvl` returns.
|
|
117
|
+
let outcome = *unsafe { Box::from_raw(result as *mut Outcome<R>) };
|
|
118
|
+
match outcome {
|
|
119
|
+
Outcome::Value(value) => value,
|
|
120
|
+
Outcome::Panic(payload) => panic::resume_unwind(payload),
|
|
121
|
+
}
|
|
122
|
+
}
|
|
@@ -0,0 +1,50 @@
|
|
|
1
|
+
use gigatoken_rs::load_tokenizer::hf::{self, HfTokenizer};
|
|
2
|
+
use magnus::{Error, Module, RString, Ruby, Value, function};
|
|
3
|
+
|
|
4
|
+
// XZM-WORKAROUND: macOS 26's xzm malloc zone SIGTRAPs on multi-GB Rust chunk
|
|
5
|
+
// frees (`_xzm_reclaim_mark_used_locked` assertion); routing Rust allocations
|
|
6
|
+
// through mimalloc avoids the xzm zone entirely.
|
|
7
|
+
#[global_allocator]
|
|
8
|
+
static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
|
|
9
|
+
|
|
10
|
+
mod error;
|
|
11
|
+
mod gvl;
|
|
12
|
+
mod sentencepiece;
|
|
13
|
+
mod sources;
|
|
14
|
+
mod tokenizer;
|
|
15
|
+
|
|
16
|
+
use error::raise;
|
|
17
|
+
use sentencepiece::SentencePieceTokenizer;
|
|
18
|
+
use tokenizer::BPETokenizer;
|
|
19
|
+
|
|
20
|
+
// The gigatoken core crate exposes no version constant of its own, so this
|
|
21
|
+
// is the ext crate's (gigatoken-rb's) version — see the builder report.
|
|
22
|
+
fn crate_version() -> String {
|
|
23
|
+
env!("CARGO_PKG_VERSION").to_string()
|
|
24
|
+
}
|
|
25
|
+
|
|
26
|
+
/// Load a tokenizer from in-memory HuggingFace `tokenizer.json` contents.
|
|
27
|
+
/// Returns a `SentencePieceTokenizer` when the model uses `byte_fallback`, a
|
|
28
|
+
/// `BPETokenizer` otherwise — the same split as pyo3's `load_hf_json` and
|
|
29
|
+
/// the two classes' own `from_hf_json` constructors.
|
|
30
|
+
fn load_hf_json(ruby: &Ruby, data: RString) -> Result<Value, Error> {
|
|
31
|
+
// SAFETY: read-only, for the duration of this synchronous call.
|
|
32
|
+
let bytes = unsafe { data.as_slice() };
|
|
33
|
+
match hf::load_hf_slice(bytes) {
|
|
34
|
+
Ok(HfTokenizer::Bpe(tokenizer)) => Ok(ruby.into_value(BPETokenizer::from_tokenizer(tokenizer))),
|
|
35
|
+
Ok(HfTokenizer::SentencePiece(tokenizer)) => Ok(ruby.into_value(SentencePieceTokenizer::from_tokenizer(tokenizer))),
|
|
36
|
+
Err(e) => Err(raise(ruby, e.to_string())),
|
|
37
|
+
}
|
|
38
|
+
}
|
|
39
|
+
|
|
40
|
+
#[magnus::init]
|
|
41
|
+
fn init(ruby: &Ruby) -> Result<(), Error> {
|
|
42
|
+
let gigatoken = ruby.define_module("Gigatoken")?;
|
|
43
|
+
let native = gigatoken.define_module("Native")?;
|
|
44
|
+
native.define_module_function("crate_version", function!(crate_version, 0))?;
|
|
45
|
+
native.define_module_function("load_hf_json", function!(load_hf_json, 1))?;
|
|
46
|
+
sources::init(ruby, native)?;
|
|
47
|
+
tokenizer::init(ruby, native)?;
|
|
48
|
+
sentencepiece::init(ruby, native)?;
|
|
49
|
+
Ok(())
|
|
50
|
+
}
|