From adf52b1ca984039e4daca482ca82b926276944ac Mon Sep 17 00:00:00 2001 From: Arthur Date: Wed, 15 Jul 2026 17:40:44 +0200 Subject: [PATCH 1/8] build: add size-optimized 'thin' release profile opt-level=z + codegen-units=1 + panic=abort + strip (inherits fat LTO). For on-device footprint testing: cargo build --profile thin. Minimal encode binary: 1.90 MB -> 1.27 MB (-33%) vs release. --- tokenizers/Cargo.toml | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/tokenizers/Cargo.toml b/tokenizers/Cargo.toml index affb24158..7ff9d4b5f 100644 --- a/tokenizers/Cargo.toml +++ b/tokenizers/Cargo.toml @@ -107,6 +107,16 @@ lto = "fat" inherits = "release" debug = true +# Size-optimized "thin" release for on-device footprint experiments: opt-for-size, +# single codegen unit (better cross-fn dedup), abort-on-panic (drops unwind tables), +# and stripped. Inherits `lto = "fat"` from release. Build: `cargo build --profile thin`. +[profile.thin] +inherits = "release" +opt-level = "z" +codegen-units = 1 +panic = "abort" +strip = true + [[example]] name = "encode_batch" required-features = ["http"] From f0dc52616f6ef1122542fce2592c68a72a6b85ea Mon Sep 17 00:00:00 2001 From: Arthur Date: Wed, 15 Jul 2026 17:59:23 +0200 Subject: [PATCH 2/8] build(tk-encode): vendor slimmed ptr_hash (drop colored/mem_dbg/tempfile) Upstream ptr_hash 2.0.1 forces colored, mem_dbg (+ its mem_dbg-derive -> syn/quote/ proc-macro2 chain) and tempfile as NON-optional deps, though they only serve CLI/ debug/on-disk-sharding paths we never hit. Vendor a copy under vendor/ptr_hash with those three removed (MemSize derives stripped, colored .bold() dropped, a std TempDir shim for the unused disk-sharding path) and [patch.crates-io] to it. tk-encode normal deps 85 -> 74. MPHF build/query unchanged: vocab tests 25/25 pass. --- tokenizers/Cargo.lock | 65 - tokenizers/Cargo.toml | 5 + tokenizers/vendor/ptr_hash/.cargo-ok | 1 + .../vendor/ptr_hash/.cargo_vcs_info.json | 6 + tokenizers/vendor/ptr_hash/Cargo.toml | 164 +++ tokenizers/vendor/ptr_hash/readme.md | 166 +++ tokenizers/vendor/ptr_hash/src/bucket_fn.rs | 197 +++ tokenizers/vendor/ptr_hash/src/bucket_idx.rs | 62 + tokenizers/vendor/ptr_hash/src/build.rs | 435 +++++++ tokenizers/vendor/ptr_hash/src/fastmod.rs | 111 ++ tokenizers/vendor/ptr_hash/src/hash.rs | 212 +++ tokenizers/vendor/ptr_hash/src/lib.rs | 1141 +++++++++++++++++ tokenizers/vendor/ptr_hash/src/pack.rs | 172 +++ tokenizers/vendor/ptr_hash/src/reduce.rs | 61 + tokenizers/vendor/ptr_hash/src/shard.rs | 282 ++++ .../vendor/ptr_hash/src/sort_buckets.rs | 176 +++ tokenizers/vendor/ptr_hash/src/stats.rs | 156 +++ tokenizers/vendor/ptr_hash/src/test.rs | 374 ++++++ tokenizers/vendor/ptr_hash/src/util.rs | 66 + 19 files changed, 3787 insertions(+), 65 deletions(-) create mode 100644 tokenizers/vendor/ptr_hash/.cargo-ok create mode 100644 tokenizers/vendor/ptr_hash/.cargo_vcs_info.json create mode 100644 tokenizers/vendor/ptr_hash/Cargo.toml create mode 100644 tokenizers/vendor/ptr_hash/readme.md create mode 100644 tokenizers/vendor/ptr_hash/src/bucket_fn.rs create mode 100644 tokenizers/vendor/ptr_hash/src/bucket_idx.rs create mode 100644 tokenizers/vendor/ptr_hash/src/build.rs create mode 100644 tokenizers/vendor/ptr_hash/src/fastmod.rs create mode 100644 tokenizers/vendor/ptr_hash/src/hash.rs create mode 100644 tokenizers/vendor/ptr_hash/src/lib.rs create mode 100644 tokenizers/vendor/ptr_hash/src/pack.rs create mode 100644 tokenizers/vendor/ptr_hash/src/reduce.rs create mode 100644 tokenizers/vendor/ptr_hash/src/shard.rs create mode 100644 tokenizers/vendor/ptr_hash/src/sort_buckets.rs create mode 100644 tokenizers/vendor/ptr_hash/src/stats.rs create mode 100644 tokenizers/vendor/ptr_hash/src/test.rs create mode 100644 tokenizers/vendor/ptr_hash/src/util.rs diff --git a/tokenizers/Cargo.lock b/tokenizers/Cargo.lock index 4e5a39e84..9dfbe508a 100644 --- a/tokenizers/Cargo.lock +++ b/tokenizers/Cargo.lock @@ -31,12 +31,6 @@ dependencies = [ "memchr", ] -[[package]] -name = "allocator-api2" -version = "0.2.21" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" - [[package]] name = "anes" version = "0.1.6" @@ -336,15 +330,6 @@ version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" -[[package]] -name = "colored" -version = "3.1.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "faf9468729b8cbcea668e36183cb69d317348c2e08e994829fb56ebfdfbaac34" -dependencies = [ - "windows-sys 0.61.2", -] - [[package]] name = "compact_str" version = "0.9.1" @@ -624,12 +609,6 @@ version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "34aa73646ffb006b8f5147f3dc182bd4bcb190227ce861fc4a4844bf8e3cb2c0" -[[package]] -name = "equivalent" -version = "1.0.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" - [[package]] name = "errno" version = "0.3.14" @@ -699,12 +678,6 @@ version = "1.0.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" -[[package]] -name = "foldhash" -version = "0.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb" - [[package]] name = "form_urlencoded" version = "1.2.2" @@ -845,17 +818,6 @@ dependencies = [ "zerocopy", ] -[[package]] -name = "hashbrown" -version = "0.16.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100" -dependencies = [ - "allocator-api2", - "equivalent", - "foldhash", -] - [[package]] name = "hermit-abi" version = "0.5.2" @@ -1300,28 +1262,6 @@ version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "670fdfda89751bc4a84ac13eaa63e205cf0fd22b4c9a5fbfa085b63c1f1d3a30" -[[package]] -name = "mem_dbg" -version = "0.4.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f4ef2d80bfa14894b6d5a3ff537e7e9a908dbf4c95de8a5b8ad2a473301676e6" -dependencies = [ - "bitflags", - "hashbrown", - "mem_dbg-derive", -] - -[[package]] -name = "mem_dbg-derive" -version = "0.3.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "73acd151c6ce84a41d8d6fb0958d9a3d5a18d649ad5a85ad5b719439af8ad257" -dependencies = [ - "proc-macro2", - "quote", - "syn", -] - [[package]] name = "memchr" version = "2.8.2" @@ -1590,23 +1530,18 @@ dependencies = [ [[package]] name = "ptr_hash" version = "2.0.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a847c2cc746ab2aeba36aad3e75fc417b47539603298c12d8373e388890aad3c" dependencies = [ "bitvec", - "colored", "fastrand", "fxhash", "itertools 0.15.0", "log", - "mem_dbg", "prefetch-index", "rand 0.10.2", "rand_chacha 0.10.0", "rayon", "rdst", "serde", - "tempfile", "xxhash-rust", ] diff --git a/tokenizers/Cargo.toml b/tokenizers/Cargo.toml index 7ff9d4b5f..fd127c120 100644 --- a/tokenizers/Cargo.toml +++ b/tokenizers/Cargo.toml @@ -117,6 +117,11 @@ codegen-units = 1 panic = "abort" strip = true +# Slimmed vendored ptr_hash: identical MPHF, minus the non-optional colored/mem_dbg/tempfile deps +# (and their syn/quote/proc-macro2 chain) that upstream forces. See vendor/ptr_hash/Cargo.toml. +[patch.crates-io] +ptr_hash = { path = "vendor/ptr_hash" } + [[example]] name = "encode_batch" required-features = ["http"] diff --git a/tokenizers/vendor/ptr_hash/.cargo-ok b/tokenizers/vendor/ptr_hash/.cargo-ok new file mode 100644 index 000000000..5f8b79583 --- /dev/null +++ b/tokenizers/vendor/ptr_hash/.cargo-ok @@ -0,0 +1 @@ +{"v":1} \ No newline at end of file diff --git a/tokenizers/vendor/ptr_hash/.cargo_vcs_info.json b/tokenizers/vendor/ptr_hash/.cargo_vcs_info.json new file mode 100644 index 000000000..c85eae24b --- /dev/null +++ b/tokenizers/vendor/ptr_hash/.cargo_vcs_info.json @@ -0,0 +1,6 @@ +{ + "git": { + "sha1": "60a379911e09ec6a046a8118ddb6d4c12856751c" + }, + "path_in_vcs": "" +} \ No newline at end of file diff --git a/tokenizers/vendor/ptr_hash/Cargo.toml b/tokenizers/vendor/ptr_hash/Cargo.toml new file mode 100644 index 000000000..fc1a4ccf4 --- /dev/null +++ b/tokenizers/vendor/ptr_hash/Cargo.toml @@ -0,0 +1,164 @@ +# Vendored + slimmed copy of ptr_hash 2.0.1 (github.com/RagnarGrootKoerkamp/ptrhash). +# The upstream crate makes `colored`, `mem_dbg` (+ its syn/quote/proc-macro2 derive chain) and +# `tempfile` NON-optional even though they are only used for CLI/debug/on-disk-sharding paths we +# never hit. This copy removes those three deps (and the example/dev-dep bloat) so `tk-encode`'s +# dependency graph shrinks. The MPHF build/query path is untouched — byte-exact with upstream. +# Applied via `[patch.crates-io] ptr_hash = { path = "vendor/ptr_hash" }` in the workspace root. + +[package] +edition = "2021" +name = "ptr_hash" +version = "2.0.1" +authors = ["Ragnar Groot Koerkamp"] +build = false +autolib = false +autobins = false +autoexamples = false +autotests = false +autobenches = false +description = "A high-throughput minimal perfect hash function" +homepage = "https://curiouscoding.nl/posts/ptrhash" +keywords = [ + "hashing", + "minimal", + "perfect", + "mphf", +] +categories = [ + "data-structures", + "science", +] +license = "MIT" +repository = "https://github.com/RagnarGrootKoerkamp/ptrhash" + +[features] +cacheline-ef = ["dep:cacheline-ef"] +clap = ["dep:clap"] +default = [ + "elias-fano", + "cacheline-ef", + "gxhash", +] +elias-fano = ["dep:sucds"] +epserde = [ + "dep:epserde", + "dep:epserde-derive", + "cacheline-ef/epserde", +] +gxhash = ["dep:gxhash"] +hashers = [ + "dep:cityhash-102-rs", + "dep:fastmurmur3", + "dep:hashers", + "dep:highway", + "dep:metrohash", + "dep:murmur2", + "dep:murmur3", + "dep:wyhash", +] +unstable = [] + +[lib] +name = "ptr_hash" +path = "src/lib.rs" + +[dependencies.bitvec] +version = "1.0.1" + +[dependencies.cacheline-ef] +version = "1.1.0" +optional = true + +[dependencies.cityhash-102-rs] +version = "0.1.0" +optional = true + +[dependencies.clap] +version = "4.4.6" +features = ["derive"] +optional = true + +[dependencies.epserde] +version = "0.13.0" +optional = true + +[dependencies.epserde-derive] +version = "0.13.0" +optional = true + +[dependencies.fastmurmur3] +version = "0.2.0" +optional = true + +[dependencies.fastrand] +version = "2.0.1" + +[dependencies.fxhash] +version = "0.2.1" + +[dependencies.gxhash] +version = "3.5.0" +optional = true + +[dependencies.hashers] +version = "1.0.1" +optional = true + +[dependencies.highway] +version = "1.1.0" +optional = true + +[dependencies.itertools] +version = "0.15.0" + +[dependencies.log] +version = "0.4.27" + +[dependencies.metrohash] +version = "1.0.6" +optional = true + +[dependencies.murmur2] +version = "0.1.0" +optional = true + +[dependencies.murmur3] +version = "0.5.2" +optional = true + +[dependencies.prefetch-index] +version = "0.2.1" + +[dependencies.rand] +version = "0.10.2" + +[dependencies.rand_chacha] +version = "0.10.0" + +[dependencies.rayon] +version = "1.8.0" + +[dependencies.rdst] +version = "0.20.11" + +[dependencies.serde] +version = "1.0.215" +features = ["derive"] + +[dependencies.sucds] +version = "0.8.0" +optional = true + +[dependencies.wyhash] +version = "0.6.0" +optional = true + +[dependencies.xxhash-rust] +version = "0.8.7" +features = [ + "xxh64", + "xxh3", +] + +# Opt out of the parent tokenizers workspace so this vendored crate is its own root. +[workspace] diff --git a/tokenizers/vendor/ptr_hash/readme.md b/tokenizers/vendor/ptr_hash/readme.md new file mode 100644 index 000000000..4644ab442 --- /dev/null +++ b/tokenizers/vendor/ptr_hash/readme.md @@ -0,0 +1,166 @@ +# PtrHash: Minimal Perfect Hashing at RAM Throughput + +[![crates.io](https://img.shields.io/crates/v/ptr_hash.svg)](https://crates.io/crates/ptr_hash) +[![docs.rs](https://img.shields.io/docsrs/ptr_hash.svg)](https://docs.rs/ptr_hash) + +PtrHash is a fast and space efficient *minimal perfect hash function* that maps +a list of `n` distinct keys into `{0,...,n-1}`. +It is based on/inspired by [PTHash](https://github.com/jermp/pthash) (and much +more than just a Rust rewrite). + +**Paper.** + +*Ragnar Groot Koerkamp*. PtrHash: Minimal Perfect Hashing at RAM Throughput. +SEA (2025). [doi.org/10.4230/LIPIcs.SEA.2025.21](https://doi.org/10.4230/LIPIcs.SEA.2025.21) + +**Evals.** Source code for the paper evals can be found in +[examples/evals.rs](examples/evals.rs), and analysis is [evals.py](evals.py). +Plots can be found [in the blog](https://github.com/RagnarGrootKoerkamp/research/blob/master/posts/ptrhash/). +The paper evals were done on the `evals` branch (which is v1.0 with GxHash added +for string hashing) and my [fork](https://github.com/ragnargrootkoerkamp/MPHF-Experiments) of [mphf-experiments](https://github.com/ByteHamster/MPHF-Experiments). + +For changes since then, see [CHANGELOG.md](./CHANGELOG.md). + +**Contact.** + +In case you run into any kind of issue or things are unclear, +please make issues and/or PRs, or reach out on [twitter]((https://twitter.com/curious_coding))/[bsky](https://bsky.app/profile/curiouscoding.nl). +I'm more than happy to help out with integrating PtrHash. + +## Performance on small inputs + +Space usage and query throughput of a for-loop in ns/key as measured by `examples/query_bench.rs`: + +| method | bits/key | hash | 100k | 1M | 10M | 100 M | +|----------------|----------|--------|------|-----|-----|-------| +| FastPtrHash | 2.67 | NoHash | 1.5 | 2.2 | 3.1 | 7.6 | +| | | FxHash | 1.7 | 2.5 | 3.2 | 8.2 | +| DefaultPtrHash | 3.00 | NoHash | 1.7 | 2.6 | 3.3 | 8.5 | +| | | FxHash | 2.0 | 2.7 | 3.7 | 9.1 | +| CompactPtrHash | 2.15 | NoHash | 4.4 | 5.5 | 6.8 | 15.6 | +| | | FxHash | 4.8 | 5.8 | 7.2 | 15.7 | + + +## Performance on large input + +PtrHash supports up to `2^40` keys (and probably more). +For `n=10^9` integer keys with `FxHash`, we get the following on my `i7-10750H` +at `3.6GHz` with 6 cores (see `examples/large_bench.rs`). +The construction uses multi-threading only for `CompactPtrHash`, and remapping +is skipped for `FastPtrHash`. +The last two columns indicate the query throughput for streaming queries with +prefetching 32 iterations ahead, and for prefetching in batches of 32. + +| method | pilots bits/key | remap bits/key | space bits/key | construction ns/key | loop ns/key | stream ns/key | batch ns/key | +|----------------|-----------------|-----------------|----------------|---------------------|-------------|---------------|--------------| +| FastPtrHash | 2.67 | - | 2.67 | 343.5 | 9.7 | 7.5 | 9.0 | +| DefaultPtrHash | 2.67 | 0.33 | 3.00 | 349.9 | 12.0 | 8.3 | 9.4 | +| CompactPtrHash | 2.05 | 0.10 | 2.15 | 49.0 (12t) | 19.8 | 8.5 | 10.1 | + +Streaming query throughput per thread fully saturates the memory bandwidth of each +core (around 7.5 ns/cache line), and with multi-threading the full DDR4 memory bandwidth is saturated. + +## Input + +PtrHash is primarily intended to be used on large sets of keys, say of size at +least 1 million. Nevertheless, it can also be used for sets as small as e.g. 10 +keys. In this case, there will be a relatively large constant space overhead, +and other methods may be smaller and/or faster. +(PtrHash should work fine and be reasonably fast, but for such small inputs the space-efficient +design of PtrHash makes little sense and faster queries might be possible.) + + +## Usage + +See [docs.rs](https://docs.rs/ptr_hash) for the different variants and parameters. +Below, we use `PtrHashParams::default()` for a reasonable trade-off between size +(2.4 bits/key) and speed. +Slightly smaller size is possible using `PtrHashParams::default_compact()`, +at the cost of significantly slower construction time (2x) and lowered reliability. + +There is also `PtrHashParams::default_fast()`, which takes 25% more space but +can be almost 2x faster when querying integer keys in tight loops. Nevertheless, +for large inputs, maximum query throughput is achieved with `index_stream` with default parameters. + +```rust +use ptr_hash::{PtrHash, PtrHashParams}; + +// Generate some random keys. +let n = 1_000_000_000; +let keys = ptr_hash::util::generate_keys(n); + +// Build the datastructure. +let mphf = ::new(&keys, PtrHashParams::default()); + +// Get the minimal index of a key. +let key = 0; +let idx = mphf.index(&key); +assert!(idx < n); + +// Get the non-minimal index of a key. Slightly faster, but can be >=n. +let _idx = mphf.index(&key); + +// An iterator over the indices of the keys. +// 32: number of iterations ahead to prefetch. +// true: remap to a minimal key in [0, n). +let indices = mphf.index_stream::<32, _>(&keys); +assert_eq!(indices.sum::(), (n * (n - 1)) / 2); + +// Test that all items map to different indices +let mut taken = vec![false; n]; +for key in &keys { + let idx = mphf.index(&key); + assert!(!taken[idx]); + taken[idx] = true; +} + +// In case you want maximum query throughput at the cost of returning non-minimal values, +// use `FastPtrHash`. `phf.max_index()` will be roughly `1.01*n`. +let phf = ::new(&keys, PtrHashParams::default()); + +for key in &keys { + let idx = phf.index(&key); + assert!(idx < phf.max_index()); +} + +// To enable parallel construction and optimize for space, use `CompactPtrHash`. +let mphf = ::new(&keys, PtrHashParams::default_compact()); +``` + +## Epserde + +The `PtrHash` datastructure can be (de)serialized to/from disk using +[epserde](https://github.com/vigna/epserde-rs) when the `epserde` feature is set. +This also allows convenient deserialization using `mmap`. +See [examples/epserde.rs](examples/epserde.rs) for an example. + +## Sharding + +In order to build PtrHash on large sets of keys that do not fit in ram, the keys +can be sharded and constructed one shard at a time. +See `fn sharding()` in [examples/evals.rs](examples/epserde.rs) for an example. + +## Compared to PTHash + +PtrHash extends PTHash in a few ways: + +- **8-bit pilots:** Instead of allowing pilots to take any integer value, we + restrict them to `[0, 256)` and store them as `Vec` directly. + This avoids the need for a compact or dictionary encoding. +- **Evicting:** To get all pilots to be small, we use *evictions*, similar + to *cuckoo hashing*: Whenever we cannot find a collision-free pilot for a + bucket, we find the pilot with the fewest collisions and *evict* all + colliding buckets, which are pushed on a queue after which they will search + for a new pilot. +- **Partitioning:** To speed up construction, we partition all keys/hashes + into parts such that each part contains `S=2^k` *slots*. + This significantly speeds up + construction since all reads of the `taken` bitvector are now very local. + + This brings the benefit that the only global memory needed is to store the + hashes for each part. The sorting, bucketing, and slot filling is per-part + and needs comparatively little memory. +- **Remap encoding:** We use the `CachelineEF` partitioned Elias-Fano encoding that stores + chunks of `44` integers into a single cacheline. This takes `~30%` more + space for remapping, but replaces the three reads needed by (global) + Elias-Fano encoding by a single read. diff --git a/tokenizers/vendor/ptr_hash/src/bucket_fn.rs b/tokenizers/vendor/ptr_hash/src/bucket_fn.rs new file mode 100644 index 000000000..ea6257cdc --- /dev/null +++ b/tokenizers/vendor/ptr_hash/src/bucket_fn.rs @@ -0,0 +1,197 @@ +//! Various "bucket functions" are implemented here. +//! +//! These functions map a uniform `u64` hash to a new `u64`, to skew the distribution of bucket sizes. +//! +//! By default, `CubicEps` is used. `Linear` is simplest and can be faster at the cost of requiring more space. +//! The remaining ones are only kept for benchmarking. + + +use crate::util::mul_high; +use std::fmt::Debug; + +pub trait BucketFn: Clone + Copy + Sync + Debug { + const LINEAR: bool = false; + const B_OUTPUT: bool = false; + fn set_buckets_per_part(&mut self, _b: u64) {} + fn call(&self, x: u64) -> u64; +} + +/// The function simply returns `x` itself. +#[derive(Clone, Copy, Debug, Default)] +#[cfg_attr(feature = "epserde", derive(epserde::prelude::Epserde))] +#[cfg_attr(feature = "epserde", epserde(deep_copy))] +pub struct Linear; + +impl BucketFn for Linear { + const LINEAR: bool = true; + fn call(&self, x: u64) -> u64 { + x + } +} + +/// A 2-piece-wise linear function, as used in FCH and PTHash. +/// +/// | . +/// | . +/// | ....---< gamma +/// | ..... | +/// |.... | +/// +------------^-- +/// beta +/// +/// line1: y = x * (gamma / beta) +/// ~~~ slope1 ~~~ +/// line2: y = x * ((1 - gamma) / (1 - beta)) + (gamma - beta) / (1 - beta) +/// ~~~~~~~~~ slope2 ~~~~~~~~~ ~~~~~~~~~~ offset ~~~~~~~~~ +#[derive(Clone, Copy, Debug)] +#[cfg_attr(feature = "epserde", derive(epserde::prelude::Epserde))] +#[cfg_attr(feature = "epserde", epserde(deep_copy))] +pub struct Skewed { + beta_f: f64, + gamma_f: f64, + /// buckets per part + b: u64, + beta: u64, + slope1: u64, + slope2: u64, + neg_offset: u64, +} + +impl Default for Skewed { + fn default() -> Self { + Skewed::new(0.6, 0.3) + } +} + +impl Skewed { + // Map the first beta% of hashes to the first gamma% of buckets. + pub fn new(beta: f64, gamma: f64) -> Self { + assert!( + beta > gamma, + "Beta={beta} must be larger than gamma={gamma}" + ); + Self { + beta_f: beta, + gamma_f: gamma, + b: 0, + beta: 0, + slope1: 0, + slope2: 0, + neg_offset: 0, + } + } +} + +impl BucketFn for Skewed { + const B_OUTPUT: bool = true; + fn set_buckets_per_part(&mut self, b: u64) { + let beta = self.beta_f; + let gamma = self.gamma_f; + self.b = b; + let as_u64 = |x: f64| (x * u64::MAX as f64) as u64; + self.slope1 = mul_high(as_u64(gamma / beta), self.b); + self.slope2 = mul_high(as_u64((1. - gamma) / (1. - beta) / 8.), self.b << 3); + self.neg_offset = mul_high(as_u64((beta - gamma) / (1. - beta) / 8.), self.b << 3); + self.beta = as_u64(beta); + } + fn call(&self, x: u64) -> u64 { + // NOTE: There is a lot of MOV/CMOV going on here. + let is_large = x >= self.beta; + let slope = if is_large { self.slope2 } else { self.slope1 }; + mul_high(x, slope) - is_large as u64 * self.neg_offset + // debug_assert!(!is_large || self.p2 <= b, "p2 {} <= b {}", self.p2, b); + // debug_assert!(!is_large || b < self.b, "b {} < p2 {}", b, self.b); + // debug_assert!(is_large || b < self.p2, "b {} < p2 {}", b, self.p2); + } +} + +/// The optimal bucket function of PHOBIC, with a variable `eps`. +#[derive(Clone, Copy, Debug)] +#[cfg_attr(feature = "epserde", derive(epserde::prelude::Epserde))] +#[cfg_attr(feature = "epserde", epserde(deep_copy))] +pub struct Optimal { + pub eps: f64, +} + +impl BucketFn for Optimal { + fn call(&self, x: u64) -> u64 { + let p32 = (1u64 << 32) as f64; + let p64 = p32 * p32; + let p64inv = 1. / p64; + let x = (x as f64) * p64inv; + let y = x + (1. - self.eps) * (1. - x) * (1. - x).ln(); + + (y * p64) as u64 + } +} + +/// `x*x` +#[derive(Clone, Copy, Debug, Default)] +#[cfg_attr(feature = "epserde", derive(epserde::prelude::Epserde))] +#[cfg_attr(feature = "epserde", epserde(deep_copy))] +pub struct Square; + +impl BucketFn for Square { + fn call(&self, x: u64) -> u64 { + mul_high(x, x) + } +} + +/// `x*x * 255/256 + x/256` +#[derive(Clone, Copy, Debug, Default)] +#[cfg_attr(feature = "epserde", derive(epserde::prelude::Epserde))] +#[cfg_attr(feature = "epserde", epserde(deep_copy))] +pub struct SquareEps; + +impl BucketFn for SquareEps { + fn call(&self, x: u64) -> u64 { + mul_high(x, x) / 256 * 255 + x / 256 + } +} + +/// `x * x * (1 + x)/2` +#[derive(Clone, Copy, Debug, Default)] +#[cfg_attr(feature = "epserde", derive(epserde::prelude::Epserde))] +#[cfg_attr(feature = "epserde", epserde(deep_copy))] +pub struct Cubic; + +impl BucketFn for Cubic { + fn call(&self, x: u64) -> u64 { + // x * x * (1 + x)/2 + mul_high(mul_high(x, x), (x >> 1) | (1 << 63)) + } +} + +/// `x * x * (1 + x)/2 * 255/256 + x/256` +#[derive(Clone, Copy, Debug, Default)] +#[cfg_attr(feature = "epserde", derive(epserde::prelude::Epserde))] +#[cfg_attr(feature = "epserde", epserde(deep_copy))] +pub struct CubicEps; + +impl BucketFn for CubicEps { + fn call(&self, x: u64) -> u64 { + // x * x * (1 + x)/2 + mul_high(mul_high(x, x), (x >> 1) | (1 << 63)) / 256 * 255 + x / 256 + } +} + +#[cfg(test)] +mod test { + use crate::bucket_fn::BucketFn; + + #[test] + fn test_skewed() { + use super::Skewed; + let mut skewed = Skewed::new(0.6, 0.3); + skewed.set_buckets_per_part(1000000000); + + let mut last_y = 0; + let n = 100; + for i in 0..100 { + let x = u64::MAX / n * i; + let y = skewed.call(x); + assert!(y >= last_y); + last_y = y; + } + } +} diff --git a/tokenizers/vendor/ptr_hash/src/bucket_idx.rs b/tokenizers/vendor/ptr_hash/src/bucket_idx.rs new file mode 100644 index 000000000..b9245fb22 --- /dev/null +++ b/tokenizers/vendor/ptr_hash/src/bucket_idx.rs @@ -0,0 +1,62 @@ +use std::ops::{Add, Index, IndexMut, Sub}; + +// Ord so we can sort them. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +pub struct BucketIdx(pub u32); + +impl Add for BucketIdx { + type Output = Self; + + fn add(self, rhs: usize) -> Self::Output { + Self(self.0 + rhs as u32) + } +} + +impl Sub for BucketIdx { + type Output = Self; + + fn sub(self, rhs: usize) -> Self::Output { + Self(self.0 - rhs as u32) + } +} + +impl BucketIdx { + pub const NONE: BucketIdx = BucketIdx(u32::MAX); + pub fn range(num_buckets: usize) -> impl Iterator + Clone { + (0..num_buckets as u32).map(Self) + } + pub fn is_some(&self) -> bool { + self.0 != u32::MAX + } + pub fn is_none(&self) -> bool { + self.0 == u32::MAX + } +} + +impl Index for [T] { + type Output = T; + + fn index(&self, index: BucketIdx) -> &Self::Output { + unsafe { self.get_unchecked(index.0 as usize) } + } +} + +impl IndexMut for [T] { + fn index_mut(&mut self, index: BucketIdx) -> &mut Self::Output { + unsafe { self.get_unchecked_mut(index.0 as usize) } + } +} + +impl Index for Vec { + type Output = T; + + fn index(&self, index: BucketIdx) -> &Self::Output { + unsafe { self.get_unchecked(index.0 as usize) } + } +} + +impl IndexMut for Vec { + fn index_mut(&mut self, index: BucketIdx) -> &mut Self::Output { + unsafe { self.get_unchecked_mut(index.0 as usize) } + } +} diff --git a/tokenizers/vendor/ptr_hash/src/build.rs b/tokenizers/vendor/ptr_hash/src/build.rs new file mode 100644 index 000000000..fbd826edf --- /dev/null +++ b/tokenizers/vendor/ptr_hash/src/build.rs @@ -0,0 +1,435 @@ +use super::*; +use crate::{bucket_idx::BucketIdx, stats::BucketStats}; +use bitvec::{slice::BitSlice, vec::BitVec}; +use log::warn; +use rayon::prelude::*; +use std::{ + collections::BinaryHeap, + iter::zip, + sync::{ + atomic::{AtomicUsize, Ordering}, + Mutex, + }, +}; + +impl, const SINGLE_PART: bool, const REMAP: bool> + PtrHash, SINGLE_PART, REMAP> +{ + pub(super) fn build_shard( + &self, + shard: usize, + hashes: &[Hx::H], + part_starts: &[u32], + pilots: &mut [u8], + taken: &mut [BitVec], + ) -> Option { + let pilots_per_part = pilots.par_chunks_exact_mut(self.buckets); + + let iter = pilots_per_part.zip(taken).enumerate(); + + // let total_evictions = AtomicUsize::new(0); + let parts_done = AtomicUsize::new(shard * self.parts_per_shard); + let stats = Mutex::new(BucketStats::new()); + + let ok = iter.try_for_each(|(part_in_shard, (pilots, taken))| { + let part = shard * self.parts_per_shard + part_in_shard; + let hashes = &hashes + [part_starts[part_in_shard] as usize..part_starts[part_in_shard + 1] as usize]; + let _cnt = self.build_part(part, hashes, pilots, taken, &stats)?; + let _parts_done = parts_done.fetch_add(1, Ordering::Relaxed); + // total_evictions.fetch_add(cnt, Ordering::Relaxed); + + // if self.params.print_stats { + // eprint!( + // "parts done: {parts_done:>6}/{:>6} ({:>4.1}%)\r", + // self.parts, + // 100. * parts_done as f32 / self.parts as f32 + // ); + // } + + Some(()) + }); + + if ok.is_none() { + return None; + } + + assert_eq!( + parts_done.load(Ordering::Relaxed), + (shard + 1) * self.parts_per_shard + ); + + // let total_evictions: usize = total_evictions.load(Ordering::Relaxed); + // let sum_pilots = pilots.iter().map(|&k| k as Pilot).sum::(); + + // // Clear the last \r line. + // if self.params.print_stats { + // eprint!("\x1b[K"); + // eprintln!( + // " displ./bkt: {:>14.3}", + // total_evictions as f32 / (self.buckets * self.parts_per_shard) as f32 + // ); + // eprintln!( + // " avg pilot: {:>14.3}", + // sum_pilots as f32 / (self.buckets * self.parts_per_shard) as f32 + // ); + // } + + // if self.params.print_stats { + // stats.lock().unwrap().print(); + // } + + Some(stats.into_inner().unwrap()) + } + + fn build_part( + &self, + part: usize, + hashes: &[Hx::H], + pilots: &mut [u8], + taken: &mut BitSlice, + _stats: &Mutex, + ) -> Option { + let (starts, bucket_order) = self.sort_buckets(part, hashes); + + let kmax = 256; + + let mut slots = vec![BucketIdx::NONE; self.slots]; + let bucket_len = |b: BucketIdx| (starts[b + 1] - starts[b]) as usize; + + let max_bucket_len = bucket_len(bucket_order[0]); + + // First process larger buckets. + // TODO: Use bucket queue instead? + // NOTE: I tried 'rattle-kicking' where we prefer evicting buckets with a small pilot, + // but in practice this ends up slower, even though it saves ~15% of evictions. + let mut stack = BinaryHeap::new(); + + let slots_for_bucket = |b: BucketIdx, p: Pilot| unsafe { + let hp = self.hash_pilot(p); + hashes + .get_unchecked(starts[b] as usize..starts[b + 1] as usize) + .iter() + .map(move |&hx| self.slot_in_part_hp(hx, hp)) + }; + let mut duplicate_slots = { + let mut slots_tmp = vec![0; max_bucket_len]; + move |b: BucketIdx, p: Pilot| { + slots_tmp.clear(); + slots_tmp.extend(slots_for_bucket(b, p)); + slots_tmp.sort_unstable(); + slots_tmp.iter().tuple_windows().any(|(a, b)| a == b) + } + }; + + let mut recent = [BucketIdx::NONE; 16]; + let mut total_evictions = 0; + + let mut rng = fastrand::Rng::new(); + + // let mut eviction_counts: Vec = vec![]; + + for (i, &new_b) in bucket_order.iter().enumerate() { + let new_bucket = &hashes[starts[new_b] as usize..starts[new_b + 1] as usize]; + if new_bucket.is_empty() { + pilots[new_b] = 0; + continue; + } + let new_b_len = new_bucket.len(); + + let mut evictions = 0usize; + + stack.push((new_b_len, new_b)); + recent.fill(BucketIdx::NONE); + let mut recent_idx = 0; + recent[0] = new_b; + + 'b: while let Some((_b_len, b)) = stack.pop() { + if evictions > self.slots && evictions.is_power_of_two() { + // log = true; + let num_taken_slots = taken.count_ones(); + // if self.params.print_stats { + // eprintln!( + // "part {part:>6} alpha {:>5.2}% bucket size {} ({}/{}, {:>5.2}%) slots filled {}/{} ({:>5.2}%) chain: {evictions:>9}", + // 100. * hashes.len() as f32 / slots.len() as f32, + // new_b_len, + // i, self.buckets, + // 100. * i as f32 / self.buckets as f32, + // num_taken_slots, + // taken.len(), + // 100. * num_taken_slots as f32 / taken.len() as f32, + // ); + // } + if evictions >= 10 * self.slots { + warn!( + "\ +Too many evictions. Aborting! +When the current bucket has size >=2, try decreasing lambda to use fewer elements per buckets. +When the current bucket has size 1 (or maybe 2), try decreasing alpha to have more empty slots for the last few buckets. + +Current part: {part:>6} with load factor alpha={:>5.2}% +Current bucket: size {} ({}/{}, {:>5.2}%) +Slots filled so far: {}/{} ({:>5.2}%) +Eviction chain length: {evictions:>9} +", + 100. * hashes.len() as f32 / slots.len() as f32, + new_b_len, + i, self.buckets, + 100. * i as f32 / self.buckets as f32, + num_taken_slots, + taken.len(), + 100. * num_taken_slots as f32 / taken.len() as f32, + ); + return None; + } + } + + // 1a) Check for a solution without collisions. + + let bucket = + unsafe { hashes.get_unchecked(starts[b] as usize..starts[b + 1] as usize) }; + let b_slots = + |hp: PilotHash| bucket.iter().map(move |&hx| self.slot_in_part_hp(hx, hp)); + + // 1b) Hot-path for when there are no collisions, which is most of the buckets. + if let Some((p, hp)) = self.find_pilot(kmax, bucket, taken) { + // HOT: Many branch misses here. + pilots[b] = p as u8; + for p in b_slots(hp) { + unsafe { + // Taken is already filled by find_pilot. + // HOT: This is a hot instruction; takes as much time as finding the pilot. + *slots.get_unchecked_mut(p) = b; + } + } + continue 'b; + } + + // 2) Search for a pilot with minimal number of collisions. + + // Start at a random pilot to prevent eviction cycles. + let p0 = rng.u8(..) as u64; + // (worst colliding bucket size, p) + let mut best = (usize::MAX, u64::MAX); + + 'p: for delta in 0u64..kmax { + // HOT: This code is slow and full of branch-misses. + // But also, it's only 20% of build_part() time, since the + // hot-path above covers most. + let p = (p0 + delta) % kmax; + let hp = self.hash_pilot(p); + let mut collision_score = 0; + for p in b_slots(hp) { + let s = unsafe { *slots.get_unchecked(p) }; + // HOT: many branches + let new_score = if s.is_none() { + continue; + } else if recent.contains(&s) { + continue 'p; + } else { + // HOT: cache misses. + bucket_len(s).pow(2) + }; + collision_score += new_score; + if collision_score >= best.0 { + continue 'p; + } + } + + // This check takes 2% of time even though it almost + // always passes. Can we delay it to filling of the + // slots table, and backtrack if needed. + if !duplicate_slots(b, p) { + best = (collision_score, p); + // Since we already checked for a collision-free solution, + // the next best is a single collision of size b_len. + if collision_score == new_b_len * new_b_len { + break; + } + } + } + + if best == (usize::MAX, u64::MAX) { + let slots = b_slots(0); + let len = bucket.len(); + let num_slots = self.slots; + eprintln!( + "part {part}: bucket of size {len} with {num_slots} slots: Indistinguishable hashes in bucket!" + ); + for (hx, slot) in zip(bucket, slots) { + eprintln!("{:x?} -> slot {slot}", hx); + } + eprintln!( + "part {part}: bucket of size {len} with {num_slots} slots: Indistinguishable hashes in bucket!" + ); + return None; + } + + let (_collision_score, p) = best; + // if self.params.print_stats { + // eprintln!( + // "{evictions:>7} | pilots[{:>7}] = {:>3} len: {} stack: {} score: {:>3}", + // b.0, + // p, + // bucket_len(b), + // stack.len(), + // _collision_score + // ); + // } + pilots[b] = p as u8; + let hp = self.hash_pilot(p); + + // Drop the collisions and set the new pilot. + for slot in b_slots(hp) { + // THIS IS A HOT INSTRUCTION. + let b2 = slots[slot]; + if b2.is_some() { + assert!(b2 != b); + // DROP BUCKET b + // if self.params.print_stats { + // eprintln!( + // "{evictions:>7} | Push {:>7} len: {}", + // b2.0, + // bucket_len(b2) + // ); + // } + stack.push((bucket_len(b2), b2)); + evictions += 1; + for p2 in slots_for_bucket(b2, pilots[b2] as Pilot) { + unsafe { + *slots.get_unchecked_mut(p2) = BucketIdx::NONE; + taken.set_unchecked(p2, false); + } + } + } + unsafe { + *slots.get_unchecked_mut(slot) = b; + taken.set_unchecked(slot, true); + } + } + + recent_idx += 1; + recent_idx %= recent.len(); + recent[recent_idx] = b; + } + total_evictions += evictions; + // if self.params.print_stats { + // eviction_counts.push(evictions); + // } + } + + // if self.params.print_stats { + // let mut stats = stats.lock().unwrap(); + // for (i, &b) in bucket_order.iter().enumerate() { + // stats.add( + // i, + // bucket_order.len(), + // bucket_len(b), + // pilots[b] as Pilot, + // *eviction_counts.get(i).unwrap_or(&0), + // ); + // } + // } + + Some(total_evictions) + } + + fn find_pilot( + &self, + kmax: u64, + bucket: &[Hx::H], + taken: &mut BitSlice, + ) -> Option<(Pilot, PilotHash)> { + // This gives ~10% speedup. + match bucket.len() { + 1 => self.find_pilot_array::<1>(kmax, bucket.try_into().unwrap(), taken), + 2 => self.find_pilot_array::<2>(kmax, bucket.try_into().unwrap(), taken), + 3 => self.find_pilot_array::<3>(kmax, bucket.try_into().unwrap(), taken), + 4 => self.find_pilot_array::<4>(kmax, bucket.try_into().unwrap(), taken), + 5 => self.find_pilot_array::<5>(kmax, bucket.try_into().unwrap(), taken), + 6 => self.find_pilot_array::<6>(kmax, bucket.try_into().unwrap(), taken), + 7 => self.find_pilot_array::<7>(kmax, bucket.try_into().unwrap(), taken), + 8 => self.find_pilot_array::<8>(kmax, bucket.try_into().unwrap(), taken), + _ => self.find_pilot_slice(kmax, bucket, taken), + } + } + fn find_pilot_array( + &self, + kmax: u64, + bucket: &[Hx::H; L], + taken: &mut BitSlice, + ) -> Option<(Pilot, PilotHash)> { + self.find_pilot_slice(kmax, bucket, taken) + } + + // Note: Prefetching on `taken` is not needed because we use parts that fit in L1 cache anyway. + // + // Note: Tried looping over multiple pilots in parallel, but the additional + // lookups this does aren't worth it. + #[inline(always)] + fn find_pilot_slice( + &self, + kmax: u64, + bucket: &[Hx::H], + taken: &mut BitSlice, + ) -> Option<(Pilot, PilotHash)> { + let r = bucket.len() / 4 * 4; + 'p: for p in 0u64..kmax { + let hp = self.hash_pilot(p); + // True when the slot for hx is already taken. + let check = |hx| unsafe { *taken.get_unchecked(self.slot_in_part_hp(hx, hp)) }; + + // Process chunks of 4 bucket elements at a time. + // This reduces branch-misses (of all of build_part) 3-fold, giving 20% speedup. + for i in (0..r).step_by(4) { + // Check all 4 elements of the chunk without early break. + // NOTE: It's hard to SIMD vectorize the `slot` computation + // here because it uses 64x64->128bit multiplies. + let checks: [bool; 4] = unsafe { + [ + check(*bucket.get_unchecked(i)), + check(*bucket.get_unchecked(i + 1)), + check(*bucket.get_unchecked(i + 2)), + check(*bucket.get_unchecked(i + 3)), + ] + }; + if checks.iter().any(|&bad| bad) { + continue 'p; + } + } + // Check remaining elements. + let mut bad = false; + for &hx in &bucket[r..] { + bad |= check(hx); + } + if bad { + continue 'p; + } + + if self.try_take_pilot(bucket, hp, taken) { + return Some((p, hp)); + } + } + None + } + + /// Fill `taken` with the slots for `hp`, but backtrack as soon as a + /// collision within the bucket is found. + /// + /// Returns true on success. + fn try_take_pilot(&self, bucket: &[Hx::H], hp: PilotHash, taken: &mut BitSlice) -> bool { + // This bucket does not collide with previous buckets, but it may still collide with itself. + for (i, &hx) in bucket.iter().enumerate() { + let slot = self.slot_in_part_hp(hx, hp); + if unsafe { *taken.get_unchecked(slot) } { + // Collision within the bucket. Clean already set entries. + for &hx in unsafe { bucket.get_unchecked(..i) } { + unsafe { taken.set_unchecked(self.slot_in_part_hp(hx, hp), false) }; + } + return false; + } + unsafe { taken.set_unchecked(slot, true) }; + } + true + } +} diff --git a/tokenizers/vendor/ptr_hash/src/fastmod.rs b/tokenizers/vendor/ptr_hash/src/fastmod.rs new file mode 100644 index 000000000..4480e0cb5 --- /dev/null +++ b/tokenizers/vendor/ptr_hash/src/fastmod.rs @@ -0,0 +1,111 @@ + +use crate::reduce::Reduce; + +// Multiply a u128 by u64 and return the upper 64 bits of the result. +// ((lowbits * d as u128) >> 128) as u64 +#[allow(unused)] +fn mul128_u64(lowbits: u128, d: u64) -> u64 { + let bot_half = ((lowbits & u64::MAX as u128) * d as u128) >> 64; // Won't overflow + let top_half = (lowbits >> 64) * d as u128; + let both_halves = bot_half + top_half; // Both halves are already shifted down by 64 + (both_halves >> 64) as u64 +} + +/// FastMod64 +/// Taken from https://github.com/lemire/fastmod/blob/master/include/fastmod.h +#[derive(Copy, Clone, Debug)] +#[cfg_attr(feature = "epserde", derive(epserde::prelude::Epserde))] +#[cfg_attr(feature = "epserde", epserde(deep_copy))] +#[allow(unused)] +pub struct FM64 { + d: u64, + m: u128, +} +impl Reduce for FM64 { + fn new(d: usize) -> Self { + Self { + d: d as u64, + m: (u128::MAX / d as u128).wrapping_add(1), + } + } + fn reduce(self, h: u64) -> usize { + let lowbits = self.m.wrapping_mul(h as u128); + mul128_u64(lowbits, self.d) as usize + } +} + +/// FastMod32, using the low 32 bits of the hash. +/// Taken from https://github.com/lemire/fastmod/blob/master/include/fastmod.h +#[derive(Copy, Clone, Debug)] +#[cfg_attr(feature = "epserde", derive(epserde::prelude::Epserde))] +#[cfg_attr(feature = "epserde", epserde(deep_copy))] +pub struct FM32 { + d: u64, + m: u64, +} +impl Reduce for FM32 { + fn new(d: usize) -> Self { + // assert!(d <= u32::MAX as usize); + Self { + d: d as u64, + m: (u64::MAX / d as u64).wrapping_add(1), + } + } + // NOTE: h must be < 2^32 for correct modulo results!. + // Otherwise, a 'random' integer up to d is returned. + fn reduce(self, h: u64) -> usize { + let lowbits = self.m.wrapping_mul(h); + ((lowbits as u128 * self.d as u128) >> 64) as usize + } +} + +#[test] +fn test_fastmod64() { + #[cfg(debug_assertions)] + let cnt = 100; + #[cfg(not(debug_assertions))] + let cnt = 10000; + + for _ in 1..cnt { + let d = rand::random::() as usize + 1; + let fm = FM64::new(d); + for _ in 0..cnt { + let x = rand::random::() as u64; + assert_eq!(fm.reduce(x), x as usize % d, "failure for d = {d}, x = {x}",); + } + } +} + +#[test] +fn test_fastmod32_equals_modulo() { + #[cfg(debug_assertions)] + let cnt = 100; + #[cfg(not(debug_assertions))] + let cnt = 10000; + + for _ in 1..cnt { + let d = rand::random::() as usize + 1; + let fm = FM32::new(d); + for _ in 0..cnt { + let x = rand::random::() as u64; + assert_eq!(fm.reduce(x), x as usize % d, "failure for d = {d}, x = {x}",); + } + } +} + +#[test] +fn test_fastmod32_doesnt_overflow() { + #[cfg(debug_assertions)] + let cnt = 100; + #[cfg(not(debug_assertions))] + let cnt = 10000; + + for _ in 1..cnt { + let d = rand::random::().saturating_add(1) as usize; + let fm = FM32::new(d); + for _ in 0..cnt { + let x = rand::random::(); + assert!(fm.reduce(x) < d, "failure for d = {d}, x = {x}",); + } + } +} diff --git a/tokenizers/vendor/ptr_hash/src/hash.rs b/tokenizers/vendor/ptr_hash/src/hash.rs new file mode 100644 index 000000000..7da890055 --- /dev/null +++ b/tokenizers/vendor/ptr_hash/src/hash.rs @@ -0,0 +1,212 @@ +//! Implementations of various hashers to use with PtrHash. +//! +//! ## Integer keys +//! +//! We provide: +//! - [`NoHash`]: does nothing -- only use on truly random keys. +//! - [`FastIntHash`] = [`FxHash`], which does a single wrapping multiplication and should be good enough most of the time. +//! - [`StrongerIntHash`]: Use this when the input keys are very regular, eg `0..1000`. (But then why do you need an MPHF anyway?) +//! Does a `u128` multiplication, and xors the high and low word together, like xxh3. Then does one more multiplication. Not very scientific but 'it works'. +//! - [`GxInt`]: GxHash, but with the type 'inlined' so that it optimized better. +//! - [`Xxh3Int`]: Xxh3, but with the type 'inlined' so that it optimized better. +//! +//! In practice, prefer [`FastIntHash`] if it's good enough. +//! Otherwise, fall back to [`StrongerIntHash`]. +//! If that still fails (which probably shouldn't happen) fall back to one of the two remaining +//! +//! If [`StrongerIntHash`] lacks sufficient randomness, use [`Xxh3Int`] instead, which is considerably slower but much stronger. +//! +//! ## String keys +//! +//! For string keys, use [`StringHash`] for 64-bit hashes and [`StringHash128`] for 128-bit hashes. +//! These are aliases for 64bit and 128bit versions of gxhash, respectively. +//! +//! Another option is to use [`FxHash`] instead. +//! +//! In general any type implementing `Hasher` can be used, but it may be more +//! efficient to implement [`KeyHasher`] yourself for your key type, to directly +//! call specialized functions rather than going through the generic `Hasher` +//! interface. +//! +#[cfg(feature = "gxhash")] +use gxhash::GxBuildHasher; + +use crate::KeyT; +use std::fmt::Debug; + +/// The [`KeyHasher`] trait returns a 64 or 128-bit `Hash`. From this, two `u64` values are extracted. +/// +/// When 64-bit hashes are enough, we simply return the same hash (the `u64` +/// `Self` value) as the low and high part. +/// +/// When 128-bit hashes are needed, the two functions return the low/high half of bits. +/// +/// Our method never needs the full hash value, and instead uses the two hashes +/// in different places to extract sufficient entropy. +pub trait Hash: Copy + Debug + Default + Send + Sync + Eq + rdst::RadixKey { + /// Returns the low 64bits of the hash. + fn low(&self) -> u64; + /// Returns the high 64bits of the hash. + fn high(&self) -> u64; +} + +impl Hash for u64 { + fn low(&self) -> u64 { + *self + } + fn high(&self) -> u64 { + *self + } +} + +impl Hash for u128 { + fn low(&self) -> u64 { + *self as u64 + } + fn high(&self) -> u64 { + (*self >> 64) as u64 + } +} + +/// Wrapper trait for various hash functions. +pub trait KeyHasher: Clone + Sync { + type H: Hash; + fn hash(x: &Key, seed: u64) -> Self::H; +} + +/// All external hashers work. +impl KeyHasher for H { + type H = u64; + #[inline(always)] + fn hash(x: &Key, seed: u64) -> u64 { + let mut hasher = H::default(); + Key::hash(x, &mut hasher); + hasher.finish() ^ seed + } +} + +// Aliases + +/// A slightly faster but weaker hash for sufficiently random integers. Uses [`fxhash::FxHasher64`]. +pub type FastIntHash = fxhash::FxHasher64; +pub type FxHash = fxhash::FxHasher64; +/// Type alias for xxhash (XXH3) hasher. +/// +/// Prefer [`Xxh3Int`] for integers, which avoids some overhead of the default hasher. +pub type Xxh3 = xxhash_rust::xxh3::Xxh3Default; +#[cfg(feature = "gxhash")] +pub type Gx = gxhash::GxHasher; + +/// Use gxhash for 64-bit string hashing. +#[cfg(feature = "gxhash")] +pub type StringHash = Gx; +/// Use gxhash for 128-bit string hashing. +#[cfg(feature = "gxhash")] +pub type StringHash128 = Gx128; + +// Implementations + +/// 128-bit version of XXH3. +#[cfg_attr(feature = "epserde", derive(epserde::prelude::Epserde))] +#[derive(Clone)] +pub struct Xxh3_128; +impl KeyHasher for Xxh3_128 { + type H = u128; + #[inline(always)] + fn hash(x: &Key, seed: u64) -> u128 { + let mut hasher = xxhash_rust::xxh3::Xxh3Default::default(); + x.hash(&mut hasher); + hasher.digest128() ^ (seed as u128 | (seed as u128) << 64) + } +} + +/// 128-bit version of XXH3. +#[cfg(feature = "gxhash")] +#[cfg_attr(feature = "epserde", derive(epserde::prelude::Epserde))] +#[derive(Clone)] +pub struct Gx128; +#[cfg(feature = "gxhash")] +impl KeyHasher for Gx128 { + type H = u128; + #[inline(always)] + fn hash(x: &Key, seed: u64) -> u128 { + use std::hash::BuildHasher; + let mut hasher = GxBuildHasher::with_seed(seed as i64).build_hasher(); + x.hash(&mut hasher); + hasher.finish_u128() + } +} + +/// A sufficiently good hash for non-random integers. Inspired by Xxh3, with one extra multiplication: +/// FIXME: IS THAT NEEDED? +/// +/// ```ignore +/// let (hi, lo) = (value ^ seed) as u128 * C as u128; +/// return (hi ^ lo).wrapping_mul(C); +/// ``` +#[cfg_attr(feature = "epserde", derive(epserde::prelude::Epserde))] +#[derive(Clone)] +pub struct StrongerIntHash; + +/// Mixing constant. +pub const C: u64 = 0x517cc1b727220a95; + +/// No hash at all; just `value ^ seed`. Use with caution. Mostly for benchmarking. +#[cfg_attr(feature = "epserde", derive(epserde::prelude::Epserde))] +#[derive(Clone)] +pub struct NoHash; + +/// Inlined version of Xxh3 for integer keys. +#[cfg_attr(feature = "epserde", derive(epserde::prelude::Epserde))] +#[derive(Clone)] +pub struct Xxh3Int; + +/// Inlined version of Xxh3 for integer keys. +#[cfg(feature = "gxhash")] +#[cfg_attr(feature = "epserde", derive(epserde::prelude::Epserde))] +#[derive(Clone)] +pub struct GxInt; + +// Macro to implement hashes for all integer types. +macro_rules! int_hashers { + ($($t:ty),*) => { + $( + impl KeyHasher<$t> for NoHash { + type H = u64; + #[inline(always)] + fn hash(x: &$t, seed: u64) -> u64 { + *x as u64 ^ seed + } + } + + impl KeyHasher<$t> for StrongerIntHash { + type H = u64; + #[inline(always)] + fn hash(x: &$t, seed: u64) -> u64 { + let r = (*x as u64 ^ seed) as u128 * C as u128; + let low = r as u64; + let high = (r >> 64) as u64; + (low ^ high).wrapping_mul(C) + } + } + + impl KeyHasher<$t> for Xxh3Int { + type H = u64; + #[inline(always)] + fn hash(x: &$t, seed: u64) -> u64 { + xxhash_rust::xxh3::xxh3_64_with_seed(&(*x as u64).to_le_bytes(), seed) + } + } + + #[cfg(feature = "gxhash")] + impl KeyHasher<$t> for GxInt { + type H = u64; + #[inline(always)] + fn hash(x: &$t, seed: u64) -> u64 { + gxhash::gxhash64(&(*x as u64).to_le_bytes(), seed as i64) + } + } + )* + }; +} +int_hashers!(u8, u16, u32, u64, usize, i8, i16, i32, i64, isize); diff --git a/tokenizers/vendor/ptr_hash/src/lib.rs b/tokenizers/vendor/ptr_hash/src/lib.rs new file mode 100644 index 000000000..d7bd9d7df --- /dev/null +++ b/tokenizers/vendor/ptr_hash/src/lib.rs @@ -0,0 +1,1141 @@ +#![cfg_attr(feature = "unstable", feature(iter_array_chunks))] +//! # PtrHash: Minimal Perfect Hashing at RAM Throughput +//! +//! PtrHash builds a _minimal perfect hash function_, that is, +//! a hash function that maps a fixed set of keys to `{0, ..., n-1}`. +//! +//! PtrHash was developed for large key sets of at least 1 million keys, and has been tested up to `10^11` keys. +//! In the default configuration, it uses 3.0 bits per key. +//! Nevertheless, it can also be used for arbitrary small sets. +//! +//! See the GitHub [readme](https://github.com/ragnargrootkoerkamp/ptrhash) +//! or paper ([arXiv](https://arxiv.org/abs/2502.15539), [blog version](https://curiouscoding.nl/posts/ptrhash/)) +//! for details on the algorithm and performance. +//! +//! Usage example: +//! ```rust +//! use ptr_hash::PtrHashParams; +//! +//! // Enable logging. +//! env_logger::init(); +//! +//! // Generate some random keys. +//! let n = 1_000_000; +//! let keys = ptr_hash::util::generate_keys(n); +//! +//! // Build the default variant of the datastructure. +//! // See `FastPtrHash` and `CompactPtrHash` below as alternatives. +//! let mphf = ::new(&keys, PtrHashParams::default()); +//! +//! // Get the index of a key. +//! let key = 0; +//! let idx = mphf.index(&key); +//! assert!(idx < n); +//! +//! // An iterator over the indices of the keys. +//! // 32: number of iterations ahead to prefetch. +//! // _: placeholder to infer the type of keys being iterated. +//! let indices = mphf.index_stream::<32, _>(&keys); +//! assert_eq!(indices.sum::(), (n * (n - 1)) / 2); +//! +//! // Query a batch of keys. +//! let query_keys = [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15]; +//! let mut indices = mphf.index_batch::<16, _>(query_keys); +//! indices.sort(); +//! for i in 0..indices.len()-1 { +//! assert!(indices[i] != indices[i+1]); +//! } +//! +//! // Test that all items map to different indices +//! let mut taken = vec![false; n]; +//! for key in &keys { +//! let idx = mphf.index(&key); +//! assert!(!taken[idx]); +//! taken[idx] = true; +//! } +//! +//! // `FastPtrHash` skips remapping and returns values up to `PtrHash::max_index()`, +//! // which is around 1.01*n. +//! let phf = ::new(&keys, PtrHashParams::default()); +//! +//! for key in &keys { +//! let idx = mphf.index(&key); +//! // NOTE: Not `n` but `phf.max_index()` here! +//! assert!(idx < phf.max_index()); +//! } +//! +//! // `CompactPtrHash` uses multi-threaded construction and is more space-efficient, +//! // but slightly slower to query. +//! let phf = ::new(&keys, PtrHashParams::default_compact()); +//! +//! for key in &keys { +//! let idx = mphf.index(&key); +//! assert!(idx < n); +//! } +//! ``` +//! +//! ## Default configurations +//! +//! - For fastest query throughput, use [`FastPtrHash`] (2.67 bits/key), which is a non-minimal PHF that skips remapping values into `[0, n)`.67 bits/key. +//! - If you want a minimal PHF that returns values in `[0, n)`, use [`DefaultPtrHash`] (3.0 bits/key). +//! - If you have many keys, use [`CompactPtrHash`] (2.15 bits/key), which allows for multi-threaded construction +//! by splitting the input into multiple parts and uses a more space-efficient bucket function. +//! Queries will be a bit slower because the part and index inside the part are computed separately. +//! `PtrHashParams::default_balanced()` is smaller and still fast to construct, while +//! `PtrHashParams::default_compact()` is even smaller and around 2x slower to construct. +//! +//! ## Hash functions +//! +//! PtrHash benefits from using an as-fast-as-possible hash function. +//! +//! - If your keys are already random integers, use [`hash::NoHash`]. +//! - For integers, use [`hash::FastIntHash`], which aliases the fast-but-weak [`hash::FxHash`]. Otherwise, try [`hash::StrongerIntHash`]. +//! - For strings, use [`hash::StringHash`] when the number of keys is at most `10^9`, and use [`hash::StringHash128`] for more keys. These alias [`hash::Gx`] and [`hash::Gx128`]. +//! +//! See the [`hash`] module documentation for better hashes in case these cause hash collisions. +//! +//! ``` +//! // Hashing strings +//! # #[cfg(feature = "gxhash")] { +//! use ptr_hash::{DefaultPtrHash, PtrHashParams, hash::StringHash}; +//! +//! let keys = vec!["abc", "def"]; +//! let mphf = >::new(&keys, PtrHashParams::default()); +//! +//! let idx = mphf.index(&"def"); +//! # } +//! ``` +//! +//! ## Partitioning +//! +//! By default, PtrHash builds all keys as a single part. +//! +//! Faster multi-threaded construction is possible using `SINGLE_PART=false` (via [`CompactPtrHash`]), +//! which splits the keys over multiple parts. +//! Additionally, having fewer keys per part improves the cache-locality of the construction. +//! Query time is slightly slower though, since computing the part and index +//! inside the part are two separate steps. +//! +//! ## Sharding +//! +//! When the keys and/or their hashes do not all fit in memory at once, use sharding. +//! This requires `SINGLE_PART=false`, e.g. via [`CompactPtrHash`]. +//! See [`shard::Sharding`] for details of different sharding methods. +//! ``` +//! use ptr_hash::{CompactPtrHash, PtrHashParams, Sharding}; +//! +//! let mut params = PtrHashParams::default_compact(); +//! // The default value. For ~16GB of u64 hashes or ~32GB of u128 hashes. +//! // Make sure to also leave space for the data structure itself. +//! params.keys_per_shard = 1<<31; +//! params.sharding = Sharding::Disk; +//! +//! let keys = vec![1,2,3]; // 10^12 or who knows how many keys. +//! let mphf = ::new(&keys, params); +//! ``` +//! +//! ## Reducing space usage +//! +//! The default parameters are chosen for reliability, construction speed, and query speed, and give around 3 bits per keys. +//! To achieve smaller sizes, consider using [`cacheline_ef::CachelineEfVec`] or [`pack::EliasFano`] as 'remap' structure, instead of `Vec`. +//! +//! Additionally, one can use [`CompactPtrHash`] with [`PtrHashParams::default_balanced()`] parameters, which use the `CubicEps` bucket function instead of `Linear`, and increase `lambda` from the default of `3.0` to `3.5`. +//! [`PtrHashParams::default_compact()`] is even smaller, but slower to construct, and generally slightly less reliable. +//! +//! ``` +//! # #[cfg(feature = "elias-fano")] { +//! use ptr_hash::{PtrHash, PtrHashParams}; +//! +//! let params = PtrHashParams::default_balanced(); +//! let keys = vec![1u64, 2, 3]; +//! let mphf = >::new(&keys, params); +//! # } +//! ``` + +/// Customizable Hasher trait. +pub mod hash; +/// Extendable backing storage trait and types. +pub mod pack; +/// Some internal logging and testing utilities. +pub mod util; + +pub mod bucket_fn; +mod bucket_idx; +mod build; +mod fastmod; +mod reduce; +mod shard; +mod sort_buckets; +#[doc(hidden)] +pub mod stats; +#[cfg(test)] +mod test; + +use bitvec::{bitvec, vec::BitVec}; +use bucket_fn::BucketFn; +use bucket_fn::CubicEps; +use bucket_fn::Linear; +use bucket_fn::SquareEps; +use itertools::izip; +use itertools::Itertools; +use log::debug; +use log::trace; +use pack::MutPacked; +use rand::{RngExt, SeedableRng}; +use rand_chacha::ChaCha8Rng; +use rayon::prelude::*; +pub use shard::Sharding; +use stats::BucketStats; +use std::array::from_fn; +use std::{borrow::Borrow, default::Default, marker::PhantomData, time::Instant}; + +use crate::{hash::*, pack::Packed, reduce::*, util::log_duration}; + +/// Parameters for PtrHash construction. +/// +/// While all fields are public, prefer one of the default functions, +/// [`PtrHashParams::default()`], [`PtrHashParams::default_fast()`], or +/// [`PtrHashParams::default_compact()`]. +#[derive(Clone, Copy, Debug)] +#[cfg_attr(feature = "epserde", derive(epserde::prelude::Epserde))] +#[cfg_attr(feature = "epserde", epserde(deep_copy))] +pub struct PtrHashParams { + /// Use `n/alpha` slots approximately. + pub alpha: f64, + /// Use average bucket size lambda. + pub lambda: f64, + /// Bucket function + pub bucket_fn: BF, + /// Upper bound on number of keys per shard. + /// Default is 2^31, or 16GB of u64 hashes per shard. + pub keys_per_shard: usize, + /// When true, write each shard to a file instead of iterating multiple + /// times. + pub sharding: Sharding, +} + +impl PtrHashParams { + /// Parameters for fast construction and queries. Use these by default. + /// + /// Takes `3.0` bits/key, and can be up to 2x faster to query than the balanced or compact versions. + /// - `alpha=0.99` + /// - `lambda=3.0` + /// - `bucket_fn=Linear` + /// + pub fn default_fast() -> Self { + Self { + alpha: 0.99, + lambda: 3.0, + bucket_fn: Linear, + keys_per_shard: 1 << 31, + sharding: Sharding::None, + } + } +} + +#[doc(hidden)] +impl PtrHashParams { + pub fn default_square() -> Self { + Self { + alpha: 0.99, + lambda: 3.5, + bucket_fn: SquareEps, + keys_per_shard: 1 << 31, + sharding: Sharding::None, + } + } +} + +impl PtrHashParams { + /// Balanced parameters, which saves some space for larger inputs. + /// This is the 'Default' from the paper. + /// + /// Takes `2.4` bits/key, and trades off space and speed. + /// - `alpha=0.99` + /// - `lambda=3.5` + /// - `bucket_fn=CubicEps` + pub fn default_balanced() -> Self { + Self { + alpha: 0.99, + lambda: 3.5, + bucket_fn: CubicEps, + keys_per_shard: 1 << 31, + sharding: Sharding::None, + } + } + + /// Default 'compact' parameters. + /// + /// Takes `2.15` bits/key, but is typically 2x slower to construct than the default version. + /// If construction fails, try again with decreased `lambda`. + /// - `alpha=0.99` + /// - `lambda=3.9` + /// - `bucket_fn=CubicEps` + pub fn default_compact() -> Self { + Self { + alpha: 0.99, + lambda: 3.9, + bucket_fn: CubicEps, + keys_per_shard: 1 << 31, + sharding: Sharding::None, + } + } +} + +/// By default, use [`PtrHashParams::default_fast()`]. +impl Default for PtrHashParams { + fn default() -> Self { + Self::default_fast() + } +} + +/// Fastest variant: a **non-minimal** PHF that **may return values slightly larger than n**. 2.67 bits/key. +/// +/// This skips remapping values into `[0, n)`, and instead returns values up to [`PtrHash::max_index()`]. +/// +/// Use as [`FastPtrHash::::new(&keys, PtrHashParams::default_fast())`]. +pub type FastPtrHash = + PtrHash, Hx, Vec, true, false>; + +/// Default variant: a minimal PHF that returns values in `[0, n)`. 3.0 bits/key. +/// +/// Use as [`>::new(&keys, PtrHashParams::default())`]. +pub type DefaultPtrHash = + PtrHash, Hx, Vec, true, true>; + +/// Variant for large data sets with multi-threaded construction. 2.15 bits/key. +/// +/// This version optimizes space usage rathen than query speed: +/// - Uses multiple parts to allow multi-threaded construction. +/// - Uses the `CubicEps` bucket function to reduce space usage. +/// - Remaps keys into `[0, n)` to achieve a minimal PHF. +/// +/// This uses 2.15 bits/key with EliasFano, and 2.35 bits/key otherwise. +/// +/// Use as [`>::new(&keys, PtrHashParams::default_compact())`]. +#[cfg(feature = "elias-fano")] +pub type CompactPtrHash = + PtrHash, false, true>; +#[cfg(not(feature = "elias-fano"))] +pub type CompactPtrHash = + PtrHash, Hx, Vec, false, true>; + +/// Trait that keys must satisfy. +pub trait KeyT: Send + Sync + std::hash::Hash {} +impl KeyT for T {} + +// Some fixed algorithmic decisions. +type Rp = FastReduce; +type Rb = FastReduce; +// NOTE: This is not a faithful modulo for 64-bit values, but either way only returns values up to `d`. +type RemSlots = fastmod::FM32; +type Pilot = u64; +type PilotHash = u64; + +/// PtrHash datastructure. +/// It is recommended to use PtrHash with default type aliases: +/// - [`FastPtrHash`]: *non-minimal* PHF that skips remapping values into `[0, n)`. +/// - [`DefaultPtrHash`]: MPHF that returns values in `[0, n)`. +/// - [`CompactPtrHash`]: MPHF that allows for multi-threaded construction. +/// +/// - `Key`: The type of keys to hash. +/// - `BF`: The bucket function to use. Inferred from `PtrHashParams` when calling `PtrHash::new()`. +/// - `F`: The packing to use for remapping free slots, default `CachelineEf`. +/// - `Hx`: The hasher to use for keys, default `FxHash` for integers, but consider +/// `hash::StringHash` (using `gxhash`) for strings, or `hash::StringHash128` when the number of string keys is very +/// large. +/// - `V`: The pilots type. Usually `Vec`, or `&[u8]` for Epserde. +/// - `SINGLE_PART`: using a single part gives faster queries, but prevents multi-threaded construction (default: true). +/// - `REMAP`: remapping results in a minimal PHF, but is slightly slower (default: true). +#[cfg_attr(feature = "epserde", derive(epserde::prelude::Epserde))] +#[derive(Clone)] +pub struct PtrHash< + Key: KeyT + ?Sized = u64, + BF: BucketFn = bucket_fn::Linear, + F: Packed = Vec, + Hx: KeyHasher = hash::FastIntHash, + V: AsRef<[u8]> = Vec, + const SINGLE_PART: bool = true, + const REMAP: bool = true, +> { + params: PtrHashParams, + + /// The number of keys. + n: usize, + /// The total number of parts. + parts: usize, + /// The number of shards. + shards: usize, + /// The maximal number of parts per shard. + /// The last shard may have fewer parts. + parts_per_shard: usize, + /// The total number of slots. + slots_total: usize, + /// The total number of buckets. + buckets_total: usize, + /// The number of slots per part, always a power of 2. + slots: usize, + /// The number of buckets per part. + buckets: usize, + + // Precomputed fast modulo operations. + /// Fast %shards. + rem_shards: Rp, + /// Fast %parts. + rem_parts: Rp, + /// Fast %b. + rem_buckets: Rb, + /// Fast %b_total. + rem_buckets_total: Rb, + + /// Fast %s when there is only a single part. + rem_slots: RemSlots, + + // Computed state. + /// The global seed. + seed: u64, + /// The pilots. + pilots: V, + /// Remap the out-of-bound slots to free slots. + remap: F, + _key: PhantomData, + _hx: PhantomData, +} + +/// An empty PtrHash instance. Mostly useless, but may be convenient. +impl< + Key: KeyT, + BF: BucketFn, + F: MutPacked, + Hx: KeyHasher, + const SINGLE_PART: bool, + const REMAP: bool, + > Default for PtrHash, SINGLE_PART, REMAP> +where + PtrHashParams: Default, +{ + fn default() -> Self { + PtrHash { + params: as Default>::default(), + + n: 0, + parts: 0, + shards: 0, + parts_per_shard: 0, + slots_total: 0, + buckets_total: 0, + slots: 0, + buckets: 0, + rem_shards: FastReduce::new(0), + rem_parts: FastReduce::new(0), + rem_buckets: FastReduce::new(0), + rem_buckets_total: FastReduce::new(0), + rem_slots: RemSlots::new(0), + seed: 0, + pilots: vec![], + remap: F::default(), + _key: PhantomData, + _hx: PhantomData, + } + } +} + +/// Construction methods taking a list of keys. +impl< + Key: KeyT, + BF: BucketFn, + F: MutPacked, + Hx: KeyHasher, + const SINGLE_PART: bool, + const REMAP: bool, + > PtrHash, SINGLE_PART, REMAP> +{ + /// Create a new PtrHash instance from the given keys. + /// + /// Use `::new()` or `DefaultPtrHash::new()` instead of simply `PtrHash::new()` to + /// get the default values for the generics. + /// + /// NOTE: This panics when construction fails after 10 attempts. + /// This should be rare, but can happen if we are unlucky with the rng seeds. + /// Consider calling [`PtrHash::try_new()`] instead. + /// + /// NOTE: Only up to 2^40 keys are supported. + pub fn new(keys: &[Key], params: PtrHashParams) -> Self { + let mut ptr_hash = Self::init(keys.len(), params); + ptr_hash + .compute_pilots(keys.par_iter()) + .expect("Unable to construct PtrHash after 10 tries. Try using a better hash or decreasing lambda."); + ptr_hash + } + + /// Version that returns build statistics. + #[doc(hidden)] + pub fn new_with_stats(keys: &[Key], params: PtrHashParams) -> (Self, BucketStats) { + let mut ptr_hash = Self::init(keys.len(), params); + let stats = ptr_hash + .compute_pilots(keys.par_iter()) + .expect("Unable to construct PtrHash after 10 tries. Try using a better hash or decreasing lambda."); + (ptr_hash, stats) + } + + /// Fallible version of `new` that returns `None` if construction fails. + /// This can happen when `lambda` is too larger (e.g. for `default_compact` + /// parameters) and the eviction chains become too long. + pub fn try_new(keys: &[Key], params: PtrHashParams) -> Option { + let mut ptr_hash = Self::init(keys.len(), params); + ptr_hash.compute_pilots(keys.par_iter())?; + Some(ptr_hash) + } +} + +/// Construction (helper) methods working with unsized keys. +impl< + Key: KeyT + ?Sized, + BF: BucketFn, + F: MutPacked, + Hx: KeyHasher, + const SINGLE_PART: bool, + const REMAP: bool, + > PtrHash, SINGLE_PART, REMAP> +{ + /// Same as `new` above, but takes a `ParallelIterator` over keys instead of a slice. + /// + /// The iterator must be cloneable, since construction can fail for the + /// first seed (e.g. due to duplicate hashes), in which case a new pass over + /// keys is need. + pub fn new_from_par_iter<'a>( + n: usize, + keys: impl ParallelIterator> + Clone + 'a, + params: PtrHashParams, + ) -> Self { + let mut ptr_hash = Self::init(n, params); + ptr_hash.compute_pilots(keys); + ptr_hash + } + + /// Only initialize the parameters; do not compute the pilots yet. + fn init(n: usize, mut params: PtrHashParams) -> Self { + // assert!(n < (1 << 40), "Number of keys must be less than 2^40."); + + let shards = match (SINGLE_PART, params.sharding) { + (_, Sharding::None) => 1, + (true, _) => panic!("Sharding does not work in combination with SINGLE_PART=true."), + _ => n.div_ceil(params.keys_per_shard), + }; + + // Formula of Vigna, eps-cost-sharding: https://arxiv.org/abs/2503.18397 + // (1-alpha)/2, so that on average we still have some room to play with. + let parts = if SINGLE_PART { + 1 + } else { + let eps = (1.0 - params.alpha) / 2.0; + let x = n as f64 * eps * eps / 2.0; + // Half the number of target parts for some safety margin. + // Otherwise, the largest part is often still too large for small inputs. + let target_parts = x / x.ln() / 2.0; + let parts_per_shard = (target_parts.floor() as usize) / shards; + parts_per_shard.max(1) * shards + }; + + let keys_per_part = n / parts; + let parts_per_shard = parts / shards; + let mut slots_per_part = (keys_per_part as f64 / params.alpha) as usize; + // Avoid powers of two, since then %S does not depend on all bits. + if slots_per_part.is_power_of_two() { + slots_per_part += 1; + } + let slots_total = parts * slots_per_part; + // Add a few extra buckets to avoid collisions for small n. + let buckets_per_part = (keys_per_part as f64 / params.lambda).ceil() as usize + 3; + let buckets_total = parts * buckets_per_part; + + trace!(" keys: {n:>10}"); + trace!(" shards: {shards:>10}"); + trace!(" parts: {parts:>10}"); + trace!(" slots/prt: {slots_per_part:>10}"); + trace!(" slots tot: {slots_total:>10}"); + trace!(" real alpha: {:>10.4}", n as f64 / slots_total as f64); + trace!(" buckets/prt: {buckets_per_part:>10}"); + trace!(" buckets tot: {buckets_total:>10}"); + trace!("keys/ bucket: {:>13.2}", n as f64 / buckets_total as f64); + + params + .bucket_fn + .set_buckets_per_part(buckets_per_part as u64); + + Self { + params, + n, + parts, + shards, + parts_per_shard, + slots_total, + slots: slots_per_part, + buckets_total, + buckets: buckets_per_part, + rem_shards: Rp::new(shards), + rem_parts: Rp::new(parts), + rem_buckets: Rb::new(buckets_per_part), + rem_buckets_total: Rb::new(buckets_total), + rem_slots: RemSlots::new(slots_per_part.max(1)), // fix for n=0 + seed: 0, + pilots: Default::default(), + remap: F::default(), + _key: PhantomData, + _hx: PhantomData, + } + } + + fn compute_pilots<'a>( + &mut self, + keys: impl ParallelIterator> + Clone + 'a, + ) -> Option { + let overall_start = std::time::Instant::now(); + // Initialize arrays; + let mut taken: Vec = vec![]; + let mut pilots: Vec = vec![]; + + let mut tries = 0; + const MAX_TRIES: usize = 10; + + let mut rng = ChaCha8Rng::seed_from_u64(31415); + + // Loop over global seeds `s`. + let stats = 's: loop { + tries += 1; + if tries > MAX_TRIES { + log::error!("PtrHash failed to find a global seed after {MAX_TRIES} tries."); + return None; + } + + let old_seed = self.seed; + + // Choose a global seed s. + self.seed = rng.random(); + if tries == 1 { + log::trace!("First seed tried: {}", self.seed); + } else { + log::warn!("Previous seed {old_seed} failed."); + log::warn!("Trying seed number {tries}: {}.", self.seed); + } + + // Reset output-memory. + pilots.clear(); + pilots.resize(self.buckets_total, 0); + + // TODO: Compress taken on the fly, instead of pre-allocating the entire thing. + for taken in taken.iter_mut() { + taken.clear(); + taken.resize(self.slots, false); + } + taken.resize_with(self.parts, || bitvec![0; self.slots]); + + // Iterate over shards. + let shard_hashes = self.shards(keys.clone()); + // Avoid chunks_mut(0) when n=0. + let shard_pilots = pilots.chunks_mut((self.buckets * self.parts_per_shard).max(1)); + let shard_taken = taken.chunks_mut(self.parts_per_shard); + let mut stats = BucketStats::default(); + // eprintln!("Num shards (keys) {}", shard_keys.()); + for (shard, (hashes, pilots, taken)) in + izip!(shard_hashes, shard_pilots, shard_taken).enumerate() + { + // Determine the buckets. + let start = std::time::Instant::now(); + let Some((hashes, part_starts)) = self.sort_parts(shard, hashes) else { + trace!("Found duplicate hashes"); + // Found duplicate hashes. + continue 's; + }; + let start = log_duration("sort buckets", start); + + // Compute pilots. + if let Some(shard_stats) = + self.build_shard(shard, &hashes, &part_starts, pilots, taken) + { + stats.merge(shard_stats); + log_duration("find pilots", start); + } else { + trace!("Could not find pilots"); + continue 's; + } + } + + let start = std::time::Instant::now(); + let remap = self.remap_free_slots(&taken); + log_duration("remap free", start); + if remap.is_err() { + trace!("Failed to construct CachelineEF"); + continue 's; + } + + break 's stats; + }; + + // Pack the data. + self.pilots = pilots; + + let (p, r) = self.bits_per_element(); + debug!("bits/element: {}", p + r); + log_duration("total build", overall_start); + Some(stats) + } + + fn remap_free_slots(&mut self, taken: &Vec) -> Result<(), ()> { + assert_eq!( + taken.iter().map(|t| t.count_zeros()).sum::(), + self.slots_total - self.n, + "Not the right number of free slots left!\n total slots {} - n {}", + self.slots_total, + self.n + ); + + if !REMAP || self.slots_total == self.n { + return Ok(()); + } + + // Compute the free spots. + let mut v = Vec::with_capacity(self.slots_total - self.n); + let get = |t: &Vec, idx: usize| t[idx / self.slots][idx % self.slots]; + for i in taken + .iter() + .enumerate() + .flat_map(|(p, t)| { + let offset = p * self.slots; + t.iter_zeros().map(move |i| offset + i) + }) + .take_while(|&i| i < self.n) + { + while !get(&taken, self.n + v.len()) { + v.push(i as u64); + } + v.push(i as u64); + } + self.remap = MutPacked::try_new(v).ok_or(())?; + Ok(()) + } +} + +/// Indexing methods. +impl< + Key: KeyT + ?Sized, + BF: BucketFn, + F: Packed, + Hx: KeyHasher, + V: AsRef<[u8]>, + const SINGLE_PART: bool, + const REMAP: bool, + > PtrHash +{ + /// Return the number of bits per element used for the pilots (`.0`) and the + /// remapping (`.1`). + pub fn bits_per_element(&self) -> (f64, f64) { + let pilots = self.pilots.as_ref().size_in_bytes() as f64 / self.n as f64; + let remap = self.remap.size_in_bytes() as f64 / self.n as f64; + (8. * pilots, 8. * remap) + } + + pub fn n(&self) -> usize { + self.n + } + + /// `self.index()` always returns below this bound. + /// Around `n/alpha ~ 1.01*n` when `REMAP=false`, or `n` when `REMAP=true`. + pub fn max_index(&self) -> usize { + if REMAP { + self.n + } else { + self.slots_total + } + } + + pub fn slots_per_part(&self) -> usize { + self.slots + } + + /// Get the index for `key`. + /// + /// When `REMAP` is true, returns an index in `[0, n)`. + /// When `REMAP` is false, returns a non-minimal index in `[0, self.max_index())`. + /// + /// When `SINGLE_PART` is true, this uses slightly faster single-part hashing. + #[inline(always)] + pub fn index(&self, key: &Key) -> usize { + let slot = self.index_no_remap(key); + + if REMAP { + if slot < self.n { + slot + } else { + self.remap.index(slot - self.n) as usize + } + } else { + slot + } + } + + /// Get the non-remapped index for `key`, regardless of `REMAP`. + /// + /// When `SINGLE_PART` is true, uses slightly faster single-part hashing. + #[doc(hidden)] + #[inline(always)] + pub fn index_no_remap(&self, key: &Key) -> usize { + if SINGLE_PART { + debug_assert_eq!(self.parts, 1); + } + + let hx = self.hash_key(key); + let slot = if SINGLE_PART { + let b = self.bucket_in_part(hx.high()); + let pilot = self.pilots.as_ref().index(b); + self.slot_in_part(hx, pilot) + } else { + let b = self.bucket(hx); + let pilot = self.pilots.as_ref().index(b); + self.slot(hx, pilot) + }; + + slot + } + + /// Takes an iterator over keys and returns an iterator over the indices of the keys. + /// + /// Uses a buffer of size `B` for prefetching ahead. `B=32` should be a good choice. + /// When `REMAP` is true (the default), returns minimal indices in `[0, n)`. + /// When `REMAP` is false, returns non-minimal indices in `[0, n/alpha)`. + /// The iterator can return either `Q=Key` or `Q=&Key`. + /// + /// See the module-level documentation for an example. + // NOTE: It would be cool to use SIMD to determine buckets/positions in + // parallel, but this is complicated, since SIMD doesn't support the + // 64x64->128 multiplications needed in bucket/slot computations. + #[inline] + pub fn index_stream<'a, const B: usize, Q: Borrow + 'a>( + &'a self, + keys: impl IntoIterator + 'a, + ) -> impl Iterator + 'a { + self.index_stream_maybe_remap::(keys) + } + + /// Version that allows disabling remapping. + #[doc(hidden)] + #[inline] + pub fn index_stream_maybe_remap< + 'a, + const B: usize, + const QUERY_REMAP: bool, + Q: Borrow + 'a, + >( + &'a self, + keys: impl IntoIterator + 'a, + ) -> impl Iterator + 'a { + let mut keys = keys.into_iter(); + + // Ring buffers to cache the hash and bucket of upcoming queries. + let mut next_hashes: [Hx::H; B] = [Hx::H::default(); B]; + let mut next_buckets: [usize; B] = [0; B]; + + // Initialize and prefetch first B values. + let mut leftover = B; + for idx in 0..B { + let hx = keys + .next() + .map(|k| { + leftover -= 1; + self.hash_key(k.borrow()) + }) + .unwrap_or_default(); + next_hashes[idx] = hx; + + next_buckets[idx] = self.bucket(next_hashes[idx]); + prefetch_index::prefetch_index(self.pilots.as_ref(), next_buckets[idx]); + } + + // Manual iterator implementation so we avoid the overhead and + // non-inlining of Chain, and instead have a manual fold. + struct It< + 'a, + const B: usize, + Key: KeyT + ?Sized, + Q: Borrow + 'a, + KeyIt: Iterator + 'a, + BF: BucketFn, + F: Packed, + Hx: KeyHasher, + V: AsRef<[u8]>, + const SINGLE_PART: bool, + const REMAP: bool, + const QUERY_REMAP: bool, + > { + ph: &'a PtrHash, + keys: KeyIt, + next_hashes: [Hx::H; B], + next_buckets: [usize; B], + leftover: usize, + } + + impl< + 'a, + const B: usize, + Key: KeyT + ?Sized, + Q: Borrow + 'a, + KeyIt: Iterator + 'a, + BF: BucketFn, + F: Packed, + Hx: KeyHasher, + V: AsRef<[u8]>, + const SINGLE_PART: bool, + const REMAP: bool, + const QUERY_REMAP: bool, + > Iterator for It<'a, B, Key, Q, KeyIt, BF, F, Hx, V, SINGLE_PART, REMAP, QUERY_REMAP> + { + type Item = usize; + fn next(&mut self) -> Option { + unimplemented!("Use a method that calls `fold()` instead."); + } + + #[inline(always)] + fn fold(mut self, init: BB, mut f: FF) -> BB + where + Self: Sized, + FF: FnMut(BB, Self::Item) -> BB, + { + let mut accum = init; + let mut i = 0; + + for key in self.keys { + let next_hash = self.ph.hash_key(key.borrow()); + let idx = i % B; + let cur_hash = self.next_hashes[idx]; + let cur_bucket = self.next_buckets[idx]; + self.next_hashes[idx] = next_hash; + self.next_buckets[idx] = self.ph.bucket(self.next_hashes[idx]); + prefetch_index::prefetch_index(self.ph.pilots.as_ref(), self.next_buckets[idx]); + let pilot = self.ph.pilots.as_ref().index(cur_bucket); + let slot = self.ph.slot(cur_hash, pilot); + + let slot = if QUERY_REMAP && slot >= self.ph.n { + self.ph.remap.index(slot - self.ph.n) as usize + } else { + slot + }; + + accum = f(accum, slot); + i += 1; + } + + for _ in 0..B - self.leftover { + let idx = i % B; + let cur_hash = self.next_hashes[idx]; + let cur_bucket = self.next_buckets[idx]; + let pilot = self.ph.pilots.as_ref().index(cur_bucket); + let slot = self.ph.slot(cur_hash, pilot); + + let slot = if REMAP && slot >= self.ph.n { + self.ph.remap.index(slot - self.ph.n) as usize + } else { + slot + }; + + accum = f(accum, slot); + i += 1; + } + + accum + } + } + It:: { + ph: self, + keys, + next_hashes, + next_buckets, + leftover, + } + } + + /// Query a batch of `K` keys at once. + /// + /// Input can be either `[Key; K]` or `[&Key; K]`. + #[inline] + pub fn index_batch<'a, const K: usize, Q: Borrow + 'a>( + &'a self, + xs: [Q; K], + ) -> [usize; K] { + let hashes = xs.map(|x| self.hash_key(x.borrow())); + let mut buckets: [usize; K] = [0; K]; + + // Prefetch. + for idx in 0..K { + buckets[idx] = self.bucket(hashes[idx]); + prefetch_index::prefetch_index(self.pilots.as_ref(), buckets[idx]); + } + // Query. + from_fn( + #[inline(always)] + move |idx| { + let pilot = self.pilots.as_ref().index(buckets[idx]); + let slot = self.slot(hashes[idx], pilot); + if REMAP && slot >= self.n { + self.remap.index(slot - self.n) as usize + } else { + slot + } + }, + ) + } + + /// Takes an iterator over keys and returns an iterator over the indices of the keys. + /// + /// Queries in batches of size K. + /// + /// NOTE: Does not process the remainder + #[doc(hidden)] + #[cfg(feature = "unstable")] + #[inline] + pub fn index_batch_exact<'a, const K: usize>( + &'a self, + xs: impl IntoIterator + 'a, + ) -> impl Iterator + 'a { + let mut buckets: [usize; K] = [0; K]; + + // Work on chunks of size K. + let mut f = { + #[inline(always)] + move |hx: [Hx::H; K]| { + // Prefetch. + for idx in 0..K { + buckets[idx] = self.bucket(hx[idx]); + crate::util::prefetch_index(self.pilots.as_ref(), buckets[idx]); + } + // Query. + (0..K).map( + #[inline(always)] + move |idx| { + let pilot = self.pilots.as_ref().index(buckets[idx]); + let slot = self.slot(hx[idx], pilot); + if REMAP && slot >= self.n { + self.remap.index(slot - self.n) as usize + } else { + slot + } + }, + ) + } + }; + let array_chunks = xs.into_iter().map(|x| self.hash_key(x)).array_chunks::(); + array_chunks.into_iter().flat_map( + #[inline(always)] + move |chunk| f(chunk), + ) + // .chain(f(&array_chunks + // .into_remainder() + // .unwrap_or_default() + // .into_iter())) + } + + /// A variant of index_batch_exact that scales better with K. + /// Somehow the version above has pretty constant speed regardless of K. + #[doc(hidden)] + #[inline] + pub fn index_batch_exact2<'a, const K: usize, const QUERY_REMAP: bool>( + &'a self, + xs: impl IntoIterator + 'a, + ) -> impl Iterator + 'a { + let mut buckets: [usize; K] = [0; K]; + let mut hs: [Hx::H; K] = [Hx::H::default(); K]; + + let mut xs = xs + .into_iter() + .map(|x| self.hash_key(x)) + .chain([Default::default(); K]); + for i in 0..K { + hs[i] = xs.next().unwrap(); + } + let mut idx = K; + xs.map(move |hx| { + if idx == K { + idx = 0; + // Prefetch. + for idx in 0..K { + buckets[idx] = self.bucket(hs[idx]); + prefetch_index::prefetch_index(self.pilots.as_ref(), buckets[idx]); + } + } + + // Query. + let pilot = self.pilots.as_ref().index(buckets[idx]); + let slot = self.slot(hs[idx], pilot); + + // Update hash in current pos and increment. + hs[idx] = hx; + idx += 1; + + // Remap? + if QUERY_REMAP && slot >= self.n { + self.remap.index(slot - self.n) as usize + } else { + slot + } + }) + } + + fn hash_key(&self, x: &Key) -> Hx::H { + Hx::hash(x, self.seed) + } + + fn hash_pilot(&self, p: Pilot) -> PilotHash { + hash::C.wrapping_mul(p ^ self.seed) + } + + fn shard(&self, hx: Hx::H) -> usize { + self.rem_shards.reduce(hx.high()) + } + + fn part(&self, hx: Hx::H) -> usize { + self.rem_parts.reduce(hx.high()) + } + + /// Map `hx_remainder` to a bucket in the range [0, self.b). + /// Hashes =self.p1 are mapped to small [self.p2, self.b). + /// + /// (Unless SPLIT_BUCKETS is false, in which case all hashes are mapped to [0, self.b).) + fn bucket_in_part(&self, x: u64) -> usize { + if BF::LINEAR { + self.rem_buckets.reduce(x) + } else if BF::B_OUTPUT { + self.params.bucket_fn.call(x) as usize + } else { + self.rem_buckets.reduce(self.params.bucket_fn.call(x)) + } + } + + /// See bucket.rs for additional implementations. + /// Returns the offset in the slots array for the current part and the bucket index. + fn bucket(&self, hx: Hx::H) -> usize { + if BF::LINEAR { + return self.rem_buckets_total.reduce(hx.high()); + } + + // Extract the high bits for part selection; do normal bucket + // computation within the part using the remaining bits. + // NOTE: This is somewhat slow, but doing better is hard. + let (part, hx) = self.rem_parts.reduce_with_remainder(hx.high()); + let bucket = self.bucket_in_part(hx); + part * self.buckets + bucket + } + + /// Slot uses the 64 low bits of the hash. + fn slot(&self, hx: Hx::H, pilot: u64) -> usize { + (self.part(hx) * self.slots) + self.slot_in_part(hx, pilot) + } + + fn slot_in_part(&self, hx: Hx::H, pilot: Pilot) -> usize { + self.slot_in_part_hp(hx, self.hash_pilot(pilot)) + } + + /// Slot uses the 64 low bits of the hash. + fn slot_in_part_hp(&self, hx: Hx::H, hp: PilotHash) -> usize { + self.rem_slots.reduce(hx.low() ^ hp) + } +} diff --git a/tokenizers/vendor/ptr_hash/src/pack.rs b/tokenizers/vendor/ptr_hash/src/pack.rs new file mode 100644 index 000000000..857aaa4df --- /dev/null +++ b/tokenizers/vendor/ptr_hash/src/pack.rs @@ -0,0 +1,172 @@ +//! The `Packed` and `MutPacked` traits are used for the underlying storage of +//! the remap vector. +//! +//! This is implemented for `Vec`, `CachelineEfVec`, and `EliasFano` from `sucds`. +//! `Packed` is also implemented for respective non-owning (slice) types to support epserde. + +#[cfg(feature = "elias-fano")] +use sucds::mii_sequences::EliasFanoBuilder; + +#[cfg(feature = "cacheline-ef")] +use cacheline_ef::{CachelineEf, CachelineEfVec}; + +/// A trait for backing storage types. +pub trait Packed: Sync { + /// This uses get_unchecked internally, so you must ensure that index is within bounds. + fn index(&self, index: usize) -> u64; + /// Prefetch the element at the given index. + fn prefetch(&self, _index: usize) {} + /// Size in bytes. + fn size_in_bytes(&self) -> usize; +} + +/// An extension of Packed that can be used during construction. +pub trait MutPacked: Packed + Sized { + fn default() -> Self; + fn try_new(vals: Vec) -> Option; + fn name() -> String; +} + +macro_rules! vec_impl { + ($t:ty) => { + impl MutPacked for Vec<$t> { + fn default() -> Self { + Default::default() + } + fn try_new(vals: Vec) -> Option { + Some( + vals.into_iter() + .map(|x| { + x.try_into() + .expect(&format!("Value {x} is larger than backing type can hold.")) + }) + .collect(), + ) + } + fn name() -> String { + stringify!(Vec<$t>).to_string() + } + } + impl Packed for Vec<$t> { + fn index(&self, index: usize) -> u64 { + unsafe { (*self.get_unchecked(index)) as u64 } + } + fn prefetch(&self, index: usize) { + prefetch_index::prefetch_index(self, index); + } + fn size_in_bytes(&self) -> usize { + std::mem::size_of_val(self.as_slice()) + } + } + }; +} + +vec_impl!(u8); +vec_impl!(u16); +vec_impl!(u32); +vec_impl!(u64); + +macro_rules! slice_impl { + ($t:ty) => { + impl Packed for [$t] { + fn index(&self, index: usize) -> u64 { + unsafe { (*self.get_unchecked(index)) as u64 } + } + fn prefetch(&self, index: usize) { + prefetch_index::prefetch_index(self, index); + } + fn size_in_bytes(&self) -> usize { + std::mem::size_of_val(self) + } + } + impl Packed for &[$t] { + fn index(&self, index: usize) -> u64 { + unsafe { (*self.get_unchecked(index)) as u64 } + } + fn prefetch(&self, index: usize) { + prefetch_index::prefetch_index(self, index); + } + fn size_in_bytes(&self) -> usize { + std::mem::size_of_val(*self) + } + } + impl Packed for Box<[$t]> { + fn index(&self, index: usize) -> u64 { + unsafe { (*self.get_unchecked(index)) as u64 } + } + fn prefetch(&self, index: usize) { + prefetch_index::prefetch_index(self, index); + } + fn size_in_bytes(&self) -> usize { + std::mem::size_of_val(&*self) + } + } + }; +} + +slice_impl!(u8); +slice_impl!(u16); +slice_impl!(u32); +slice_impl!(u64); + +#[cfg(feature = "cacheline-ef")] +impl MutPacked for CachelineEfVec> { + fn default() -> Self { + Default::default() + } + fn try_new(vals: Vec) -> Option { + Self::try_new(&vals) + } + fn name() -> String { + "CacheLineEF".to_string() + } +} + +#[cfg(feature = "cacheline-ef")] +impl + Sync> Packed for CachelineEfVec { + fn index(&self, index: usize) -> u64 { + unsafe { self.index_unchecked(index) } + } + fn prefetch(&self, index: usize) { + self.prefetch(index) + } + fn size_in_bytes(&self) -> usize { + self.size_in_bytes() + } +} + +/// Wrapper around the Sucds implementation. +#[cfg(feature = "elias-fano")] +pub struct EliasFano(sucds::mii_sequences::EliasFano); + +#[cfg(feature = "elias-fano")] +impl MutPacked for EliasFano { + fn default() -> Self { + EliasFano(Default::default()) + } + + fn try_new(vals: Vec) -> Option { + if vals.is_empty() { + Some(Self::default()) + } else { + let mut builder = + EliasFanoBuilder::new(*vals.last().unwrap() as usize + 1, vals.len()).unwrap(); + builder.extend(vals.iter().map(|&x| x as usize)).unwrap(); + Some(EliasFano(builder.build())) + } + } + fn name() -> String { + "EF".to_string() + } +} + +#[cfg(feature = "elias-fano")] +impl Packed for EliasFano { + fn index(&self, index: usize) -> u64 { + self.0.select(index as _).unwrap() as u64 + } + + fn size_in_bytes(&self) -> usize { + sucds::Serializable::size_in_bytes(&self.0) + } +} diff --git a/tokenizers/vendor/ptr_hash/src/reduce.rs b/tokenizers/vendor/ptr_hash/src/reduce.rs new file mode 100644 index 000000000..e7b5c2986 --- /dev/null +++ b/tokenizers/vendor/ptr_hash/src/reduce.rs @@ -0,0 +1,61 @@ + +use crate::{ + hash::{self}, + util::mul_high, +}; + +pub trait Reduce: Copy + Sync + std::fmt::Debug { + /// Reduce into the range [0, d). + fn new(d: usize) -> Self; + /// Reduce a (uniform random 64 bit) number into the range [0, d). + fn reduce(self, h: u64) -> usize; + /// Reduce a (uniform random 64 bit) number into the range [0, d), + /// and also return a remainder that can be used for further reductions. + fn reduce_with_remainder(self, _h: u64) -> (usize, u64) { + unimplemented!(); + } +} + +/// FastReduce64 +/// Taken from https://lemire.me/blog/2016/06/27/a-fast-alternative-to-the-modulo-reduction/ +/// NOTE: This only uses the lg(n) high-order bits of entropy from the hash. +#[derive(Copy, Clone, Debug)] +#[cfg_attr(feature = "epserde", derive(epserde::prelude::Epserde))] +#[cfg_attr(feature = "epserde", epserde(deep_copy))] +pub struct FastReduce { + d: u64, +} +impl Reduce for FastReduce { + fn new(d: usize) -> Self { + Self { d: d as u64 } + } + fn reduce(self, h: u64) -> usize { + mul_high(self.d, h) as usize + } + fn reduce_with_remainder(self, h: u64) -> (usize, u64) { + let r = self.d as u128 * h as u128; + ((r >> 64) as usize, r as u64) + } +} + +/// Multiply-Reduce 64 +/// Multiply by mixing constant C and take the required number of bits. +/// Only works when the modulus is a power of 2. +#[derive(Copy, Clone, Debug)] +#[cfg_attr(feature = "epserde", derive(epserde::prelude::Epserde))] +#[cfg_attr(feature = "epserde", epserde(deep_copy))] +#[allow(unused)] +pub struct MulReduce { + mask: u64, +} +impl Reduce for MulReduce { + fn new(d: usize) -> Self { + assert!(d.is_power_of_two(), "{d} is not a power of 2"); + Self { + mask: (d - 1) as u64, + } + } + fn reduce(self, h: u64) -> usize { + (mul_high(hash::C, h) & self.mask) as usize + } +} diff --git a/tokenizers/vendor/ptr_hash/src/shard.rs b/tokenizers/vendor/ptr_hash/src/shard.rs new file mode 100644 index 000000000..676911d7b --- /dev/null +++ b/tokenizers/vendor/ptr_hash/src/shard.rs @@ -0,0 +1,282 @@ +use std::{ + fs::File, + io::{BufReader, BufWriter, Read, Seek, SeekFrom, Write}, + sync::Mutex, +}; + +#[cfg(feature = "clap")] +use clap::builder::PossibleValue; + +use log::{info, trace}; + +use super::*; + +/// Minimal stand-in for `tempfile::TempDir`, used only by the on-disk sharding path (never +/// entered by the in-memory `Sharding::None` / `default_fast` build). Creates a uniquely-named +/// directory under the system temp dir and removes it on drop. Avoids the `tempfile` dependency. +struct TempDir { + path: std::path::PathBuf, +} +impl TempDir { + fn new() -> std::io::Result { + let mut path = std::env::temp_dir(); + let nanos = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_nanos()) + .unwrap_or(0); + path.push(format!("ptrhash-shard-{}-{}", std::process::id(), nanos)); + std::fs::create_dir_all(&path)?; + Ok(Self { path }) + } + fn path(&self) -> &std::path::Path { + &self.path + } +} +impl Drop for TempDir { + fn drop(&mut self) { + let _ = std::fs::remove_dir_all(&self.path); + } +} + +/// The sharding method to use. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)] +#[cfg_attr(feature = "epserde", derive(epserde::prelude::Epserde))] +#[cfg_attr(feature = "epserde", epserde(deep_copy))] +pub enum Sharding { + /// Process all hashes as a single Vec in memory. + #[default] + None, + /// Repeatedly hash all elements, and each time only process a chunk of 2^31 of them. + Memory, + /// Hash everything once, write shards of up to 2^31 hashes to disk. + Disk, + /// Hybrid that repeatedly fills the given amount (in bytes) of disk space with hashes. + Hybrid(usize), +} + +#[cfg(feature = "clap")] +impl clap::ValueEnum for Sharding { + fn value_variants<'a>() -> &'a [Self] { + // 128 GiB for Hybrid & usize::MAX for 32-bit systems. + #[cfg(target_pointer_width = "64")] + const HYBRID_BYTES: usize = 1 << 37; + #[cfg(not(target_pointer_width = "64"))] + const HYBRID_BYTES: usize = usize::MAX; + &[ + Sharding::None, + Sharding::Memory, + Sharding::Disk, + Sharding::Hybrid(HYBRID_BYTES), + ] + } + fn to_possible_value<'a>(&self) -> Option { + Some(match self { + Sharding::None => PossibleValue::new("none"), + Sharding::Memory => PossibleValue::new("memory"), + Sharding::Disk => PossibleValue::new("disk"), + Sharding::Hybrid(_) => PossibleValue::new("hybrid"), + }) + } +} + +impl< + Key: KeyT + ?Sized, + BF: BucketFn, + F: Packed, + Hx: KeyHasher, + const SINGLE_PART: bool, + const REMAP: bool, + > PtrHash, SINGLE_PART, REMAP> +{ + /// Return an iterator over the Vec of hashes of each shard. + pub(crate) fn shards<'a>( + &'a self, + keys: impl ParallelIterator> + Clone + 'a, + ) -> Box> + 'a> { + match self.params.sharding { + Sharding::None => self.no_sharding(keys.clone()), + Sharding::Memory => self.shard_keys_in_memory(keys.clone()), + Sharding::Disk => self.shard_keys_hybrid(usize::MAX, keys.clone()), + Sharding::Hybrid(mem) => self.shard_keys_hybrid(mem, keys.clone()), + } + } + + /// Collect all hashes to a Vec directly and return it. + fn no_sharding<'a>( + &'a self, + keys: impl ParallelIterator> + Clone + 'a, + ) -> Box> + 'a> { + trace!("No sharding: collecting all {} hashes in memory.", self.n); + let start = std::time::Instant::now(); + let hashes = keys.map(|key| self.hash_key(key.borrow())).collect(); + log_duration("collect hash", start); + Box::new(std::iter::once(hashes)) + } + + /// Loop over the keys once per shard. + /// Return an iterator over shards. + /// For each shard, a filtered copy of the ParallelIterator is returned. + fn shard_keys_in_memory<'a>( + &'a self, + keys: impl ParallelIterator> + Clone + 'a, + ) -> Box> + 'a> { + trace!( + "In-memory sharding: iterate keys once for each of {} shards, each of ~{} keys.", + self.shards, + self.n / self.shards + ); + let it = (0..self.shards).map(move |shard| { + trace!("Shard {shard:>3}/{:3}\r", self.shards); + let start = std::time::Instant::now(); + let hashes: Vec<_> = keys + .clone() + .map(|key| self.hash_key(key.borrow())) + .filter(move |h| self.shard(*h) == shard) + .collect(); + trace!("Shard {shard:>3}/{:3}: {} keys", self.shards, hashes.len()); + log_duration("collect shrd", start); + hashes + }); + Box::new(it) + } + + /// Loop over the keys and write each keys hash to the corresponding shard. + /// Returns an iterator over shards. + /// Files are written to /tmp by default, but this can be changed using the + /// TMPDIR environment variable. + /// + /// This is based on `SigStore` in `sux-rs`, but simplified for the specific use case here. + /// https://github.com/vigna/sux-rs/blob/main/src/utils/sig_store.rs + fn shard_keys_hybrid<'a>( + &'a self, + mem: usize, + keys: impl ParallelIterator> + Clone + 'a, + ) -> Box> + 'a> { + let total_shards = self.shards; + let keys_per_shard = self.n / total_shards; + let shards_on_disk = mem / std::mem::size_of::() / keys_per_shard; + assert!( + shards_on_disk > 0, + "Each shard takes more than the provided memory." + ); + if mem < usize::MAX { + info!("Hybrid sharding: writing hashes to disk for {shards_on_disk} shards at a time, for total {} shards each of ~{} keys.", self.shards, self.n / self.shards); + } else { + info!( + "On-disk sharding: writing hashes to disk for all {} shards at a time, each of ~{} keys.", + self.shards, self.n / self.shards + + ); + } + + let it = (0..self.shards) + .step_by(shards_on_disk) + .flat_map(move |first_shard| { + let temp_dir = TempDir::new().unwrap(); + info!("TMP PATH: {:?}", temp_dir.path()); + + let shard_range = first_shard..(first_shard + shards_on_disk).min(self.shards); + info!("Writing keys for shards {shard_range:?}/{}", self.shards); + + let start = std::time::Instant::now(); + + // Create a file writer and count for each shard. + let writers = shard_range + .clone() + .map(|shard| { + Mutex::new(( + BufWriter::new( + File::options() + .read(true) + .write(true) + .create(true) + .open(temp_dir.path().join(format!("{}.tmp", shard))) + .unwrap(), + ), + 0, + )) + }) + .collect_vec(); + + // Each thread has a local buffer per shard. + let init = || writers.iter().map(ThreadLocalBuf::new).collect_vec(); + // Iterate over keys. + keys.clone() + .map(|key| self.hash_key(key.borrow())) + .for_each_init(init, |bufs, h| { + let shard = self.shard(h); + if shard_range.contains(&shard) { + bufs[shard - shard_range.start].push(h); + } + }); + let start = log_duration("Writing files", start); + + // Flush writers and convert to files. + let files = writers + .into_iter() + .map(|w| { + let (mut w, cnt) = w.into_inner().unwrap(); + w.flush().unwrap(); + let mut file = w.into_inner().unwrap(); + file.seek(SeekFrom::Start(0)).unwrap(); + (file, cnt) + }) + .collect_vec(); + log_duration("Flushing writers", start); + + files + .into_iter() + .zip(shard_range) + .map(move |((f, cnt), _shard)| { + let start = std::time::Instant::now(); + let mut v = vec![Hx::H::default(); cnt]; + let mut reader = BufReader::new(f); + let (pre, data, post) = unsafe { v.align_to_mut::() }; + assert!(pre.is_empty()); + assert!(post.is_empty()); + Read::read_exact(&mut reader, data).unwrap(); + log_duration("Read shard", start); + v + }) + + // Files are cleaned up automatically when tmpdir goes out of scope. + }); + Box::new(it) + } +} + +struct ThreadLocalBuf<'a, H> { + buf: Vec, + file: &'a Mutex<(BufWriter, usize)>, +} + +impl<'a, H> ThreadLocalBuf<'a, H> { + fn new(file: &'a Mutex<(BufWriter, usize)>) -> Self { + Self { + // buffer 1GB of data at a time. + buf: Vec::with_capacity(1 << 28), + file, + } + } + fn push(&mut self, h: H) { + self.buf.push(h); + if self.buf.len() == self.buf.capacity() { + self.flush(); + } + } + fn flush(&mut self) { + let mut file = self.file.lock().unwrap(); + let (pre, bytes, post) = unsafe { self.buf.align_to::() }; + assert!(pre.is_empty()); + assert!(post.is_empty()); + file.0.write_all(bytes).unwrap(); + file.1 += self.buf.len(); + self.buf.clear(); + } +} + +impl<'a, H> Drop for ThreadLocalBuf<'a, H> { + fn drop(&mut self) { + self.flush(); + } +} diff --git a/tokenizers/vendor/ptr_hash/src/sort_buckets.rs b/tokenizers/vendor/ptr_hash/src/sort_buckets.rs new file mode 100644 index 000000000..df03b08b5 --- /dev/null +++ b/tokenizers/vendor/ptr_hash/src/sort_buckets.rs @@ -0,0 +1,176 @@ +use super::*; +use crate::bucket_idx::BucketIdx; +use rdst::RadixSort; +use std::time::Instant; + +impl, const SINGLE_PART: bool, const REMAP: bool> + PtrHash, SINGLE_PART, REMAP> +{ + /// Returns: + /// 1. Hashes + /// 2. Start indices of each bucket. + /// 3. Order of the buckets within each part. + /// + /// This returns None if duplicate hashes are found. + #[must_use] + pub(super) fn sort_parts( + &self, + shard: usize, + mut hashes: Vec, + ) -> Option<(Vec, Vec)> { + // For FastReduce methods, we can just sort by hash directly + // instead of sorting by bucket id: For FR32L, first partition by those + // =self.p1, and then sort each group using the low + // 32 bits. + // NOTE: This does not work for other reduction methods. + + let start = Instant::now(); + // 2. Radix sort hashes. + // HOT: This takes half the time for 128bit hashes. + // TODO: Just append each hash to its part directly, where each part has + // space for exactly its number of slots. + // + // TODO: Write robinhood sort that inserts in the right place directly. + // A) Sort L1 sized ranges. + // B) Splat the front of each range to the next part of the target interval. + hashes.radix_sort_unstable(); + let start = log_duration("┌ radix sort", start); + + // 3. Check duplicates. + let distinct = hashes.par_windows(2).all(|w| w[0] != w[1]); + let start = log_duration("├ check dups", start); + if !distinct { + eprintln!("Hashes are not distinct!"); + return None; + } + + // 4. Find the start of each part using binary search. + if !hashes.is_empty() { + assert!(shard * self.parts_per_shard <= self.part(hashes[0])); + assert!(self.part(*hashes.last().unwrap()) < (shard + 1) * self.parts_per_shard); + } + let mut part_starts = vec![0u32; self.parts_per_shard + 1]; + for part_in_shard in 1..=self.parts_per_shard { + part_starts[part_in_shard] = hashes + .binary_search_by(|h| { + if self.part(*h) < shard * self.parts_per_shard + part_in_shard { + std::cmp::Ordering::Less + } else { + std::cmp::Ordering::Greater + } + }) + .unwrap_err() as u32; + } + + // Check max part len. + let mut max_part_len = 0; + for (start, end) in part_starts.iter().tuple_windows() { + let len = (end - start) as usize; + max_part_len = max_part_len.max(len); + } + let exp = self.n / self.parts; + let stddev = exp.isqrt(); + + // https://math.stackexchange.com/a/89147/91741: + // expected max of N (here #parts) samples of a random variable is + // exp + sigma * sqrt(2 * ln N). + let exp_max = exp + stddev * ((self.parts as f32).ln() * 2.).sqrt() as usize; + trace!("exp key/part: {exp:>10} stddev {stddev:>10}"); + trace!( + "exp max k/pt: {exp_max:>10} {:>10} {:>8.2}", + exp_max - exp, + (exp_max - exp) as f32 / stddev as f32 + ); + trace!( + " max k/pt: {max_part_len:>10} {:>10} {:>8.2}", + max_part_len - exp, + (max_part_len - exp) as f32 / stddev as f32 + ); + trace!( + " slots/pt: {:>10} {:>10} {:>8.2}", + self.slots, + self.slots - exp, + (self.slots - exp) as f32 / stddev as f32 + ); + trace!("exp alpha: {:>13.2}%", 100. * self.params.alpha); + trace!( + "max alpha: {:>13.2}%", + 100. * max_part_len as f32 / self.slots as f32 + ); + + if max_part_len as usize > self.slots { + trace!( + "Shard {shard}: Part has more elements than slots! elements {max_part_len} > {} slots", + self.slots + ); + return None; + } + + log_duration("├part starts", start); + + Some((hashes, part_starts)) + } + + // Sort the buckets in the given part and corresponding range of hashes. + pub(super) fn sort_buckets(&self, part: usize, hashes: &[Hx::H]) -> (Vec, Vec) { + // Where each bucket starts in hashes. + let mut bucket_starts = Vec::with_capacity(self.buckets + 1); + + // The order of buckets, from large to small. + let mut order: Vec = vec![BucketIdx::NONE; self.buckets]; + + // The number of buckets of each length. + let mut bucket_len_cnt = vec![0; 32]; + + let mut end = 0; + bucket_starts.push(end as u32); + + // Loop over buckets in part, setting start positions and counting # buckets of each size. + for b in 0..self.buckets { + let start = end; + // NOTE: Many branch misses here. + while end < hashes.len() && self.bucket(hashes[end]) == part * self.buckets + b { + end += 1; + } + + let l = end - start; + if l >= bucket_len_cnt.len() { + bucket_len_cnt.resize(l + 1, 0); + } + bucket_len_cnt[l] += 1; + bucket_starts.push(end as u32); + } + + assert_eq!(end, hashes.len()); + + let max_bucket_size = bucket_len_cnt.len() - 1; + // This assert is disabled, because it only holds when using uniform buckets. + if false { + let expected_bucket_size = self.slots as f32 / self.buckets as f32; + assert!(max_bucket_size <= (20. * expected_bucket_size) as usize, "Part {part}: Bucket size {max_bucket_size} is too much larger than the expected size of {expected_bucket_size}." ); + } + + // Compute start positions of each range of buckets of equal size. + let mut acc = 0; + for i in (0..=max_bucket_size).rev() { + let tmp = bucket_len_cnt[i]; + bucket_len_cnt[i] = acc; + acc += tmp; + } + + // Write buckets to their right location. + for b in BucketIdx::range(self.buckets) { + let l = (bucket_starts[b + 1] - bucket_starts[b]) as usize; + order[bucket_len_cnt[l]] = b; + bucket_len_cnt[l] += 1; + } + + // for i in 0..max_bucket_size { + // if bucket_len_cnt[i] > 0 { + // eprintln!("Bucket size {:>2}: {:>6}", i, bucket_len_cnt[i]); + // } + // } + + (bucket_starts, order) + } +} diff --git a/tokenizers/vendor/ptr_hash/src/stats.rs b/tokenizers/vendor/ptr_hash/src/stats.rs new file mode 100644 index 000000000..a14caa040 --- /dev/null +++ b/tokenizers/vendor/ptr_hash/src/stats.rs @@ -0,0 +1,156 @@ +use crate::Pilot; + +#[derive(Default, Clone, serde::Serialize, Debug)] +struct Row { + buckets: usize, + elements: usize, + elements_max: usize, + pilot_sum: Pilot, + pilot_max: Pilot, + evictions: usize, + evictions_max: usize, +} + +impl Row { + fn add(&mut self, bucket_len: usize, pilot: Pilot, evictions: usize) { + self.buckets += 1; + self.elements += bucket_len; + self.elements_max = self.elements_max.max(bucket_len); + self.pilot_sum += pilot; + self.pilot_max = self.pilot_max.max(pilot); + self.evictions += evictions; + self.evictions_max = self.evictions_max.max(evictions); + } +} + +#[derive(Default, serde::Serialize, Debug)] +pub struct BucketStats { + by_pct: Vec, + by_bucket_len: Vec, +} + +impl BucketStats { + pub fn new() -> Self { + Self { + by_pct: vec![Row::default(); 100], + by_bucket_len: vec![Row::default(); 100], + } + } + + pub fn merge(&mut self, other: Self) { + self.by_pct.resize(100, Row::default()); + self.by_bucket_len.resize( + self.by_bucket_len.len().max(other.by_bucket_len.len()), + Row::default(), + ); + for (a, b) in self.by_pct.iter_mut().zip(other.by_pct.iter()) { + a.buckets += b.buckets; + a.elements += b.elements; + a.elements_max = a.elements_max.max(b.elements_max); + a.pilot_sum += b.pilot_sum; + a.pilot_max = a.pilot_max.max(b.pilot_max); + a.evictions += b.evictions; + a.evictions_max = a.evictions_max.max(b.evictions_max); + } + for (a, b) in self + .by_bucket_len + .iter_mut() + .zip(other.by_bucket_len.iter()) + { + a.buckets += b.buckets; + a.elements += b.elements; + a.elements_max = a.elements_max.max(b.elements_max); + a.pilot_sum += b.pilot_sum; + a.pilot_max = a.pilot_max.max(b.pilot_max); + a.evictions += b.evictions; + a.evictions_max = a.evictions_max.max(b.evictions_max); + } + } + + pub fn add( + &mut self, + bucket_id: usize, + buckets_total: usize, + bucket_len: usize, + pilot: Pilot, + evictions: usize, + ) { + let pct = bucket_id * 100 / buckets_total; + self.by_pct[pct].add(bucket_len, pilot, evictions); + if self.by_bucket_len.len() <= bucket_len { + self.by_bucket_len.resize(bucket_len + 1, Row::default()); + } + self.by_bucket_len[bucket_len].add(bucket_len, pilot, evictions); + } + + pub fn print(&self) { + eprintln!(); + Self::print_rows(&self.by_pct, false); + // eprintln!(); + // Self::print_rows(&self.by_bucket_len, true); + eprintln!(); + } + + fn print_rows(rows: &[Row], reverse: bool) { + let b_total = rows.iter().map(|r| r.buckets).sum::(); + let n = rows.iter().map(|r| r.elements).sum::(); + + eprintln!( + "{:>4} {:>11} {:>7} {:>6} {:>6} {:>6} {:>10} {:>10} {:>10} {:>10}", + "sz", + "cnt", + "bucket%", + "cuml%", + "elem%", + "cuml%", + "avg p", + "max p", + "avg evict", + "max evict" + ); + let mut bucket_cuml = 0; + let mut elem_cuml = 0; + let process_row = |row: &Row| { + if row.buckets == 0 { + return; + } + bucket_cuml += row.buckets; + elem_cuml += row.elements; + eprintln!( + "{:>4}: {:>11} {:>7.2} {:>6.2} {:>6.2} {:>6.2} {:>10.1} {:>10} {:>10.5} {:>10}", + row.elements_max, + row.buckets, + row.buckets as f32 / b_total as f32 * 100., + bucket_cuml as f32 / b_total as f32 * 100., + row.elements as f32 / n as f32 * 100., + elem_cuml as f32 / n as f32 * 100., + row.pilot_sum as f32 / row.buckets as f32, + row.pilot_max, + row.evictions as f32 / row.buckets as f32, + row.evictions_max + ); + }; + if reverse { + rows.iter().rev().for_each(process_row); + } else { + rows.iter().for_each(process_row); + } + let sum_pilots = rows.iter().map(|r| r.pilot_sum).sum::(); + let max_pilot = rows.iter().map(|r| r.pilot_max).max().unwrap(); + let sum_evictions = rows.iter().map(|r| r.evictions).sum::(); + let max_evictions = rows.iter().map(|r| r.evictions_max).max().unwrap(); + eprintln!( + "{:>4}: {:>11} {:>7.2} {:>6.2} {:>6.2} {:>6.2} {:>10.1} {:>10} {:>10.5} {:>10}", + "", + b_total, + 100., + 100., + 100., + 100., + sum_pilots as f32 / b_total as f32, + max_pilot, + sum_evictions as f32 / b_total as f32, + max_evictions + ); + } +} diff --git a/tokenizers/vendor/ptr_hash/src/test.rs b/tokenizers/vendor/ptr_hash/src/test.rs new file mode 100644 index 000000000..12be173f1 --- /dev/null +++ b/tokenizers/vendor/ptr_hash/src/test.rs @@ -0,0 +1,374 @@ +use std::hint::black_box; + +use super::*; +use crate::util::generate_keys; + +/// Construct the MPHF and test all keys are mapped to unique indices. +#[test] +fn construct_random() { + #[allow(unused_mut)] + let mut ns = vec![ + 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 30, 100, 300, 1000, 3000, 10_000, 30_000, 100_000, + ]; + #[cfg(not(debug_assertions))] + ns.extend_from_slice(&[1_000_000, 3_000_000, 10_000_000, 20_000_000, 30_000_000]); + + for n in (0..10).chain(ns) { + eprintln!("RANDOM Testing n = {}", n); + let keys = generate_keys(n); + let ptr_hash = DefaultPtrHash::::new(&keys, PtrHashParams::default_fast()); + let mut done = bitvec![0; n]; + for key in &keys { + let idx = ptr_hash.index(&key); + assert!(!done[idx]); + done.set(idx, true); + } + let ptr_hash = FastPtrHash::::new(&keys, PtrHashParams::default_fast()); + let mut done = bitvec![0; ptr_hash.max_index()]; + for key in &keys { + let idx = ptr_hash.index(key); + assert!(!done[idx]); + done.set(idx, true); + } + let ptr_hash = + CompactPtrHash::::new(&keys, PtrHashParams::default_compact()); + let mut done = bitvec![0; n]; + for key in &keys { + let idx = ptr_hash.index(key); + assert!(!done[idx]); + done.set(idx, true); + } + } +} + +/// Construct the MPHF and test all keys are mapped to unique indices. +#[test] +#[ignore = "large"] +fn test_1e9() { + env_logger::init(); + let n = 1_000_000_000; + eprintln!("RANDOM Testing n = {}", n); + let keys = generate_keys(n); + + let ptr_hash = DefaultPtrHash::::new(&keys, PtrHashParams::default_fast()); + let mut done = bitvec![0; n]; + for key in &keys { + let idx = ptr_hash.index(key); + assert!(!done[idx]); + done.set(idx, true); + } + + let ptr_hash = FastPtrHash::::new(&keys, PtrHashParams::default_fast()); + let mut done = bitvec![0; ptr_hash.max_index()]; + for key in &keys { + let idx = ptr_hash.index(key); + assert!(!done[idx]); + done.set(idx, true); + } + + let ptr_hash = CompactPtrHash::::new(&keys, PtrHashParams::default_compact()); + let mut done = bitvec![0; n]; + for key in &keys { + let idx = ptr_hash.index(key); + assert!(!done[idx]); + done.set(idx, true); + } +} + +#[test] +#[ignore = "for benchmarking only"] +fn int_hash_speed() { + let n = 10000000; + let keys = (0..n as u64).map(|i| hash::C * i).collect::>(); + let seed = black_box(132); + + let start = std::time::Instant::now(); + for k in &keys { + black_box(FastIntHash::hash(k, seed)); + } + eprintln!("Time {:?}", start.elapsed()); + + let start = std::time::Instant::now(); + for k in &keys { + black_box(StrongerIntHash::hash(k, seed)); + } + eprintln!("Time {:?}", start.elapsed()); + + #[cfg(feature = "gxhash")] + { + let start = std::time::Instant::now(); + for k in &keys { + black_box(GxInt::hash(k, seed)); + } + eprintln!("Time {:?}", start.elapsed()); + } + + let start = std::time::Instant::now(); + for k in &keys { + black_box(Xxh3Int::hash(k, seed)); + } + eprintln!("Time {:?}", start.elapsed()); +} + +/// Keys are multiples of 1, 2^40, and 10^12 +#[test] +fn construct_multiples() { + env_logger::init(); + for m in [1, 1 << 40, 1_000_000_000_000, 3u64.pow(23)] { + #[allow(unused_mut)] + let mut ns = vec![ + 0usize, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 30, 100, 300, 1000, 3000, 10_000, 30_000, + 100_000, + ]; + #[cfg(not(debug_assertions))] + ns.extend_from_slice(&[1_000_000, 3_000_000, 10_000_000, 20_000_000, 30_000_000]); + + for n in ns { + eprintln!("MULTIPLES OF {m} Testing n = {}", n); + let mut keys = (0..n as u64).map(|i| m * i).collect::>(); + keys.sort(); + for i in 0..n.saturating_sub(1) { + if keys[i] >= keys[i + 1] { + keys[i + 1] = keys[i] + 1; + } + } + let ptr_hash = + DefaultPtrHash::::new(&keys, PtrHashParams::default_fast()); + let mut done = bitvec![0; n]; + for key in keys { + let idx = ptr_hash.index(&key); + assert!(!done[idx]); + done.set(idx, true); + } + } + } +} + +#[test] +fn index_stream() { + for n in [2, 10, 100, 1000, 10_000, 100_000, 1_000_000] { + let keys = generate_keys(n); + let ptr_hash = ::new(&keys, Default::default()); + let sum = ptr_hash.index_stream::<32, _>(&keys).sum::(); + assert_eq!(sum, (n * (n - 1)) / 2, "Failure for n = {n}"); + } +} + +#[cfg(feature = "unstable")] +#[test] +fn index_batch() { + for n in [10usize, 100, 1000, 10_000, 100_000, 1_000_000] { + let n = n.next_multiple_of(32); + let keys = generate_keys(n); + let ptr_hash = ::new(&keys, Default::default()); + let sum = ptr_hash.index_batch_exact::<32>(&keys).sum::(); + assert_eq!(sum, (n * (n - 1)) / 2); + } +} + +#[test] +fn new_par_iter() { + let n = 100_000; + let keys = generate_keys(n); + ::new_from_par_iter(n, keys.par_iter(), Default::default()); +} + +#[test] +#[cfg_attr(debug_assertions, ignore = "only run in release mode")] +fn in_memory_sharding() { + let n = 1 << 25; + let range = 0..n as u64; + let keys = range.clone().into_par_iter(); + let ptr_hash = >::new_from_par_iter( + n, + keys.clone(), + PtrHashParams { + keys_per_shard: 1 << 22, + sharding: Sharding::Memory, + ..PtrHashParams::default_compact() + }, + ); + eprintln!("Checking duplicates..."); + let mut done = bitvec![0; n]; + for key in range { + let idx = ptr_hash.index(&key); + assert!(!done[idx]); + done.set(idx, true); + } +} + +#[test] +#[cfg_attr(debug_assertions, ignore = "only run in release mode")] +fn on_disk_sharding() { + let n = 1 << 25; + let range = 0..n as u64; + let keys = range.clone().into_par_iter(); + let ptr_hash = >::new_from_par_iter( + n, + keys.clone(), + PtrHashParams { + keys_per_shard: 1 << 22, + sharding: Sharding::Disk, + ..PtrHashParams::default_compact() + }, + ); + eprintln!("Checking duplicates..."); + let mut done = bitvec![0; n]; + for key in range { + let idx = ptr_hash.index(&key); + assert!(!done[idx]); + done.set(idx, true); + } +} + +/// Test that sharded construction and queries work with more than 2^32 keys. +#[test] +#[ignore = "very slow"] +fn many_keys_memory() { + let n = 1 << 33; + let n_query = 1 << 27; + let range = 0..n as u64; + let keys = range.clone().into_par_iter(); + let ptr_hash = >::new_from_par_iter( + n, + keys.clone(), + PtrHashParams { + keys_per_shard: 1 << 30, + sharding: Sharding::Memory, + ..PtrHashParams::default_compact() + }, + ); + // Since running all queries is super slow, we only check a subset of them. + // Although this doesn't completely check that there are no duplicate + // mappings, by the birthday paradox we can be quite sure there are none + // since we check way more than sqrt(n) of them. + eprintln!("Checking duplicates..."); + let mut done = bitvec![0; n]; + for key in 0..n_query { + let idx = ptr_hash.index(&key); + assert!(!done[idx]); + done.set(idx, true); + } +} + +/// Test that sharded construction and queries work with more than 2^32 keys. +#[test] +#[ignore = "very slow; writes 64GB to disk"] +fn many_keys_disk() { + let n = 1 << 33; + let n_query = 1 << 27; + let range = 0..n as u64; + let keys = range.clone().into_par_iter(); + let ptr_hash = >::new_from_par_iter( + n, + keys.clone(), + PtrHashParams { + keys_per_shard: 1 << 30, + sharding: Sharding::Disk, + ..PtrHashParams::default_compact() + }, + ); + // Since running all queries is super slow, we only check a subset of them. + // Although this doesn't completely check that there are no duplicate + // mappings, by the birthday paradox we can be quite sure there are none + // since we check way more than sqrt(n) of them. + eprintln!("Checking duplicates..."); + let mut done = bitvec![0; n]; + for key in 0..n_query { + let idx = ptr_hash.index(&key); + assert!(!done[idx]); + done.set(idx, true); + } +} + +#[test] +fn ptr_hash_can_clone() { + let ptr_hash = PtrHash::<_>::new(&[0, 1], PtrHashParams::default()); + + // test succeeds if this compiles + let _y = ptr_hash.clone(); +} + +#[test] +fn integer_key_types() { + let h = PtrHash::<_>::new(&[0u8], PtrHashParams::default()); + h.index(&0u8); + let h = PtrHash::<_>::new(&[0u16], PtrHashParams::default()); + h.index(&0u16); + let h = PtrHash::<_>::new(&[0u32], PtrHashParams::default()); + h.index(&0u32); + let h = PtrHash::<_>::new(&[0u64], PtrHashParams::default()); + h.index(&0u64); + let h = PtrHash::<_>::new(&[0usize], PtrHashParams::default()); + h.index(&0usize); + let h = PtrHash::<_>::new(&[0i8], PtrHashParams::default()); + h.index(&0i8); + let h = PtrHash::<_>::new(&[0i16], PtrHashParams::default()); + h.index(&0i16); + let h = PtrHash::<_>::new(&[0i32], PtrHashParams::default()); + h.index(&0i32); + let h = PtrHash::<_>::new(&[0i64], PtrHashParams::default()); + h.index(&0i64); + let h = PtrHash::<_>::new(&[0isize], PtrHashParams::default()); + h.index(&0isize); +} + +#[test] +#[cfg(feature = "gxhash")] +fn string_key_types() { + let h = DefaultPtrHash::::new(&["a"], PtrHashParams::default()); + + // h.index("a"); + h.index(&"a"); + h.index(&"a".to_string().as_str()); + h.index(&Box::new("a")); + + // TODO: The below don't work yet. + // See https://github.com/beling/bsuccinct-rs/issues/9 for some comments. + + // let h = DefaultPtrHash::::new(&["a".to_string()], PtrHashParams::default()); + + // h.index(&&"a"); + // h.index(&"a"); + // h.index("a".to_string()); + // h.index(&"a".to_string()); + // h.index(Box::new("a")); + // h.index(&Box::new("a")); + // h.index(Box::new("a".to_string())); + // h.index(&Box::new("a".to_string())); + + // let h = DefaultPtrHash::::new(&[Box::new("a")], PtrHashParams::default()); + // h.index("a"); + // h.index(&"a"); + // h.index("a".to_string()); + // h.index(&"a".to_string()); + // h.index(Box::new("a")); + // h.index(&Box::new("a")); + // h.index(Box::new("a".to_string())); + // h.index(&Box::new("a".to_string())); + + // let h = DefaultPtrHash::::new( + // &[Box::new("a".to_string())], + // PtrHashParams::default(), + // ); + // h.index("a"); + // h.index(&"a"); + // h.index("a".to_string()); + // h.index(&"a".to_string()); + // h.index(Box::new("a")); + // h.index(&Box::new("a")); + // h.index(Box::new("a".to_string())); + // h.index(&Box::new("a".to_string())); +} + +#[test] +fn single_part() { + let n = 1_000_000; + let keys = util::generate_keys(n); + + let params = PtrHashParams::default(); + + let mphf = , hash::FastIntHash, _, true, true>>::new(&keys, params); + + mphf.index(&0); +} diff --git a/tokenizers/vendor/ptr_hash/src/util.rs b/tokenizers/vendor/ptr_hash/src/util.rs new file mode 100644 index 000000000..53e3742ec --- /dev/null +++ b/tokenizers/vendor/ptr_hash/src/util.rs @@ -0,0 +1,66 @@ +//! Internal utilities that are only exposed for testing/benchmarking purposes. +//! Do not use externally. +use super::*; +use log::{trace, warn}; +use rand::{rng, RngExt}; +use rayon::prelude::*; +use rdst::RadixSort; + +pub(crate) fn mul_high(a: u64, b: u64) -> u64 { + ((a as u128 * b as u128) >> 64) as u64 +} + +thread_local! { + /// TODO: Use trace! instead. + static LOG: std::cell::Cell = std::cell::Cell::new(true); +} + +pub(crate) fn log_duration(name: &str, start: Instant) -> Instant { + if !LOG.with(|log| log.get()) { + return start; + } + trace!( + "{}", + format!("{name:>12}: {:>13.2?}s", start.elapsed().as_secs_f32()) + ); + Instant::now() +} + +pub fn generate_keys(n: usize) -> Vec { + // TODO: Deterministic key generation. + let start = Instant::now(); + let keys = loop { + let start = Instant::now(); + let keys: Vec<_> = (0..n) + .into_par_iter() + .map_init(rng, |rng, _| rng.random()) + .collect(); + let start = log_duration("┌ gen keys", start); + let mut keys2: Vec<_> = keys.par_iter().copied().collect(); + let start = log_duration("├ clone", start); + keys2.radix_sort_unstable(); + let start = log_duration("├ sort", start); + let distinct = keys2.par_windows(2).all(|w| w[0] < w[1]); + log_duration("├ duplicates", start); + if distinct { + break keys; + } + warn!("DUPLICATE KEYS GENERATED"); + }; + log_duration("generatekeys", start); + keys +} + +pub fn generate_string_keys(n: usize) -> Vec> { + let start = Instant::now(); + // let start = Instant::now(); + let keys: Vec<_> = (0..n) + .into_par_iter() + .map_init(rng, |rng, _| { + let len = rng.random_range(10..=50); + (0..len).map(|_| rng.random_range(1..=255)).collect() + }) + .collect(); + log_duration("generatekeys", start); + keys +} From b25b035edc6dd2e6e5958801051fc377920e53b7 Mon Sep 17 00:00:00 2001 From: Arthur Date: Wed, 15 Jul 2026 15:44:09 +0200 Subject: [PATCH 3/8] feat(tk-train): pretokenize for training via the pipeline Add tk_encode pipeline::TrainingPretokenizer (normalize + atomsplit split, with the byte-level alphabet remap the encode-path TryFrom drops). Retarget TokenizerTrainExt from the legacy generic TokenizerImpl (which drove training through the legacy Normalizer/PreTokenizer trait runtime) onto the concrete config wrappers, and feed the trainer from TrainingPretokenizer instead of pre_tokenize_for_training. --- .../tk-encode/src/tokenizer/pipeline.rs | 118 ++++++++++++++++++ tokenizers/tk-train/src/train_ext.rs | 27 ++-- 2 files changed, 135 insertions(+), 10 deletions(-) diff --git a/tokenizers/tk-encode/src/tokenizer/pipeline.rs b/tokenizers/tk-encode/src/tokenizer/pipeline.rs index 9a8ee0abc..46d3a3759 100644 --- a/tokenizers/tk-encode/src/tokenizer/pipeline.rs +++ b/tokenizers/tk-encode/src/tokenizer/pipeline.rs @@ -149,6 +149,124 @@ impl TryFrom for PipelinePreTokenizer { } } +/// Normalize `seq` with the pipeline `Cow` normalizer (identity when there is no normalizer). +fn normalize_for_training<'a>( + normalizer: Option<&NormalizerWrapper>, + seq: &'a str, +) -> Result> { + match normalizer { + Some(n) => Normalizer::normalize(n, seq), + None => Ok(Cow::Borrowed(seq)), + } +} + +/// Remap a word's bytes to the GPT-2 byte-level alphabet (one printable char per byte), the +/// content rewrite the legacy [`crate::utils::byte_level::byte_level_transform`] applied before +/// counting. A byte-level model must train on these symbols, not raw UTF-8. +fn byte_level_word(word: &str) -> String { + word.bytes() + .map(|b| crate::utils::byte_level::BYTES_CHAR_LOOKUP[b as usize]) + .collect() +} + +/// Pipeline-based pre-tokenizer for **training**: built once from a tokenizer's normalizer + +/// pre-tokenizer config, then applied per sequence to yield the normalized words a `Trainer` +/// counts. This is the pipeline replacement for the legacy +/// `TokenizerImpl::pre_tokenize_for_training` — `tk-train` builds one before feeding the trainer. +pub struct TrainingPretokenizer { + normalizer: Option, + engine: TrainingSplit, +} + +enum TrainingSplit { + /// atomsplit-backed pipeline pre-tokenizer over the normalized text. + Pipeline(PipelinePreTokenizer), + /// ByteLevel: split on the GPT-2 regex (or not), then remap each word to the byte-level + /// alphabet — the alphabet remap that `PipelinePreTokenizer`'s encode-path `TryFrom` drops. + ByteLevel { + split: Option, + add_prefix_space: bool, + }, +} + +impl TrainingPretokenizer { + /// Prepare a training pre-tokenizer from a tokenizer's normalizer + pre-tokenizer config. + /// + /// ponytail: Metaspace (and any other pretokenizer the pipeline doesn't cover) surfaces the + /// `PipelinePreTokenizer::try_from` error rather than silently mis-training. Add a Metaspace + /// variant to `PipelinePreTokenizer` when training those is needed. + pub fn new( + normalizer: Option<&NormalizerWrapper>, + pre_tokenizer: Option<&PreTokenizerWrapper>, + ) -> Result { + let engine = match pre_tokenizer { + Some(PreTokenizerWrapper::ByteLevel(bl)) => { + let split = if bl.use_regex { + Some(SplitPretok::new( + SplitPattern::Regex(GPT2_REGEX_STR.to_owned()), + SplitDelimiterBehavior::Isolated, + false, + )?) + } else { + None + }; + TrainingSplit::ByteLevel { + split, + add_prefix_space: bl.add_prefix_space, + } + } + Some(w) => TrainingSplit::Pipeline(w.clone().try_into()?), + None => TrainingSplit::Pipeline(PipelinePreTokenizer::None), + }; + Ok(Self { + normalizer: normalizer.cloned(), + engine, + }) + } + + /// The normalized words of `seq`, ready for the trainer to count. + pub fn pretokenize(&self, seq: &str) -> Result> { + let normalized = normalize_for_training(self.normalizer.as_ref(), seq)?; + match &self.engine { + TrainingSplit::Pipeline(pretok) => { + let mut spans = Vec::new(); + pretok.pre_tokenize(&normalized, &mut spans)?; + Ok(spans + .iter() + .map(|s| normalized[s.start as usize..s.end as usize].to_owned()) + .collect()) + } + TrainingSplit::ByteLevel { + split, + add_prefix_space, + } => { + let prefixed; + let text: &str = if *add_prefix_space + && !normalized.is_empty() + && !normalized.starts_with(' ') + { + prefixed = format!(" {normalized}"); + &prefixed + } else { + &normalized + }; + let mut spans = Vec::new(); + match split { + Some(split) => split.pre_tokenize(text, &mut spans)?, + None => spans.push(Span { + start: 0, + end: text.len() as u32, + }), + } + Ok(spans + .iter() + .map(|s| byte_level_word(&text[s.start as usize..s.end as usize])) + .collect()) + } + } + } +} + /// An output token. Carries only the vocabulary `id` — offsets and the token /// string are dropped, which is all an encode-only caller needs. #[derive(Debug, Clone, Copy)] diff --git a/tokenizers/tk-train/src/train_ext.rs b/tokenizers/tk-train/src/train_ext.rs index 6dfa3cbde..70e88e358 100644 --- a/tokenizers/tk-train/src/train_ext.rs +++ b/tokenizers/tk-train/src/train_ext.rs @@ -1,9 +1,12 @@ use std::fs::File; use std::io::BufReader; -use tk_encode::tokenizer::{ - Decoder, Model, Normalizer, PostProcessor, PreTokenizer, TokenizerImpl, -}; +use tk_encode::decoders::DecoderWrapper; +use tk_encode::normalizers::NormalizerWrapper; +use tk_encode::pre_tokenizers::PreTokenizerWrapper; +use tk_encode::processors::PostProcessorWrapper; +use tk_encode::tokenizer::pipeline::TrainingPretokenizer; +use tk_encode::tokenizer::{Model, TokenizerImpl}; use tk_encode::utils::iter::ResultShunt; use tk_encode::utils::progress::{ProgressBar, ProgressStyle}; use tk_encode::{LinesWithEnding, Result}; @@ -35,13 +38,14 @@ pub trait TokenizerTrainExt { S: AsRef + Send; } -impl TokenizerTrainExt for TokenizerImpl +// Retargeted from the legacy generic `TokenizerImpl` (which drove training through the +// legacy `Normalizer`/`PreTokenizer` trait runtime) to the concrete config wrappers, so training +// pre-tokenizes through the fast `TrainingPretokenizer` (atomsplit) pipeline instead. `M` stays +// generic so both `ModelWrapper` (via `TrainerWrapper`) and single-model trainers still apply. +impl TokenizerTrainExt + for TokenizerImpl where M: Model + Send + Sync, - N: Normalizer + Send + Sync, - PT: PreTokenizer + Send + Sync, - PP: PostProcessor + Send + Sync, - D: Decoder + Send + Sync, { fn train_from_files(&mut self, trainer: &mut T, files: Vec) -> Result<&mut Self> where @@ -84,13 +88,15 @@ where None }; + let pretok = + TrainingPretokenizer::new(self.get_normalizer(), self.get_pre_tokenizer())?; trainer.feed( sequences.inspect(|s| { if let Some(progress) = &progress { progress.inc(s.len() as u64) } }), - |seq| self.pre_tokenize_for_training(seq), + |seq| pretok.pretokenize(seq), )?; if let Some(pbar) = progress { @@ -126,13 +132,14 @@ where None }; + let pretok = TrainingPretokenizer::new(self.get_normalizer(), self.get_pre_tokenizer())?; trainer.feed( sequences.inspect(|_s| { if let Some(progress) = &progress { progress.inc(1) } }), - |seq| self.pre_tokenize_for_training(seq), + |seq| pretok.pretokenize(seq), )?; if let Some(pbar) = progress { pbar.finish(); From 8c19d80dc0188af4ca00f839347420129b1e09bd Mon Sep 17 00:00:00 2001 From: Arthur Date: Wed, 15 Jul 2026 16:16:59 +0200 Subject: [PATCH 4/8] refactor(tk-encode): strip legacy encode/decode engine, pipeline-only Remove the legacy inference path so tk-encode is pipeline-first: - delete decoders/ and processors/ modules, tokenizer/encoding.rs, utils/{padding,truncation}.rs - collapse TokenizerImpl -> ; Tokenizer keeps only loading (from_file/from_bytes/from_pretrained/from_str), config accessors, add_tokens, get_vocab and save; drop encode*/decode*/DecodeStream/post_process/padding/truncation - slim AddedVocabulary to a data store: drop the daachorse matcher + regex word-boundary logic (find_matches/split_with_indices/extract_and_normalize); keep add_tokens id assignment, get_added_tokens_decoder, is_special_token and serde - drop the legacy Decoder/PostProcessor traits + their component impls (ByteLevel, Metaspace, Replace); keep the pipeline (Cow/span) impls - serialization: emit/parse the pipeline fields only; ignore truncation/padding/ post_processor/decoder sections on load - tk-train: target the collapsed TokenizerImpl Umbrella crate + python/node bindings intentionally break here (follow-up). --- tokenizers/tk-encode/src/decoders/bpe.rs | 38 - .../tk-encode/src/decoders/byte_fallback.rs | 109 -- tokenizers/tk-encode/src/decoders/ctc.rs | 120 -- tokenizers/tk-encode/src/decoders/fuse.rs | 43 - tokenizers/tk-encode/src/decoders/mod.rs | 233 --- tokenizers/tk-encode/src/decoders/sequence.rs | 55 - tokenizers/tk-encode/src/decoders/strip.rs | 80 -- .../tk-encode/src/decoders/wordpiece.rs | 86 -- tokenizers/tk-encode/src/lib.rs | 27 +- .../tk-encode/src/normalizers/replace.rs | 32 - .../src/pre_tokenizers/byte_level.rs | 533 +------ .../tk-encode/src/pre_tokenizers/metaspace.rs | 27 +- tokenizers/tk-encode/src/processors/bert.rs | 296 ---- tokenizers/tk-encode/src/processors/mod.rs | 128 -- .../tk-encode/src/processors/roberta.rs | 341 ----- .../tk-encode/src/processors/sequence.rs | 172 --- .../tk-encode/src/processors/template.rs | 1156 --------------- .../src/tokenizer/added_vocabulary.rs | 803 +---------- .../tk-encode/src/tokenizer/encoding.rs | 909 ------------ tokenizers/tk-encode/src/tokenizer/mod.rs | 1261 +---------------- .../tk-encode/src/tokenizer/pipeline.rs | 4 +- .../tk-encode/src/tokenizer/pre_tokenizer.rs | 140 +- .../tk-encode/src/tokenizer/serialization.rs | 68 +- tokenizers/tk-encode/src/utils/mod.rs | 19 - tokenizers/tk-encode/src/utils/padding.rs | 142 -- tokenizers/tk-encode/src/utils/truncation.rs | 326 ----- tokenizers/tk-train/src/train_ext.rs | 5 +- 27 files changed, 56 insertions(+), 7097 deletions(-) delete mode 100644 tokenizers/tk-encode/src/decoders/bpe.rs delete mode 100644 tokenizers/tk-encode/src/decoders/byte_fallback.rs delete mode 100644 tokenizers/tk-encode/src/decoders/ctc.rs delete mode 100644 tokenizers/tk-encode/src/decoders/fuse.rs delete mode 100644 tokenizers/tk-encode/src/decoders/mod.rs delete mode 100644 tokenizers/tk-encode/src/decoders/sequence.rs delete mode 100644 tokenizers/tk-encode/src/decoders/strip.rs delete mode 100644 tokenizers/tk-encode/src/decoders/wordpiece.rs delete mode 100644 tokenizers/tk-encode/src/processors/bert.rs delete mode 100644 tokenizers/tk-encode/src/processors/mod.rs delete mode 100644 tokenizers/tk-encode/src/processors/roberta.rs delete mode 100644 tokenizers/tk-encode/src/processors/sequence.rs delete mode 100644 tokenizers/tk-encode/src/processors/template.rs delete mode 100644 tokenizers/tk-encode/src/tokenizer/encoding.rs delete mode 100644 tokenizers/tk-encode/src/utils/padding.rs delete mode 100644 tokenizers/tk-encode/src/utils/truncation.rs diff --git a/tokenizers/tk-encode/src/decoders/bpe.rs b/tokenizers/tk-encode/src/decoders/bpe.rs deleted file mode 100644 index 813dc7083..000000000 --- a/tokenizers/tk-encode/src/decoders/bpe.rs +++ /dev/null @@ -1,38 +0,0 @@ -use crate::tokenizer::{Decoder, Result}; - -use serde::{Deserialize, Serialize}; - -#[derive(Deserialize, Clone, Debug, Serialize)] -/// Allows decoding Original BPE by joining all the tokens and then replacing -/// the suffix used to identify end-of-words by whitespaces -#[serde(tag = "type")] -#[non_exhaustive] -pub struct BPEDecoder { - pub suffix: String, -} - -impl BPEDecoder { - pub fn new(suffix: String) -> Self { - Self { suffix } - } -} - -impl Default for BPEDecoder { - fn default() -> Self { - Self::new("".into()) - } -} - -impl Decoder for BPEDecoder { - fn decode_chain(&self, tokens: Vec) -> Result> { - let n = tokens.len() - 1; - Ok(tokens - .into_iter() - .enumerate() - .map(|(i, token)| { - let replacement = if i == n { "" } else { " " }; - token.replace(&self.suffix, replacement) - }) - .collect()) - } -} diff --git a/tokenizers/tk-encode/src/decoders/byte_fallback.rs b/tokenizers/tk-encode/src/decoders/byte_fallback.rs deleted file mode 100644 index 57b7b63cd..000000000 --- a/tokenizers/tk-encode/src/decoders/byte_fallback.rs +++ /dev/null @@ -1,109 +0,0 @@ -use crate::tokenizer::{Decoder, Result}; -use monostate::MustBe; - -use serde::{Deserialize, Serialize}; - -#[derive(Deserialize, Clone, Debug, Serialize, Default)] -/// ByteFallback is a simple trick which converts tokens looking like `<0x61>` -/// to pure bytes, and attempts to make them into a string. If the tokens -/// cannot be decoded you will get � instead for each inconvertible byte token -#[non_exhaustive] -pub struct ByteFallback { - #[serde(rename = "type")] - type_: MustBe!("ByteFallback"), -} - -impl ByteFallback { - pub fn new() -> Self { - Self { - type_: MustBe!("ByteFallback"), - } - } -} - -impl Decoder for ByteFallback { - fn decode_chain(&self, tokens: Vec) -> Result> { - let mut new_tokens: Vec = vec![]; - let mut previous_byte_tokens: Vec = vec![]; - - for token in tokens { - let bytes = if token.len() == 6 && token.starts_with("<0x") && token.ends_with('>') { - u8::from_str_radix(&token[3..5], 16).ok() - } else { - None - }; - if let Some(bytes) = bytes { - previous_byte_tokens.push(bytes); - } else { - if !previous_byte_tokens.is_empty() { - if let Ok(string) = String::from_utf8(previous_byte_tokens.clone()) { - new_tokens.push(string); - } else { - for _ in 0..previous_byte_tokens.len() { - new_tokens.push("�".into()); - } - } - previous_byte_tokens.clear(); - } - new_tokens.push(token); - } - } - if !previous_byte_tokens.is_empty() { - if let Ok(string) = String::from_utf8(previous_byte_tokens.clone()) { - new_tokens.push(string); - } else { - for _ in 0..previous_byte_tokens.len() { - new_tokens.push("�".into()); - } - } - } - - Ok(new_tokens) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn decode() { - let decoder = ByteFallback::new(); - let res = decoder - .decode_chain(vec!["Hey".into(), "friend!".into()]) - .unwrap(); - assert_eq!(res, vec!["Hey", "friend!"]); - - let res = decoder.decode_chain(vec!["<0x61>".into()]).unwrap(); - assert_eq!(res, vec!["a"]); - - let res = decoder.decode_chain(vec!["<0xE5>".into()]).unwrap(); - assert_eq!(res, vec!["�"]); - - let res = decoder - .decode_chain(vec!["<0xE5>".into(), "<0x8f>".into()]) - .unwrap(); - assert_eq!(res, vec!["�", "�"]); - - // 叫 - let res = decoder - .decode_chain(vec!["<0xE5>".into(), "<0x8f>".into(), "<0xab>".into()]) - .unwrap(); - assert_eq!(res, vec!["叫"]); - - let res = decoder - .decode_chain(vec![ - "<0xE5>".into(), - "<0x8f>".into(), - "<0xab>".into(), - "a".into(), - ]) - .unwrap(); - assert_eq!(res, vec!["叫", "a"]); - - let res = decoder - .decode_chain(vec!["<0xE5>".into(), "<0x8f>".into(), "a".into()]) - .unwrap(); - assert_eq!(res, vec!["�", "�", "a"]); - } -} diff --git a/tokenizers/tk-encode/src/decoders/ctc.rs b/tokenizers/tk-encode/src/decoders/ctc.rs deleted file mode 100644 index 9d5a57188..000000000 --- a/tokenizers/tk-encode/src/decoders/ctc.rs +++ /dev/null @@ -1,120 +0,0 @@ -use crate::decoders::wordpiece; -use crate::tokenizer::{Decoder, Result}; - -use itertools::Itertools; -use serde::{Deserialize, Serialize}; - -#[derive(Debug, Clone, Serialize, Deserialize)] -/// The CTC (Connectionist Temporal Classification) decoder takes care -/// of sanitizing a list of inputs token. -/// Due to some alignment problem the output of some models can come -/// with duplicated token. -#[serde(tag = "type")] -#[non_exhaustive] -pub struct CTC { - /// The pad token used by CTC to delimit a new token. - pub pad_token: String, - /// The word delimiter token. It will be replaced by a ``. - pub word_delimiter_token: String, - /// Whether to cleanup some tokenization artifacts. - /// Mainly spaces before punctuation, and some abbreviated english forms. - pub cleanup: bool, -} - -impl CTC { - pub fn new(pad_token: String, word_delimiter_token: String, cleanup: bool) -> Self { - Self { - pad_token, - word_delimiter_token, - cleanup, - } - } -} - -impl Default for CTC { - fn default() -> Self { - Self { - pad_token: "".to_string(), - word_delimiter_token: "|".to_string(), - cleanup: true, - } - } -} - -impl Decoder for CTC { - fn decode_chain(&self, tokens: Vec) -> Result> { - Ok(tokens - .into_iter() - .dedup() - .filter_map(|token| { - let mut replaced = token.replace(&self.pad_token, ""); - if self.cleanup { - replaced = - wordpiece::cleanup(&replaced).replace(&self.word_delimiter_token, " "); - } - if replaced.is_empty() { - None - } else { - Some(replaced) - } - }) - .collect()) - } -} - -#[cfg(test)] -mod tests { - use super::*; - #[test] - fn handmade_sample() { - let ctc_decoder = CTC::default(); - let id_to_string_result = " h e e l l l o o o " - .split(' ') - .map(|s| s.to_string()) - .collect(); - assert_eq!( - ctc_decoder.decode_chain(id_to_string_result).unwrap(), - vec!["h", "e", "l", "l", "o"] - ); - } - #[test] - fn handmade_with_delimiter_sample() { - let ctc_decoder = CTC::default(); - let id_to_string_result = " h e e l l l o o o | w o o o r l l d " - .split(' ') - .map(|s| s.to_string()) - .collect(); - assert_eq!( - ctc_decoder.decode_chain(id_to_string_result).unwrap(), - vec!["h", "e", "l", "l", "o", " ", "w", "o", "r", "l", "d"] - ); - } - #[test] - fn librispeech_sample() { - let ctc_decoder = CTC::default(); - let id_to_string_result = " A | | M A N | | | S A I D D | | T T O | | T H E E | | | U U N N I V E R R S E E | | S S I R R | | | I | E X I S T | | ".split(' ').map(|s| s.to_string()).collect(); - assert_eq!( - ctc_decoder.decode_chain(id_to_string_result).unwrap(), - vec![ - "A", " ", "M", "A", "N", " ", "S", "A", "I", "D", " ", "T", "O", " ", "T", "H", - "E", " ", "U", "N", "I", "V", "E", "R", "S", "E", " ", "S", "I", "R", " ", "I", - " ", "E", "X", "I", "S", "T", " " - ] - ); - } - #[test] - fn another_librispeech_sample() { - let ctc_decoder = CTC::default(); - let id_to_string_result = " H I S S | | I N S T T A N C C T | | | | | P A N N N I C | | W A S | | F O L L L O O W E E D | | B Y | | | A | | S S S M M A L L L | | | S H H A R R P | B L L O W W | | | H I G H H | | O N | | H I S S | | C H H E S S T T | | | ".split(' ').map(|s| s.to_string()).collect(); - assert_eq!( - ctc_decoder.decode_chain(id_to_string_result).unwrap(), - vec![ - "H", "I", "S", " ", "I", "N", "S", "T", "A", "N", "C", "T", " ", "P", "A", "N", - "I", "C", " ", "W", "A", "S", " ", "F", "O", "L", "L", "O", "W", "E", "D", " ", - "B", "Y", " ", "A", " ", "S", "M", "A", "L", "L", " ", "S", "H", "A", "R", "P", - " ", "B", "L", "O", "W", " ", "H", "I", "G", "H", " ", "O", "N", " ", "H", "I", - "S", " ", "C", "H", "E", "S", "T", " " - ] - ); - } -} diff --git a/tokenizers/tk-encode/src/decoders/fuse.rs b/tokenizers/tk-encode/src/decoders/fuse.rs deleted file mode 100644 index 5e4a1c119..000000000 --- a/tokenizers/tk-encode/src/decoders/fuse.rs +++ /dev/null @@ -1,43 +0,0 @@ -use crate::tokenizer::{Decoder, Result}; -use monostate::MustBe; -use serde::{Deserialize, Serialize}; - -#[derive(Clone, Debug, Serialize, Deserialize, Default)] -/// Fuse simply fuses all tokens into one big string. -/// It's usually the last decoding step anyway, but this -/// decoder exists incase some decoders need to happen after that -/// step -#[non_exhaustive] -pub struct Fuse { - #[serde(rename = "type")] - type_: MustBe!("Fuse"), -} - -impl Fuse { - pub fn new() -> Self { - Self { - type_: MustBe!("Fuse"), - } - } -} - -impl Decoder for Fuse { - fn decode_chain(&self, tokens: Vec) -> Result> { - let new_string = tokens.join(""); - Ok(vec![new_string]) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn decode() { - let decoder = Fuse::new(); - let res = decoder - .decode_chain(vec!["Hey".into(), " friend!".into()]) - .unwrap(); - assert_eq!(res, vec!["Hey friend!"]); - } -} diff --git a/tokenizers/tk-encode/src/decoders/mod.rs b/tokenizers/tk-encode/src/decoders/mod.rs deleted file mode 100644 index 6e79e7029..000000000 --- a/tokenizers/tk-encode/src/decoders/mod.rs +++ /dev/null @@ -1,233 +0,0 @@ -pub mod bpe; -pub mod byte_fallback; -pub mod ctc; -pub mod fuse; -pub mod sequence; -pub mod strip; -pub mod wordpiece; - -// Re-export these as decoders -pub use super::pre_tokenizers::byte_level; -pub use super::pre_tokenizers::metaspace; - -use serde::{Deserialize, Deserializer, Serialize}; - -use crate::decoders::bpe::BPEDecoder; -use crate::decoders::byte_fallback::ByteFallback; -use crate::decoders::ctc::CTC; -use crate::decoders::fuse::Fuse; -use crate::decoders::sequence::Sequence; -use crate::decoders::strip::Strip; -use crate::decoders::wordpiece::WordPiece; -use crate::normalizers::replace::Replace; -use crate::pre_tokenizers::byte_level::ByteLevel; -use crate::pre_tokenizers::metaspace::Metaspace; -use crate::{Decoder, Result}; - -#[derive(Serialize, Clone, Debug)] -#[serde(untagged)] -pub enum DecoderWrapper { - BPE(BPEDecoder), - ByteLevel(ByteLevel), - WordPiece(WordPiece), - Metaspace(Metaspace), - CTC(CTC), - Sequence(Sequence), - Replace(Replace), - Fuse(Fuse), - Strip(Strip), - ByteFallback(ByteFallback), -} - -impl<'de> Deserialize<'de> for DecoderWrapper { - fn deserialize(deserializer: D) -> std::result::Result - where - D: Deserializer<'de>, - { - #[derive(Deserialize)] - pub struct Tagged { - #[serde(rename = "type")] - variant: EnumType, - #[serde(flatten)] - rest: serde_json::Value, - } - #[derive(Serialize, Deserialize)] - pub enum EnumType { - BPEDecoder, - ByteLevel, - WordPiece, - Metaspace, - CTC, - Sequence, - Replace, - Fuse, - Strip, - ByteFallback, - } - - #[derive(Deserialize)] - #[serde(untagged)] - pub enum DecoderHelper { - Tagged(Tagged), - Legacy(serde_json::Value), - } - - #[derive(Deserialize)] - #[serde(untagged)] - pub enum DecoderUntagged { - BPE(BPEDecoder), - ByteLevel(ByteLevel), - WordPiece(WordPiece), - Metaspace(Metaspace), - CTC(CTC), - Sequence(Sequence), - Replace(Replace), - Fuse(Fuse), - Strip(Strip), - ByteFallback(ByteFallback), - } - - let helper = DecoderHelper::deserialize(deserializer).expect("Helper"); - Ok(match helper { - DecoderHelper::Tagged(model) => { - let mut values: serde_json::Map = - serde_json::from_value(model.rest).map_err(serde::de::Error::custom)?; - values.insert( - "type".to_string(), - serde_json::to_value(&model.variant).map_err(serde::de::Error::custom)?, - ); - let values = serde_json::Value::Object(values); - match model.variant { - EnumType::BPEDecoder => DecoderWrapper::BPE( - serde_json::from_value(values).map_err(serde::de::Error::custom)?, - ), - EnumType::ByteLevel => DecoderWrapper::ByteLevel( - serde_json::from_value(values).map_err(serde::de::Error::custom)?, - ), - EnumType::WordPiece => DecoderWrapper::WordPiece( - serde_json::from_value(values).map_err(serde::de::Error::custom)?, - ), - EnumType::Metaspace => DecoderWrapper::Metaspace( - serde_json::from_value(values).map_err(serde::de::Error::custom)?, - ), - EnumType::CTC => DecoderWrapper::CTC( - serde_json::from_value(values).map_err(serde::de::Error::custom)?, - ), - EnumType::Sequence => DecoderWrapper::Sequence( - serde_json::from_value(values).map_err(serde::de::Error::custom)?, - ), - EnumType::Replace => DecoderWrapper::Replace( - serde_json::from_value(values).map_err(serde::de::Error::custom)?, - ), - EnumType::Fuse => DecoderWrapper::Fuse( - serde_json::from_value(values).map_err(serde::de::Error::custom)?, - ), - EnumType::Strip => DecoderWrapper::Strip( - serde_json::from_value(values).map_err(serde::de::Error::custom)?, - ), - EnumType::ByteFallback => DecoderWrapper::ByteFallback( - serde_json::from_value(values).map_err(serde::de::Error::custom)?, - ), - } - } - DecoderHelper::Legacy(value) => { - let untagged = serde_json::from_value(value).map_err(serde::de::Error::custom)?; - match untagged { - DecoderUntagged::BPE(dec) => DecoderWrapper::BPE(dec), - DecoderUntagged::ByteLevel(dec) => DecoderWrapper::ByteLevel(dec), - DecoderUntagged::WordPiece(dec) => DecoderWrapper::WordPiece(dec), - DecoderUntagged::Metaspace(dec) => DecoderWrapper::Metaspace(dec), - DecoderUntagged::CTC(dec) => DecoderWrapper::CTC(dec), - DecoderUntagged::Sequence(dec) => DecoderWrapper::Sequence(dec), - DecoderUntagged::Replace(dec) => DecoderWrapper::Replace(dec), - DecoderUntagged::Fuse(dec) => DecoderWrapper::Fuse(dec), - DecoderUntagged::Strip(dec) => DecoderWrapper::Strip(dec), - DecoderUntagged::ByteFallback(dec) => DecoderWrapper::ByteFallback(dec), - } - } - }) - } -} - -impl Decoder for DecoderWrapper { - fn decode_chain(&self, tokens: Vec) -> Result> { - match self { - Self::BPE(bpe) => bpe.decode_chain(tokens), - Self::ByteLevel(bl) => bl.decode_chain(tokens), - Self::Metaspace(ms) => ms.decode_chain(tokens), - Self::WordPiece(wp) => wp.decode_chain(tokens), - Self::CTC(ctc) => ctc.decode_chain(tokens), - Self::Sequence(seq) => seq.decode_chain(tokens), - Self::Replace(seq) => seq.decode_chain(tokens), - Self::ByteFallback(bf) => bf.decode_chain(tokens), - Self::Strip(bf) => bf.decode_chain(tokens), - Self::Fuse(bf) => bf.decode_chain(tokens), - } - } -} - -impl_enum_from!(BPEDecoder, DecoderWrapper, BPE); -impl_enum_from!(ByteLevel, DecoderWrapper, ByteLevel); -impl_enum_from!(ByteFallback, DecoderWrapper, ByteFallback); -impl_enum_from!(Fuse, DecoderWrapper, Fuse); -impl_enum_from!(Strip, DecoderWrapper, Strip); -impl_enum_from!(Metaspace, DecoderWrapper, Metaspace); -impl_enum_from!(WordPiece, DecoderWrapper, WordPiece); -impl_enum_from!(CTC, DecoderWrapper, CTC); -impl_enum_from!(Sequence, DecoderWrapper, Sequence); -impl_enum_from!(Replace, DecoderWrapper, Replace); - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn decoder_serialization() { - let oldjson = r#"{"type":"Sequence","decoders":[{"type":"ByteFallback"},{"type":"Metaspace","replacement":"▁","add_prefix_space":true,"prepend_scheme":"always"}]}"#; - let olddecoder: DecoderWrapper = serde_json::from_str(oldjson).unwrap(); - let oldserialized = serde_json::to_string(&olddecoder).unwrap(); - let json = r#"{"type":"Sequence","decoders":[{"type":"ByteFallback"},{"type":"Metaspace","replacement":"▁","prepend_scheme":"always","split":true}]}"#; - assert_eq!(oldserialized, json); - - let decoder: DecoderWrapper = serde_json::from_str(json).unwrap(); - let serialized = serde_json::to_string(&decoder).unwrap(); - assert_eq!(serialized, json); - } - #[test] - fn decoder_serialization_other_no_arg() { - let json = r#"{"type":"Sequence","decoders":[{"type":"Fuse"},{"type":"Metaspace","replacement":"▁","prepend_scheme":"always","split":true}]}"#; - let decoder: DecoderWrapper = serde_json::from_str(json).unwrap(); - let serialized = serde_json::to_string(&decoder).unwrap(); - assert_eq!(serialized, json); - } - - #[test] - fn decoder_serialization_no_decode() { - let json = r#"{"type":"Sequence","decoders":[{},{"type":"Metaspace","replacement":"▁","prepend_scheme":"always"}]}"#; - let parse = serde_json::from_str::(json); - match parse { - Err(err) => assert_eq!( - format!("{err}"), - "data did not match any variant of untagged enum DecoderUntagged" - ), - _ => panic!("Expected error"), - } - - let json = r#"{"replacement":"▁","prepend_scheme":"always"}"#; - let parse = serde_json::from_str::(json); - match parse { - Err(err) => assert_eq!( - format!("{err}"), - "data did not match any variant of untagged enum DecoderUntagged" - ), - _ => panic!("Expected error"), - } - - let json = r#"{"type":"Sequence","prepend_scheme":"always"}"#; - let parse = serde_json::from_str::(json); - match parse { - Err(err) => assert_eq!(format!("{err}"), "missing field `decoders`"), - _ => panic!("Expected error"), - } - } -} diff --git a/tokenizers/tk-encode/src/decoders/sequence.rs b/tokenizers/tk-encode/src/decoders/sequence.rs deleted file mode 100644 index 73169b695..000000000 --- a/tokenizers/tk-encode/src/decoders/sequence.rs +++ /dev/null @@ -1,55 +0,0 @@ -use crate::decoders::DecoderWrapper; -use crate::tokenizer::{Decoder, Result}; -use crate::utils::macro_rules_attribute; -use serde::{Deserialize, Serialize}; - -#[derive(Clone, Debug)] -#[macro_rules_attribute(impl_serde_type!)] -pub struct Sequence { - decoders: Vec, -} - -impl Sequence { - pub fn new(decoders: Vec) -> Self { - Self { decoders } - } - - pub fn get_decoders(&self) -> &[DecoderWrapper] { - &self.decoders - } - - pub fn get_decoders_mut(&mut self) -> &mut [DecoderWrapper] { - &mut self.decoders - } -} - -impl Decoder for Sequence { - fn decode_chain(&self, mut tokens: Vec) -> Result> { - for decoder in &self.decoders { - tokens = decoder.decode_chain(tokens)?; - } - Ok(tokens) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::decoders::ctc::CTC; - use crate::pre_tokenizers::metaspace::Metaspace; - - #[test] - fn sequence_basic() { - let decoders = vec![ - DecoderWrapper::CTC(CTC::default()), - DecoderWrapper::Metaspace(Metaspace::default()), - ]; - let decoder = Sequence::new(decoders); - let tokens: Vec = vec!["▁", "▁", "H", "H", "i", "i", "▁", "y", "o", "u"] - .into_iter() - .map(|s| s.to_string()) - .collect(); - let out_tokens = decoder.decode(tokens).unwrap(); - assert_eq!(out_tokens, "Hi you"); - } -} diff --git a/tokenizers/tk-encode/src/decoders/strip.rs b/tokenizers/tk-encode/src/decoders/strip.rs deleted file mode 100644 index 9aeffec64..000000000 --- a/tokenizers/tk-encode/src/decoders/strip.rs +++ /dev/null @@ -1,80 +0,0 @@ -use crate::tokenizer::{Decoder, Result}; - -use serde::{Deserialize, Serialize}; - -#[derive(Deserialize, Clone, Debug, Serialize, Default)] -/// Strip is a simple trick which converts tokens looking like `<0x61>` -/// to pure bytes, and attempts to make them into a string. If the tokens -/// cannot be decoded you will get � instead for each inconvertible byte token -#[serde(tag = "type")] -#[non_exhaustive] -pub struct Strip { - pub content: char, - pub start: usize, - pub stop: usize, -} - -impl Strip { - pub fn new(content: char, start: usize, stop: usize) -> Self { - Self { - content, - start, - stop, - } - } -} - -impl Decoder for Strip { - fn decode_chain(&self, tokens: Vec) -> Result> { - Ok(tokens - .into_iter() - .map(|token| { - let chars: Vec = token.chars().collect(); - - let mut start_cut = 0; - for (i, &c) in chars.iter().enumerate().take(self.start) { - if c == self.content { - start_cut = i + 1; - continue; - } else { - break; - } - } - - let mut stop_cut = chars.len(); - for i in 0..self.stop { - let index = chars.len() - i - 1; - if chars[index] == self.content { - stop_cut = index; - continue; - } else { - break; - } - } - - let new_token: String = chars[start_cut..stop_cut].iter().collect(); - new_token - }) - .collect()) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn decode() { - let decoder = Strip::new('H', 1, 0); - let res = decoder - .decode_chain(vec!["Hey".into(), " friend!".into(), "HHH".into()]) - .unwrap(); - assert_eq!(res, vec!["ey", " friend!", "HH"]); - - let decoder = Strip::new('y', 0, 1); - let res = decoder - .decode_chain(vec!["Hey".into(), " friend!".into()]) - .unwrap(); - assert_eq!(res, vec!["He", " friend!"]); - } -} diff --git a/tokenizers/tk-encode/src/decoders/wordpiece.rs b/tokenizers/tk-encode/src/decoders/wordpiece.rs deleted file mode 100644 index a2da414c0..000000000 --- a/tokenizers/tk-encode/src/decoders/wordpiece.rs +++ /dev/null @@ -1,86 +0,0 @@ -use crate::tokenizer::{Decoder, Result}; - -use serde::{Deserialize, Serialize}; - -#[derive(Deserialize, Clone, Debug, Serialize)] -/// The WordPiece decoder takes care of decoding a list of wordpiece tokens -/// back into a readable string. -#[serde(tag = "type")] -#[non_exhaustive] -pub struct WordPiece { - /// The prefix to be used for continuing subwords - pub prefix: String, - /// Whether to cleanup some tokenization artifacts (spaces before punctuation, ...) - pub cleanup: bool, -} - -impl WordPiece { - pub fn new(prefix: String, cleanup: bool) -> Self { - Self { prefix, cleanup } - } -} - -impl Default for WordPiece { - fn default() -> Self { - Self { - prefix: "##".to_owned(), - cleanup: true, - } - } -} -pub fn cleanup(dirty_input: &str) -> String { - dirty_input - .replace(" .", ".") - .replace(" ?", "?") - .replace(" !", "!") - .replace(" ,", ",") - .replace(" ' ", "'") - .replace(" n't", "n't") - .replace(" 'm", "'m") - .replace(" do not", " don't") - .replace(" 's", "'s") - .replace(" 've", "'ve") - .replace(" 're", "'re") -} - -impl Decoder for WordPiece { - fn decode_chain(&self, mut tokens: Vec) -> Result> { - for (i, token) in tokens.iter_mut().enumerate() { - if i != 0 { - if let Some(tk) = token.strip_prefix(&self.prefix) { - *token = tk.to_string(); - } else { - *token = format!(" {token}"); - } - } - if self.cleanup { - *token = cleanup(token); - } - } - Ok(tokens) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn wordpiece_decoder() { - let decoder = WordPiece::new("##".to_string(), false); - - assert_eq!( - decoder - .decode(vec![ - "##uelo".to_string(), - "Ara".to_string(), - "##új".to_string(), - "##o".to_string(), - "No".to_string(), - "##guera".to_string() - ]) - .unwrap(), - "##uelo Araújo Noguera" - ); - } -} diff --git a/tokenizers/tk-encode/src/lib.rs b/tokenizers/tk-encode/src/lib.rs index 1f876685b..ead898916 100644 --- a/tokenizers/tk-encode/src/lib.rs +++ b/tokenizers/tk-encode/src/lib.rs @@ -25,17 +25,19 @@ //! that, for example, a language model would need, such as special tokens. //! //! ## Loading a pretrained tokenizer from the Hub -//! ``` +//! ```no_run //! use tk_encode::tokenizer::{Result, Tokenizer}; +//! use tk_encode::tokenizer::pipeline::PipelineTokenizer; +//! use std::convert::TryFrom; //! //! fn main() -> Result<()> { //! # #[cfg(feature = "http")] //! # { //! // needs http feature enabled //! let tokenizer = Tokenizer::from_pretrained("bert-base-cased", None)?; -//! -//! let encoding = tokenizer.encode("Hey there!", false)?; -//! println!("{:?}", encoding.get_tokens()); +//! let pipeline = PipelineTokenizer::try_from(&tokenizer)?; +//! let ids = pipeline.encode("Hey there!", false)?; +//! println!("{:?}", ids); //! # } //! Ok(()) //! } @@ -44,8 +46,10 @@ //! ## Deserialization and tokenization example //! //! ```no_run -//! use tk_encode::tokenizer::{Result, Tokenizer, EncodeInput}; +//! use tk_encode::tokenizer::{Result, Tokenizer}; +//! use tk_encode::tokenizer::pipeline::PipelineTokenizer; //! use tk_encode::models::bpe::BPE; +//! use std::convert::TryFrom; //! //! fn main() -> Result<()> { //! let bpe_builder = BPE::from_file("./path/to/vocab.json", "./path/to/merges.txt"); @@ -54,10 +58,10 @@ //! .unk_token("[UNK]".into()) //! .build()?; //! -//! let mut tokenizer = Tokenizer::new(bpe); -//! -//! let encoding = tokenizer.encode("Hey there!", false)?; -//! println!("{:?}", encoding.get_tokens()); +//! let tokenizer = Tokenizer::new(bpe); +//! let pipeline = PipelineTokenizer::try_from(&tokenizer)?; +//! let ids = pipeline.encode("Hey there!", false)?; +//! println!("{:?}", ids); //! //! Ok(()) //! } @@ -85,16 +89,11 @@ #[macro_use] extern crate log; -#[macro_use] -extern crate derive_builder; - #[macro_use] pub mod utils; -pub mod decoders; pub mod models; pub mod normalizers; pub mod pre_tokenizers; -pub mod processors; pub mod tokenizer; pub mod vocab; /// Legacy map-backed token store; drop-in twin of the fast `bucket_vocab_store::BucketVocabStore`. diff --git a/tokenizers/tk-encode/src/normalizers/replace.rs b/tokenizers/tk-encode/src/normalizers/replace.rs index 6237a88aa..88e8451a7 100644 --- a/tokenizers/tk-encode/src/normalizers/replace.rs +++ b/tokenizers/tk-encode/src/normalizers/replace.rs @@ -1,8 +1,6 @@ use std::borrow::Cow; use crate::pipeline; -use crate::tokenizer::pattern::Pattern; -use crate::tokenizer::Decoder; use crate::tokenizer::{NormalizedString, Normalizer, Result}; use crate::utils::SysRegex; use serde::{Deserialize, Serialize}; @@ -111,26 +109,6 @@ impl pipeline::Normalizer for Replace { } } -impl Decoder for Replace { - fn decode_chain(&self, tokens: Vec) -> Result> { - tokens - .into_iter() - .map(|token| -> Result { - let mut new_token = "".to_string(); - - for ((start, stop), is_match) in (&self.regex).find_matches(&token)? { - if is_match { - new_token.push_str(&self.content); - } else { - new_token.push_str(&token[start..stop]); - } - } - Ok(new_token) - }) - .collect() - } -} - // `Replace` needs a system-regex backend (SysRegex) for every test here. #[cfg(all(test, feature = "fancy-regex"))] mod tests { @@ -174,16 +152,6 @@ mod tests { assert_eq!(serde_json::from_str::(replace_s).unwrap(), replace); } - #[test] - fn test_replace_decode() { - let original = vec!["hello".to_string(), "_hello".to_string()]; - let replace = Replace::new("_", " ").unwrap(); - assert_eq!( - replace.decode_chain(original).unwrap(), - vec!["hello", " hello"] - ); - } - #[test] fn pipeline_replace_matches_legacy() { let n = Replace::new("''", "\"").unwrap(); diff --git a/tokenizers/tk-encode/src/pre_tokenizers/byte_level.rs b/tokenizers/tk-encode/src/pre_tokenizers/byte_level.rs index ee4695fd8..c74b41aa5 100644 --- a/tokenizers/tk-encode/src/pre_tokenizers/byte_level.rs +++ b/tokenizers/tk-encode/src/pre_tokenizers/byte_level.rs @@ -1,11 +1,8 @@ -use crate::utils::byte_level::{byte_level_transform, BYTES_CHAR_LOOKUP, CHAR_BYTES_LOOKUP}; +use crate::utils::byte_level::{byte_level_transform, BYTES_CHAR_LOOKUP}; use crate::utils::{GptFsm, GptFsmPattern}; use serde::{Deserialize, Serialize}; -use crate::tokenizer::{ - Decoder, Encoding, PostProcessor, PreTokenizedString, PreTokenizer, Result, - SplitDelimiterBehavior, -}; +use crate::tokenizer::{PreTokenizedString, PreTokenizer, Result, SplitDelimiterBehavior}; use crate::utils::macro_rules_attribute; #[derive(Copy, Clone, Debug, PartialEq, Eq)] @@ -100,529 +97,3 @@ impl PreTokenizer for ByteLevel { } } -/// As a `Decoder`, `ByteLevel` is in charge of converting any byte-level characters to their -/// unicode counterpart, before merging everything back into a single String. -/// This decoder will consume the tokens and merge them in one step to alleviate -/// the fact that single token decoded might be a byte not representable as -/// as String. -impl Decoder for ByteLevel { - fn decode_chain(&self, tokens: Vec) -> Result> { - let toks = tokens - .into_iter() - .flat_map(|t| { - t.chars() - .try_fold(vec![], |mut acc, c| { - CHAR_BYTES_LOOKUP.get(&c).map(|b| { - acc.push(*b); - acc - }) - }) - .unwrap_or_else(|| t.as_bytes().to_vec()) - }) - .collect::>(); - Ok(vec![String::from_utf8_lossy(&toks).to_string()]) - } -} - -/// As a `PostProcessor`, `ByteLevel` is in charge of trimming the offsets if necessary. -impl PostProcessor for ByteLevel { - fn added_tokens(&self, _is_pair: bool) -> usize { - 0 - } - - fn process_encodings( - &self, - mut encodings: Vec, - _add_special_tokens: bool, - ) -> Result> { - if self.trim_offsets { - for encoding in encodings.iter_mut() { - process_offsets(encoding, self.add_prefix_space); - encoding - .get_overflowing_mut() - .iter_mut() - .for_each(|encoding| process_offsets(encoding, self.add_prefix_space)); - } - } - for (i, encoding) in encodings.iter_mut().enumerate() { - encoding.set_sequence_id(i); - } - Ok(encodings) - //::default_process(encodings, add_special_tokens) - } -} - -pub fn process_offsets(encoding: &mut Encoding, add_prefix_space: bool) { - encoding.process_tokens_with_offsets_mut(|(i, (token, offsets))| { - let mut leading_spaces = token - .chars() - .take_while(|c| *c == BYTES_CHAR_LOOKUP[b' ' as usize] || c.is_whitespace()) - .count(); - let trailing_spaces = token - .chars() - .rev() - .take_while(|c| *c == BYTES_CHAR_LOOKUP[b' ' as usize] || c.is_whitespace()) - .count(); - - if leading_spaces > 0 || trailing_spaces > 0 { - if leading_spaces > 0 { - // If user uses `is_pretokenized=True` we might have - // offsets that might begin at the start of the string but are - // NOT the first token. - let is_first = i == 0 || offsets.0 == 0; - if is_first && add_prefix_space && leading_spaces == 1 { - // If we are processing the first pair of offsets, with `add_prefix_space`, - // then we shouldn't remove anything we added. If there are more than one - // leading spaces though, it means we didn't add them, and they should be - // removed. - leading_spaces = 0; - } - offsets.0 = std::cmp::min(offsets.0 + leading_spaces, offsets.1); - } - if trailing_spaces > 0 && offsets.1 >= trailing_spaces { - offsets.1 = std::cmp::max(offsets.1 - trailing_spaces, offsets.0); - } - } - }); -} - -#[cfg(test)] -mod tests { - - use super::*; - use crate::tokenizer::{ - Decoder, Encoding, OffsetReferential, OffsetType, PostProcessor, PreTokenizedString, - PreTokenizer, - }; - use ahash::AHashMap; - use std::iter::FromIterator; - - #[test] - fn pre_tokenization() { - let bytelevel = ByteLevel::default().add_prefix_space(false); - let mut pretokenized: PreTokenizedString = "Hello my friend, how is your day going?".into(); - bytelevel.pre_tokenize(&mut pretokenized).unwrap(); - assert_eq!( - pretokenized - .get_splits(OffsetReferential::Original, OffsetType::Byte) - .into_iter() - .map(|(s, o, _)| (s, o)) - .collect::>(), - vec![ - ("Hello", (0, 5)), - ("Ġmy", (5, 8)), - ("Ġfriend", (8, 15)), - (",", (15, 16)), - ("Ġhow", (16, 20)), - ("Ġis", (20, 23)), - ("Ġyour", (23, 28)), - ("Ġday", (28, 32)), - ("Ġgoing", (32, 38)), - ("?", (38, 39)) - ] - ); - } - - #[test] - fn pre_tokenization_no_regex() { - let bytelevel = ByteLevel::default().use_regex(false); - let mut pretokenized: PreTokenizedString = "Hello my friend, how is your day going?".into(); - bytelevel.pre_tokenize(&mut pretokenized).unwrap(); - assert_eq!( - pretokenized - .get_splits(OffsetReferential::Original, OffsetType::Byte) - .into_iter() - .map(|(s, o, _)| (s, o)) - .collect::>(), - vec![("ĠHelloĠmyĠfriend,ĠhowĠisĠyourĠdayĠgoing?", (0, 39))] - ); - } - - #[test] - fn decoding() { - let bytelevel = ByteLevel::default().add_prefix_space(false); - assert_eq!( - bytelevel - .decode_chain( - vec![ - "Hello", "Ġmy", "Ġfriend", ",", "Ġhow", "Ġis", "Ġyour", "Ġday", "Ġgoing", - "?" - ] - .into_iter() - .map(|s| s.into()) - .collect::>() - ) - .unwrap(), - vec!["Hello my friend, how is your day going?"] - ); - } - - #[test] - fn add_prefix_space() { - let bytelevel = ByteLevel::default().add_prefix_space(true); - for s in &[ - " Hello my friend, how is your day going?", - "Hello my friend, how is your day going?", - ] { - let mut pretokenized = PreTokenizedString::from(*s); - bytelevel.pre_tokenize(&mut pretokenized).unwrap(); - assert_eq!( - pretokenized - .get_splits(OffsetReferential::Normalized, OffsetType::Byte) - .into_iter() - .map(|(s, o, _)| (s, o)) - .collect::>(), - vec![ - ("ĠHello", (0, 7)), - ("Ġmy", (7, 11)), - ("Ġfriend", (11, 19)), - (",", (19, 20)), - ("Ġhow", (20, 25)), - ("Ġis", (25, 29)), - ("Ġyour", (29, 35)), - ("Ġday", (35, 40)), - ("Ġgoing", (40, 47)), - ("?", (47, 48)) - ] - ); - } - } - - #[test] - fn decode_works_on_separated_tokens() { - let samples = vec![ - "A Nuskhuri abbreviation of იესუ ქრისტე ( iesu kriste ) \" Jesus Christ \"", - "An equal number have descenders , like p or q in English \ - : გ , დ , ე , ვ , კ , ლ , ჟ , ტ , უ , ფ , ღ , ყ , ც", - ]; - - let bytelevel = ByteLevel::default().add_prefix_space(false); - for sample in samples { - let mut pretokenized = PreTokenizedString::from(sample); - bytelevel.pre_tokenize(&mut pretokenized).unwrap(); - let separated_tokens = pretokenized - .get_splits(OffsetReferential::Original, OffsetType::Byte) - .iter() - .flat_map(|(s, _, _)| s.split("").map(|t| t.into())) - .collect::>(); - assert_eq!( - sample, - bytelevel.decode_chain(separated_tokens).unwrap().join("") - ); - } - } - - #[test] - fn handling_of_newlines() { - let mut pretokenized = PreTokenizedString::from("Hello there\nHello there"); - let bytelevel = ByteLevel::default().add_prefix_space(false); - bytelevel.pre_tokenize(&mut pretokenized).unwrap(); - - assert_eq!( - pretokenized - .get_splits(OffsetReferential::Original, OffsetType::Byte) - .into_iter() - .map(|(s, o, _)| (s, o)) - .collect::>(), - vec![ - ("Hello", (0, 5)), - ("Ġthere", (5, 11)), - ("Ċ", (11, 12)), - ("Hello", (12, 17)), - ("Ġthere", (17, 23)) - ] - ); - } - - #[test] - fn handling_of_multiple_whitespaces() { - let mut pretokenized = PreTokenizedString::from("Hello there dear"); - let bytelevel = ByteLevel::default().add_prefix_space(false); - bytelevel.pre_tokenize(&mut pretokenized).unwrap(); - - assert_eq!( - pretokenized - .get_splits(OffsetReferential::Original, OffsetType::Byte) - .into_iter() - .map(|(s, o, _)| (s, o)) - .collect::>(), - vec![ - ("Hello", (0, 5)), - ("Ġthere", (5, 11)), - ("ĠĠĠĠĠĠ", (11, 17)), - ("Ġdear", (17, 22)) - ] - ); - } - - #[test] - fn offsets_when_char_split_up() { - let input = "i⭢j"; - let mut pretokenized = PreTokenizedString::from(input); - let bytelevel = ByteLevel::default().add_prefix_space(false); - bytelevel.pre_tokenize(&mut pretokenized).unwrap(); - - assert_eq!( - pretokenized - .get_splits(OffsetReferential::Original, OffsetType::Byte) - .into_iter() - .map(|(s, o, _)| (s, o)) - .collect::>(), - vec![("i", (0, 1)), ("âŃ¢", (1, 4)), ("j", (4, 5))] - ); - assert_eq!( - pretokenized - .get_splits(OffsetReferential::Normalized, OffsetType::Byte) - .into_iter() - .map(|(s, o, _)| (s, o)) - .collect::>(), - vec![("i", (0, 1)), ("âŃ¢", (1, 7)), ("j", (7, 8))] - ); - assert_eq!( - pretokenized - .get_splits(OffsetReferential::Original, OffsetType::Byte) - .into_iter() - .map(|(_, o, _)| &input[o.0..o.1]) - .collect::>(), - vec!["i", "⭢", "j"] - ); - } - - #[test] - fn processor_trims_offsets_pre_tokenized() { - // If user uses `is_pretokenized=True` we might have - // offsets that might begin at the start of the string but are - // NOT the first token. - let mut encoding = Encoding::new( - vec![0; 5], - vec![], - vec!["Ġl".into(), "ove".into(), "Ġl".into(), "ove".into()], - vec![], - vec![(0, 1), (1, 4), (0, 1), (1, 4)], - vec![], - vec![], - vec![], - AHashMap::new(), - ); - process_offsets(&mut encoding, true); - assert_eq!( - encoding, - Encoding::new( - vec![0; 5], - vec![], - vec!["Ġl".into(), "ove".into(), "Ġl".into(), "ove".into()], - vec![], - vec![(0, 1), (1, 4), (0, 1), (1, 4)], - vec![], - vec![], - vec![], - AHashMap::new(), - ) - ); - } - - #[test] - fn processor_trims_offsets() { - let start = Encoding::new( - vec![0; 5], - vec![], - vec![ - "Ġ".into(), - "ĠĠĠĠHelloĠĠ".into(), - "ĠĠHello".into(), - "HelloĠĠ".into(), - "ĠĠĠĠ".into(), - ], - vec![], - vec![(0, 1), (0, 11), (11, 18), (18, 25), (25, 29)], - vec![], - vec![], - vec![], - AHashMap::new(), - ); - let expected = Encoding::new( - vec![0; 5], - vec![0; 5], - vec![ - "Ġ".into(), - "ĠĠĠĠHelloĠĠ".into(), - "ĠĠHello".into(), - "HelloĠĠ".into(), - "ĠĠĠĠ".into(), - ], - vec![], - vec![(0, 0), (4, 9), (13, 18), (18, 23), (29, 29)], - vec![], - vec![], - vec![], - AHashMap::from_iter(vec![(0, 0..5)]), - ); - - let bytelevel = ByteLevel::default().trim_offsets(true); - assert_eq!( - expected, - bytelevel.process(start.clone(), None, false).unwrap() - ); - - let pair_expected = Encoding::new( - vec![0; 10], - vec![0, 0, 0, 0, 0, 1, 1, 1, 1, 1], - vec![ - "Ġ".into(), - "ĠĠĠĠHelloĠĠ".into(), - "ĠĠHello".into(), - "HelloĠĠ".into(), - "ĠĠĠĠ".into(), - "Ġ".into(), - "ĠĠĠĠHelloĠĠ".into(), - "ĠĠHello".into(), - "HelloĠĠ".into(), - "ĠĠĠĠ".into(), - ], - vec![], - vec![ - (0, 0), - (4, 9), - (13, 18), - (18, 23), - (29, 29), - (0, 0), - (4, 9), - (13, 18), - (18, 23), - (29, 29), - ], - vec![], - vec![], - vec![], - AHashMap::from_iter(vec![(0, 0..5), (1, 5..10)]), - ); - assert_eq!( - pair_expected, - bytelevel - .process(start.clone(), Some(start), false) - .unwrap() - ); - } - - #[test] - fn decode_unknown_characters() { - let byte_level = ByteLevel::default(); - assert_eq!( - byte_level - .decode_chain(vec![ - "Hello".into(), - "Ġthere".into(), - "Ġdear".into(), - "Ġfriend!".into(), - "Ġ".into(), - "[PA D]".into() - ]) - .unwrap(), - vec!["Hello there dear friend! [PA D]"] - ); - } - - /// Splits from the pipeline conversion of `byte_level`, with the raw text of each - /// range transformed to the byte-level alphabet so it's comparable with the legacy - /// oracle's output strings. - fn pipeline_splits(byte_level: ByteLevel, text: &str) -> Vec<(String, (usize, usize))> { - use std::convert::TryFrom; - let converted = crate::pipeline::PipelinePreTokenizer::try_from( - crate::PreTokenizerWrapper::ByteLevel(byte_level), - ) - .unwrap(); - let mut out = Vec::new(); - crate::pipeline::PreTokenizer::pre_tokenize(&converted, text, &mut out).unwrap(); - out.iter() - .map(|s| { - let transformed = text[s.range()] - .bytes() - .map(|b| BYTES_CHAR_LOOKUP[b as usize]) - .collect(); - (transformed, (s.start as usize, s.end as usize)) - }) - .collect() - } - - fn legacy_splits(byte_level: ByteLevel, text: &str) -> Vec<(String, (usize, usize))> { - let mut pre = PreTokenizedString::from(text); - byte_level.pre_tokenize(&mut pre).unwrap(); - pre.get_splits(OffsetReferential::Original, OffsetType::Byte) - .into_iter() - .map(|(s, o, _)| (s.to_string(), o)) - .collect() - } - - #[test] - fn pipeline_conversion_matches_legacy_splits() { - let byte_level = ByteLevel::default().add_prefix_space(false); - for text in [ - "Hello my friend, how is your day going?", - "Hello there\nHello there", - "Hello there dear", - " leading space", - "trailing space ", - "i⭢j", - "中文 text 123, mixed! 🤗 emoji", - "I'm can't we've they'll it's", - "tabs\tand\r\nnewlines", - "café über naïve", - "!!!???...", - "single", - ] { - assert_eq!( - pipeline_splits(byte_level, text), - legacy_splits(byte_level, text), - "diverged on {text:?}", - ); - } - } - - #[test] - fn pipeline_conversion_no_regex_is_identity_split() { - let byte_level = ByteLevel::default() - .add_prefix_space(false) - .use_regex(false); - let text = "Hello my friend, how is your day going?"; - assert_eq!( - pipeline_splits(byte_level, text), - legacy_splits(byte_level, text), - ); - } - - #[test] - fn pipeline_conversion_rejects_add_prefix_space() { - // The range-based pipeline can't prepend text; converting must fail loudly - // rather than silently produce different splits than the legacy path. - use std::convert::TryFrom; - let byte_level = ByteLevel::default().add_prefix_space(true); - assert!(crate::pipeline::PipelinePreTokenizer::try_from( - crate::PreTokenizerWrapper::ByteLevel(byte_level) - ) - .is_err()); - } - - #[test] - fn deserialization() { - // Before use_regex - let byte_level: ByteLevel = serde_json::from_str( - r#"{"type": "ByteLevel", "add_prefix_space": true, "trim_offsets": false}"#, - ) - .unwrap(); - assert!(byte_level.use_regex); - - // Loading works, new future BC test. - let byte_level: ByteLevel = serde_json::from_str( - r#"{"type": "ByteLevel", "add_prefix_space": true, "trim_offsets": false, "use_regex": true}"#, - ) - .unwrap(); - assert!(byte_level.use_regex); - - let byte_level: ByteLevel = serde_json::from_str( - r#"{"type": "ByteLevel", "add_prefix_space": true, "trim_offsets": false, "use_regex": false}"#, - ) - .unwrap(); - assert!(!byte_level.use_regex); - } -} diff --git a/tokenizers/tk-encode/src/pre_tokenizers/metaspace.rs b/tokenizers/tk-encode/src/pre_tokenizers/metaspace.rs index d821f1184..d4e3a8db7 100644 --- a/tokenizers/tk-encode/src/pre_tokenizers/metaspace.rs +++ b/tokenizers/tk-encode/src/pre_tokenizers/metaspace.rs @@ -1,4 +1,4 @@ -use crate::tokenizer::{Decoder, PreTokenizedString, PreTokenizer, Result, SplitDelimiterBehavior}; +use crate::tokenizer::{PreTokenizedString, PreTokenizer, Result, SplitDelimiterBehavior}; use serde::{de, Deserialize, Deserializer, Serialize}; /// Enum representing options for the metaspace prepending scheme. @@ -147,31 +147,6 @@ impl PreTokenizer for Metaspace { } } -impl Decoder for Metaspace { - fn decode_chain(&self, tokens: Vec) -> Result> { - Ok(tokens - .iter() - .enumerate() - .map(|(i, token)| { - token - .chars() - .flat_map(|c| { - if c == self.replacement { - if i == 0 && self.prepend_scheme != PrependScheme::Never { - None - } else { - Some(' ') - } - } else { - Some(c) - } - }) - .collect::() - }) - .collect()) - } -} - #[cfg(test)] mod tests { use regex::Regex; diff --git a/tokenizers/tk-encode/src/processors/bert.rs b/tokenizers/tk-encode/src/processors/bert.rs deleted file mode 100644 index a1cab8abd..000000000 --- a/tokenizers/tk-encode/src/processors/bert.rs +++ /dev/null @@ -1,296 +0,0 @@ -use crate::tokenizer::{Encoding, PostProcessor, Result}; -use ahash::AHashMap; -use serde::{Deserialize, Serialize}; -use std::iter::FromIterator; - -#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)] -#[serde(tag = "type")] -pub struct BertProcessing { - pub sep: (String, u32), - pub cls: (String, u32), -} - -impl Default for BertProcessing { - fn default() -> Self { - Self { - sep: ("[SEP]".into(), 102), - cls: ("[CLS]".into(), 101), - } - } -} - -impl BertProcessing { - pub fn new(sep: (String, u32), cls: (String, u32)) -> Self { - Self { sep, cls } - } - - pub fn get_sep_copy(&self) -> (String, u32) { - (self.sep.0.clone(), self.sep.1) - } - - pub fn get_cls_copy(&self) -> (String, u32) { - (self.cls.0.clone(), self.cls.1) - } -} - -#[derive(thiserror::Error, Debug)] -pub enum BertProcessorError { - #[error("encodings vector length must be either 1 or 2")] - InvalidEncodingsVecLength, -} - -impl PostProcessor for BertProcessing { - fn added_tokens(&self, is_pair: bool) -> usize { - if is_pair { - 3 - } else { - 2 - } - } - - fn process_encodings( - &self, - mut encodings: Vec, - add_special_tokens: bool, - ) -> Result> { - if !add_special_tokens { - return Ok(encodings); - } - - let encodings: Vec = encodings - .iter_mut() - .enumerate() - .map(|(i, encoding)| { - if i == 0 { - let ids = [&[self.cls.1], encoding.get_ids(), &[self.sep.1]].concat(); - let type_ids = [&[0], encoding.get_type_ids(), &[0]].concat(); - let tokens = [ - std::slice::from_ref(&self.cls.0), - encoding.get_tokens(), - std::slice::from_ref(&self.sep.0), - ] - .concat(); - let words = [&[None], encoding.get_word_ids(), &[None]].concat(); - let offsets = [&[(0, 0)], encoding.get_offsets(), &[(0, 0)]].concat(); - let special_tokens = - [&[1u32], &vec![0; encoding.get_ids().len()][..], &[1]].concat(); - let attention_mask = vec![1; ids.len()]; - - // For compatibility with `TemplateProcessing`, the sequence_ranges shouldn't contain - // the special tokens. - let sequence_ranges = AHashMap::from_iter(vec![(0, 1..ids.len() - 1)]); - Encoding::new( - ids, - type_ids, - tokens, - words, - offsets, - special_tokens, - attention_mask, - encoding - .take_overflowing() - .into_iter() - .map(|encoding| { - let ids = - [&[self.cls.1], encoding.get_ids(), &[self.sep.1]].concat(); - let type_ids = [&[0], encoding.get_type_ids(), &[0]].concat(); - let tokens = [ - std::slice::from_ref(&self.cls.0), - encoding.get_tokens(), - std::slice::from_ref(&self.sep.0), - ] - .concat(); - let words = [&[None], encoding.get_word_ids(), &[None]].concat(); - let offsets = - [&[(0, 0)], encoding.get_offsets(), &[(0, 0)]].concat(); - let special_tokens = - [&[1u32], &vec![0; encoding.get_ids().len()][..], &[1]] - .concat(); - let attention_mask = vec![1; ids.len()]; - - // For compatibility with `TemplateProcessing`, the sequence_ranges shouldn't - // contain the special tokens. - let sequence_ranges = - AHashMap::from_iter(vec![(0, 1..ids.len() - 1)]); - Encoding::new( - ids, - type_ids, - tokens, - words, - offsets, - special_tokens, - attention_mask, - vec![], - sequence_ranges, - ) - }) - .collect(), - sequence_ranges, - ) - } else { - let pair_ids = [encoding.get_ids(), &[self.sep.1]].concat(); - let pair_type_ids = [encoding.get_type_ids(), &[1]].concat(); - let pair_tokens = - [encoding.get_tokens(), std::slice::from_ref(&self.sep.0)].concat(); - let pair_words = [encoding.get_word_ids(), &[None]].concat(); - let pair_offsets = [encoding.get_offsets(), &[(0, 0)]].concat(); - let pair_special_tokens = - [&vec![0u32; encoding.get_type_ids().len()][..], &[1]].concat(); - let pair_attention_mask = vec![1; pair_ids.len()]; - - // For compatibility with `TemplateProcessing`, the sequence_ranges shouldn't contain - // the special tokens. - let pair_sequence_ranges = - AHashMap::from_iter(vec![(1, 0..pair_ids.len() - 1)]); - Encoding::new( - pair_ids, - pair_type_ids, - pair_tokens, - pair_words, - pair_offsets, - pair_special_tokens, - pair_attention_mask, - encoding - .take_overflowing() - .into_iter() - .map(|encoding| { - let pair_ids = [encoding.get_ids(), &[self.sep.1]].concat(); - let pair_type_ids = [encoding.get_type_ids(), &[1]].concat(); - let pair_tokens = - [encoding.get_tokens(), std::slice::from_ref(&self.sep.0)] - .concat(); - let pair_words = [encoding.get_word_ids(), &[None]].concat(); - let pair_offsets = [encoding.get_offsets(), &[(0, 0)]].concat(); - let pair_special_tokens = - [&vec![0u32; encoding.get_type_ids().len()][..], &[1]].concat(); - let pair_attention_mask = vec![1; pair_ids.len()]; - - // For compatibility with `TemplateProcessing`, the sequence_ranges - // shouldn't contain the special tokens. - let pair_sequence_ranges = - AHashMap::from_iter(vec![(1, 0..pair_ids.len() - 1)]); - Encoding::new( - pair_ids, - pair_type_ids, - pair_tokens, - pair_words, - pair_offsets, - pair_special_tokens, - pair_attention_mask, - vec![], - pair_sequence_ranges, - ) - }) - .collect(), - pair_sequence_ranges, - ) - } - }) - .collect(); - - Ok(encodings) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn serde() { - let bert = BertProcessing::default(); - let bert_r = r#"{"type":"BertProcessing","sep":["[SEP]",102],"cls":["[CLS]",101]}"#; - assert_eq!(serde_json::to_string(&bert).unwrap(), bert_r); - assert_eq!( - serde_json::from_str::(bert_r).unwrap(), - bert - ); - } - - #[test] - fn bert_processing() { - let processor = BertProcessing::default(); - assert_eq!(processor.added_tokens(false), 2); - assert_eq!(processor.added_tokens(true), 3); - - use crate::Token; - let encoding = Encoding::from_tokens( - vec![ - Token::new(12, "Hello".into(), (0, 5)), - Token::new(14, "there".into(), (6, 11)), - ], - 0, - ); - let pair = Encoding::from_tokens(vec![Token::new(15, "pair".into(), (0, 4))], 0); - let single_encoding = processor.process(encoding.clone(), None, true).unwrap(); - assert_eq!( - single_encoding, - Encoding::new( - vec![101, 12, 14, 102], - vec![0, 0, 0, 0], - vec![ - "[CLS]".into(), - "Hello".into(), - "there".into(), - "[SEP]".into() - ], - vec![None, None, None, None], - vec![(0, 0), (0, 5), (6, 11), (0, 0)], - vec![1, 0, 0, 1], - vec![1, 1, 1, 1], - vec![], - AHashMap::from_iter(vec![(0, 1..3)]), - ) - ); - assert_eq!(single_encoding.token_to_sequence(2), Some(0)); - assert_eq!(single_encoding.token_to_sequence(3), None); - let pair_encoding = processor - .process(encoding.clone(), Some(pair.clone()), true) - .unwrap(); - assert_eq!( - pair_encoding, - Encoding::new( - vec![101, 12, 14, 102, 15, 102], - vec![0, 0, 0, 0, 1, 1], - vec![ - "[CLS]".into(), - "Hello".into(), - "there".into(), - "[SEP]".into(), - "pair".into(), - "[SEP]".into() - ], - vec![None, None, None, None, None, None], - vec![(0, 0), (0, 5), (6, 11), (0, 0), (0, 4), (0, 0)], - vec![1, 0, 0, 1, 0, 1], - vec![1, 1, 1, 1, 1, 1], - vec![], - AHashMap::from_iter(vec![(0, 1..3), (1, 4..5)]), - ) - ); - assert_eq!(pair_encoding.token_to_sequence(2), Some(0)); - assert_eq!(pair_encoding.token_to_sequence(3), None); - assert_eq!(pair_encoding.token_to_sequence(4), Some(1)); - assert_eq!(pair_encoding.token_to_sequence(5), None); - - // No special tokens - let pair_encoding = processor.process(encoding, Some(pair), false).unwrap(); - assert_eq!( - pair_encoding, - Encoding::new( - vec![12, 14, 15], - vec![0, 0, 1], - vec!["Hello".into(), "there".into(), "pair".into(),], - vec![None, None, None], - vec![(0, 5), (6, 11), (0, 4)], - vec![0, 0, 0], - vec![1, 1, 1], - vec![], - AHashMap::from_iter(vec![(0, 0..2), (1, 2..3)]), - ) - ); - assert_eq!(pair_encoding.token_to_sequence(0), Some(0)); - assert_eq!(pair_encoding.token_to_sequence(1), Some(0)); - assert_eq!(pair_encoding.token_to_sequence(2), Some(1)); - } -} diff --git a/tokenizers/tk-encode/src/processors/mod.rs b/tokenizers/tk-encode/src/processors/mod.rs deleted file mode 100644 index 869cc6891..000000000 --- a/tokenizers/tk-encode/src/processors/mod.rs +++ /dev/null @@ -1,128 +0,0 @@ -pub mod bert; -pub mod roberta; -pub mod sequence; -pub mod template; - -// Re-export these as processors -pub use super::pre_tokenizers::byte_level; - -use serde::{Deserialize, Serialize}; - -use crate::pre_tokenizers::byte_level::ByteLevel; -use crate::processors::bert::BertProcessing; -use crate::processors::roberta::RobertaProcessing; -use crate::processors::sequence::Sequence; -use crate::processors::template::TemplateProcessing; -use crate::{Encoding, PostProcessor, Result}; - -#[derive(Serialize, Deserialize, PartialEq, Debug, Clone, Eq)] -#[serde(untagged)] -pub enum PostProcessorWrapper { - // Roberta must be before Bert for deserialization (serde does not validate tags) - Roberta(RobertaProcessing), - Bert(BertProcessing), - ByteLevel(ByteLevel), - Template(TemplateProcessing), - Sequence(Sequence), -} - -impl PostProcessor for PostProcessorWrapper { - fn added_tokens(&self, is_pair: bool) -> usize { - match self { - Self::Bert(bert) => bert.added_tokens(is_pair), - Self::ByteLevel(bl) => bl.added_tokens(is_pair), - Self::Roberta(roberta) => roberta.added_tokens(is_pair), - Self::Template(template) => template.added_tokens(is_pair), - Self::Sequence(bl) => bl.added_tokens(is_pair), - } - } - - fn process_encodings( - &self, - encodings: Vec, - add_special_tokens: bool, - ) -> Result> { - match self { - Self::Bert(bert) => bert.process_encodings(encodings, add_special_tokens), - Self::ByteLevel(bl) => bl.process_encodings(encodings, add_special_tokens), - Self::Roberta(roberta) => roberta.process_encodings(encodings, add_special_tokens), - Self::Template(template) => template.process_encodings(encodings, add_special_tokens), - Self::Sequence(bl) => bl.process_encodings(encodings, add_special_tokens), - } - } -} - -impl_enum_from!(BertProcessing, PostProcessorWrapper, Bert); -impl_enum_from!(ByteLevel, PostProcessorWrapper, ByteLevel); -impl_enum_from!(RobertaProcessing, PostProcessorWrapper, Roberta); -impl_enum_from!(TemplateProcessing, PostProcessorWrapper, Template); -impl_enum_from!(Sequence, PostProcessorWrapper, Sequence); - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn deserialize_bert_roberta_correctly() { - let roberta = RobertaProcessing::default(); - let roberta_r = r#"{ - "type":"RobertaProcessing", - "sep":["",2], - "cls":["",0], - "trim_offsets":true, - "add_prefix_space":true - }"# - .replace(char::is_whitespace, ""); - assert_eq!(serde_json::to_string(&roberta).unwrap(), roberta_r); - assert_eq!( - serde_json::from_str::(&roberta_r).unwrap(), - PostProcessorWrapper::Roberta(roberta) - ); - - let bert = BertProcessing::default(); - let bert_r = r#"{"type":"BertProcessing","sep":["[SEP]",102],"cls":["[CLS]",101]}"#; - assert_eq!(serde_json::to_string(&bert).unwrap(), bert_r); - assert_eq!( - serde_json::from_str::(bert_r).unwrap(), - PostProcessorWrapper::Bert(bert) - ); - } - - #[test] - fn post_processor_deserialization_no_type() { - let json = r#"{"add_prefix_space": true, "trim_offsets": false, "use_regex": false}"#; - let reconstructed = serde_json::from_str::(json); - match reconstructed { - Err(err) => assert_eq!( - err.to_string(), - "data did not match any variant of untagged enum PostProcessorWrapper" - ), - _ => panic!("Expected an error here"), - } - - let json = r#"{"sep":["[SEP]",102],"cls":["[CLS]",101]}"#; - let reconstructed = serde_json::from_str::(json); - assert!(matches!( - reconstructed.unwrap(), - PostProcessorWrapper::Bert(_) - )); - - let json = - r#"{"sep":["",2], "cls":["",0], "trim_offsets":true, "add_prefix_space":true}"#; - let reconstructed = serde_json::from_str::(json); - assert!(matches!( - reconstructed.unwrap(), - PostProcessorWrapper::Roberta(_) - )); - - let json = r#"{"type":"RobertaProcessing", "sep":["",2] }"#; - let reconstructed = serde_json::from_str::(json); - match reconstructed { - Err(err) => assert_eq!( - err.to_string(), - "data did not match any variant of untagged enum PostProcessorWrapper" - ), - _ => panic!("Expected an error here"), - } - } -} diff --git a/tokenizers/tk-encode/src/processors/roberta.rs b/tokenizers/tk-encode/src/processors/roberta.rs deleted file mode 100644 index f2a47a9d3..000000000 --- a/tokenizers/tk-encode/src/processors/roberta.rs +++ /dev/null @@ -1,341 +0,0 @@ -use crate::processors::byte_level::process_offsets; -use crate::tokenizer::{Encoding, PostProcessor, Result}; -use ahash::AHashMap; -use serde::{Deserialize, Serialize}; -use std::iter::FromIterator; - -#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)] -#[serde(tag = "type")] -pub struct RobertaProcessing { - pub sep: (String, u32), - pub cls: (String, u32), - pub trim_offsets: bool, - pub add_prefix_space: bool, -} - -impl Default for RobertaProcessing { - fn default() -> Self { - Self { - sep: ("".into(), 2), - cls: ("".into(), 0), - trim_offsets: true, - add_prefix_space: true, - } - } -} - -impl RobertaProcessing { - pub fn new(sep: (String, u32), cls: (String, u32)) -> Self { - Self { - sep, - cls, - ..Default::default() - } - } - - #[must_use] - pub fn trim_offsets(mut self, v: bool) -> Self { - self.trim_offsets = v; - self - } - - #[must_use] - pub fn add_prefix_space(mut self, v: bool) -> Self { - self.add_prefix_space = v; - self - } - - pub fn get_sep_copy(&self) -> (String, u32) { - (self.sep.0.clone(), self.sep.1) - } - - pub fn get_cls_copy(&self) -> (String, u32) { - (self.cls.0.clone(), self.cls.1) - } -} - -impl PostProcessor for RobertaProcessing { - fn added_tokens(&self, is_pair: bool) -> usize { - if is_pair { - 4 - } else { - 2 - } - } - - fn process_encodings( - &self, - mut encodings: Vec, - add_special_tokens: bool, - ) -> Result> { - if self.trim_offsets { - for encoding in encodings.iter_mut() { - process_offsets(encoding, self.add_prefix_space); - encoding - .get_overflowing_mut() - .iter_mut() - .for_each(|encoding| process_offsets(encoding, self.add_prefix_space)); - } - } - - // Roberta is weird, and every encoding is type_id=0. - encodings - .iter_mut() - .for_each(|encoding| encoding.set_type_ids(vec![0; encoding.len()])); - - if !add_special_tokens { - return Ok(encodings); - } - - let encodings: Vec = encodings - .iter_mut() - .enumerate() - .map(|(i, encoding)| { - if i == 0 { - let ids = [&[self.cls.1], encoding.get_ids(), &[self.sep.1]].concat(); - let type_ids = [&[0], encoding.get_type_ids(), &[0]].concat(); - let tokens = [ - std::slice::from_ref(&self.cls.0), - encoding.get_tokens(), - std::slice::from_ref(&self.sep.0), - ] - .concat(); - let words = [&[None], encoding.get_word_ids(), &[None]].concat(); - let offsets = [&[(0, 0)], encoding.get_offsets(), &[(0, 0)]].concat(); - let special_tokens = - [&[1u32], &vec![0; encoding.get_ids().len()][..], &[1]].concat(); - let attention_mask = vec![1; ids.len()]; - - // For compatibility with `TemplateProcessing`, the sequence_ranges shouldn't contain - // the special tokens. - let sequence_ranges = AHashMap::from_iter(vec![(0, 1..ids.len() - 1)]); - Encoding::new( - ids, - type_ids, - tokens, - words, - offsets, - special_tokens, - attention_mask, - encoding - .take_overflowing() - .into_iter() - .map(|encoding| { - let ids = - [&[self.cls.1], encoding.get_ids(), &[self.sep.1]].concat(); - let type_ids = vec![0; encoding.get_ids().len() + 2]; - let tokens = [ - std::slice::from_ref(&self.cls.0), - encoding.get_tokens(), - std::slice::from_ref(&self.sep.0), - ] - .concat(); - let words = [&[None], encoding.get_word_ids(), &[None]].concat(); - let offsets = - [&[(0, 0)], encoding.get_offsets(), &[(0, 0)]].concat(); - let special_tokens = - [&[1u32], &vec![0; encoding.get_ids().len()][..], &[1]] - .concat(); - let attention_mask = vec![1; ids.len()]; - - // For compatibility with `TemplateProcessing`, the sequence_ranges shouldn't - // contain the special tokens. - let sequence_ranges = - AHashMap::from_iter(vec![(0, 1..ids.len() - 1)]); - Encoding::new( - ids, - type_ids, - tokens, - words, - offsets, - special_tokens, - attention_mask, - vec![], - sequence_ranges, - ) - }) - .collect(), - sequence_ranges, - ) - } else { - let pair_ids = [&[self.sep.1], encoding.get_ids(), &[self.sep.1]].concat(); - let pair_type_ids = vec![0; encoding.get_ids().len() + 2]; - let pair_tokens = [ - std::slice::from_ref(&self.sep.0), - encoding.get_tokens(), - std::slice::from_ref(&self.sep.0), - ] - .concat(); - let pair_words = [&[None], encoding.get_word_ids(), &[None]].concat(); - let pair_offsets = [&[(0, 0)], encoding.get_offsets(), &[(0, 0)]].concat(); - let pair_special_tokens = - [&[1], &vec![0u32; encoding.get_type_ids().len()][..], &[1]].concat(); - let pair_attention_mask = vec![1; pair_ids.len()]; - - // For compatibility with `TemplateProcessing`, the sequence_ranges shouldn't contain - // the special tokens. - let pair_sequence_ranges = - AHashMap::from_iter(vec![(1, 1..pair_ids.len() - 1)]); - Encoding::new( - pair_ids, - pair_type_ids, - pair_tokens, - pair_words, - pair_offsets, - pair_special_tokens, - pair_attention_mask, - encoding - .take_overflowing() - .into_iter() - .map(|encoding| { - let pair_ids = - [&[self.sep.1], encoding.get_ids(), &[self.sep.1]].concat(); - let pair_type_ids = vec![0; encoding.get_ids().len() + 2]; - let pair_tokens = [ - std::slice::from_ref(&self.sep.0), - encoding.get_tokens(), - std::slice::from_ref(&self.sep.0), - ] - .concat(); - let pair_words = - [&[None], encoding.get_word_ids(), &[None]].concat(); - let pair_offsets = - [&[(0, 0)], encoding.get_offsets(), &[(0, 0)]].concat(); - let pair_special_tokens = - [&[1], &vec![0u32; encoding.get_type_ids().len()][..], &[1]] - .concat(); - let pair_attention_mask = vec![1; pair_ids.len()]; - - // For compatibility with `TemplateProcessing`, the sequence_ranges - // shouldn't contain the special tokens. - let pair_sequence_ranges = - AHashMap::from_iter(vec![(1, 1..pair_ids.len() - 1)]); - Encoding::new( - pair_ids, - pair_type_ids, - pair_tokens, - pair_words, - pair_offsets, - pair_special_tokens, - pair_attention_mask, - vec![], - pair_sequence_ranges, - ) - }) - .collect(), - pair_sequence_ranges, - ) - } - }) - .collect(); - - Ok(encodings) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn serde() { - let roberta = RobertaProcessing::default(); - let roberta_r = r#"{ - "type":"RobertaProcessing", - "sep":["",2], - "cls":["",0], - "trim_offsets":true, - "add_prefix_space":true - }"# - .replace(char::is_whitespace, ""); - assert_eq!(serde_json::to_string(&roberta).unwrap(), roberta_r); - assert_eq!( - serde_json::from_str::(&roberta_r).unwrap(), - roberta - ); - } - - #[test] - fn roberta_processing() { - let processor = RobertaProcessing::default(); - assert_eq!(processor.added_tokens(false), 2); - assert_eq!(processor.added_tokens(true), 4); - - use crate::Token; - let encoding = Encoding::from_tokens( - vec![ - Token::new(12, "Hello".into(), (0, 5)), - Token::new(14, "there".into(), (6, 11)), - ], - 0, - ); - let pair = Encoding::from_tokens(vec![Token::new(15, "pair".into(), (0, 4))], 0); - let single_encoding = processor.process(encoding.clone(), None, true).unwrap(); - assert_eq!( - single_encoding, - Encoding::new( - vec![0, 12, 14, 2], - vec![0, 0, 0, 0], - vec!["".into(), "Hello".into(), "there".into(), "".into()], - vec![None, None, None, None], - vec![(0, 0), (0, 5), (6, 11), (0, 0)], - vec![1, 0, 0, 1], - vec![1, 1, 1, 1], - vec![], - AHashMap::from_iter(vec![(0, 1..3)]), - ) - ); - assert_eq!(single_encoding.token_to_sequence(2), Some(0)); - assert_eq!(single_encoding.token_to_sequence(3), None); - let pair_encoding = processor - .process(encoding.clone(), Some(pair.clone()), true) - .unwrap(); - assert_eq!( - pair_encoding, - Encoding::new( - vec![0, 12, 14, 2, 2, 15, 2], - vec![0, 0, 0, 0, 0, 0, 0], - vec![ - "".into(), - "Hello".into(), - "there".into(), - "".into(), - "".into(), - "pair".into(), - "".into() - ], - vec![None, None, None, None, None, None, None], - vec![(0, 0), (0, 5), (6, 11), (0, 0), (0, 0), (0, 4), (0, 0)], - vec![1, 0, 0, 1, 1, 0, 1], - vec![1, 1, 1, 1, 1, 1, 1], - vec![], - AHashMap::from_iter(vec![(0, 1..3), (1, 5..6)]), - ) - ); - assert_eq!(pair_encoding.token_to_sequence(2), Some(0)); - assert_eq!(pair_encoding.token_to_sequence(3), None); - assert_eq!(pair_encoding.token_to_sequence(4), None); - assert_eq!(pair_encoding.token_to_sequence(5), Some(1)); - assert_eq!(pair_encoding.token_to_sequence(6), None); - - // No special tokens - let pair_encoding = processor.process(encoding, Some(pair), false).unwrap(); - assert_eq!( - pair_encoding, - Encoding::new( - vec![12, 14, 15], - vec![0, 0, 0], - vec!["Hello".into(), "there".into(), "pair".into(),], - vec![None, None, None], - vec![(0, 5), (6, 11), (0, 4)], - vec![0, 0, 0], - vec![1, 1, 1], - vec![], - AHashMap::from_iter(vec![(0, 0..2), (1, 2..3)]), - ) - ); - assert_eq!(pair_encoding.token_to_sequence(0), Some(0)); - assert_eq!(pair_encoding.token_to_sequence(1), Some(0)); - assert_eq!(pair_encoding.token_to_sequence(2), Some(1)); - } -} diff --git a/tokenizers/tk-encode/src/processors/sequence.rs b/tokenizers/tk-encode/src/processors/sequence.rs deleted file mode 100644 index f44cf54ac..000000000 --- a/tokenizers/tk-encode/src/processors/sequence.rs +++ /dev/null @@ -1,172 +0,0 @@ -use crate::processors::PostProcessorWrapper; -use crate::tokenizer::{Encoding, PostProcessor, Result}; -use crate::utils::macro_rules_attribute; -use serde::{Deserialize, Serialize}; - -#[derive(Clone, Debug, PartialEq, Eq)] -#[macro_rules_attribute(impl_serde_type!)] -pub struct Sequence { - processors: Vec, -} - -impl Sequence { - pub fn new(processors: Vec) -> Self { - Self { processors } - } - - pub fn get(&self, index: usize) -> Option<&PostProcessorWrapper> { - self.processors.get(index) - } - - pub fn get_mut(&mut self, index: usize) -> Option<&mut PostProcessorWrapper> { - self.processors.get_mut(index) - } - - pub fn set_mut(&mut self, index: usize, post_proc: PostProcessorWrapper) { - self.processors[index] = post_proc; - } -} - -impl AsRef<[PostProcessorWrapper]> for Sequence { - fn as_ref(&self) -> &[PostProcessorWrapper] { - &self.processors - } -} - -impl AsMut<[PostProcessorWrapper]> for Sequence { - fn as_mut(&mut self) -> &mut [PostProcessorWrapper] { - &mut self.processors - } -} - -impl IntoIterator for Sequence { - type Item = PostProcessorWrapper; - type IntoIter = std::vec::IntoIter; - - fn into_iter(self) -> Self::IntoIter { - self.processors.into_iter() - } -} - -impl PostProcessor for Sequence { - fn added_tokens(&self, is_pair: bool) -> usize { - self.processors - .iter() - .map(|p| p.added_tokens(is_pair)) - .sum::() - } - - fn process_encodings( - &self, - mut encodings: Vec, - add_special_tokens: bool, - ) -> Result> { - for processor in &self.processors { - encodings = processor.process_encodings(encodings, add_special_tokens)?; - } - Ok(encodings) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::processors::{ByteLevel, PostProcessorWrapper}; - use crate::tokenizer::{Encoding, PostProcessor}; - use ahash::AHashMap; - use std::iter::FromIterator; - - #[test] - fn process_chain() { - let start = Encoding::new( - vec![0; 5], - vec![0; 5], - vec![ - "Ġ".into(), - "ĠĠĠĠHelloĠĠ".into(), - "ĠĠHello".into(), - "HelloĠĠ".into(), - "ĠĠĠĠ".into(), - ], - vec![], - vec![(0, 1), (0, 11), (11, 18), (18, 25), (25, 29)], - vec![], - vec![], - vec![], - AHashMap::new(), - ); - - let bytelevel = ByteLevel::default().trim_offsets(true); - let sequence = Sequence::new(vec![PostProcessorWrapper::ByteLevel(bytelevel)]); - let expected = Encoding::new( - vec![0; 5], - vec![0; 5], - vec![ - "Ġ".into(), - "ĠĠĠĠHelloĠĠ".into(), - "ĠĠHello".into(), - "HelloĠĠ".into(), - "ĠĠĠĠ".into(), - ], - vec![], - vec![(0, 0), (4, 9), (13, 18), (18, 23), (29, 29)], - vec![], - vec![], - vec![], - AHashMap::from_iter(vec![(0, 0..5)]), - ); - - assert_eq!( - expected, - bytelevel.process(start.clone(), None, false).unwrap() - ); - assert_eq!( - expected, - sequence.process(start.clone(), None, false).unwrap() - ); - - let pair_expected = Encoding::new( - vec![0; 10], - vec![0, 0, 0, 0, 0, 1, 1, 1, 1, 1], - vec![ - "Ġ".into(), - "ĠĠĠĠHelloĠĠ".into(), - "ĠĠHello".into(), - "HelloĠĠ".into(), - "ĠĠĠĠ".into(), - "Ġ".into(), - "ĠĠĠĠHelloĠĠ".into(), - "ĠĠHello".into(), - "HelloĠĠ".into(), - "ĠĠĠĠ".into(), - ], - vec![], - vec![ - (0, 0), - (4, 9), - (13, 18), - (18, 23), - (29, 29), - (0, 0), - (4, 9), - (13, 18), - (18, 23), - (29, 29), - ], - vec![], - vec![], - vec![], - AHashMap::from_iter(vec![(0, 0..5), (1, 5..10)]), - ); - assert_eq!( - pair_expected, - bytelevel - .process(start.clone(), Some(start.clone()), false) - .unwrap() - ); - assert_eq!( - pair_expected, - sequence.process(start.clone(), Some(start), false).unwrap() - ); - } -} diff --git a/tokenizers/tk-encode/src/processors/template.rs b/tokenizers/tk-encode/src/processors/template.rs deleted file mode 100644 index 4ca93bc6a..000000000 --- a/tokenizers/tk-encode/src/processors/template.rs +++ /dev/null @@ -1,1156 +0,0 @@ -//! # Template Processing -//! -//! Provides a way to specify templates in order to add the special tokens to each -//! input sequence as relevant. -//! -//! ## Example -//! -//! Let's take `BERT` tokenizer as an example. It uses two special tokens, used to -//! delimitate each sequence. `[CLS]` is always used at the beginning of the first -//! sequence, and `[SEP]` is added at the end of both the first, and the pair -//! sequences. The final result looks like this: -//! - Single sequence: `[CLS] Hello there [SEP]` -//! - Pair sequences: `[CLS] My name is Anthony [SEP] What is my name? [SEP]` -//! -//! With the type ids as following: -//! ```markdown -//! [CLS] ... [SEP] ... [SEP] -//! 0 0 0 1 1 -//! ``` -//! -//! So, we can define a [`TemplateProcessing`] that will achieve this result: -//! ``` -//! # use tk_encode::processors::template::TemplateProcessing; -//! let template = TemplateProcessing::builder() -//! // The template when we only have a single sequence: -//! .try_single(vec!["[CLS]", "$0", "[SEP]"]).unwrap() -//! // Same as: -//! .try_single("[CLS] $0 [SEP]").unwrap() -//! -//! // The template when we have both sequences: -//! .try_pair(vec!["[CLS]:0", "$A:0", "[SEP]:0", "$B:1", "[SEP]:1"]).unwrap() -//! // Same as: -//! .try_pair("[CLS]:0 $A:0 [SEP]:0 $B:1 [SEP]:1").unwrap() -//! // Or: -//! .try_pair("[CLS] $0 [SEP] $B:1 [SEP]:1").unwrap() -//! -//! // The list of special tokens used by each sequences -//! .special_tokens(vec![("[CLS]", 1), ("[SEP]", 0)]) -//! .build() -//! .unwrap(); -//! ``` -//! -//! In this example, each input sequence is identified using a `$` construct. This identifier -//! lets us specify each input sequence, and the type_id to use. When nothing is specified, -//! it uses the default values. Here are the different ways to specify it: -//! - Specifying the sequence, with default `type_id == 0`: `$A` or `$B` -//! - Specifying the `type_id` with default `sequence == A`: `$0`, `$1`, `$2`, ... -//! - Specifying both: `$A:0`, `$B:1`, ... -//! -//! The same construct is used for special tokens: `(:)?`. -//! -//! **Warning**: You must ensure that you are giving the correct tokens/ids as these will -//! be added to the `Encoding` without any further check. If the given ids correspond to -//! something totally different in a `Tokenizer` using this `PostProcessor`, it might lead -//! to unexpected results. -//! -//! [`TemplateProcessing`]: struct.TemplateProcessing.html -//! -use crate::{Encoding, PostProcessor, Result}; -use ahash::{AHashMap, AHashSet}; -use itertools::Itertools; -use serde::{Deserialize, Serialize}; -use std::convert::{TryFrom, TryInto}; -use std::result::Result as StdResult; - -/// Represents any sequences received as input of the PostProcessor -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Eq)] -pub enum Sequence { - /// This is the first sequence, the one that is always specified - A, - /// This is the pair sequence, that is optional - B, -} - -/// Represents the different kind of pieces that constitute a template. -/// It can be either the input sequence or a [`SpecialToken`]: -/// -/// - The `Sequence` has an associated `type_id` which is used by default -/// for any token inside this sequence. The `Sequence` corresponds to one -/// of the input sequence given as input of the `PostProcessor`. -/// -/// - The `SpecialToken` has an associated `id`. It corresponds to a [`SpecialToken`]. -/// -/// The easiest way to build a `Piece` is actually by converting it from a string: -/// ``` -/// # use tk_encode::processors::template::Piece; -/// # use std::convert::TryFrom; -/// let sequence_with_type_id_0 = Piece::try_from("$0").unwrap(); -/// let sequence_with_type_id_1 = Piece::try_from("$1").unwrap(); -/// let special_token_cls = Piece::try_from("[CLS]").unwrap(); -/// ``` -/// -/// [`SpecialToken`]: struct.SpecialToken.html -/// -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Eq)] -pub enum Piece { - Sequence { id: Sequence, type_id: u32 }, - SpecialToken { id: String, type_id: u32 }, -} - -impl Piece { - fn extract_id(s: &str) -> Option { - if s.starts_with('$') { - let rest = &s['$'.len_utf8()..]; - - // If the id is just `$`, we use 0 as type_id, and Sequence A - match rest { - "" => Some(Self::Sequence { - id: Sequence::A, - type_id: 0, - }), - "A" | "a" => Some(Self::Sequence { - id: Sequence::A, - type_id: 0, - }), - "B" | "b" => Some(Self::Sequence { - id: Sequence::B, - type_id: 0, - }), - n => { - if let Ok(type_id) = n.parse::() { - Some(Self::Sequence { - id: Sequence::A, - type_id, - }) - } else { - None - } - } - } - } else { - Some(Self::SpecialToken { - id: s.to_owned(), - type_id: 0, - }) - } - } - - fn with_type_id(self, type_id: u32) -> Self { - match self { - Self::Sequence { id, .. } => Self::Sequence { id, type_id }, - Self::SpecialToken { id, .. } => Self::SpecialToken { id, type_id }, - } - } -} - -impl TryFrom for Piece { - type Error = String; - - fn try_from(s: String) -> StdResult { - let parts = s.split(':').collect::>(); - - let err = || format!("Cannot build Piece from string \"{s}\""); - match parts.as_slice() { - [id, type_id] => { - let type_id: u32 = type_id.parse().map_err(|_| err())?; - let piece = Self::extract_id(id).ok_or_else(err)?; - Ok(piece.with_type_id(type_id)) - } - [id] => Self::extract_id(id).ok_or_else(err), - _ => Err(err()), - } - } -} - -impl TryFrom<&str> for Piece { - type Error = String; - - fn try_from(s: &str) -> StdResult { - Piece::try_from(s.to_owned()) - } -} - -/// Represents a bunch of tokens to be used in a template. -/// Usually, special tokens have only one associated id/token but in -/// some cases, it might be interesting to have multiple ids/tokens. -/// -/// # Examples -/// ``` -/// # use tk_encode::processors::template::SpecialToken; -/// // Simple cases, where a single id/token is necessary: -/// let cls = SpecialToken::from(("[CLS]", 1)); -/// let sep = SpecialToken::from((0, "[SEP]")); // The order in the tuple is not important -/// -/// // More complex case with multiple values: -/// let complex = SpecialToken::new( -/// "A complex special token:".into(), -/// vec![0, 1, 2, 3, 4], -/// vec!["A".into(), "complex".into(), "special".into(), "token".into(), ":".into()] -/// ).unwrap(); -/// ``` -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Eq)] -pub struct SpecialToken { - /// A unique id used to identify this SpecialToken in the template - id: String, - /// The list of associated ids - ids: Vec, - /// The list of associated tokens - tokens: Vec, -} - -impl From<(String, u32)> for SpecialToken { - fn from(v: (String, u32)) -> Self { - Self { - id: v.0.clone(), - ids: vec![v.1], - tokens: vec![v.0], - } - } -} -impl From<(&str, u32)> for SpecialToken { - fn from(v: (&str, u32)) -> Self { - Self::from((v.0.to_owned(), v.1)) - } -} -impl From<(u32, String)> for SpecialToken { - fn from(v: (u32, String)) -> Self { - Self::from((v.1, v.0)) - } -} -impl From<(u32, &str)> for SpecialToken { - fn from(v: (u32, &str)) -> Self { - Self::from((v.1.to_owned(), v.0)) - } -} - -impl SpecialToken { - pub fn new(id: String, ids: Vec, tokens: Vec) -> Result { - if ids.len() != tokens.len() { - Err("SpecialToken: ids and tokens must be of the same length".into()) - } else { - Ok(Self { id, ids, tokens }) - } - } -} - -/// A Template represents a Vec<[`Piece`]>. -/// -/// We can easily build one as follows -/// ``` -/// # use tk_encode::processors::template::Template; -/// # use std::convert::TryFrom; -/// // By providing a `String` or `&str`, we just split on whitespaces: -/// let template = Template::try_from("[CLS] $0 [SEP]").unwrap(); -/// -/// // By providing pieces directly: -/// let template = Template::try_from(vec!["[CLS]", "$0", "[SEP]"]).unwrap(); -/// ``` -/// Both of these methods give the same result. -/// -/// [`Piece`]: enum.Piece.html -/// -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Eq)] -#[serde(transparent)] -pub struct Template(Vec); - -impl TryFrom> for Template -where - T: TryInto, -{ - type Error = String; - - fn try_from(v: Vec) -> StdResult { - Ok(Self( - v.into_iter() - .map(|p| p.try_into()) - .collect::, Self::Error>>()?, - )) - } -} - -impl TryFrom for Template { - type Error = String; - - fn try_from(s: String) -> StdResult { - Self::try_from(s.as_ref()) - } -} - -impl TryFrom<&str> for Template { - type Error = String; - - fn try_from(s: &str) -> StdResult { - Self::try_from(s.split(' ').collect::>()) - } -} - -/// A bunch of [`SpecialToken`] represented by their ID. -/// Internally, `Tokens` is a `HashMap` and can be built -/// from a HashMap or a Vec<[`SpecialToken`]>. -/// -/// [`SpecialToken`]: struct.SpecialToken.html -#[derive(Debug, Clone, PartialEq, Default, Serialize, Deserialize, Eq)] -#[serde(transparent)] -pub struct Tokens( - #[serde(serialize_with = "crate::utils::ordered_map")] pub AHashMap, -); - -impl> From> for Tokens { - fn from(v: Vec) -> Self { - Self( - v.into_iter() - .map(|t| { - let token: SpecialToken = t.into(); - (token.id.clone(), token) - }) - .collect(), - ) - } -} - -impl From> for Tokens { - fn from(v: AHashMap) -> Self { - Self(v) - } -} - -/// This PostProcessor takes care of processing each input `Encoding` by applying -/// the corresponding template, before merging them in the final Encoding. -/// -/// A `Template` is actually a sequence of `Piece` that will be -/// concatenated together in the given order. Each `Piece` represents either -/// one of the input `Encoding` or a `SpecialToken`. -/// -/// ## Example -/// ``` -/// # use tk_encode::processors::template::TemplateProcessing; -/// let template = TemplateProcessing::builder() -/// .try_single("[CLS] $A [SEP]").unwrap() -/// .try_pair("[CLS] $A [SEP] $B:1 [SEP]:1").unwrap() -/// .special_tokens(vec![("[CLS]", 1), ("[SEP]", 0)]) -/// .build() -/// .unwrap(); -/// ``` -/// -#[derive(Debug, Clone, PartialEq, Builder, Serialize, Deserialize, Eq)] -#[serde(tag = "type", from = "TemplateProcessingDeserializer")] -#[builder(build_fn(validate = "Self::validate"))] -pub struct TemplateProcessing { - #[builder(try_setter, default = "\"$0\".try_into().unwrap()")] - pub single: Template, - #[builder(try_setter, default = "\"$A:0 $B:1\".try_into().unwrap()")] - pair: Template, - #[builder(setter(skip), default = "self.default_added(true)")] - #[serde(skip)] - added_single: usize, - #[builder(setter(skip), default = "self.default_added(false)")] - #[serde(skip)] - added_pair: usize, - #[builder(setter(into), default)] - special_tokens: Tokens, -} - -impl TemplateProcessing { - // Getter for `single` - pub fn get_single(&self) -> String { - format!("{:?}", self.single) - } - - // Setter for `single` - pub fn set_single(&mut self, single: Template) { - self.single = single; - } - - // Getter for `pair` - pub fn get_pair(&self) -> &Template { - &self.pair - } - - // Setter for `pair` - pub fn set_pair(&mut self, pair: Template) { - self.pair = pair; - } - - // Getter for `added_single` - pub fn get_added_single(&self) -> usize { - self.added_single - } - - // Setter for `added_single` - pub fn set_added_single(&mut self, added_single: usize) { - self.added_single = added_single; - } - - // Getter for `added_pair` - pub fn get_added_pair(&self) -> usize { - self.added_pair - } - - // Setter for `added_pair` - pub fn set_added_pair(&mut self, added_pair: usize) { - self.added_pair = added_pair; - } - - // Getter for `special_tokens` - pub fn get_special_tokens(&self) -> &Tokens { - &self.special_tokens - } - - // Setter for `special_tokens` - pub fn set_special_tokens(&mut self, special_tokens: Tokens) { - self.special_tokens = special_tokens; - } -} - -impl From<&str> for TemplateProcessingBuilderError { - fn from(e: &str) -> Self { - e.to_string().into() - } -} - -impl PartialEq for TemplateProcessingBuilderError { - fn eq(&self, other: &Self) -> bool { - self.to_string() == other.to_string() - } -} - -/// We use this custom deserializer to provided the values for `added_single` -/// and `added_pair` during deserialization, while not having to serialize them -#[doc(hidden)] -#[derive(Deserialize)] -#[serde(tag = "type")] -struct TemplateProcessingDeserializer { - single: Template, - pair: Template, - special_tokens: Tokens, -} -impl From for TemplateProcessing { - fn from(t: TemplateProcessingDeserializer) -> Self { - let added_single = count_added(&t.single, Some(&t.special_tokens)); - let added_pair = count_added(&t.pair, Some(&t.special_tokens)); - Self { - single: t.single, - pair: t.pair, - added_single, - added_pair, - special_tokens: t.special_tokens, - } - } -} - -/// Count the number of added tokens in the given template -fn count_added(container: &Template, special_tokens: Option<&Tokens>) -> usize { - container - .0 - .iter() - .map(|p| match p { - Piece::Sequence { .. } => 0, - Piece::SpecialToken { id, .. } => { - special_tokens.map_or(0, |spt| spt.0.get(id).map_or(0, |s| s.ids.len())) - } - }) - .sum() -} - -impl TemplateProcessingBuilder { - fn default_added(&self, is_single: bool) -> usize { - let container = if is_single { - self.single.as_ref() - } else { - self.pair.as_ref() - }; - container.map_or(0, |pieces| { - count_added(pieces, self.special_tokens.as_ref()) - }) - } - - fn validate(&self) -> std::result::Result<(), String> { - let pair_has_both = self.pair.as_ref().is_none_or(|pair| { - let mut has_a = false; - let mut has_b = false; - for piece in &pair.0 { - if let Piece::Sequence { - id: Sequence::A, .. - } = piece - { - has_a = true; - } - if let Piece::Sequence { - id: Sequence::B, .. - } = piece - { - has_b = true; - } - } - has_a && has_b - }); - if !pair_has_both { - return Err("Template for `pair` must use both sequences".into()); - } - - let check = |sp| { - let exist = self - .special_tokens - .as_ref() - .is_some_and(|map| map.0.contains_key(sp)); - - match exist { - false => Some(sp), - true => None, - } - }; - - let empty = []; - let missing: AHashSet<&str> = self - .single - .as_ref() - .map_or(empty.iter(), |s| s.0.iter()) - .chain(self.pair.as_ref().map_or(empty.iter(), |s| s.0.iter())) - .filter_map(|piece| match piece { - Piece::Sequence { .. } => None, - Piece::SpecialToken { id, .. } => check(id.as_ref()), - }) - .collect::>(); - - if missing.is_empty() { - Ok(()) - } else { - Err(format!( - "Missing SpecialToken(s) with id(s) `{}`", - missing.iter().join(", ") - )) - } - } -} - -impl Default for TemplateProcessing { - fn default() -> Self { - Self { - single: "$0".try_into().unwrap(), - pair: "$1".try_into().unwrap(), - added_single: 0, - added_pair: 0, - special_tokens: Tokens::default(), - } - } -} - -impl TemplateProcessing { - pub fn builder() -> TemplateProcessingBuilder { - TemplateProcessingBuilder::default() - } - - fn apply_template( - &self, - template: &[Piece], - mut encodings: Vec, - add_special_tokens: bool, - ) -> Result> { - let final_encodings: Vec = template - .iter() - .flat_map(|piece| { - match piece { - Piece::Sequence { id, type_id } => { - let i = usize::from(*id != Sequence::A); - let encoding = &mut encodings[i]; - encoding.set_type_ids(vec![*type_id; encoding.len()]); - encoding.set_sequence_id(i); - Some(encoding.clone()) - } - Piece::SpecialToken { id, type_id } => { - if add_special_tokens { - let tok = &self.special_tokens.0[id]; // We already checked existence above - let len = tok.ids.len(); - - let encoding = Encoding::new( - tok.ids.clone(), - std::iter::repeat_n(*type_id, len).collect(), - tok.tokens.clone(), - // words - std::iter::repeat_n(None, len).collect(), - // offsets - std::iter::repeat_n((0, 0), len).collect(), - // special_tokens_mask - std::iter::repeat_n(1, len).collect(), - // attention_mask - std::iter::repeat_n(1, len).collect(), - // overflowing - vec![], - // sequence_range - AHashMap::new(), - ); - Some(encoding) - } else { - None - } - } - } - }) - .collect(); - - //let mut pair = if encodings.len() > 1 { - // Some(encodings.pop().unwrap()) - //} else { - // None - //}; - //let mut encoding = encodings.pop().unwrap(); - - //let pair_overflowing = pair.as_mut().map_or(vec![], |e| e.take_overflowing()); - //let mut overflowing: Vec = encoding - // .take_overflowing() - // .iter() - // .map(|encoding| -> Result> { - // // 1. The pair itself - // let mut overflowings = self.apply_template( - // template, - // if encodings.len() > 1 { - // vec![encoding.clone(), encodings[1].clone()] - // } else { - // vec![encoding.clone()] - // }, - // add_special_tokens, - // )?; - - // // 2. Its overflowings - // for other_o in &pair_overflowing { - // overflowings.extend(self.apply_template( - // template, - // vec![encoding.clone(), other_o.clone()], - // add_special_tokens, - // )?); - // } - - // Ok(overflowings) - // }) - // .collect::>>>()? - // .into_iter() - // .flatten() - // .collect(); - //// We also need to combine the first sequence with all other overflowings - //overflowing.extend( - // pair_overflowing - // .into_iter() - // .map(|pair| { - // self.apply_template(template, vec![encoding.clone(), pair], add_special_tokens) - // }) - // .collect::>>()? - // .into_iter() - // .flatten(), - //); - - Ok(final_encodings) - } -} - -impl PostProcessor for TemplateProcessing { - fn added_tokens(&self, is_pair: bool) -> usize { - if is_pair { - self.added_pair - } else { - self.added_single - } - } - - fn process_encodings( - &self, - encodings: Vec, - add_special_tokens: bool, - ) -> Result> { - // let (encoding, pair): (Encoding, Option) = match encodings.len() { - // 1 => ( - // encodings - // .pop() - // .ok_or(ProcessorError::InvalidEncodingsVecLength)?, - // None, - // ), - // 2 => { - // let pair = encodings - // .pop() - // .ok_or(ProcessorError::InvalidEncodingsVecLength)?; - // let encoding = encodings - // .pop() - // .ok_or(ProcessorError::InvalidEncodingsVecLength)?; - // (encoding, Some(pair)) - // } - // _ => return Err(Box::new(ProcessorError::InvalidEncodingsVecLength)), - // }; - let template = match encodings.len() { - 2 => &self.pair.0, - 1 => &self.single.0, - _ => todo!(), - }; - let encodings = self.apply_template(template, encodings, add_special_tokens)?; - Ok(encodings) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use std::convert::TryInto; - use std::iter::FromIterator; - - #[test] - fn piece_serde() { - let seq_0 = Piece::Sequence { - id: Sequence::A, - type_id: 0, - }; - let seq_0_s = r#"{"Sequence":{"id":"A","type_id":0}}"#; - - assert_eq!(serde_json::to_string(&seq_0).unwrap(), seq_0_s); - assert_eq!(serde_json::from_str::(seq_0_s).unwrap(), seq_0); - - let seq_1 = Piece::Sequence { - id: Sequence::B, - type_id: 1, - }; - let seq_1_s = r#"{"Sequence":{"id":"B","type_id":1}}"#; - assert_eq!(serde_json::to_string(&seq_1).unwrap(), seq_1_s); - assert_eq!(serde_json::from_str::(seq_1_s).unwrap(), seq_1); - - let spe = Piece::SpecialToken { - id: "[CLS]".into(), - type_id: 0, - }; - let spe_s = r#"{"SpecialToken":{"id":"[CLS]","type_id":0}}"#; - assert_eq!(serde_json::to_string(&spe).unwrap(), spe_s); - assert_eq!(serde_json::from_str::(spe_s).unwrap(), spe); - } - - #[test] - fn piece() { - assert_eq!( - Ok(Piece::Sequence { - id: Sequence::A, - type_id: 0 - }), - "$".try_into() - ); - assert_eq!( - Ok(Piece::Sequence { - id: Sequence::B, - type_id: 0 - }), - "$B".try_into() - ); - assert_eq!( - Ok(Piece::Sequence { - id: Sequence::A, - type_id: 1 - }), - "$1".try_into() - ); - assert_eq!( - Ok(Piece::Sequence { - id: Sequence::B, - type_id: 2 - }), - "$B:2".try_into() - ); - assert_eq!( - Ok(Piece::Sequence { - id: Sequence::A, - type_id: 1 - }), - "$:1".try_into() - ); - assert!(Piece::try_from("$C:1").is_err()); - assert!(Piece::try_from("$A:").is_err()); - } - - #[test] - fn special_token_serde() { - let simple = SpecialToken::from(("[CLS]", 0)); - let simple_s = r#"{"id":"[CLS]","ids":[0],"tokens":["[CLS]"]}"#; - assert_eq!(serde_json::to_string(&simple).unwrap(), simple_s); - assert_eq!( - serde_json::from_str::(simple_s).unwrap(), - simple - ); - - let complete = SpecialToken::new( - "[2FR]".into(), - vec![1, 2, 3], - vec!["convert".into(), "to".into(), "FR".into()], - ) - .unwrap(); - let complete_s = r#"{"id":"[2FR]","ids":[1,2,3],"tokens":["convert","to","FR"]}"#; - assert_eq!(serde_json::to_string(&complete).unwrap(), complete_s); - assert_eq!( - serde_json::from_str::(complete_s).unwrap(), - complete - ); - - let malformed = SpecialToken::new( - "[2FR]".into(), - vec![1, 2], - vec!["convert".into(), "to".into(), "FR".into()], - ); - assert!(malformed.is_err()); - let malformed = SpecialToken::new( - "[2FR]".into(), - vec![1, 2, 3], - vec!["convert".into(), "FR".into()], - ); - assert!(malformed.is_err()); - } - - #[test] - fn template_serde() { - let template = Template(vec![ - Piece::Sequence { - id: Sequence::A, - type_id: 0, - }, - Piece::SpecialToken { - id: "[CLS]".into(), - type_id: 0, - }, - ]); - let template_s = - r#"[{"Sequence":{"id":"A","type_id":0}},{"SpecialToken":{"id":"[CLS]","type_id":0}}]"#; - assert_eq!(serde_json::to_string(&template).unwrap(), template_s); - assert_eq!( - serde_json::from_str::