gigatoken 0.2.2 → 0.4.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 CHANGED
@@ -1,7 +1,7 @@
1
1
  ---
2
2
  SHA256:
3
- metadata.gz: ae6f31ee2b6996844f0f7460616ed7a3c4802993256fc68b231c7aeabf870df6
4
- data.tar.gz: d5e9eb5dc28c10f26b4dbb29f04f5f2156ae1a50697687251fb2a1e1ad5f8d1b
3
+ metadata.gz: 1fd5d74bb0d8cb5be9f7cf9fa077cc9c695ab5279436f3ecf3db4374952be586
4
+ data.tar.gz: 5953850c0513e14bb6aceccf4205b774c1db027f6d4d0f4508d90107f0d483ef
5
5
  SHA512:
6
- metadata.gz: 49ad46adee00b833739bd523905fc3f44d59374267aeff242be355246cbc173e6efdfefe8f509c287d429929df36413ccbdac487ec8a9dc6b9c5f3d758d06779
7
- data.tar.gz: 80d49b5c3246c91028cb669eb4ff534f8a49af393fe974e6537b6e3fdd33c916e6c0070eaabbae8014df6d9ece96318dae1c07352568d5108426714d93d57f54
6
+ metadata.gz: 9b85351d2af130d4f3b0b96c1f654b94ca5a865bd8865b6bd7d29ce2c5d297f052e8ee4a7ee40391573b88b3d81df5ba3aff8fd6edf64e6e03dee596644f1ae9
7
+ data.tar.gz: c8aca888f7b6477e406fa1234633e8e334a3242d9056951153c2306a6cc5b5eda8d46636c00bfd2c515ecf309130ffaf3543ced5305cc01e00473d8cf2966a73
data/Cargo.lock CHANGED
@@ -771,7 +771,7 @@ dependencies = [
771
771
 
772
772
  [[package]]
773
773
  name = "gigatoken-rb"
774
- version = "0.2.2"
774
+ version = "0.4.0"
775
775
  dependencies = [
776
776
  "gigatoken",
777
777
  "magnus",
data/Cargo.toml CHANGED
@@ -70,6 +70,13 @@ debug = true
70
70
  strip = "none"
71
71
  split-debuginfo = "packed"
72
72
 
73
+ # Frame pointers make Instruments' stack unwinding reliable, but a per-profile
74
+ # `rustflags` here is nightly-only and would need a `.cargo/config.toml` opt-in
75
+ # that the gem does not ship — which made the source gem unbuildable. The
76
+ # profile keeps the debug info; scripts/profile-cpu.fish exports RUSTFLAGS for
77
+ # its own build instead (this profile has its own target subdirectory, so that
78
+ # never reaches a release build). Mangling stays at the toolchain default (v0);
79
+ # `cargo instruments` demangles v0 symbols itself.
73
80
  [profile.profiling]
74
81
  inherits = "release"
75
82
  debug = 2
@@ -77,9 +84,6 @@ lto = false
77
84
  codegen-units = 1
78
85
  strip = "none"
79
86
  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
87
 
84
88
  [profile.profiling.package."*"]
85
89
  debug = 2
data/README.md CHANGED
@@ -15,7 +15,7 @@ Zero-copy Ruby bindings for [marcelroed/gigatoken](https://github.com/marcelroed
15
15
 
16
16
  Mac Studio M4 Max, OpenWebText, GPT-2 tokenizer; every library produces the same tokenization, gigatoken just does it faster. The Python row includes a fix for [marcelroed/gigatoken#38](https://github.com/marcelroed/gigatoken/issues/38) — a hidden memcpy in `shrink_to_fit()` we found while chasing the last of the Ruby–Python gap and sent upstream (one-line fix; without it the wheel lands around 7.4 GB/s).
17
17
 
18
- **340x faster** than the fastest existing Ruby gem (tiktoken_ruby) and **1,050x faster** than the tokenizers gem. Full methodology & exact counts: [docs/rb/benchmarks.md](docs/rb/benchmarks.md).
18
+ **340x faster** than the fastest existing Ruby gem (tiktoken_ruby) and **1,050x faster** than the tokenizers gem. Full methodology & exact counts: [docs/explanation/benchmarks.md](docs/explanation/benchmarks.md).
19
19
 
20
20
  ## Install
21
21
 
@@ -57,9 +57,9 @@ Gigatoken::Tokenizer.from_tiktoken("cl100k_base.tiktoken", pretokenizer: "gpt4",
57
57
  Gigatoken::Tokenizer.from_json(File.binread("tokenizer.json"))
58
58
  ```
59
59
 
60
- A `.tiktoken` file holds mergeable ranks only — its pretokenization scheme and special tokens live in the code that defines the encoding, not the file — so `pretokenizer:` is a required keyword (one of `Gigatoken::Native.pretokenizer_names`: `gpt2`/`r50k`, `gpt4`/`cl100k`, `qwen2`, `qwen35`, `olmo3`, `deepseek_v3`, `o200k`, `nemotron`, `kimi`) and `special_tokens:` defaults to none. Nothing is guessed: an unknown scheme raises `Gigatoken::Error` naming the valid ones, and `Tokenizer.load` on a `.tiktoken` path with no `pretokenizer:` raises rather than silently picking one.
60
+ A `.tiktoken` file holds mergeable ranks only — its pretokenization scheme and special tokens live in the code that defines the encoding, not the file — so `pretokenizer:` is a required keyword (one of `Gigatoken::Native.pretokenizer_names`: `gpt2`/`r50k`, `gpt4`/`cl100k`, `qwen2`, `qwen35`, `olmo3`, `deepseek_v3`, `o200k`, `nemotron`, `kimi`) and `special_tokens:` defaults to none. Nothing is guessed: an unknown scheme raises `Gigatoken::ModelError` naming the valid ones, and `Tokenizer.load` on a `.tiktoken` path with no `pretokenizer:` raises rather than silently picking one.
61
61
 
62
- 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.
62
+ 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::InputError` on invalid UTF-8 instead of guessing.
63
63
 
64
64
  ### Packaged tiktoken encodings
65
65
 
@@ -71,7 +71,7 @@ Gigatoken::Tokenizer.load("cl100k_base") # same result — packaged nam
71
71
  # checked before the Hub-repo-id shape
72
72
  ```
73
73
 
74
- `p50k_base` and `p50k_edit` are deliberately not packaged: both load the same non-dense ranks (id 50256 is left free for `<|endoftext|>`), and the rank loader rejects non-dense ranks. Both entry points raise `Gigatoken::Error` explaining that, rather than `load` falling through to the Hub for a name that happens to look like a legacy repo id.
74
+ `p50k_base` and `p50k_edit` are deliberately not packaged: both load the same non-dense ranks (id 50256 is left free for `<|endoftext|>`), and the rank loader rejects non-dense ranks. Both entry points raise `Gigatoken::ModelError` explaining that, rather than `load` falling through to the Hub for a name that happens to look like a legacy repo id.
75
75
 
76
76
  `encode` on a packaged tokenizer honours its special-token table: text containing `<|endoftext|>` (or any other literal special-token string) is tokenized as that special token, not as ordinary text. That matches [`tiktoken`](https://github.com/openai/tiktoken)'s `encode_with_special_tokens`, not its plain `encode`, which treats the same literal as ordinary text — a difference worth knowing if you're tokenizing untrusted input. To get tiktoken's non-honouring default instead, build a tokenizer from the same rank file with an empty special-token table:
77
77
 
@@ -117,9 +117,15 @@ packed.token_count # => total tokens
117
117
  packed[3] # => document 3's ids as an Array, on demand
118
118
  ```
119
119
 
120
+ ### Errors
121
+
122
+ Everything the library raises is a `Gigatoken::Error`, never a raw Rust panic — and, under it, one of three: `Gigatoken::ModelError` when a tokenizer can't be loaded (bad or hostile JSON, a missing file or directory, an unknown or unpackable encoding name, a malformed `.tiktoken`), `Gigatoken::InputError` when a document or an id can't be taken (a String that won't transcode, invalid UTF-8 on the SentencePiece path, an id outside the vocabulary in `decode`), and `Gigatoken::HubError` for everything `Gigatoken::Hub` raises (HTTP status, transport, timeout, the repo-id and header checks). `rescue Gigatoken::Error` catches all three.
123
+
124
+ `encode` and `encode_batch` honour the String's encoding tag: UTF-8, US-ASCII and binary go through byte-wise, and a real non-UTF-8 encoding (ISO-8859-1, UTF-16LE, …) is transcoded first, so it gives the same ids as the same text read as UTF-8. See [the reference](docs/reference/tokenizer.md#input-encodings).
125
+
120
126
  ### Async
121
127
 
122
- `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).
128
+ `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/how-to/run-under-async.md](docs/how-to/run-under-async.md).
123
129
 
124
130
  ## CLI
125
131
 
@@ -138,6 +144,16 @@ gigatoken bench lib/gigatoken/encodings/cl100k_base.tiktoken README.md --pretoke
138
144
 
139
145
  Leave it off against a `.tiktoken` TOKENIZER and both commands raise `Gigatoken::Error` naming the valid schemes instead of crashing; for every other TOKENIZER shape (`tokenizer.json`, a packaged name, a Hub repo id) `--pretokenizer` is accepted but ignored.
140
146
 
147
+ ## Documentation
148
+
149
+ In-depth docs live under [`docs/`](docs/README.md), organized by
150
+ [Diátaxis](https://diataxis.fr/):
151
+
152
+ - **Tutorial:** [Getting started](docs/tutorials/getting-started.md)
153
+ - **How-to:** [Load a tokenizer](docs/how-to/load-a-tokenizer.md), [Tokenize files](docs/how-to/tokenize-files.md), [Packed results](docs/how-to/use-packed-results.md), [Cache budget](docs/how-to/tune-the-cache-budget.md), [Async](docs/how-to/run-under-async.md), [Measure](docs/how-to/measure-time-and-allocations.md)
154
+ - **Reference:** [`Tokenizer`](docs/reference/tokenizer.md), [`PackedResult`](docs/reference/packed-result.md), [File sources](docs/reference/file-sources.md), [Encodings and settings](docs/reference/encodings-and-settings.md), [CLI](docs/reference/cli.md)
155
+ - **Explanation:** [Benchmarks](docs/explanation/benchmarks.md), [Allocations](docs/explanation/allocations.md), [Async design](docs/explanation/async-design.md)
156
+
141
157
  ## Development
142
158
 
143
159
  ```bash
@@ -145,6 +161,7 @@ bundle install
145
161
  bundle exec rake compile # native extension (Rust nightly, via rust-toolchain.toml)
146
162
  bundle exec rspec
147
163
  bundle exec standardrb
164
+ ruby -Ilib bench/operations.rb # every operation: i/s, objects and malloc per call
148
165
  ```
149
166
 
150
167
  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).
@@ -0,0 +1,105 @@
1
+ //! Shared input loading for the bench targets. Lives in a subdirectory so
2
+ //! cargo does not treat it as a bench target; each bench pulls it in with
3
+ //! `mod common;`.
4
+
5
+ use std::time::Instant;
6
+
7
+ /// Re-enable transparent huge pages for this process. The encode paths
8
+ /// madvise their big tables and buffers to 2 MiB pages (they far exceed
9
+ /// 4 KiB dTLB coverage, and Zen drops software prefetches that miss the
10
+ /// TLB), but some session managers launch children with
11
+ /// PR_SET_THP_DISABLE, which silently vetoes MADV_HUGEPAGE; clear it so
12
+ /// the bench measures the tokenizer, not the launcher's memory policy.
13
+ /// No-op off Linux.
14
+ pub fn allow_thp() {
15
+ #[cfg(target_os = "linux")]
16
+ // SAFETY: prctl(PR_SET_THP_DISABLE, 0) only clears a per-process flag.
17
+ unsafe {
18
+ libc::prctl(libc::PR_SET_THP_DISABLE, 0, 0, 0, 0);
19
+ }
20
+ }
21
+
22
+ /// Hint 2 MiB pages for a buffer's reserved capacity, BEFORE it is first
23
+ /// written (the ordering is what makes the hint effective: pages fault in
24
+ /// huge only if the madvise precedes the first touch). Multi-GB bench
25
+ /// buffers otherwise saturate the dTLB alongside the encode's own tables.
26
+ /// No-op off Linux.
27
+ #[allow(unused_variables, clippy::missing_safety_doc)]
28
+ pub fn madvise_hugepage_capacity<T>(v: &mut Vec<T>) {
29
+ #[cfg(target_os = "linux")]
30
+ if v.capacity() > 0 {
31
+ // Align the start inward: malloc's mmap chunks carry a 16-byte
32
+ // header, and an unaligned madvise start is EINVAL (silent no-op).
33
+ const PAGE: usize = 4096;
34
+ let addr = v.as_mut_ptr() as usize;
35
+ let start = (addr + PAGE - 1) & !(PAGE - 1);
36
+ let end = addr + v.capacity() * std::mem::size_of::<T>();
37
+ if end > start {
38
+ // SAFETY: the range is one live allocation; the hint neither
39
+ // reads nor writes it.
40
+ unsafe {
41
+ libc::madvise(start as *mut libc::c_void, end - start, libc::MADV_HUGEPAGE);
42
+ }
43
+ }
44
+ }
45
+ }
46
+
47
+ /// Load the benchmark input from `~/data/owt_train.txt`, truncated to a
48
+ /// UTF-8 character boundary.
49
+ ///
50
+ /// ENCODE_MB caps the input for fast profiling iterations (only that many
51
+ /// bytes are read from disk, so the read does not dominate a profile of the
52
+ /// encode loop). When it is unset, `default_mb` applies; `None` reads the
53
+ /// whole file.
54
+ pub fn load_owt_input(default_mb: Option<usize>) -> Vec<u8> {
55
+ let owt_path = std::env::home_dir().unwrap().join("data/owt_train.txt");
56
+ eprintln!("Reading {owt_path:?}...");
57
+ let t0 = Instant::now();
58
+
59
+ let cap_mb = std::env::var("ENCODE_MB")
60
+ .ok()
61
+ .map(|mb| {
62
+ mb.trim()
63
+ .parse::<usize>()
64
+ .expect("ENCODE_MB must be an integer")
65
+ })
66
+ .or(default_mb);
67
+ let mut data = match cap_mb {
68
+ Some(mb) => {
69
+ use std::io::Read;
70
+ let max_bytes = mb * 1_000_000;
71
+ let file =
72
+ std::fs::File::open(&owt_path).expect("Could not open ~/data/owt_train.txt");
73
+ let mut data = Vec::with_capacity(max_bytes);
74
+ madvise_hugepage_capacity(&mut data);
75
+ file.take(max_bytes as u64)
76
+ .read_to_end(&mut data)
77
+ .expect("read failed");
78
+ data
79
+ }
80
+ None => {
81
+ let len = std::fs::metadata(&owt_path)
82
+ .expect("Could not stat ~/data/owt_train.txt")
83
+ .len() as usize;
84
+ use std::io::Read;
85
+ let mut data = Vec::with_capacity(len + 1);
86
+ madvise_hugepage_capacity(&mut data);
87
+ std::fs::File::open(&owt_path)
88
+ .expect("Could not open ~/data/owt_train.txt")
89
+ .read_to_end(&mut data)
90
+ .expect("read failed");
91
+ data
92
+ }
93
+ };
94
+ // Back up to a UTF-8 character boundary (a byte cap can split a
95
+ // multibyte character).
96
+ if let Err(e) = std::str::from_utf8(&data) {
97
+ data.truncate(e.valid_up_to());
98
+ }
99
+ eprintln!(
100
+ "Read {:.2} GB in {:.1}s",
101
+ data.len() as f64 / 1e9,
102
+ t0.elapsed().as_secs_f64()
103
+ );
104
+ data
105
+ }
data/benches/encode.rs ADDED
@@ -0,0 +1,79 @@
1
+ //! Whole-file parallel encode benchmark. For both backends the entire input
2
+ //! is ONE document handed to the library's parallel encode path
3
+ //! (`encode_docs_ragged` / `sp_encode_docs_ragged`) — the same chunking
4
+ //! policy, safe-boundary document splitting, and (for BPE) persistent worker
5
+ //! pool as `encode_batch` / `encode_files` — so this measures gigatoken's
6
+ //! own parallelism, not a bench-local split.
7
+ //!
8
+ //! Run with: cargo bench --bench encode (full OWT)
9
+ //! ENCODE_MB=500 cargo bench --bench encode
10
+ //! TOKENIZER_JSON=data/qwen3_5_tokenizer.json cargo bench --bench encode
11
+ //! TOKENIZER_JSON=meta-llama/Llama-3.1-8B cargo bench --bench encode
12
+ //! (a HuggingFace repo id: resolved via the standard HF cache,
13
+ //! downloaded into it on a miss)
14
+
15
+ use gigatoken_rs::load_tokenizer::hf::{HfTokenizer, load_hf_slice};
16
+ use gigatoken_rs::load_tokenizer::hub;
17
+ use gigatoken_rs::{WorkerPool, encode_docs_ragged, sp_encode_docs_ragged};
18
+ use std::hint::black_box;
19
+ use std::path::PathBuf;
20
+ use std::time::Instant;
21
+
22
+ mod common;
23
+
24
+ fn main() {
25
+ common::allow_thp();
26
+ let tokenizer_json =
27
+ std::env::var("TOKENIZER_JSON").unwrap_or_else(|_| "data/gpt2_tokenizer.json".to_string());
28
+ let tokenizer_path = {
29
+ let path = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join(&tokenizer_json);
30
+ if !path.exists() && hub::looks_like_repo_id(&tokenizer_json) {
31
+ hub::hub_file(&tokenizer_json, "tokenizer.json", "main")
32
+ .expect("Could not fetch tokenizer.json from the HuggingFace Hub")
33
+ } else {
34
+ path
35
+ }
36
+ };
37
+ eprintln!("Loading tokenizer from {tokenizer_path:?}...");
38
+ let data = std::fs::read(&tokenizer_path).expect("Could not read tokenizer.json");
39
+ let tokenizer = load_hf_slice(&data).expect("Could not load tokenizer");
40
+
41
+ let input = common::load_owt_input(None);
42
+ let size_gb = input.len() as f64 / 1e9;
43
+
44
+ let workers = WorkerPool::new();
45
+ let (n_tokens, elapsed) = match &tokenizer {
46
+ HfTokenizer::Bpe(tokenizer) => {
47
+ eprintln!(
48
+ "Encoding (ByteLevel BPE, 1 document, {} threads)...",
49
+ rayon::current_num_threads()
50
+ );
51
+ let start = Instant::now();
52
+ let (ids, lens) = encode_docs_ragged(&workers, tokenizer, &[&input]);
53
+ black_box((&ids, lens));
54
+ (ids.len(), start.elapsed().as_secs_f64())
55
+ }
56
+ HfTokenizer::SentencePiece(tokenizer) => {
57
+ // Trailing partial UTF-8 char (from the ENCODE_MB byte cap) is
58
+ // dropped; SP encoding takes str input.
59
+ let text = match std::str::from_utf8(&input) {
60
+ Ok(text) => text,
61
+ Err(e) => std::str::from_utf8(&input[..e.valid_up_to()]).unwrap(),
62
+ };
63
+ eprintln!(
64
+ "Encoding (SentencePiece, 1 document, {} threads)...",
65
+ rayon::current_num_threads()
66
+ );
67
+ let start = Instant::now();
68
+ let (ids, lens) = sp_encode_docs_ragged(tokenizer, &[text]);
69
+ black_box((&ids, lens));
70
+ (ids.len(), start.elapsed().as_secs_f64())
71
+ }
72
+ };
73
+ let throughput_gb = size_gb / elapsed;
74
+
75
+ eprintln!(
76
+ "{n_tokens} tokens in {elapsed:.2}s — {throughput_gb:.2} GB/s ({:.0} MB/s)",
77
+ throughput_gb * 1000.0
78
+ );
79
+ }
@@ -0,0 +1,60 @@
1
+ //! Whole-document multithreaded encode benchmark, mirroring
2
+ //! `BPETokenizer.encode_files` on a single plain-text file: the entire input
3
+ //! is ONE document handed to the library's parallel encode path
4
+ //! (`encode_docs_ragged`), which splits it at pretoken-safe boundaries
5
+ //! (token-identical to a serial pass), encodes with a persistent worker
6
+ //! pool, and gathers one flat id buffer.
7
+ //!
8
+ //! One round per process by default: a first pass over a fresh dataset is
9
+ //! the workload that matters, and everything a second in-process round
10
+ //! reuses — worker pretoken caches, allocator arenas, faulted pages, grown
11
+ //! buffer capacities — makes later rounds unrealistically fast. Restart the
12
+ //! binary to collect independent samples; ENCODE_ROUNDS>1 remains available
13
+ //! for cache-behavior experiments.
14
+ //!
15
+ //! Run with: cargo bench --bench encode_doc (2 GB default)
16
+ //! ENCODE_MB=500 cargo bench --bench encode_doc
17
+ //! TOKENIZER_JSON=data/qwen3_5_tokenizer.json cargo bench --bench encode_doc
18
+
19
+ use gigatoken_rs::load_tokenizer::hf::load_hf_bpe;
20
+ use gigatoken_rs::{WorkerPool, encode_docs_ragged};
21
+ use std::hint::black_box;
22
+ use std::path::PathBuf;
23
+ use std::time::Instant;
24
+
25
+ mod common;
26
+
27
+ const DEFAULT_MB: usize = 2000;
28
+
29
+ fn main() {
30
+ common::allow_thp();
31
+ let tokenizer_json = std::env::var("TOKENIZER_JSON")
32
+ .unwrap_or_else(|_| "data/olmo3_tokenizer.json".to_string());
33
+ let tokenizer_path = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join(&tokenizer_json);
34
+ eprintln!("Loading tokenizer from {tokenizer_path:?}...");
35
+ let tokenizer = load_hf_bpe(&tokenizer_path).expect("Could not load tokenizer");
36
+
37
+ let input = common::load_owt_input(Some(DEFAULT_MB));
38
+ let size_mb = input.len() as f64 / 1e6;
39
+ eprintln!("1 document, {} threads\n", rayon::current_num_threads());
40
+
41
+ let rounds: usize = std::env::var("ENCODE_ROUNDS")
42
+ .ok()
43
+ .and_then(|v| v.trim().parse().ok())
44
+ .unwrap_or(1);
45
+ for round in 0..rounds {
46
+ // Fresh worker pool per round so extra rounds at least start with
47
+ // cold pretoken caches (the pool retains one forked tokenizer per
48
+ // rayon thread).
49
+ let workers = WorkerPool::new();
50
+ let t0 = Instant::now();
51
+ let (flat, lens) = encode_docs_ragged(&workers, &tokenizer, &[&input]);
52
+ black_box((&flat, lens));
53
+ let elapsed = t0.elapsed().as_secs_f64();
54
+ eprintln!(
55
+ "round {round}: {} tokens in {elapsed:.3}s — {:.0} MB/s",
56
+ flat.len(),
57
+ size_mb / elapsed
58
+ );
59
+ }
60
+ }
@@ -0,0 +1,105 @@
1
+ use gigatoken_rs::load_tokenizer::hf::{HfTokenizer, load_hf_slice};
2
+ use gigatoken_rs::pretokenize::FastR50kPretokenizer;
3
+ use std::path::PathBuf;
4
+ use std::time::Instant;
5
+
6
+ mod common;
7
+ fn main() {
8
+ common::allow_thp();
9
+ // ENCODE_TOKENIZER overrides the tokenizer.json: a local path (e.g.
10
+ // data/qwen3_tokenizer.json to bench the qwen2-scheme encode path) or a
11
+ // HuggingFace repo id (e.g. Qwen/Qwen2-1.5B-Instruct), served from the
12
+ // standard HF cache and downloaded into it on a miss. Encoding then runs
13
+ // through the scheme dispatch instead of the hardcoded r50k pretokenizer.
14
+ let tokenizer_override = std::env::var("ENCODE_TOKENIZER").ok().map(|value| {
15
+ let path = PathBuf::from(&value);
16
+ if !path.exists() && gigatoken_rs::load_tokenizer::hub::looks_like_repo_id(&value) {
17
+ gigatoken_rs::load_tokenizer::hub::hub_file(&value, "tokenizer.json", "main")
18
+ .expect("Could not fetch tokenizer.json from the HuggingFace Hub")
19
+ } else {
20
+ path
21
+ }
22
+ });
23
+ let tokenizer_path = tokenizer_override.clone().unwrap_or_else(|| {
24
+ PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("data/gpt2_tokenizer.json")
25
+ });
26
+ eprintln!("Loading tokenizer from {tokenizer_path:?}...");
27
+ let data = std::fs::read(&tokenizer_path).expect("Could not read tokenizer.json");
28
+ // byte_fallback configs (Llama/gemma-style SentencePiece) dispatch to the
29
+ // SP encode path; everything else to ByteLevel BPE, as in benches/encode.
30
+ let tokenizer = load_hf_slice(&data).expect("Could not load tokenizer");
31
+
32
+ let input = common::load_owt_input(None);
33
+ let size_gb = input.len() as f64 / 1e9;
34
+ // Encode the whole buffer in one pass (matches real usage; the pretokenizer
35
+ // handles newlines itself, so pre-splitting into lines is unnecessary).
36
+ let buf: &[u8] = &input;
37
+
38
+ // ENCODE_PASSES=N re-encodes the same buffer N times; passes after the
39
+ // first run with a fully warm pretoken cache, isolating the hit path.
40
+ let passes: usize = std::env::var("ENCODE_PASSES")
41
+ .ok()
42
+ .and_then(|v| v.trim().parse().ok())
43
+ .unwrap_or(1);
44
+ eprintln!("Encoding (single-threaded)...");
45
+ // Flat output buffer, reserved once from the batch engine's
46
+ // bytes-per-token estimate (see batch::encode_chunk) and reused across
47
+ // passes; the token count is its length.
48
+ let mut out: Vec<u32> = Vec::with_capacity(input.len() / 4 + 16);
49
+ common::madvise_hugepage_capacity(&mut out);
50
+ // ENCODE_DUMP=path writes the final pass's tokens as little-endian u32s,
51
+ // for byte-exact A/B comparison across encode-path changes.
52
+ let dump_path = std::env::var("ENCODE_DUMP").ok();
53
+ match tokenizer {
54
+ HfTokenizer::Bpe(mut tokenizer) => {
55
+ for pass in 0..passes {
56
+ out.clear();
57
+ let start = Instant::now();
58
+ if tokenizer_override.is_some() {
59
+ let pretokens = tokenizer.pretokenizer_type().pretokenize(buf);
60
+ tokenizer.memoized_encode_flat(pretokens, &mut out);
61
+ } else {
62
+ tokenizer.memoized_encode_flat(FastR50kPretokenizer::new(buf), &mut out);
63
+ }
64
+ report_pass(pass, out.len(), size_gb, start);
65
+ }
66
+ let (sl, sc, ll, lc, lkb, al, ac) = tokenizer.cache_mem_stats();
67
+ eprintln!(
68
+ "cache: short {sl} entries (cap {sc}), long {ll} (cap {lc}, {lkb} key bytes), arena {al} tokens (cap {ac})"
69
+ );
70
+ }
71
+ HfTokenizer::SentencePiece(tokenizer) => {
72
+ // Trailing partial UTF-8 char (from the ENCODE_MB byte cap) is
73
+ // dropped; SP encoding takes str input.
74
+ let text = match std::str::from_utf8(buf) {
75
+ Ok(text) => text,
76
+ Err(e) => std::str::from_utf8(&buf[..e.valid_up_to()]).unwrap(),
77
+ };
78
+ let mut state = gigatoken_rs::EncodeState::new();
79
+ for pass in 0..passes {
80
+ out.clear();
81
+ let start = Instant::now();
82
+ let out_ref = &mut out;
83
+ tokenizer.encode_raw_cb(&mut state, text, &mut |tokens| {
84
+ out_ref.extend(tokens.iter().map(|t| t.0))
85
+ });
86
+ report_pass(pass, out.len(), size_gb, start);
87
+ }
88
+ eprintln!("cache: {} units", state.cache_size());
89
+ }
90
+ }
91
+ if let Some(path) = dump_path {
92
+ let bytes: Vec<u8> = out.iter().flat_map(|t| t.to_le_bytes()).collect();
93
+ std::fs::write(&path, bytes).expect("Could not write ENCODE_DUMP");
94
+ eprintln!("dumped {} tokens to {path}", out.len());
95
+ }
96
+ }
97
+
98
+ fn report_pass(pass: usize, total_tokens: usize, size_gb: f64, start: Instant) {
99
+ let elapsed = start.elapsed().as_secs_f64();
100
+ let throughput_gb = size_gb / elapsed;
101
+ eprintln!(
102
+ "pass {pass}: {total_tokens} tokens in {elapsed:.2}s — {throughput_gb:.2} GB/s ({:.0} MB/s)",
103
+ throughput_gb * 1000.0
104
+ );
105
+ }
@@ -0,0 +1,51 @@
1
+ //! Single-threaded SentencePiece encode throughput on OWT, mirroring
2
+ //! `encode_st` (GPT-2 byte-level) for a like-for-like comparison.
3
+ //!
4
+ //! Select the tokenizer with SP_TOKENIZER=tinyllama (default) or sp4096;
5
+ //! cap the input with ENCODE_MB like `encode_st`.
6
+
7
+ use gigatoken_rs::load_tokenizer::hf::load_hf_sentencepiece;
8
+ use std::path::PathBuf;
9
+ use std::time::Instant;
10
+
11
+ mod common;
12
+
13
+ fn main() {
14
+ let which = std::env::var("SP_TOKENIZER").unwrap_or_else(|_| "tinyllama".to_string());
15
+ let file = match which.as_str() {
16
+ "tinyllama" => "data/tinyllama_tokenizer.json",
17
+ "sp4096" => "data/fineweb_4096_bpe_tokenizer.json",
18
+ other => panic!("unknown SP_TOKENIZER {other:?}"),
19
+ };
20
+ let tokenizer_path = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join(file);
21
+ eprintln!("Loading {which} tokenizer from {tokenizer_path:?}...");
22
+ let tokenizer = load_hf_sentencepiece(&tokenizer_path).expect("Could not load tokenizer");
23
+
24
+ let input = common::load_owt_input(None);
25
+ let size_gb = input.len() as f64 / 1e9;
26
+ let text = std::str::from_utf8(&input).expect("input must be UTF-8");
27
+
28
+ eprintln!("Encoding (single-threaded)...");
29
+ // Count-only callback, mirroring encode_st's measurement of the GPT-2
30
+ // path (the full encode runs; only output materialization is skipped).
31
+ let passes: usize = std::env::var("SP_PASSES")
32
+ .ok()
33
+ .and_then(|v| v.parse().ok())
34
+ .unwrap_or(2);
35
+ let mut state = gigatoken_rs::EncodeState::new();
36
+ for pass in 1..=passes {
37
+ let mut total_tokens: usize = 0;
38
+ let start = Instant::now();
39
+ tokenizer.encode_raw_cb(&mut state, text, &mut |tokens: &[_]| {
40
+ total_tokens += tokens.len();
41
+ });
42
+ let elapsed = start.elapsed().as_secs_f64();
43
+ let throughput_gb = size_gb / elapsed;
44
+ eprintln!(
45
+ "pass {pass} (cache {}): {total_tokens} tokens in {elapsed:.2}s — {throughput_gb:.2} GB/s ({:.0} MB/s), {} cached units",
46
+ if pass == 1 { "cold" } else { "warm" },
47
+ throughput_gb * 1000.0,
48
+ state.cache_size()
49
+ );
50
+ }
51
+ }