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
data/src/lib.rs ADDED
@@ -0,0 +1,578 @@
1
+ #![feature(portable_simd)]
2
+
3
+ pub(crate) mod batch;
4
+ #[cfg(feature = "python")]
5
+ pub(crate) mod bindings;
6
+ pub(crate) mod bpe;
7
+ pub(crate) mod bpe_train;
8
+ pub mod input;
9
+ pub mod pretokenize;
10
+ #[cfg(test)]
11
+ pub(crate) mod test_hub;
12
+ pub(crate) mod token;
13
+ pub use crate::batch::{
14
+ GatherBuf, GatherOutcome, WorkerPool, encode_docs_into, encode_docs_ragged,
15
+ encode_files_docs, encode_files_docs_serial, sp_encode_docs_ragged,
16
+ sp_encode_docs_ragged_serial, sp_encode_files_docs, sp_encode_files_docs_serial,
17
+ };
18
+ pub use crate::bpe::SentencePieceBPE;
19
+ pub use crate::bpe::Tokenizer;
20
+ pub use crate::bpe::sentencepiece::EncodeState;
21
+ pub mod load_tokenizer;
22
+
23
+ #[cfg(feature = "python")]
24
+ use crate::batch::encode_into;
25
+ #[cfg(feature = "python")]
26
+ use crate::bindings::bridge::{
27
+ EncodeInput, encode_batch_pylist, encode_batch_ragged, extract_doc, extract_token_ids,
28
+ merges_to_pylist, vocab_to_pydict,
29
+ };
30
+ #[cfg(feature = "python")]
31
+ use crate::bindings::matcher::{SpecialTokenFound, SubstringMatcher};
32
+ #[cfg(feature = "python")]
33
+ use crate::bindings::padding;
34
+ #[cfg(feature = "python")]
35
+ use crate::bindings::pretokenize::{PretokenizerIter, pretokenized_counts, pretokenizer};
36
+ #[cfg(feature = "python")]
37
+ use crate::bindings::sources::{
38
+ BytesSource, FileSource, JsonlFileSource, ParquetFileSource, TextFileSource,
39
+ encode_files_ragged,
40
+ };
41
+ #[cfg(feature = "python")]
42
+ use crate::bindings::train::train_bpe;
43
+ #[cfg(feature = "python")]
44
+ use crate::input::file_source::DocFormat;
45
+ #[cfg(feature = "python")]
46
+ use numpy::{IntoPyArray, PyArray1};
47
+ #[cfg(feature = "python")]
48
+ use pyo3::prelude::*;
49
+ #[cfg(feature = "python")]
50
+ use pyo3::pybacked::{PyBackedBytes, PyBackedStr};
51
+ #[cfg(feature = "python")]
52
+ use pyo3::types::{PyBytes, PyDict, PyList};
53
+ #[cfg(feature = "python")]
54
+ use std::path::PathBuf;
55
+
56
+ #[cfg(feature = "python")]
57
+ #[pyclass]
58
+ struct BPETokenizer {
59
+ tokenizer: Tokenizer,
60
+ workers: WorkerPool,
61
+ }
62
+
63
+ #[cfg(feature = "python")]
64
+ impl BPETokenizer {
65
+ /// See `batch::encode_files_docs` / `batch::encode_files_docs_serial`:
66
+ /// each region splits into documents per `format` during the encode —
67
+ /// one document per region for plain batches, separator pieces for a
68
+ /// BytesSource. Call with the GIL released.
69
+ fn encode_regions_ragged(
70
+ &self,
71
+ regions: &[&[u8]],
72
+ format: &DocFormat,
73
+ parallel: bool,
74
+ ) -> (Vec<u32>, Vec<i64>) {
75
+ if parallel {
76
+ encode_files_docs(&self.workers, &self.tokenizer, regions, format)
77
+ } else {
78
+ encode_files_docs_serial(&self.workers, &self.tokenizer, regions, format)
79
+ }
80
+ }
81
+ }
82
+
83
+ #[cfg(feature = "python")]
84
+ #[pymethods]
85
+ impl BPETokenizer {
86
+ #[staticmethod]
87
+ fn from_tiktoken(path: PathBuf) -> PyResult<Self> {
88
+ Ok(Self {
89
+ tokenizer: load_tokenizer::tiktoken::load_tiktoken(&path)?,
90
+ workers: WorkerPool::new(),
91
+ })
92
+ }
93
+ /// Load from a tiktoken rank file plus a tokenizer_config.json carrying
94
+ /// the special tokens — the layout of repos that ship no tokenizer.json
95
+ /// (e.g. the moonshotai Kimi line) — with the named pretokenizer scheme.
96
+ #[staticmethod]
97
+ fn from_tiktoken_model(
98
+ model_path: PathBuf,
99
+ config_path: PathBuf,
100
+ pretokenizer: &str,
101
+ ) -> PyResult<Self> {
102
+ let scheme = pretokenize::PretokenizerType::from_name(pretokenizer).ok_or_else(|| {
103
+ pyo3::exceptions::PyValueError::new_err(format!(
104
+ "unknown pretokenizer scheme {pretokenizer:?}; expected one of \
105
+ gpt2, gpt4, qwen2, qwen35, olmo3, deepseek_v3, o200k, nemotron, kimi"
106
+ ))
107
+ })?;
108
+ Ok(Self {
109
+ tokenizer: load_tokenizer::tiktoken::load_tiktoken_model(&model_path, &config_path, scheme)?,
110
+ workers: WorkerPool::new(),
111
+ })
112
+ }
113
+
114
+ #[staticmethod]
115
+ fn from_hf(path: PathBuf) -> PyResult<Self> {
116
+ Ok(Self {
117
+ tokenizer: load_tokenizer::hf::load_hf_bpe(&path)?,
118
+ workers: WorkerPool::new(),
119
+ })
120
+ }
121
+
122
+ /// Encode a single document (str or bytes) with the main tokenizer,
123
+ /// whose pretoken cache persists across calls.
124
+ fn encode<'py>(
125
+ &mut self,
126
+ py: Python<'py>,
127
+ input: Bound<'py, PyAny>,
128
+ ) -> PyResult<Bound<'py, PyArray1<u32>>> {
129
+ let input = extract_doc(&input)?;
130
+ let (mut ids, mut lens) = (Vec::new(), Vec::new());
131
+ encode_into(&mut self.tokenizer, input.as_bytes(), &mut ids, &mut lens);
132
+ Ok(ids.into_pyarray(py))
133
+ }
134
+
135
+ /// Encode a batch of documents in parallel with rayon, releasing the GIL.
136
+ /// Takes a list of str or a list of bytes (all elements of the same
137
+ /// type), an awkward Array of strings/bytestrings — whose flat buffers
138
+ /// are used directly, with no per-document Python objects — or a
139
+ /// BytesSource, whose buffers are split into documents on its separator
140
+ /// during the encode itself (pass pre-read corpora that way instead of
141
+ /// pre-splitting them). For files, use encode_files. Returns an
142
+ /// awkward.Array with one row of token ids per document (a single flat
143
+ /// buffer plus offsets, not one numpy array per document).
144
+ ///
145
+ /// Documents are grouped into chunks of at least MIN_CHUNK_BYTES (small
146
+ /// batches are encoded serially), and a document larger than a chunk is
147
+ /// split at pretoken-safe boundaries and reassembled with identical
148
+ /// tokens — a single huge document still uses all cores. Chunks are
149
+ /// encoded by pooled workers whose pretoken caches persist across calls.
150
+ ///
151
+ /// `parallel=False` encodes everything on the calling thread instead,
152
+ /// with identical output, never touching the process-global thread pool
153
+ /// — for calls inside multiprocessing worker processes (the
154
+ /// gigatoken.Tokenizer wrapper detects those and passes it
155
+ /// automatically).
156
+ #[pyo3(signature = (inputs, *, parallel = true))]
157
+ fn encode_batch<'py>(
158
+ &self,
159
+ py: Python<'py>,
160
+ inputs: Bound<'py, PyAny>,
161
+ parallel: bool,
162
+ ) -> PyResult<Bound<'py, PyAny>> {
163
+ encode_batch_ragged(py, &inputs, |docs, format| {
164
+ Ok(self.encode_regions_ragged(docs, format, parallel))
165
+ })
166
+ }
167
+
168
+ /// encode_batch returned as plain Python lists (one list of ints per
169
+ /// document), assembled in Rust for callers that need lists — same
170
+ /// inputs and `parallel` keyword as encode_batch, which would otherwise
171
+ /// convert the awkward result one Python object at a time.
172
+ #[pyo3(signature = (inputs, *, parallel = true))]
173
+ fn encode_batch_list<'py>(
174
+ &self,
175
+ py: Python<'py>,
176
+ inputs: Bound<'py, PyAny>,
177
+ parallel: bool,
178
+ ) -> PyResult<Bound<'py, PyList>> {
179
+ encode_batch_pylist(py, &inputs, None, parallel, |docs, format| {
180
+ Ok(self.encode_regions_ragged(docs, format, parallel))
181
+ })
182
+ }
183
+
184
+ /// Non-public entrypoint for the compat wrappers: encode_batch_list
185
+ /// with row assembly options (`options` is a _WrapTruncate —
186
+ /// prefix/suffix wrapping, max_tokens truncation, and the fused
187
+ /// forbidden-specials scan, which raises SpecialTokenFound on any hit).
188
+ #[pyo3(signature = (inputs, options, *, parallel = true))]
189
+ fn _encode_batch_list_compat<'py>(
190
+ &self,
191
+ py: Python<'py>,
192
+ inputs: Bound<'py, PyAny>,
193
+ options: Py<bindings::bridge::WrapTruncate>,
194
+ parallel: bool,
195
+ ) -> PyResult<Bound<'py, PyList>> {
196
+ encode_batch_pylist(py, &inputs, Some(options.get()), parallel, |docs, format| {
197
+ Ok(self.encode_regions_ragged(docs, format, parallel))
198
+ })
199
+ }
200
+
201
+ /// encode_batch assembled into one padded/truncated (rows x width)
202
+ /// uint32 matrix plus each row's real length, serving the compatibility
203
+ /// APIs — see src/bindings/padding.rs for the semantics (`options` is a
204
+ /// PadTruncate) and gigatoken.Tokenizer.encode_batch_padded for the
205
+ /// friendly keyword signature.
206
+ #[pyo3(signature = (inputs, options, *, parallel = true))]
207
+ fn encode_batch_padded<'py>(
208
+ &self,
209
+ py: Python<'py>,
210
+ inputs: Bound<'py, PyAny>,
211
+ options: padding::PadTruncate,
212
+ parallel: bool,
213
+ ) -> PyResult<padding::PaddedMatrix<'py>> {
214
+ padding::encode_batch_matrix(py, &inputs, options, parallel, |docs, format| {
215
+ Ok(self.encode_regions_ragged(docs, format, parallel))
216
+ })
217
+ }
218
+
219
+ /// Encode all documents from files in parallel, releasing the GIL.
220
+ /// Returns an awkward.Array with one row of token ids per document.
221
+ ///
222
+ /// `source` is a TextFileSource / JsonlFileSource, a single path, or a
223
+ /// list of paths (defaults per extension: .jsonl → JSONL with field
224
+ /// "text", anything else → plain text with each file as one document).
225
+ /// Everything happens in Rust: files are mmapped (or decompressed into
226
+ /// memory for .gz/.zst) and cut into chunks at document boundaries; a
227
+ /// file that is one huge document is split at pretoken-safe boundaries
228
+ /// and reassembled with identical tokens, so it still uses all cores.
229
+ /// Chunks are encoded by pooled workers whose pretoken caches persist
230
+ /// across calls.
231
+ ///
232
+ /// `parallel=False` loads and encodes everything on the calling thread
233
+ /// instead, with identical output, never touching the process-global
234
+ /// thread pool — for calls inside multiprocessing worker processes (the
235
+ /// gigatoken.Tokenizer wrapper detects those and passes it
236
+ /// automatically).
237
+ #[pyo3(signature = (source, *, parallel = true))]
238
+ fn encode_files<'py>(
239
+ &self,
240
+ py: Python<'py>,
241
+ source: Bound<'py, PyAny>,
242
+ parallel: bool,
243
+ ) -> PyResult<Bound<'py, PyAny>> {
244
+ encode_files_ragged(py, &source, parallel, |files, format| {
245
+ if parallel {
246
+ encode_files_docs(&self.workers, &self.tokenizer, files, format)
247
+ } else {
248
+ encode_files_docs_serial(&self.workers, &self.tokenizer, files, format)
249
+ }
250
+ })
251
+ }
252
+
253
+ /// Size of the vocabulary: one greater than the largest token ID,
254
+ /// including added tokens.
255
+ #[getter]
256
+ fn vocab_size(&self) -> usize {
257
+ self.tokenizer.vocab_size()
258
+ }
259
+
260
+ /// The vocabulary as a freshly built dict mapping token ID to token
261
+ /// bytes, in ID order, including added tokens.
262
+ #[getter]
263
+ fn vocab<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyDict>> {
264
+ vocab_to_pydict(py, self.tokenizer.vocab_entries())
265
+ }
266
+
267
+ /// The merge rules as a freshly built list of `(left, right)` byte
268
+ /// pairs in merge-priority order.
269
+ #[getter]
270
+ fn merges<'py>(&self, py: Python<'py>) -> Vec<(Bound<'py, PyBytes>, Bound<'py, PyBytes>)> {
271
+ merges_to_pylist(py, self.tokenizer.merge_entries())
272
+ }
273
+
274
+ fn decode(&self, tokens: Bound<'_, PyAny>) -> PyResult<Vec<u8>> {
275
+ let ids = extract_token_ids(&tokens)?;
276
+ Ok(self.tokenizer.decode(ids.as_slice()?).collect())
277
+ }
278
+
279
+ fn __repr__(&self) -> PyResult<String> {
280
+ Ok(format!("{:?}", self.tokenizer))
281
+ }
282
+ }
283
+
284
+ #[cfg(feature = "python")]
285
+ #[pyclass]
286
+ struct SentencePieceTokenizer {
287
+ tokenizer: bpe::SentencePieceBPE,
288
+ /// Pretoken cache + scratch for single-document `encode`, persisting
289
+ /// across calls (parallel paths use per-worker states instead).
290
+ state: bpe::sentencepiece::EncodeState,
291
+ }
292
+
293
+ #[cfg(feature = "python")]
294
+ impl SentencePieceTokenizer {
295
+ /// See `batch::sp_encode_docs_ragged` (`_serial` when `parallel` is
296
+ /// false); a separator format (a BytesSource input) goes through
297
+ /// `batch::sp_encode_files_docs` instead, which splits the regions into
298
+ /// documents during the encode. Regions are trusted to be valid UTF-8
299
+ /// per the documented contract — never validated, matching
300
+ /// encode_files' treatment of file contents — which is what makes the
301
+ /// unchecked str conversions here and inside sp_encode_files_docs
302
+ /// sound. Call with the GIL released.
303
+ fn encode_regions_ragged(
304
+ &self,
305
+ regions: &[&[u8]],
306
+ format: &DocFormat,
307
+ parallel: bool,
308
+ ) -> PyResult<(Vec<u32>, Vec<i64>)> {
309
+ debug_assert!(regions.iter().all(|d| std::str::from_utf8(d).is_ok()));
310
+ Ok(match format {
311
+ DocFormat::Text { separator: None } => {
312
+ // SAFETY: valid UTF-8 by the documented input contract.
313
+ let texts: Vec<&str> = regions
314
+ .iter()
315
+ .map(|d| unsafe { std::str::from_utf8_unchecked(d) })
316
+ .collect();
317
+ if parallel {
318
+ sp_encode_docs_ragged(&self.tokenizer, &texts)
319
+ } else {
320
+ sp_encode_docs_ragged_serial(&self.tokenizer, &texts)
321
+ }
322
+ }
323
+ format => {
324
+ // A separator match inside valid UTF-8 lands on char
325
+ // boundaries only when the separator is itself valid UTF-8
326
+ // (its first byte is a lead byte, and a complete match ends
327
+ // a complete character) — an arbitrary byte separator could
328
+ // cut a document mid-character and break the unchecked
329
+ // conversion inside sp_encode_files_docs. This one
330
+ // constant-time argument check stays even though document
331
+ // bytes are trusted.
332
+ if let DocFormat::Text { separator: Some(sep) } = format
333
+ && std::str::from_utf8(sep).is_err()
334
+ {
335
+ return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
336
+ "the SentencePiece backend requires a separator that is valid UTF-8",
337
+ ));
338
+ }
339
+ if parallel {
340
+ sp_encode_files_docs(&self.tokenizer, regions, format)
341
+ } else {
342
+ sp_encode_files_docs_serial(&self.tokenizer, regions, format)
343
+ }
344
+ }
345
+ })
346
+ }
347
+ }
348
+
349
+ #[cfg(feature = "python")]
350
+ #[pymethods]
351
+ impl SentencePieceTokenizer {
352
+ #[staticmethod]
353
+ fn from_hf(path: PathBuf) -> PyResult<Self> {
354
+ Ok(Self {
355
+ tokenizer: load_tokenizer::hf::load_hf_sentencepiece(&path)?,
356
+ state: bpe::sentencepiece::EncodeState::new(),
357
+ })
358
+ }
359
+
360
+ /// Encode a batch of documents in parallel, releasing the GIL. Accepts
361
+ /// the same inputs, returns the same awkward.Array shape, and honors the
362
+ /// same `parallel` keyword as BPETokenizer.encode_batch. Documents must
363
+ /// be valid UTF-8 — trusted, never validated, like encode_files.
364
+ #[pyo3(signature = (inputs, *, parallel = true))]
365
+ fn encode_batch<'py>(
366
+ &self,
367
+ py: Python<'py>,
368
+ inputs: Bound<'py, PyAny>,
369
+ parallel: bool,
370
+ ) -> PyResult<Bound<'py, PyAny>> {
371
+ encode_batch_ragged(py, &inputs, |docs, format| {
372
+ self.encode_regions_ragged(docs, format, parallel)
373
+ })
374
+ }
375
+
376
+ /// See BPETokenizer.encode_batch_list: the same lists-built-in-Rust
377
+ /// batch encode, for the SentencePiece backend.
378
+ #[pyo3(signature = (inputs, *, parallel = true))]
379
+ fn encode_batch_list<'py>(
380
+ &self,
381
+ py: Python<'py>,
382
+ inputs: Bound<'py, PyAny>,
383
+ parallel: bool,
384
+ ) -> PyResult<Bound<'py, PyList>> {
385
+ encode_batch_pylist(py, &inputs, None, parallel, |docs, format| {
386
+ self.encode_regions_ragged(docs, format, parallel)
387
+ })
388
+ }
389
+
390
+ /// See BPETokenizer._encode_batch_list_compat: the same non-public
391
+ /// compat entrypoint (row assembly options via _WrapTruncate), for the
392
+ /// SentencePiece backend.
393
+ #[pyo3(signature = (inputs, options, *, parallel = true))]
394
+ fn _encode_batch_list_compat<'py>(
395
+ &self,
396
+ py: Python<'py>,
397
+ inputs: Bound<'py, PyAny>,
398
+ options: Py<bindings::bridge::WrapTruncate>,
399
+ parallel: bool,
400
+ ) -> PyResult<Bound<'py, PyList>> {
401
+ encode_batch_pylist(py, &inputs, Some(options.get()), parallel, |docs, format| {
402
+ self.encode_regions_ragged(docs, format, parallel)
403
+ })
404
+ }
405
+
406
+ /// See BPETokenizer.encode_batch_padded: the same padded-matrix batch
407
+ /// encode, for the SentencePiece backend.
408
+ #[pyo3(signature = (inputs, options, *, parallel = true))]
409
+ fn encode_batch_padded<'py>(
410
+ &self,
411
+ py: Python<'py>,
412
+ inputs: Bound<'py, PyAny>,
413
+ options: padding::PadTruncate,
414
+ parallel: bool,
415
+ ) -> PyResult<padding::PaddedMatrix<'py>> {
416
+ padding::encode_batch_matrix(py, &inputs, options, parallel, |docs, format| {
417
+ self.encode_regions_ragged(docs, format, parallel)
418
+ })
419
+ }
420
+
421
+ /// Encode a single document (str or UTF-8 bytes), with a pretoken cache
422
+ /// that persists across calls.
423
+ fn encode<'py>(
424
+ &mut self,
425
+ py: Python<'py>,
426
+ input: Bound<'py, PyAny>,
427
+ ) -> PyResult<Bound<'py, PyArray1<u32>>> {
428
+ let input = extract_doc(&input)?;
429
+ let text: &str = match &input {
430
+ EncodeInput::Text(s) => s,
431
+ EncodeInput::Bytes(b) => {
432
+ debug_assert!(std::str::from_utf8(b).is_ok());
433
+ // SAFETY: valid UTF-8 by the documented input contract
434
+ // (trusted unchecked, like the batch and file paths).
435
+ unsafe { std::str::from_utf8_unchecked(b) }
436
+ }
437
+ };
438
+ let mut ids: Vec<u32> = Vec::new();
439
+ self.tokenizer
440
+ .encode_raw_cb(&mut self.state, text, &mut |tokens| {
441
+ ids.extend(tokens.iter().map(|&t| u32::from(t)))
442
+ });
443
+ Ok(ids.into_pyarray(py))
444
+ }
445
+
446
+ /// Encode all documents from files in parallel. Accepts the same
447
+ /// sources, applies the same chunking policy, and honors the same
448
+ /// `parallel` keyword as BPETokenizer.encode_files, and likewise
449
+ /// returns an awkward.Array with one row of token ids per document.
450
+ #[pyo3(signature = (source, *, parallel = true))]
451
+ fn encode_files<'py>(
452
+ &self,
453
+ py: Python<'py>,
454
+ source: Bound<'py, PyAny>,
455
+ parallel: bool,
456
+ ) -> PyResult<Bound<'py, PyAny>> {
457
+ encode_files_ragged(py, &source, parallel, |files, format| {
458
+ if parallel {
459
+ sp_encode_files_docs(&self.tokenizer, files, format)
460
+ } else {
461
+ sp_encode_files_docs_serial(&self.tokenizer, files, format)
462
+ }
463
+ })
464
+ }
465
+
466
+ /// Size of the vocabulary: one greater than the largest token ID,
467
+ /// including added tokens.
468
+ #[getter]
469
+ fn vocab_size(&self) -> usize {
470
+ self.tokenizer.vocab_size()
471
+ }
472
+
473
+ /// The vocabulary as a freshly built dict mapping token ID to token
474
+ /// bytes, in ID order, including added tokens.
475
+ #[getter]
476
+ fn vocab<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyDict>> {
477
+ vocab_to_pydict(py, self.tokenizer.vocab_entries())
478
+ }
479
+
480
+ /// The merge rules as a freshly built list of `(left, right)` byte
481
+ /// pairs in merge-priority order.
482
+ #[getter]
483
+ fn merges<'py>(&self, py: Python<'py>) -> Vec<(Bound<'py, PyBytes>, Bound<'py, PyBytes>)> {
484
+ merges_to_pylist(py, self.tokenizer.merge_entries())
485
+ }
486
+
487
+ fn encode_no_normalize<'py>(
488
+ &mut self,
489
+ py: Python<'py>,
490
+ input: &str,
491
+ ) -> PyResult<Bound<'py, PyArray1<u32>>> {
492
+ let mut ids: Vec<u32> = Vec::new();
493
+ self.tokenizer
494
+ .encode_normalized_cb(&mut self.state, input, &mut |tokens| {
495
+ ids.extend(tokens.iter().map(|&t| u32::from(t)))
496
+ });
497
+ Ok(ids.into_pyarray(py))
498
+ }
499
+
500
+ fn decode(&self, tokens: Bound<'_, PyAny>) -> PyResult<Vec<u8>> {
501
+ let ids = extract_token_ids(&tokens)?;
502
+ Ok(self.tokenizer.decode(ids.as_slice()?))
503
+ }
504
+
505
+ fn __repr__(&self) -> PyResult<String> {
506
+ Ok(format!("{:?}", self.tokenizer))
507
+ }
508
+ }
509
+
510
+ /// Load a tokenizer from in-memory HuggingFace `tokenizer.json` contents
511
+ /// (str or bytes). Returns a SentencePieceTokenizer when the model uses
512
+ /// byte_fallback, a BPETokenizer otherwise — the same split as the two
513
+ /// classes' from_hf constructors.
514
+ #[cfg(feature = "python")]
515
+ #[pyfunction]
516
+ fn load_hf_json(py: Python<'_>, data: Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
517
+ let backed_str;
518
+ let backed_bytes;
519
+ let bytes: &[u8] = if let Ok(s) = data.extract::<PyBackedStr>() {
520
+ backed_str = s;
521
+ backed_str.as_bytes()
522
+ } else if let Ok(b) = data.extract::<PyBackedBytes>() {
523
+ backed_bytes = b;
524
+ &backed_bytes
525
+ } else {
526
+ return Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>(format!(
527
+ "expected tokenizer.json contents as str or bytes, got {}",
528
+ data.get_type()
529
+ )));
530
+ };
531
+ match load_tokenizer::hf::load_hf_slice(bytes)? {
532
+ load_tokenizer::hf::HfTokenizer::Bpe(tokenizer) => Ok(Py::new(
533
+ py,
534
+ BPETokenizer {
535
+ tokenizer,
536
+ workers: WorkerPool::new(),
537
+ },
538
+ )?
539
+ .into_any()),
540
+ load_tokenizer::hf::HfTokenizer::SentencePiece(tokenizer) => Ok(Py::new(
541
+ py,
542
+ SentencePieceTokenizer {
543
+ tokenizer,
544
+ state: bpe::sentencepiece::EncodeState::new(),
545
+ },
546
+ )?
547
+ .into_any()),
548
+ }
549
+ }
550
+
551
+ // ---------------------------------------------------------------------------
552
+ // Module registration
553
+ // ---------------------------------------------------------------------------
554
+
555
+ #[cfg(feature = "python")]
556
+ #[pymodule]
557
+ fn gigatoken_rs<'py>(py: Python, m: &Bound<'py, PyModule>) -> PyResult<()> {
558
+ m.add("SpecialTokenFound", py.get_type::<SpecialTokenFound>())?;
559
+ m.add_function(wrap_pyfunction!(train_bpe, m)?)?;
560
+ m.add_class::<FileSource>()?;
561
+ m.add_class::<TextFileSource>()?;
562
+ m.add_class::<JsonlFileSource>()?;
563
+ m.add_class::<ParquetFileSource>()?;
564
+ m.add_class::<BytesSource>()?;
565
+ m.add_class::<PretokenizerIter>()?;
566
+ m.add_class::<SubstringMatcher>()?;
567
+ m.add_class::<bindings::bridge::WrapTruncate>()?;
568
+ m.add_class::<padding::PadTruncate>()?;
569
+ m.add_class::<BPETokenizer>()?;
570
+ m.add_class::<SentencePieceTokenizer>()?;
571
+ m.add_function(wrap_pyfunction!(pretokenizer, m)?)?;
572
+ m.add_function(wrap_pyfunction!(pretokenized_counts, m)?)?;
573
+ m.add_function(wrap_pyfunction!(load_hf_json, m)?)?;
574
+ m.add_function(wrap_pyfunction!(bindings::hub::hub_file, m)?)?;
575
+ m.add_function(wrap_pyfunction!(bindings::hub::looks_like_repo_id, m)?)?;
576
+ m.add_function(wrap_pyfunction!(bindings::hub::get_hf_token, m)?)?;
577
+ Ok(())
578
+ }