diff --git a/fast/Cargo.lock b/fast/Cargo.lock index 111723d..fd3743a 100644 --- a/fast/Cargo.lock +++ b/fast/Cargo.lock @@ -2,6 +2,20 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "ahash" +version = "0.8.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75" +dependencies = [ + "cfg-if", + "getrandom", + "once_cell", + "serde", + "version_check", + "zerocopy", +] + [[package]] name = "aho-corasick" version = "1.1.4" @@ -17,12 +31,58 @@ version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" +[[package]] +name = "base64" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e1b586273c5702936fe7b7d6896644d8be71e6314cfe09d3167c95f712589e8" + +[[package]] +name = "bitflags" +version = "2.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b588b76d00fde79687d7646a9b5bdf3cc0f655e0bbd080335a95d7e96f3587da" + +[[package]] +name = "castaway" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dec551ab6e7578819132c713a93c022a05d60159dc86e7a7050223577484c55a" +dependencies = [ + "rustversion", +] + +[[package]] +name = "cc" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c89588d05638b5b4594a3348a2d6c20277e43a7f5c5202b05cc56888475a47b8" +dependencies = [ + "find-msvc-tools", + "shlex", +] + [[package]] name = "cfg-if" version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" +[[package]] +name = "compact_str" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9dfdd1c2274d9aa354115b09dc9a901d6c5576818cdf70d14cae2bdb47df00ab" +dependencies = [ + "castaway", + "cfg-if", + "itoa", + "rustversion", + "ryu", + "serde", + "static_assertions", +] + [[package]] name = "complex-tokenization" version = "0.0.1" @@ -31,6 +91,7 @@ dependencies = [ "rayon", "regex", "rustc-hash", + "tokenizers", "unicode-segmentation", ] @@ -59,18 +120,129 @@ version = "0.8.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" +[[package]] +name = "darling" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc7f46116c46ff9ab3eb1597a45688b6715c6e628b5c133e288e709a29bcb4ee" +dependencies = [ + "darling_core", + "darling_macro", +] + +[[package]] +name = "darling_core" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d00b9596d185e565c2207a0b01f8bd1a135483d02d9b7b0a54b11da8d53412e" +dependencies = [ + "fnv", + "ident_case", + "proc-macro2", + "quote", + "strsim", + "syn 2.0.117", +] + +[[package]] +name = "darling_macro" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc34b93ccb385b40dc71c6fceac4b2ad23662c7eeb248cf10d529b7e055b6ead" +dependencies = [ + "darling_core", + "quote", + "syn 2.0.117", +] + +[[package]] +name = "dary_heap" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b1e3a325bc115f096c8b77bbf027a7c2592230e70be2d985be950d3d5e60ebe" +dependencies = [ + "serde", +] + +[[package]] +name = "derive_builder" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "507dfb09ea8b7fa618fcf76e953f4f5e192547945816d5358edffe39f6f94947" +dependencies = [ + "derive_builder_macro", +] + +[[package]] +name = "derive_builder_core" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d5bcf7b024d6835cfb3d473887cd966994907effbe9227e8c8219824d06c4e8" +dependencies = [ + "darling", + "proc-macro2", + "quote", + "syn 2.0.117", +] + +[[package]] +name = "derive_builder_macro" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab63b0e2bf4d5928aff72e83a7dace85d7bba5fe12dcc3c5a572d78caffd3f3c" +dependencies = [ + "derive_builder_core", + "syn 2.0.117", +] + [[package]] name = "either" version = "1.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719" +[[package]] +name = "esaxx-rs" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d817e038c30374a4bcb22f94d0a8a0e216958d4c3dcde369b1439fec4bdda6e6" + +[[package]] +name = "find-msvc-tools" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" + +[[package]] +name = "fnv" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" + +[[package]] +name = "getrandom" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" +dependencies = [ + "cfg-if", + "libc", + "r-efi", + "wasip2", +] + [[package]] name = "heck" version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" +[[package]] +name = "ident_case" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39" + [[package]] name = "indoc" version = "2.0.7" @@ -80,12 +252,49 @@ dependencies = [ "rustversion", ] +[[package]] +name = "itertools" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b192c782037fadd9cfa75548310488aabdbf3d2da73885b31bd0abd03351285" +dependencies = [ + "either", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + [[package]] name = "libc" version = "0.2.184" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "48f5d2a454e16a5ea0f4ced81bd44e4cfc7bd3a507b61887c99fd3538b28e4af" +[[package]] +name = "log" +version = "0.4.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" + +[[package]] +name = "macro_rules_attribute" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65049d7923698040cd0b1ddcced9b0eb14dd22c5f86ae59c3740eab64a676520" +dependencies = [ + "macro_rules_attribute-proc_macro", + "paste", +] + +[[package]] +name = "macro_rules_attribute-proc_macro" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "670fdfda89751bc4a84ac13eaa63e205cf0fd22b4c9a5fbfa085b63c1f1d3a30" + [[package]] name = "memchr" version = "2.8.0" @@ -101,18 +310,99 @@ dependencies = [ "autocfg", ] +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + +[[package]] +name = "monostate" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3341a273f6c9d5bef1908f17b7267bbab0e95c9bf69a0d4dcf8e9e1b2c76ef67" +dependencies = [ + "monostate-impl", + "serde", + "serde_core", +] + +[[package]] +name = "monostate-impl" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4db6d5580af57bf992f59068d4ea26fd518574ff48d7639b255a36f9de6e7e9" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + [[package]] name = "once_cell" version = "1.21.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" +[[package]] +name = "onig" +version = "6.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0cc3cbf698f9438986c11a880c90a6d04b9de27575afd28bbf45b154b6c709e2" +dependencies = [ + "bitflags", + "libc", + "once_cell", + "onig_sys", +] + +[[package]] +name = "onig_sys" +version = "69.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e68317604e77e53b85896388e1a803c1d21b74c899ec9e5e1112db90735edd7" +dependencies = [ + "cc", + "pkg-config", +] + +[[package]] +name = "paste" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" + +[[package]] +name = "pkg-config" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e" + [[package]] name = "portable-atomic" version = "1.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" +[[package]] +name = "ppv-lite86" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9" +dependencies = [ + "zerocopy", +] + [[package]] name = "proc-macro2" version = "1.0.106" @@ -169,7 +459,7 @@ dependencies = [ "proc-macro2", "pyo3-macros-backend", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -182,7 +472,7 @@ dependencies = [ "proc-macro2", "pyo3-build-config", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -194,6 +484,41 @@ dependencies = [ "proc-macro2", ] +[[package]] +name = "r-efi" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" + +[[package]] +name = "rand" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9ef1d0d795eb7d84685bca4f72f3649f064e6641543d3a8c415898726a57b41" +dependencies = [ + "rand_chacha", + "rand_core", +] + +[[package]] +name = "rand_chacha" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" +dependencies = [ + "ppv-lite86", + "rand_core", +] + +[[package]] +name = "rand_core" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" +dependencies = [ + "getrandom", +] + [[package]] name = "rayon" version = "1.11.0" @@ -204,6 +529,17 @@ dependencies = [ "rayon-core", ] +[[package]] +name = "rayon-cond" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2964d0cf57a3e7a06e8183d14a8b527195c706b7983549cd5462d5aa3747438f" +dependencies = [ + "either", + "itertools", + "rayon", +] + [[package]] name = "rayon-core" version = "1.13.0" @@ -255,6 +591,91 @@ version = "1.0.22" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" +[[package]] +name = "ryu" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" + +[[package]] +name = "serde" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.3", +] + +[[package]] +name = "serde_json" +version = "1.0.151" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c841b55ecdae098c80dcae9cf767f6f8a0c2cdb3416bbef72181df4d0fe73f14" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "shlex" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" + +[[package]] +name = "smallvec" +version = "1.15.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ed6a63f02c8539c91a8685a86f4099661ba3da017932f6ebbea6de3f0fa7c90" + +[[package]] +name = "spm_precompiled" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5851699c4033c63636f7ea4cf7b7c1f1bf06d0cc03cfb42e711de5a5c46cf326" +dependencies = [ + "base64", + "nom", + "serde", + "unicode-segmentation", +] + +[[package]] +name = "static_assertions" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" + +[[package]] +name = "strsim" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" + [[package]] name = "syn" version = "2.0.117" @@ -266,26 +687,152 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "syn" +version = "3.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53e9bae58849f64dfa4f5d5ae372c8341f7305f82a3868709269343628b659a3" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + [[package]] name = "target-lexicon" version = "0.13.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca" +[[package]] +name = "thiserror" +version = "2.0.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09a43598840e33d5b0331f38c5e30d13bb11c11210a4b58f0d9b18a5a5eefcd9" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43cbfe0cf76104d42a574802844187e84a305e531ed54455f11fbde0f10541cd" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.3", +] + +[[package]] +name = "tokenizers" +version = "0.22.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b238e22d44a15349529690fb07bd645cf58149a1b1e44d6cb5bd1641ff1a6223" +dependencies = [ + "ahash", + "aho-corasick", + "compact_str", + "dary_heap", + "derive_builder", + "esaxx-rs", + "getrandom", + "itertools", + "log", + "macro_rules_attribute", + "monostate", + "onig", + "paste", + "rand", + "rayon", + "rayon-cond", + "regex", + "regex-syntax", + "serde", + "serde_json", + "spm_precompiled", + "thiserror", + "unicode-normalization-alignments", + "unicode-segmentation", + "unicode_categories", +] + [[package]] name = "unicode-ident" version = "1.0.24" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" +[[package]] +name = "unicode-normalization-alignments" +version = "0.1.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43f613e4fa046e69818dd287fdc4bc78175ff20331479dab6e1b0f98d57062de" +dependencies = [ + "smallvec", +] + [[package]] name = "unicode-segmentation" version = "1.13.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9629274872b2bfaf8d66f5f15725007f635594914870f65218920345aa11aa8c" +[[package]] +name = "unicode_categories" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e" + [[package]] name = "unindent" version = "0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "wasip2" +version = "1.0.4+wasi-0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b67efb37e106e55ce722a510d6b5f9c17f083e5fc79afc2badeb12cc313d9487" +dependencies = [ + "wit-bindgen", +] + +[[package]] +name = "wit-bindgen" +version = "0.57.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" + +[[package]] +name = "zerocopy" +version = "0.8.55" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b5a105cd7b140f6eeec8acff2ea38135d3cab283ada58540f629fe51e46696eb" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.55" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fe976fb70c78cd64cccfe3a6fc142244e8a77b70959b30faf9d0ac37ee228eb" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/fast/Cargo.toml b/fast/Cargo.toml index 461e71d..7096fe7 100644 --- a/fast/Cargo.toml +++ b/fast/Cargo.toml @@ -14,3 +14,6 @@ unicode-segmentation = "1.12" regex = "1" rayon = "1" rustc-hash = "2" +# Rust-side pretokenization for the default-config batch ingestion — the same +# crate the Python `tokenizers` package wraps, so splits are identical. +tokenizers = { version = "0.22", default-features = false, features = ["onig"] } diff --git a/fast/python/complex_tokenization_fast/tokenizer.py b/fast/python/complex_tokenization_fast/tokenizer.py index 9653ff2..41ed216 100644 --- a/fast/python/complex_tokenization_fast/tokenizer.py +++ b/fast/python/complex_tokenization_fast/tokenizer.py @@ -1,6 +1,12 @@ from functools import lru_cache, reduce -from complex_tokenization_fast._rs import Node, Trainer, str_to_bytes +from complex_tokenization_fast._rs import ( + Node, + Trainer, + has_cluster_handlers_py, + str_to_bytes, + trainer_from_texts, +) from complex_tokenization_fast.graphs.settings import GraphSettings from complex_tokenization_fast.graphs.units import characters, register_script, utf8, utf8_clusters from complex_tokenization_fast.graphs.words import GPTPretokenizer, words @@ -50,8 +56,19 @@ def make_trainer(self, texts): GraphSettings.ONLY_MINIMAL_MERGES = True GraphSettings.MAX_MERGE_SIZE = self.merge_size - graphs = self._build_graphs(texts) - trainer = Trainer(graphs=graphs) + # Default configuration ingests entirely in Rust (one boundary + # crossing; same `tokenizers` crate as GPTPretokenizer, so splits are + # identical). Custom pretokenizers/units/script handlers take the + # per-document Python path. + if ( + self.units is utf8_clusters + and self.pretokenizer is GPTPretokenizer + and not has_cluster_handlers_py() + ): + trainer = trainer_from_texts(list(texts), connected=self.connected) + else: + graphs = self._build_graphs(texts) + trainer = Trainer(graphs=graphs) if self.merges: merge_list = [] diff --git a/fast/src/ingest.rs b/fast/src/ingest.rs new file mode 100644 index 0000000..4ec9461 --- /dev/null +++ b/fast/src/ingest.rs @@ -0,0 +1,158 @@ +//! Batch corpus ingestion for the default-configuration fast path: one +//! boundary crossing takes the raw text list and does pretokenization, word +//! dedup, graph building, and Trainer construction entirely in Rust. +//! Pretokenization runs the same `tokenizers` crate the Python side wraps +//! (`Split(Regex(GPT_PATTERN), behavior="isolated")`), so splits are +//! identical by construction. The Python shim gates this on the default +//! pretokenizer, `utf8_clusters` units, and no registered script handlers; +//! anything else takes the per-document Python path. + +use pyo3::prelude::*; +use rayon::prelude::*; +use rustc_hash::FxHashMap; +use std::sync::Arc; +use tokenizers::pre_tokenizers::split::{Split, SplitPattern}; +use tokenizers::tokenizer::{ + OffsetReferential, OffsetType, PreTokenizedString, PreTokenizer, SplitDelimiterBehavior, +}; +use unicode_segmentation::UnicodeSegmentation; + +use crate::graph::GraphV; +use crate::trainer::Trainer; +use crate::units::{extend_cluster_cache, snapshot_cluster_cache, utf8_inner}; + +// Mirrors GPT_PATTERN in complex_tokenization/graphs/words.py verbatim. +const GPT_PATTERN: &str = "[^\\r\\n\\p{L}\\p{N}]?[\\p{Lu}\\p{Lt}\\p{Lm}\\p{Lo}\\p{M}]*[\\p{Ll}\\p{Lm}\\p{Lo}\\p{M}]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?|[^\\r\\n\\p{L}\\p{N}]?[\\p{Lu}\\p{Lt}\\p{Lm}\\p{Lo}\\p{M}]+[\\p{Ll}\\p{Lm}\\p{Lo}\\p{M}]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n/]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+"; + +fn gpt_splitter() -> Split { + Split::new( + SplitPattern::Regex(GPT_PATTERN.into()), + SplitDelimiterBehavior::Isolated, + false, + ) + .expect("GPT_PATTERN is a valid onig regex") +} + +fn pretokenize<'a>(splitter: &Split, text: &'a str) -> Vec<&'a str> { + let mut pts = PreTokenizedString::from(text); + splitter + .pre_tokenize(&mut pts) + .expect("Split pre-tokenization is infallible for plain text"); + pts.get_splits(OffsetReferential::Original, OffsetType::Byte) + .into_iter() + .map(|(_, (start, end), _)| &text[start..end]) + .collect() +} + +/// `utf8_clusters` for one grapheme cluster with no script handlers: a +/// single-char cluster is its UTF-8 graph, a multi-char cluster is a Seq of +/// per-char UTF-8 graphs. Mirrors `resolve_cluster`'s no-handler branch. +fn cluster_graph_pure(cluster: &str) -> GraphV { + if cluster.chars().nth(1).is_none() { + return utf8_inner(cluster); + } + let char_nodes: Vec = cluster + .chars() + .map(|c| utf8_inner(c.encode_utf8(&mut [0u8; 4]))) + .collect(); + GraphV::new_seq(char_nodes) +} + +pub(crate) fn build_corpus_graph(texts: &[String], connected: bool) -> GraphV { + let splitter = gpt_splitter(); + let doc_tokens: Vec> = texts + .par_iter() + .map(|t| pretokenize(&splitter, t)) + .collect(); + + // Unique words, first-occurrence order. + let mut word_id: FxHashMap<&str, u32> = FxHashMap::default(); + let mut unique_words: Vec<&str> = Vec::new(); + for tokens in &doc_tokens { + for &w in tokens { + word_id.entry(w).or_insert_with(|| { + unique_words.push(w); + (unique_words.len() - 1) as u32 + }); + } + } + + // One graph per unique grapheme cluster, corpus-wide. Existing cluster + // cache entries are reused and new ones written back, exactly as the + // per-word `utf8_clusters` path would. + let mut cluster_of: FxHashMap<&str, GraphV> = FxHashMap::default(); + let cache = snapshot_cluster_cache(); + let mut new_clusters: Vec<(String, GraphV)> = Vec::new(); + for &w in &unique_words { + for cl in w.graphemes(true) { + if !cluster_of.contains_key(cl) { + let g = match cache.get(cl) { + Some(g) => g.clone(), + None => { + let g = cluster_graph_pure(cl); + new_clusters.push((cl.to_string(), g.clone())); + g + } + }; + cluster_of.insert(cl, g); + } + } + } + extend_cluster_cache(new_clusters); + + // One graph per unique word, shared by Arc across its occurrences + // (candidate counting is content-based, so sharing does not affect + // results; the trainer re-canonicalizes by content anyway). + let word_graphs: Vec = unique_words + .par_iter() + .map(|w| { + let mut nodes: Vec = + w.graphemes(true).map(|cl| cluster_of[cl].clone()).collect(); + if nodes.len() == 1 { + nodes.pop().unwrap() + } else { + GraphV::new_seq(nodes) + } + }) + .collect(); + + // Assemble docs the way `words()` + `Trainer(graphs=...)` would: + // connected docs are Seqs of their words (a single-word doc is the word + // itself); disconnected docs flatten into one word-occurrence list. + let subs: Vec = if connected { + doc_tokens + .par_iter() + .map(|tokens| { + let mut nodes: Vec = tokens + .iter() + .map(|w| word_graphs[word_id[w] as usize].clone()) + .collect(); + if nodes.len() == 1 { + nodes.pop().unwrap() + } else { + GraphV::new_seq(nodes) + } + }) + .collect() + } else { + doc_tokens + .iter() + .flat_map(|tokens| tokens.iter()) + .map(|w| word_graphs[word_id[w] as usize].clone()) + .collect() + }; + + GraphV::Unconn(Arc::new(subs)) +} + +#[pyfunction] +#[pyo3(signature = (texts, connected=false))] +pub fn trainer_from_texts(py: Python<'_>, texts: Vec, connected: bool) -> Trainer { + let graph = py.allow_threads(|| build_corpus_graph(&texts, connected)); + Trainer::from_graph(graph) +} + +#[pyfunction] +pub fn has_cluster_handlers_py() -> bool { + crate::units::has_cluster_handlers() +} diff --git a/fast/src/lib.rs b/fast/src/lib.rs index c025895..099b115 100644 --- a/fast/src/lib.rs +++ b/fast/src/lib.rs @@ -1,4 +1,5 @@ mod graph; +mod ingest; mod settings; mod trainer; mod units; @@ -22,6 +23,8 @@ fn _rs(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(units::set_ids_reverse_dict_py, m)?)?; m.add_function(wrap_pyfunction!(units::warm_word_cache_py, m)?)?; m.add_function(wrap_pyfunction!(units::clear_word_cache, m)?)?; + m.add_function(wrap_pyfunction!(ingest::trainer_from_texts, m)?)?; + m.add_function(wrap_pyfunction!(ingest::has_cluster_handlers_py, m)?)?; m.add_function(wrap_pyfunction!(graph::bytes_to_str_py, m)?)?; m.add_function(wrap_pyfunction!(graph::str_to_bytes_py, m)?)?; m.add_function(wrap_pyfunction!(sync_settings, m)?)?; diff --git a/fast/src/trainer.rs b/fast/src/trainer.rs index 7e2d16a..c59879f 100644 --- a/fast/src/trainer.rs +++ b/fast/src/trainer.rs @@ -1608,6 +1608,18 @@ pub struct Trainer { premerge_registry: Option, } +impl Trainer { + pub(crate) fn from_graph(graph: GraphV) -> Self { + Trainer { + graph, + merges: Vec::new(), + doc_words: None, + streaming_connected: false, + premerge_registry: None, + } + } +} + #[pymethods] impl Trainer { #[new] diff --git a/fast/src/units.rs b/fast/src/units.rs index 4bd66b1..dd7ecfd 100644 --- a/fast/src/units.rs +++ b/fast/src/units.rs @@ -52,6 +52,21 @@ pub fn clear_handlers() { CLUSTER_CACHE.lock().unwrap().clear(); } +pub(crate) fn snapshot_cluster_cache() -> HashMap { + CLUSTER_CACHE.lock().unwrap().clone() +} + +pub(crate) fn extend_cluster_cache(new_entries: Vec<(String, GraphV)>) { + let mut cache = CLUSTER_CACHE.lock().unwrap(); + for (k, v) in new_entries { + cache.entry(k).or_insert(v); + } +} + +pub(crate) fn has_cluster_handlers() -> bool { + !CLUSTER_HANDLERS.lock().unwrap().is_empty() +} + static WORD_CACHE: LazyLock>> = LazyLock::new(|| Mutex::new(HashMap::new()));