diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index d908d693..534fa320 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -113,6 +113,8 @@ jobs: toolchain: 1.97.1 components: rustfmt, clippy - run: npm ci + - name: Fetch prediction model + run: npm run prediction-model - run: cargo fmt --manifest-path src-tauri/Cargo.toml --check - run: cargo clippy --locked --manifest-path src-tauri/Cargo.toml --all-targets -- -D warnings - run: cargo test --locked --manifest-path src-tauri/Cargo.toml diff --git a/.github/workflows/release-tauri.yml b/.github/workflows/release-tauri.yml index e45f99b0..91ec6fbf 100644 --- a/.github/workflows/release-tauri.yml +++ b/.github/workflows/release-tauri.yml @@ -164,6 +164,9 @@ jobs: - name: Install dependencies run: npm ci + - name: Fetch prediction model + run: npm run prediction-model + - name: Validate Rust at release tag run: | cargo fmt --manifest-path src-tauri/Cargo.toml --check @@ -340,6 +343,9 @@ jobs: - name: Install dependencies run: npm ci + - name: Fetch prediction model + run: npm run prediction-model + - name: Validate Rust at release tag shell: powershell run: | diff --git a/.gitignore b/.gitignore index 5896deba..6182a6af 100644 --- a/.gitignore +++ b/.gitignore @@ -16,3 +16,4 @@ src-tauri/icons/android/ src-tauri/icons/ios/ src-tauri/tauri.release*.generated.json tools/prediction-data/target/ +src-tauri/resources/prediction-model/ diff --git a/AGENTS.md b/AGENTS.md index 5116d8f7..1f246ddd 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -24,6 +24,7 @@ The frozen C# application is not maintained in this repository. Do not reintrodu Use Node.js 24 and Rust 1.97.1. Run checks from the repository root: ```bash +npm run prediction-model npm run lint npm test npm run build diff --git a/docs/word-prediction.md b/docs/word-prediction.md index 9c262dd2..3962c480 100644 --- a/docs/word-prediction.md +++ b/docs/word-prediction.md @@ -10,6 +10,27 @@ A passive observer records only an activity counter and timestamp, never externa The bundled English lookup database works offline. Ranking backs off from the longest available three-word context to shorter contexts, then to word frequency. Prefix matching is case-insensitive. Duplicate words are removed. The data is fixed; typing is never used to train or update it. +## Enhanced word prediction + +Enhanced word prediction is off by default. When it is on, a small language model (SmolLM2-135M, int8 ONNX, run with ONNX Runtime on the CPU) spells suggestions from subword pieces, so it can suggest words the lookup lacks, such as WhatsApp. It sees only the buffer described above: at most its last 256 characters, starting at a word boundary, including earlier sentences still in the buffer. It runs offline in the prediction worker, never learns, and nothing it reads is logged or stored. + +- **Loading and fallback:** the model loads in the background when the worker starts, taking about 0.7 s. Until it is ready, and whenever it fails, the lookup answers as before. Turning the setting on or off restarts the worker. +- **Filters:** + - words contain only ASCII letters and apostrophes; + - a word at the start of the text must be capitalised, which rules out subword fragments; + - single letters must be in the lookup; + - blocklisted words are never suggested; + - words the lookup does not know must reach a minimum probability. +- **Spelling:** words the lookup knows take its usual spelling. Unknown words keep the model's spelling, except that a capital coming only from sentence position is removed. +- **Merging:** lookup suggestions fill any remaining positions. +- **Time limits:** the model stops exploring new words after 400 ms. If one call ever takes longer than 1 s, the worker uses the lookup for the rest of its life, so replies stay well inside the two-second deadline. +- **Clipped buffers:** when clipping leaves no complete earlier word, the lookup answers. +- **Cost:** predictions take roughly 40–80 ms on a recent laptop, and the worker needs about 300–450 MB more memory. + +The model files are too large to commit. `npm run prediction-model` downloads them from a pinned upstream revision and verifies their SHA-256. The Tauri dev and build commands and CI run it automatically. Run it once before `cargo test` or `cargo clippy` in a fresh checkout. + +## Worker process + A separate process owns the buffer, activity observer and database. Private bounded inherited pipes carry successful edits and results to native rendering. Text and suggestions are not sent to the React UI, diagnostic history or telemetry. One request is outstanding at a time with a two-second deadline. Timeout stops predictions until the keyboard is reopened; ordinary keyboard operation remains available. Closing the keyboard, ending scanning, or exiting kills and reaps the worker. Before accepting a suggestion, the worker checks its token, edit revision, foreground identity and external activity. The main process checks its generation and foreground again before injection. Verification and native insertion cannot be atomic across applications; the target can still change in that short interval. @@ -22,6 +43,12 @@ Automated tests use fixture databases, fake foreground/activity sources, fake in cargo test --manifest-path src-tauri/Cargo.toml bundled_database_benchmark -- --ignored --nocapture ``` +`cargo test` also runs the language model against the fetched files with synthetic text. For release-build suggestions and timings: + +```powershell +cargo test --release --manifest-path src-tauri/Cargo.toml bundled_model_suggests_current_words -- --nocapture +``` + For native validation, use disposable synthetic text in Notepad and a browser on Windows, and TextEdit and a browser on macOS. Launch macOS with `npm run macos:run` to retain its signed Accessibility identity. 1. Open the keyboard from the action menu and an assigned Open keyboard switch, including from idle, a paused scan and an active drag. Verify no click or focus change occurs and owned drag input is released. diff --git a/package.json b/package.json index bb9adefa..2cc25026 100644 --- a/package.json +++ b/package.json @@ -9,6 +9,7 @@ "scripts": { "dev": "vite", "build": "tsc --noEmit && vite build", + "prediction-model": "node scripts/fetch-prediction-model.mjs", "test": "vitest run && node --test scripts/create-update-feed.node-test.mjs", "lint": "tsc --noEmit", "macos:setup-signing": "./scripts/setup-macos-dev-signing.sh", diff --git a/scripts/fetch-prediction-model.mjs b/scripts/fetch-prediction-model.mjs new file mode 100644 index 00000000..68d6434e --- /dev/null +++ b/scripts/fetch-prediction-model.mjs @@ -0,0 +1,58 @@ +// Downloads the enhanced word prediction model into src-tauri/resources/prediction-model. +// The model is too large to commit, so every file is pinned to one upstream +// revision and verified by SHA-256. Verified files are not downloaded again. +import { createHash } from "node:crypto"; +import { createReadStream, createWriteStream } from "node:fs"; +import { mkdir, rename, rm, stat } from "node:fs/promises"; +import { dirname, join } from "node:path"; +import { Readable } from "node:stream"; +import { pipeline } from "node:stream/promises"; +import { fileURLToPath } from "node:url"; + +const revision = "d0ae6834f1df45e0e95b5fdae95e536f9ca7cd3f"; +const base = `https://huggingface.co/onnx-community/SmolLM2-135M-ONNX/resolve/${revision}`; +const files = [ + { name: "model.onnx", source: "onnx/model_int8.onnx", size: 135658354, sha256: "50ba80511ce74634d232a043b6c37775cca756b826b49d0a4a8eff958c4bbcc9" }, + { name: "tokenizer.json", source: "tokenizer.json", size: 2053526, sha256: "139d2f4b4919b90953bdd3c0c40c94c9b23074799a766508dc3bf5eb8ab73351" }, + { name: "config.json", source: "config.json", size: 1035, sha256: "2c5f23fddabecdf9c47d0048f555899822ec87ec4a28169393840ac7e74192c4" }, +]; +const target = join(dirname(fileURLToPath(import.meta.url)), "..", "src-tauri", "resources", "prediction-model"); + +async function digest(path) { + const hash = createHash("sha256"); + await pipeline(createReadStream(path), hash); + return hash.digest("hex"); +} + +async function verified(path, file) { + try { + return (await stat(path)).size === file.size && (await digest(path)) === file.sha256; + } catch { + return false; + } +} + +async function download(file) { + const path = join(target, file.name); + if (await verified(path, file)) return false; + const pending = `${path}.pending`; + for (let attempt = 1; ; attempt++) { + try { + const response = await fetch(`${base}/${file.source}`); + if (!response.ok) throw new Error(`HTTP ${response.status}`); + await pipeline(Readable.fromWeb(response.body), createWriteStream(pending)); + if (!(await verified(pending, file))) throw new Error("size or SHA-256 mismatch"); + await rename(pending, path); + return true; + } catch (error) { + await rm(pending, { force: true }); + if (attempt === 3) throw new Error(`Could not fetch ${file.name}: ${error.message}`); + await new Promise((resolve) => setTimeout(resolve, attempt * 2000)); + } + } +} + +await mkdir(target, { recursive: true }); +for (const file of files) { + if (await download(file)) console.log(`Fetched ${file.name}`); +} diff --git a/src-tauri/Cargo.lock b/src-tauri/Cargo.lock index 2feb2cbc..26030140 100644 --- a/src-tauri/Cargo.lock +++ b/src-tauri/Cargo.lock @@ -19,6 +19,20 @@ dependencies = [ "cpufeatures 0.2.17", ] +[[package]] +name = "ahash" +version = "0.8.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75" +dependencies = [ + "cfg-if", + "getrandom 0.3.4", + "once_cell", + "serde", + "version_check", + "zerocopy", +] + [[package]] name = "aho-corasick" version = "1.1.5" @@ -299,6 +313,12 @@ dependencies = [ "pkg-config", ] +[[package]] +name = "base64" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e1b586273c5702936fe7b7d6896644d8be71e6314cfe09d3167c95f712589e8" + [[package]] name = "base64" version = "0.21.7" @@ -311,6 +331,12 @@ version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" +[[package]] +name = "base64" +version = "0.23.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac07cdecf99051d9a5238b80f35af32cdeba5b336e55d957b318b50137e18da5" + [[package]] name = "bit-set" version = "0.8.0" @@ -511,6 +537,15 @@ dependencies = [ "toml 0.9.12+spec-1.1.0", ] +[[package]] +name = "castaway" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dec551ab6e7578819132c713a93c022a05d60159dc86e7a7050223577484c55a" +dependencies = [ + "rustversion", +] + [[package]] name = "cbc" version = "0.1.2" @@ -623,6 +658,21 @@ dependencies = [ "memchr", ] +[[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 = "concurrent-queue" version = "2.5.0" @@ -728,6 +778,25 @@ dependencies = [ "crossbeam-utils", ] +[[package]] +name = "crossbeam-deque" +version = "0.8.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "622f3fc73690be383c7214310406f28a90e6edeadc3cea882f9d71e495b9711a" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc74980687109a3b14c72fd458107bf0baa1da1a1a805e178d15501ba9b86d9d" +dependencies = [ + "crossbeam-utils", +] + [[package]] name = "crossbeam-queue" version = "0.3.13" @@ -792,14 +861,38 @@ version = "0.0.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "52560adf09603e58c9a7ee1fe1dcb95a16927b17c127f0ac02d6e768a0e25bc1" +[[package]] +name = "darling" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc7f46116c46ff9ab3eb1597a45688b6715c6e628b5c133e288e709a29bcb4ee" +dependencies = [ + "darling_core 0.20.11", + "darling_macro 0.20.11", +] + [[package]] name = "darling" version = "0.23.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "25ae13da2f202d56bd7f91c25fba009e7717a1e4a1cc98a76d844b65ae912e9d" dependencies = [ - "darling_core", - "darling_macro", + "darling_core 0.23.0", + "darling_macro 0.23.0", +] + +[[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.119", ] [[package]] @@ -815,17 +908,37 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "darling_macro" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc34b93ccb385b40dc71c6fceac4b2ad23662c7eeb248cf10d529b7e055b6ead" +dependencies = [ + "darling_core 0.20.11", + "quote", + "syn 2.0.119", +] + [[package]] name = "darling_macro" version = "0.23.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ac3984ec7bd6cfa798e62b4a642426a5be0e68f9401cfc2a01e3fa9ea2fcdb8d" dependencies = [ - "darling_core", + "darling_core 0.23.0", "quote", "syn 2.0.119", ] +[[package]] +name = "dary_heap" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b1e3a325bc115f096c8b77bbf027a7c2592230e70be2d985be950d3d5e60ebe" +dependencies = [ + "serde", +] + [[package]] name = "dbus" version = "0.9.12" @@ -857,6 +970,37 @@ dependencies = [ "syn 2.0.119", ] +[[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 0.20.11", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[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.119", +] + [[package]] name = "derive_more" version = "2.1.1" @@ -1061,6 +1205,12 @@ version = "1.0.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555" +[[package]] +name = "either" +version = "1.18.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "252afb9ae5eaa683babdc6a068b3f5726eb19e05070c731f9b2a23a7c3e8ed34" + [[package]] name = "embed-resource" version = "3.0.11" @@ -1098,7 +1248,7 @@ dependencies = [ "foreign-types-shared", "libc", "log", - "nom", + "nom 8.0.0", "objc2", "objc2-app-kit", "objc2-foundation", @@ -1156,6 +1306,12 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "esaxx-rs" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d817e038c30374a4bcb22f94d0a8a0e216958d4c3dcde369b1439fec4bdda6e6" + [[package]] name = "event-listener" version = "5.4.2" @@ -1188,6 +1344,17 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7360491ce676a36bf9bb3c56c1aa791658183a54d2744120f27285738d90465a" +[[package]] +name = "fancy-regex" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e24cb5a94bcae1e5408b0effca5cd7172ea3c5755049c5f3af4cd283a165298" +dependencies = [ + "bit-set", + "regex-automata", + "regex-syntax", +] + [[package]] name = "fastrand" version = "2.5.0" @@ -1758,6 +1925,12 @@ dependencies = [ "digest", ] +[[package]] +name = "hmac-sha256" +version = "1.1.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec9d92d097f4749b64e8cc33d924d9f40a2d4eb91402b458014b781f5733d60f" + [[package]] name = "html5ever" version = "0.38.0" @@ -2069,6 +2242,15 @@ version = "2.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6a756c3fac73139e83f14c2d742155dd2b78d3ee56597b419a0579b7bdd6dd78" +[[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" @@ -2349,6 +2531,28 @@ version = "0.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154" +[[package]] +name = "lzma-rust2" +version = "0.15.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e20f57f9918e5bd7bc58c22cdd70a6afc7375d4dd9683af5f2b34bd3d2bba619" + +[[package]] +name = "macro_rules_attribute" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b3ae8f6d608c795738406608304d30a2dfbdc8e58e44f7ba43236da5208ded3c" +dependencies = [ + "macro_rules_attribute-proc_macro", + "pastey", +] + +[[package]] +name = "macro_rules_attribute-proc_macro" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc04a4c58212d57930a24bf47d3fa87485264a3a054e9c10e042eb373573ad3c" + [[package]] name = "markup5ever" version = "0.38.0" @@ -2360,6 +2564,16 @@ dependencies = [ "web_atoms", ] +[[package]] +name = "matrixmultiply" +version = "0.3.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f607c237553f086e7043417a51df26b2eb899d3caff94e6a67592ff992fedc7" +dependencies = [ + "autocfg", + "rawpointer", +] + [[package]] name = "memchr" version = "2.8.3" @@ -2390,6 +2604,12 @@ version = "0.3.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a" +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + [[package]] name = "minisign-verify" version = "0.2.5" @@ -2417,6 +2637,28 @@ dependencies = [ "windows-sys 0.61.2", ] +[[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.119", +] + [[package]] name = "moxcms" version = "0.8.1" @@ -2448,6 +2690,21 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "ndarray" +version = "0.17.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "520080814a7a6b4a6e9070823bb24b4531daac8c4627e08ba5de8c5ef2f2752d" +dependencies = [ + "matrixmultiply", + "num-complex", + "num-integer", + "num-traits", + "portable-atomic", + "portable-atomic-util", + "rawpointer", +] + [[package]] name = "ndk" version = "0.9.0" @@ -2478,6 +2735,16 @@ version = "1.0.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "650eef8c711430f1a879fdd01d4745a7deea475becfb90269c06775983bbf086" +[[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 = "nom" version = "8.0.0" @@ -2846,6 +3113,30 @@ dependencies = [ "pin-project-lite", ] +[[package]] +name = "ort" +version = "2.0.0-rc.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4336a1e2b38848325241c72889086886004e589b7c74f335e60a8e8db5138a0b" +dependencies = [ + "ndarray", + "ort-sys", + "smallvec", + "tracing", + "ureq", +] + +[[package]] +name = "ort-sys" +version = "2.0.0-rc.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf211e3776eea6aec988552fa118dd746d70e1b1e5e244058d1c98015f3e5872" +dependencies = [ + "hmac-sha256", + "lzma-rust2", + "ureq", +] + [[package]] name = "osakit" version = "0.3.1" @@ -2914,6 +3205,18 @@ dependencies = [ "windows-link 0.2.1", ] +[[package]] +name = "paste" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" + +[[package]] +name = "pastey" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2ee67f1008b1ba2321834326597b8e186293b049a023cdef258527550b9935b4" + [[package]] name = "percent-encoding" version = "2.3.2" @@ -3049,6 +3352,21 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "portable-atomic" +version = "1.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05c8b63e8d9609db387f0324918f81d68fe27748f084ef092fb35954d0539a85" + +[[package]] +name = "portable-atomic-util" +version = "0.2.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10ab3eb7f3becc3a1cbc4f2c6f20267996cfc1a6467a873763411b136a122715" +dependencies = [ + "portable-atomic", +] + [[package]] name = "potential_utf" version = "0.1.5" @@ -3241,10 +3559,20 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "22f6172bdec972074665ed81ed53b71da00bfc44b65a753cfde883ec4c702a1a" dependencies = [ "libc", - "rand_chacha", + "rand_chacha 0.3.1", "rand_core 0.6.4", ] +[[package]] +name = "rand" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9ef1d0d795eb7d84685bca4f72f3649f064e6641543d3a8c415898726a57b41" +dependencies = [ + "rand_chacha 0.9.0", + "rand_core 0.9.5", +] + [[package]] name = "rand" version = "0.10.2" @@ -3266,6 +3594,16 @@ dependencies = [ "rand_core 0.6.4", ] +[[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 0.9.5", +] + [[package]] name = "rand_core" version = "0.6.4" @@ -3275,6 +3613,15 @@ dependencies = [ "getrandom 0.2.17", ] +[[package]] +name = "rand_core" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" +dependencies = [ + "getrandom 0.3.4", +] + [[package]] name = "rand_core" version = "0.10.1" @@ -3296,6 +3643,43 @@ version = "0.6.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "20675572f6f24e9e76ef639bc5552774ed45f1c30e2951e1e99c59888861c539" +[[package]] +name = "rawpointer" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3" + +[[package]] +name = "rayon" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d" +dependencies = [ + "either", + "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" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + [[package]] name = "redox_syscall" version = "0.5.18" @@ -3490,6 +3874,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0d41d731c7d2f962d1ccc364cec258de3c0e93b38c2fb3ba97ac74513048d634" dependencies = [ "aws-lc-rs", + "log", "once_cell", "ring", "rustls-pki-types", @@ -3565,6 +3950,12 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" +[[package]] +name = "ryu" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" + [[package]] name = "same-file" version = "1.0.6" @@ -3832,7 +4223,7 @@ version = "3.21.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "84d57bc0c8b9a17920c178daa6bb924850d54a9c97ab45194bb8c17ad66bb660" dependencies = [ - "darling", + "darling 0.23.0", "proc-macro2", "quote", "syn 2.0.119", @@ -3946,6 +4337,17 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "socks" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0c3dbbd9ae980613c6dd8e28a9407b50509d3803b57624d5dfe8315218cd58b" +dependencies = [ + "byteorder", + "libc", + "winapi", +] + [[package]] name = "softbuffer" version = "0.4.8" @@ -3994,6 +4396,18 @@ dependencies = [ "system-deps", ] +[[package]] +name = "spm_precompiled" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5851699c4033c63636f7ea4cf7b7c1f1bf06d0cc03cfb42e711de5a5c46cf326" +dependencies = [ + "base64 0.13.1", + "nom 7.1.3", + "serde", + "unicode-segmentation", +] + [[package]] name = "sqlite-wasm-rs" version = "0.5.5" @@ -4012,6 +4426,12 @@ version = "1.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" +[[package]] +name = "static_assertions" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" + [[package]] name = "strict-num" version = "0.1.1" @@ -4085,6 +4505,7 @@ dependencies = [ "objc2-core-graphics", "objc2-foundation", "objc2-quartz-core", + "ort", "rand 0.8.7", "reqwest", "rusqlite", @@ -4097,6 +4518,7 @@ dependencies = [ "tauri-plugin-single-instance", "tauri-plugin-updater", "tiny-skia", + "tokenizers", "tokio", "unicode-segmentation", "uuid", @@ -4674,6 +5096,39 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" +[[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", + "fancy-regex", + "getrandom 0.3.4", + "itertools", + "log", + "macro_rules_attribute", + "monostate", + "paste", + "rand 0.9.5", + "rayon", + "rayon-cond", + "regex", + "regex-syntax", + "serde", + "serde_json", + "spm_precompiled", + "thiserror 2.0.19", + "unicode-normalization-alignments", + "unicode-segmentation", + "unicode_categories", +] + [[package]] name = "tokio" version = "1.53.1" @@ -5017,18 +5472,62 @@ 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.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c6f5d3c3b1bf09027a88a6bc961fc00497d651009560b5463668dc81b0fa87a8" +[[package]] +name = "unicode_categories" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e" + [[package]] name = "untrusted" version = "0.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1" +[[package]] +name = "ureq" +version = "3.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a7ac20be9b7726e0bbdbf974c059676d9acb1cd414961f570a4e8231cacd7fc" +dependencies = [ + "base64 0.23.1", + "log", + "percent-encoding", + "rustls", + "rustls-pki-types", + "socks", + "ureq-proto", + "utf8-zero", + "webpki-roots", +] + +[[package]] +name = "ureq-proto" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f86fd172ccca569e458f61b6bdd6220965a9ef36e672a6852953b51a0e1583be" +dependencies = [ + "base64 0.23.1", + "http", + "httparse", + "log", +] + [[package]] name = "url" version = "2.5.8" @@ -5054,6 +5553,12 @@ dependencies = [ "url", ] +[[package]] +name = "utf8-zero" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8c0a043c9540bae7c578c88f91dda8bd82e59ae27c21baca69c8b191aaf5a6e" + [[package]] name = "utf8_iter" version = "1.0.4" @@ -5297,6 +5802,15 @@ dependencies = [ "rustls-pki-types", ] +[[package]] +name = "webpki-roots" +version = "1.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7dcd9d09a39985f5344844e66b0c530a33843579125f23e21e9f0f220850f22a" +dependencies = [ + "rustls-pki-types", +] + [[package]] name = "webview2-com" version = "0.38.2" diff --git a/src-tauri/Cargo.toml b/src-tauri/Cargo.toml index c497ba6a..01ccb6df 100644 --- a/src-tauri/Cargo.toml +++ b/src-tauri/Cargo.toml @@ -20,6 +20,7 @@ directories = "6.0.0" enigo = "=0.6.1" hmac = "0.12" keyring = "4.1.6" +ort = { version = "=2.0.0-rc.13", default-features = false, features = ["std", "download-binaries", "tls-rustls", "api-27"] } rand = "0.8" unicode-segmentation = "1.12" reqwest = { version = "0.13", default-features = false, features = ["json", "rustls"] } @@ -31,6 +32,7 @@ tauri-plugin-autostart = "2.5.1" tauri-plugin-single-instance = "2.4.3" tauri-plugin-updater = "2.10.1" tiny-skia = "0.11.4" +tokenizers = { version = "0.22", default-features = false, features = ["fancy-regex"] } tokio = { version = "1", features = ["macros", "rt", "sync", "time"] } uuid = { version = "1", features = ["v4", "serde"] } diff --git a/src-tauri/resources/PROVENANCE.md b/src-tauri/resources/PROVENANCE.md index ace5137c..4e7d5f14 100644 --- a/src-tauri/resources/PROVENANCE.md +++ b/src-tauri/resources/PROVENANCE.md @@ -20,3 +20,21 @@ cargo run --release --locked --manifest-path tools/prediction-data/Cargo.toml -- ``` Use `-- --check` before the two paths to regenerate into a temporary file and verify byte-for-byte reproducibility. CI checks the artifact on Mac and Windows. The original SQLite implementation is retained only in tests to verify identical ordered suggestions. + +# Enhanced prediction model + +`prediction-model/` is not committed. `npm run prediction-model` (`scripts/fetch-prediction-model.mjs`) downloads it from [`onnx-community/SmolLM2-135M-ONNX`](https://huggingface.co/onnx-community/SmolLM2-135M-ONNX) at revision `d0ae6834f1df45e0e95b5fdae95e536f9ca7cd3f` and verifies each file's size and SHA-256 before use: + +- `model.onnx`: upstream `onnx/model_int8.onnx`, 135,658,354 bytes, SHA-256 `50ba80511ce74634d232a043b6c37775cca756b826b49d0a4a8eff958c4bbcc9`. +- `tokenizer.json`: 2,053,526 bytes, SHA-256 `139d2f4b4919b90953bdd3c0c40c94c9b23074799a766508dc3bf5eb8ab73351`. +- `config.json`: 1,035 bytes, SHA-256 `2c5f23fddabecdf9c47d0048f555899822ec87ec4a28169393840ac7e74192c4`. + +It is an int8 ONNX conversion of [HuggingFaceTB/SmolLM2-135M](https://huggingface.co/HuggingFaceTB/SmolLM2-135M), licensed Apache-2.0. The app bundles these files unchanged and only reads them. + +# Prediction blocklist + +`prediction-blocklist.txt` is the unchanged English list from [LDNOOBW/List-of-Dirty-Naughty-Obscene-and-Otherwise-Bad-Words](https://github.com/LDNOOBW/List-of-Dirty-Naughty-Obscene-and-Otherwise-Bad-Words) (`en` at commit `4638b970cb8d9d82789564fcba1f4a1eb508ff1a`, 403 lines, SHA-256 `af851ecef1d5f212caba17339b12ac39cc2fef7d78c74876f67237644fcee8bd`). It is licensed CC BY 4.0 by its contributors. Enhanced prediction never suggests a listed word. + +# ONNX Runtime + +Enhanced prediction statically links ONNX Runtime 1.28.0 through the `ort` crate (`=2.0.0-rc.13`, MIT/Apache-2.0). At build time, `ort-sys` downloads the prebuilt archive for the target from `cdn.pyke.io` and verifies it against a SHA-256 embedded in the crate. So every build, including signed releases, depends on that download. Only the CPU provider is used, and ONNX Runtime telemetry is disabled when the worker starts it. diff --git a/src-tauri/resources/prediction-blocklist.txt b/src-tauri/resources/prediction-blocklist.txt new file mode 100644 index 00000000..a438b9ca --- /dev/null +++ b/src-tauri/resources/prediction-blocklist.txt @@ -0,0 +1,403 @@ +2g1c +2 girls 1 cup +acrotomophilia +alabama hot pocket +alaskan pipeline +anal +anilingus +anus +apeshit +arsehole +ass +asshole +assmunch +auto erotic +autoerotic +babeland +baby batter +baby juice +ball gag +ball gravy +ball kicking +ball licking +ball sack +ball sucking +bangbros +bangbus +bareback +barely legal +barenaked +bastard +bastardo +bastinado +bbw +bdsm +beaner +beaners +beaver cleaver +beaver lips +beastiality +bestiality +big black +big breasts +big knockers +big tits +bimbos +birdlock +bitch +bitches +black cock +blonde action +blonde on blonde action +blowjob +blow job +blow your load +blue waffle +blumpkin +bollocks +bondage +boner +boob +boobs +booty call +brown showers +brunette action +bukkake +bulldyke +bullet vibe +bullshit +bung hole +bunghole +busty +butt +buttcheeks +butthole +camel toe +camgirl +camslut +camwhore +carpet muncher +carpetmuncher +chocolate rosebuds +cialis +circlejerk +cleveland steamer +clit +clitoris +clover clamps +clusterfuck +cock +cocks +coprolagnia +coprophilia +cornhole +coon +coons +creampie +cum +cumming +cumshot +cumshots +cunnilingus +cunt +darkie +date rape +daterape +deep throat +deepthroat +dendrophilia +dick +dildo +dingleberry +dingleberries +dirty pillows +dirty sanchez +doggie style +doggiestyle +doggy style +doggystyle +dog style +dolcett +domination +dominatrix +dommes +donkey punch +double dong +double penetration +dp action +dry hump +dvda +eat my ass +ecchi +ejaculation +erotic +erotism +escort +eunuch +fag +faggot +fecal +felch +fellatio +feltch +female squirting +femdom +figging +fingerbang +fingering +fisting +foot fetish +footjob +frotting +fuck +fuck buttons +fuckin +fucking +fucktards +fudge packer +fudgepacker +futanari +gangbang +gang bang +gay sex +genitals +giant cock +girl on +girl on top +girls gone wild +goatcx +goatse +god damn +gokkun +golden shower +goodpoop +goo girl +goregasm +grope +group sex +g-spot +guro +hand job +handjob +hard core +hardcore +hentai +homoerotic +honkey +hooker +horny +hot carl +hot chick +how to kill +how to murder +huge fat +humping +incest +intercourse +jack off +jail bait +jailbait +jelly donut +jerk off +jigaboo +jiggaboo +jiggerboo +jizz +juggs +kike +kinbaku +kinkster +kinky +knobbing +leather restraint +leather straight jacket +lemon party +livesex +lolita +lovemaking +make me come +male squirting +masturbate +masturbating +masturbation +menage a trois +milf +missionary position +mong +motherfucker +mound of venus +mr hands +muff diver +muffdiving +nambla +nawashi +negro +neonazi +nigga +nigger +nig nog +nimphomania +nipple +nipples +nsfw +nsfw images +nude +nudity +nutten +nympho +nymphomania +octopussy +omorashi +one cup two girls +one guy one jar +orgasm +orgy +paedophile +paki +panties +panty +pedobear +pedophile +pegging +penis +phone sex +piece of shit +pikey +pissing +piss pig +pisspig +playboy +pleasure chest +pole smoker +ponyplay +poof +poon +poontang +punany +poop chute +poopchute +porn +porno +pornography +prince albert piercing +pthc +pubes +pussy +queaf +queef +quim +raghead +raging boner +rape +raping +rapist +rectum +reverse cowgirl +rimjob +rimming +rosy palm +rosy palm and her 5 sisters +rusty trombone +sadism +santorum +scat +schlong +scissoring +semen +sex +sexcam +sexo +sexy +sexual +sexually +sexuality +shaved beaver +shaved pussy +shemale +shibari +shit +shitblimp +shitty +shota +shrimping +skeet +slanteye +slut +s&m +smut +snatch +snowballing +sodomize +sodomy +spastic +spic +splooge +splooge moose +spooge +spread legs +spunk +strap on +strapon +strappado +strip club +style doggy +suck +sucks +suicide girls +sultry women +swastika +swinger +tainted love +taste my +tea bagging +threesome +throating +thumbzilla +tied up +tight white +tit +tits +titties +titty +tongue in a +topless +tosser +towelhead +tranny +tribadism +tub girl +tubgirl +tushy +twat +twink +twinkie +two girls one cup +undressing +upskirt +urethra play +urophilia +vagina +venus mound +viagra +vibrator +violet wand +vorarephilia +voyeur +voyeurweb +voyuer +vulva +wank +wetback +wet dream +white power +whore +worldsex +wrapping men +wrinkled starfish +xx +xxx +yaoi +yellow showers +yiffy +zoophilia +🖕 diff --git a/src-tauri/src/point_scan.rs b/src-tauri/src/point_scan.rs index fd2d3921..06d98cdf 100644 --- a/src-tauri/src/point_scan.rs +++ b/src-tauri/src/point_scan.rs @@ -26,6 +26,9 @@ pub struct Config { #[serde(deserialize_with = "crate::scan_preferences::deserialize_preferences")] pub scan_preferences: crate::scan_preferences::Preferences, pub word_prediction: bool, + /// Uses the on-device language model for word prediction. Off by default + /// because it needs several hundred megabytes of memory. + pub enhanced_word_prediction: bool, pub keyboard_wait_after_typing: bool, /// Moves the keyboard or mouse panel away while the pointer is over it. pub panel_avoids_pointer: bool, @@ -48,6 +51,7 @@ impl Default for Config { control_mode: ControlMode::Point, scan_preferences: Default::default(), word_prediction: true, + enhanced_word_prediction: false, keyboard_wait_after_typing: false, panel_avoids_pointer: false, scanner_color: Default::default(), @@ -109,6 +113,7 @@ impl Config { keyboard_scan: self.resolved(crate::scan_preferences::Area::Keyboard), mouse_scan: self.resolved(crate::scan_preferences::Area::Mouse), word_prediction: self.word_prediction, + enhanced_word_prediction: self.enhanced_word_prediction, keyboard_wait_after_typing: self.keyboard_wait_after_typing, panel_avoids_pointer: self.panel_avoids_pointer, scanner_color: self.resolved(crate::scan_preferences::Area::Point).color, @@ -135,6 +140,7 @@ pub struct PointSettings { pub keyboard_scan: crate::scan_preferences::Resolved, pub mouse_scan: crate::scan_preferences::Resolved, pub word_prediction: bool, + pub enhanced_word_prediction: bool, pub keyboard_wait_after_typing: bool, pub panel_avoids_pointer: bool, pub scanner_color: crate::scanning::ScannerColor, @@ -612,6 +618,14 @@ mod tests { assert!(!restored.word_prediction); } #[test] + fn enhanced_prediction_defaults_off_for_saved_settings_and_round_trips() { + let c: Config = serde_json::from_str(r#"{"wordPrediction":true}"#).unwrap(); + assert!(!c.enhanced_word_prediction); + let c: Config = serde_json::from_str(r#"{"enhancedWordPrediction":true}"#).unwrap(); + let restored: Config = serde_json::from_value(serde_json::to_value(c).unwrap()).unwrap(); + assert!(restored.enhanced_word_prediction); + } + #[test] fn keyboard_wait_defaults_off_and_round_trips() { let legacy: Config = serde_json::from_str("{}").unwrap(); assert!(!legacy.keyboard_wait_after_typing); @@ -1040,6 +1054,7 @@ mod tests { json["scannerColor"] = serde_json::json!("blue"); json["controlMode"] = serde_json::json!("point"); json["wordPrediction"] = serde_json::json!(true); + json["enhancedWordPrediction"] = serde_json::json!(false); json["keyboardWaitAfterTyping"] = serde_json::json!(false); json["panelAvoidsPointer"] = serde_json::json!(false); json["scanPreferences"] = diff --git a/src-tauri/src/point_scan_runtime.rs b/src-tauri/src/point_scan_runtime.rs index 888b2040..6b920790 100644 --- a/src-tauri/src/point_scan_runtime.rs +++ b/src-tauri/src/point_scan_runtime.rs @@ -282,7 +282,14 @@ impl Adapter for PointScan { crate::prediction::start_keyboard_activity(captured_keys); } let enabled = technique.prediction_enabled(); - crate::prediction::poll(app, technique.prediction_keyboard(), enabled, captured_keys); + let enhanced = technique.prediction_enhanced(); + crate::prediction::poll( + app, + technique.prediction_keyboard(), + enabled, + enhanced, + captured_keys, + ); } fn cleanup(_app: &AppHandle) -> Result<(), String> { crate::prediction::stop(); diff --git a/src-tauri/src/point_workflow.rs b/src-tauri/src/point_workflow.rs index d9fbb4d4..550daefa 100644 --- a/src-tauri/src/point_workflow.rs +++ b/src-tauri/src/point_workflow.rs @@ -308,6 +308,9 @@ impl Workflow { pub fn prediction_enabled(&self) -> bool { self.point.config.word_prediction } + pub fn prediction_enhanced(&self) -> bool { + self.point.config.enhanced_word_prediction + } pub fn set_keyboard_area(&mut self, area: Rect) { self.keyboard_area = area; } diff --git a/src-tauri/src/prediction/context.rs b/src-tauri/src/prediction/context.rs index d00de0ae..2b5418f3 100644 --- a/src-tauri/src/prediction/context.rs +++ b/src-tauri/src/prediction/context.rs @@ -4,8 +4,27 @@ pub struct Context { pub words: Vec, pub prefix: String, pub sentence_start: bool, + /// Buffered text before `prefix`, without a clipped leading fragment. + pub before: String, + /// Clipping left no complete earlier word, so an empty `before` is not + /// the true start of the text. + pub partial: bool, } pub fn extract(text: &str, clipped: bool) -> Context { + let context = extract_words(text, clipped); + let mut before = &text[..text.len() - context.prefix.len()]; + if clipped { + before = before + .find(char::is_whitespace) + .map_or("", |i| &before[i..]); + } + Context { + partial: clipped && before.trim().is_empty(), + before: before.to_owned(), + ..context + } +} +fn extract_words(text: &str, clipped: bool) -> Context { let boundary = text.rfind(['.', '!', '?', '\n', '\r', '。', '!', '?']); let start = boundary.map_or(0, |i| i + text[i..].chars().next().unwrap().len_utf8()); let tail = &text[start..]; @@ -31,6 +50,8 @@ pub fn extract(text: &str, clipped: bool) -> Context { .map(|(_, w)| (*w).to_owned()) .collect(), prefix, + before: String::new(), + partial: false, } } #[cfg(test)] @@ -43,5 +64,14 @@ mod tests { assert_eq!(c.prefix, "cafe\u{301}"); assert_eq!(extract("agment whole pa", true).words, ["whole"]); assert!(extract("Done! ", false).sentence_start); + assert_eq!( + extract("Hi. I would like wa", false).before, + "Hi. I would like " + ); + assert_eq!(extract("agment whole pa", true).before, " whole "); + assert_eq!(extract("agment", true).before, ""); + assert!(extract("agment", true).partial); + assert!(!extract("agment whole pa", true).partial); + assert!(!extract("wa", false).partial); } } diff --git a/src-tauri/src/prediction/database.rs b/src-tauri/src/prediction/database.rs index d240ec96..4747c7ad 100644 --- a/src-tauri/src/prediction/database.rs +++ b/src-tauri/src/prediction/database.rs @@ -1,29 +1,205 @@ -use super::{context::Context, lookup::Lookup}; -use std::path::Path; +use super::{ + context::Context, + lookup::Lookup, + model::{Model, Predict}, +}; +use std::{ + path::Path, + sync::mpsc, + time::{Duration, Instant}, +}; -pub struct Database(Lookup); +/// The model stops exploring new words after this long. +const SEARCH_BUDGET: Duration = Duration::from_millis(400); +/// A model call slower than this switches the worker to the lookup for good, +/// keeping replies well inside the main process's two-second deadline. +const SLOW_CALL: Duration = Duration::from_millis(1000); + +enum Enhanced { + Loading(mpsc::Receiver>), + Ready(Box), + Unavailable, +} + +// The lookup answers until the optional model has loaded in the background, +// and whenever the model fails or runs slowly, so it never delays typing. +pub struct Database { + lookup: Lookup, + enhanced: Enhanced, + slow_call: Duration, +} impl Database { - pub fn open(path: &Path) -> Result { - Lookup::open(path).map(Self) + pub fn open(path: &Path, enhanced: bool) -> Result { + let lookup = Lookup::open(path)?; + let enhanced = match path.parent().filter(|_| enhanced) { + Some(resources) => { + let model = resources.join("prediction-model"); + let blocklist = resources.join("prediction-blocklist.txt"); + let (tx, rx) = mpsc::sync_channel(1); + std::thread::spawn(move || { + let _ = tx.send(Model::open(&model, &blocklist)); + }); + Enhanced::Loading(rx) + } + None => Enhanced::Unavailable, + }; + Ok(Self { + lookup, + enhanced, + slow_call: SLOW_CALL, + }) } - pub fn predict(&self, context: &Context) -> Result, ()> { - Ok(self.0.predict(&context.words, &context.prefix)) + pub fn predict(&mut self, context: &Context) -> Result, ()> { + let base = self.lookup.predict(&context.words, &context.prefix); + if let Enhanced::Loading(rx) = &self.enhanced { + self.enhanced = match rx.try_recv() { + Ok(Ok(model)) => Enhanced::Ready(Box::new(model)), + Err(mpsc::TryRecvError::Empty) => return Ok(base), + _ => Enhanced::Unavailable, + }; + } + let Enhanced::Ready(model) = &mut self.enhanced else { + return Ok(base); + }; + if context.partial { + return Ok(base); + } + let lookup = &self.lookup; + let known = |w: &str| lookup.spelling(w).map(str::to_owned); + let start = Instant::now(); + let result = model.predict( + &context.before, + &context.prefix, + &known, + start + SEARCH_BUDGET, + ); + if start.elapsed() > self.slow_call { + self.enhanced = Enhanced::Unavailable; + } + match result { + Ok(mut words) => { + for word in base { + if words.len() < 5 + && !words + .iter() + .any(|w| w.to_lowercase() == word.to_lowercase()) + { + words.push(word); + } + } + words.truncate(5); + Ok(words) + } + Err(()) => Ok(base), + } } #[cfg(test)] pub fn fixture() -> Self { - Self(Lookup::fixture()) + Self { + lookup: Lookup::fixture(), + enhanced: Enhanced::Unavailable, + slow_call: SLOW_CALL, + } } } #[cfg(test)] mod tests { use super::*; + struct Fake(Result, ()>, Duration); + impl Predict for Fake { + fn predict( + &mut self, + _before: &str, + _prefix: &str, + _known: &dyn Fn(&str) -> Option, + _deadline: Instant, + ) -> Result, ()> { + std::thread::sleep(self.1); + self.0 + .clone() + .map(|w| w.into_iter().map(String::from).collect()) + } + } + fn ready(db: &mut Database, words: Result, ()>, delay: Duration) { + db.enhanced = Enhanced::Ready(Box::new(Fake(words, delay))); + } + #[test] + fn model_words_lead_and_the_lookup_fills_without_duplicates() { + let ctx = super::super::context::extract("I would like wa", false); + let mut db = Database::fixture(); + let lookup = db.predict(&ctx).unwrap(); + assert!(lookup.len() >= 2); + ready(&mut db, Ok(vec!["WhatsApp", "WATER"]), Duration::ZERO); + let words = db.predict(&ctx).unwrap(); + assert_eq!(words[..2], ["WhatsApp", "WATER"]); + assert!(!words[2..].iter().any(|w| w.eq_ignore_ascii_case("water"))); + assert!(words.len() <= 5); + ready( + &mut db, + Ok(vec!["a", "b", "c", "d", "e", "f"]), + Duration::ZERO, + ); + assert_eq!(db.predict(&ctx).unwrap(), ["a", "b", "c", "d", "e"]); + ready(&mut db, Err(()), Duration::ZERO); + assert_eq!(db.predict(&ctx).unwrap(), lookup); + assert!(matches!(db.enhanced, Enhanced::Ready(_))); + } + #[test] + fn slow_models_and_clipped_fragments_use_the_lookup() { + let mut db = Database::fixture(); + let clipped = super::super::context::extract("agmentwa", true); + let lookup = db.predict(&clipped).unwrap(); + ready(&mut db, Ok(vec!["model"]), Duration::ZERO); + assert_eq!(db.predict(&clipped).unwrap(), lookup); + let ctx = super::super::context::extract("I would like wa", false); + db.slow_call = Duration::from_millis(20); + ready(&mut db, Ok(vec!["model"]), Duration::from_millis(40)); + assert_eq!(db.predict(&ctx).unwrap()[0], "model"); + assert!(matches!(db.enhanced, Enhanced::Unavailable)); + assert_ne!(db.predict(&ctx).unwrap()[0], "model"); + } + #[test] + fn lookup_answers_while_the_model_loads_or_after_it_fails() { + let ctx = super::super::context::extract("I would like wa", false); + let mut db = Database::fixture(); + let expected = db.predict(&ctx).unwrap(); + let (tx, rx) = mpsc::sync_channel(1); + db.enhanced = Enhanced::Loading(rx); + assert_eq!(db.predict(&ctx).unwrap(), expected); + assert!(matches!(db.enhanced, Enhanced::Loading(_))); + tx.send(Err(())).unwrap(); + assert_eq!(db.predict(&ctx).unwrap(), expected); + assert!(matches!(db.enhanced, Enhanced::Unavailable)); + let (tx, rx) = mpsc::sync_channel::>(1); + db.enhanced = Enhanced::Loading(rx); + drop(tx); + assert_eq!(db.predict(&ctx).unwrap(), expected); + assert!(matches!(db.enhanced, Enhanced::Unavailable)); + } + #[test] + fn a_missing_model_leaves_the_lookup_working() { + let dir = std::env::temp_dir().join(format!("switchify-model-{}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + let lookup = dir.join("word-predictions.lookup"); + std::fs::write(&lookup, Lookup::fixture_bytes()).unwrap(); + let mut db = Database::open(&lookup, true).unwrap(); + let ctx = super::super::context::extract("I would like wa", false); + let deadline = std::time::Instant::now() + std::time::Duration::from_secs(5); + while matches!(db.enhanced, Enhanced::Loading(_)) && std::time::Instant::now() < deadline { + assert!(!db.predict(&ctx).unwrap().is_empty()); + std::thread::sleep(std::time::Duration::from_millis(10)); + } + assert!(matches!(db.enhanced, Enhanced::Unavailable)); + assert!(!db.predict(&ctx).unwrap().is_empty()); + std::fs::remove_dir_all(dir).unwrap(); + } #[test] fn bundled_lookup_matches_sqlite_corpus_and_expanded_samples() { let root = Path::new(env!("CARGO_MANIFEST_DIR")).join("resources"); let source = root.join("WordData2017051601.db"); let reference = super::super::database_reference::Database::open(&source).unwrap(); - let lookup = Database::open(&root.join("word-predictions.lookup")).unwrap(); + let mut lookup = Database::open(&root.join("word-predictions.lookup"), false).unwrap(); let c = rusqlite::Connection::open_with_flags( source, rusqlite::OpenFlags::SQLITE_OPEN_READ_ONLY, diff --git a/src-tauri/src/prediction/lookup.rs b/src-tauri/src/prediction/lookup.rs index aad8b820..13283ea0 100644 --- a/src-tauri/src/prediction/lookup.rs +++ b/src-tauri/src/prediction/lookup.rs @@ -226,6 +226,11 @@ impl Lookup { } result } + /// The most frequent spelling of `word`, compared case-insensitively. + pub fn spelling(&self, word: &str) -> Option<&str> { + let rank = *self.ids.get(&word.to_lowercase())?; + Some(&self.words[rank as usize].text) + } fn push(&self, rank: u32, prefix: &str, result: &mut Vec, seen: &mut HashSet) { let word = &self.words[rank as usize]; if result.len() < 5 @@ -246,6 +251,9 @@ impl Lookup { pub fn fixture() -> Self { Self::read(&tests::fixture_bytes()[..]).unwrap() } + pub fn fixture_bytes() -> Vec { + tests::fixture_bytes() + } } #[cfg(test)] mod tests { diff --git a/src-tauri/src/prediction/mod.rs b/src-tauri/src/prediction/mod.rs index 91569cc5..e7bcf086 100644 --- a/src-tauri/src/prediction/mod.rs +++ b/src-tauri/src/prediction/mod.rs @@ -4,6 +4,7 @@ mod database; #[cfg(test)] mod database_reference; mod lookup; +mod model; pub mod worker; use crate::scan_keyboard::{Key, Keyboard, Modifier, Page, Stroke}; @@ -65,12 +66,13 @@ struct Client { _job: Job, } impl Client { - fn start(path: &Path, ignored: Vec) -> Result { + fn start(path: &Path, ignored: Vec, enhanced: bool) -> Result { let mut command = Command::new(std::env::current_exe().map_err(|_| ())?); command .arg(worker::ARG) .arg(path) .arg(serde_json::to_string(&ignored).map_err(|_| ())?) + .arg(serde_json::to_string(&enhanced).map_err(|_| ())?) .stdin(Stdio::piped()) .stdout(Stdio::piped()) .stderr(Stdio::null()); @@ -141,6 +143,7 @@ impl Drop for Client { struct Service { client: Option, failed: bool, + received_reply: bool, generation: u64, outstanding: Option, last: Option, @@ -152,6 +155,8 @@ struct Service { accepting: bool, case: Option<(bool, bool, bool)>, tracking: bool, + /// Whether the running worker was started with the language model. + enhanced: bool, } #[derive(Clone, Copy)] pub struct InputScope { @@ -191,6 +196,18 @@ pub fn reset() { }); } impl Service { + fn response_timed_out(&self, now: Instant) -> bool { + // The worker validates and loads the bundled lookup before its first + // reply. On slower Windows machines that takes longer than a query. + let deadline = if self.received_reply { + Duration::from_secs(2) + } else { + Duration::from_secs(30) + }; + self.outstanding + .is_some_and(|sent| now.duration_since(sent) >= deadline) + } + fn queue_edit(&mut self, edit: Edit, scope: InputScope) { self.edit_revision = self.edit_revision.wrapping_add(1); if self.edit.len() >= 512 { @@ -378,7 +395,13 @@ fn database_resource( .ok_or(()) } -pub fn poll(app: &AppHandle, keyboard: Option<&mut Keyboard>, enabled: bool, ignored: &[String]) { +pub fn poll( + app: &AppHandle, + keyboard: Option<&mut Keyboard>, + enabled: bool, + enhanced: bool, + ignored: &[String], +) { let Some(keyboard) = keyboard else { stop(); return; @@ -390,6 +413,17 @@ pub fn poll(app: &AppHandle, keyboard: Option<&mut Keyboard>, enabled: bool, ign } SERVICE.with(|slot| { let mut s = slot.borrow_mut(); + if s.enhanced != enhanced { + // Changing the engine restarts the worker, like reopening the keyboard. + if s.accepting || s.accept.is_some() { + keyboard.failed(); + } + *s = Service { + enhanced, + ..Service::default() + }; + keyboard.predictions(None, false); + } let case = ( keyboard.prediction_shift(), keyboard.caps, @@ -411,7 +445,7 @@ pub fn poll(app: &AppHandle, keyboard: Option<&mut Keyboard>, enabled: bool, ign .filter_map(|name| crate::switch_input::prediction_key_code(name)) .collect(); s.client = resource(app) - .and_then(|path| Client::start(&path, ignored)) + .and_then(|path| Client::start(&path, ignored, enhanced)) .ok(); if s.client.is_none() { s.failed = true; @@ -419,9 +453,7 @@ pub fn poll(app: &AppHandle, keyboard: Option<&mut Keyboard>, enabled: bool, ign return; } } - if s.outstanding - .is_some_and(|t| t.elapsed() >= Duration::from_secs(2)) - { + if s.response_timed_out(Instant::now()) { s.fail(keyboard); return; } @@ -429,6 +461,7 @@ pub fn poll(app: &AppHandle, keyboard: Option<&mut Keyboard>, enabled: bool, ign match reply { Ok(Ok(response)) => { s.outstanding = None; + s.received_reply = true; match response { Response::Suggestions { generation, @@ -572,6 +605,7 @@ mod tests { SERVICE.with(|slot| { let s = slot.borrow(); assert!(!s.failed); + assert!(!s.received_reply); assert!(!s.accepting); assert!(s.accept.is_none()); assert!(s.outstanding.is_none()); @@ -580,6 +614,24 @@ mod tests { } } #[test] + fn first_worker_reply_has_a_startup_deadline_then_queries_use_two_seconds() { + let now = Instant::now(); + let mut service = Service { + outstanding: Some(now - Duration::from_secs(7)), + ..Default::default() + }; + assert!(!service.response_timed_out(now)); + service.outstanding = Some(now - Duration::from_secs(30)); + assert!(service.response_timed_out(now)); + service.received_reply = true; + service.outstanding = Some(now - Duration::from_secs(1)); + assert!(!service.response_timed_out(now)); + service.outstanding = Some(now - Duration::from_secs(2)); + assert!(service.response_timed_out(now)); + service.outstanding = None; + assert!(!service.response_timed_out(now)); + } + #[test] fn only_successful_supported_edits_enter_the_buffer() { stop(); let stroke = Stroke { diff --git a/src-tauri/src/prediction/model.rs b/src-tauri/src/prediction/model.rs new file mode 100644 index 00000000..04cb779f --- /dev/null +++ b/src-tauri/src/prediction/model.rs @@ -0,0 +1,633 @@ +//! Optional on-device language model for word prediction. +//! +//! The model spells candidate words from subword pieces, constrained to the +//! typed prefix, so suggestions are not limited to the bundled lookup's words. +//! A word scores log P(pieces | text) + log P(a word boundary follows), with +//! case variants merged. Words the lookup does not know must reach a minimum +//! probability, which keeps rare invented words out; blocked words never +//! appear. Runs only inside the prediction worker; text is never logged. +use ort::{ + session::{builder::GraphOptimizationLevel, Session, SessionInputValue}, + value::Tensor, +}; +use std::{ + borrow::Cow, + collections::{HashMap, HashSet}, + path::Path, + time::Instant, +}; +use tokenizers::Tokenizer; + +const BEAM: usize = 8; +const MAX_PIECES: usize = 5; +/// Pieces below this log-probability are never explored. +const PRUNE: f32 = -12.0; +/// Minimum log-probability for words the lookup does not contain. +const UNKNOWN_MIN: f32 = -10.0; +/// Most recent buffered characters the model reads. +const CONTEXT_CHARS: usize = 256; + +pub trait Scorer { + /// Next-piece log-probabilities after `text`, which becomes the base for `extend`. + fn context(&mut self, text: &str) -> Result, ()>; + /// Next-piece log-probabilities after each continuation of the last context. + fn extend(&mut self, continuations: &[Vec]) -> Result>, ()>; +} + +pub struct Vocabulary { + pieces: Vec, + /// Pieces that can end a word: space-led, punctuation or special tokens. + boundary: Vec, + /// Pieces made only of ASCII letters and apostrophes. + continues: Vec, + blocked: HashSet, +} + +impl Vocabulary { + pub fn new(pieces: Vec, blocked: HashSet) -> Self { + let wordy = |s: &str| { + !s.is_empty() + && s.chars() + .all(|c| c.is_ascii_alphabetic() || c == '\'' || c == '’') + }; + Self { + boundary: pieces + .iter() + .map(|p| { + !p.chars() + .next() + .is_some_and(|c| c.is_alphanumeric() || c == '\'' || c == '’') + }) + .collect(), + continues: pieces.iter().map(|p| wordy(p)).collect(), + pieces, + blocked, + } + } +} + +fn normalize(word: &str) -> String { + word.to_lowercase().replace('’', "'") +} + +/// Up to five words for `prefix` after `before`, most probable first. +/// `known` returns the lookup's spelling of a word it contains. +/// Expansion stops at `deadline`, returning the words finished so far. +pub fn spell( + scorer: &mut impl Scorer, + vocabulary: &Vocabulary, + before: &str, + prefix: &str, + known: impl Fn(&str) -> Option, + deadline: Instant, +) -> Result, ()> { + let typed = prefix; + let prefix = normalize(prefix); + let text = before.trim_end(); + // A word after the last one needs a leading space; at the very start of + // text a capital marks a word start, since fragments are never capitalised. + let start = text.is_empty(); + if !start && text.len() == before.len() { + return Err(()); + } + let compatible = |word: &str| { + let w = normalize(word); + w.starts_with(&prefix) || prefix.starts_with(&w) + }; + let next = scorer.context(text)?; + if next.len() != vocabulary.pieces.len() { + return Err(()); + } + let mut beam: Vec<(f32, Vec, String)> = Vec::new(); + for (id, piece) in vocabulary.pieces.iter().enumerate() { + let word = if start { + piece.as_str() + } else { + match piece.strip_prefix(' ') { + Some(word) => word, + None => continue, + } + }; + let first = word.chars().next(); + let begins = if start { + first.is_some_and(char::is_uppercase) + } else { + first.is_some_and(|c| c.is_ascii_alphabetic()) + }; + if begins + && next[id] > PRUNE + && word + .chars() + .all(|c| c.is_ascii_alphabetic() || c == '\'' || c == '’') + && compatible(word) + { + beam.push((next[id], vec![id as i64], word.to_owned())); + } + } + // Probability mass per normalized word, plus its most probable spelling. + let mut finished: HashMap = HashMap::new(); + let fifth = |finished: &HashMap| { + // Only words that can survive the final filter may prune the beam. + let mut p: Vec = finished + .iter() + .filter(|(key, f)| f.0.ln() >= UNKNOWN_MIN || known(key).is_some()) + .map(|(_, f)| f.0) + .collect(); + p.sort_by(|a, b| b.total_cmp(a)); + p.get(4).map_or(f32::NEG_INFINITY, |p| p.ln()) + }; + for _ in 0..MAX_PIECES { + if Instant::now() >= deadline { + break; + } + beam.sort_by(|a, b| b.0.total_cmp(&a.0)); + let bound = fifth(&finished); + beam.retain(|b| b.0 > bound); + beam.truncate(BEAM); + if beam.is_empty() { + break; + } + let rows = scorer.extend(&beam.iter().map(|b| b.1.clone()).collect::>())?; + if rows.len() != beam.len() { + return Err(()); + } + let mut grown = Vec::new(); + for ((logp, pieces, word), row) in beam.iter().zip(&rows) { + if row.len() != vocabulary.pieces.len() { + return Err(()); + } + let key = normalize(word); + let plausible = key.chars().count() > 1 || known(&key).is_some(); + if plausible && key.starts_with(&prefix) && !vocabulary.blocked.contains(&key) { + let boundary: f32 = row + .iter() + .zip(&vocabulary.boundary) + .filter(|(_, b)| **b) + .map(|(l, _)| l.exp()) + .sum(); + let score = logp + boundary.max(1e-12).ln(); + let entry = finished + .entry(key) + .or_insert((0.0, f32::NEG_INFINITY, String::new())); + entry.0 += score.exp(); + if score > entry.1 { + entry.1 = score; + entry.2 = word.clone(); + } + } + for (id, piece) in vocabulary.pieces.iter().enumerate() { + if vocabulary.continues[id] && row[id] > PRUNE { + let longer = format!("{word}{piece}"); + if compatible(&longer) { + let mut pieces = pieces.clone(); + pieces.push(id as i64); + grown.push((logp + row[id], pieces, longer)); + } + } + } + } + beam = grown; + } + let mut words: Vec<(f32, String, String)> = finished + .into_iter() + .filter_map(|(key, (p, _, spelling))| match known(&key) { + // Known words keep the lookup's usual spelling, so "who" is not "WHO". + Some(known) => Some((p, key, known)), + None => (p.ln() >= UNKNOWN_MIN).then_some((p, key, spelling)), + }) + .collect(); + words.sort_by(|a, b| b.0.total_cmp(&a.0).then_with(|| a.1.cmp(&b.1))); + Ok(words + .into_iter() + .take(5) + .map(|(_, key, spelling)| { + // Capitals on unknown words at a sentence start come from position. + let word = if known(&key).is_none() && (start || text.ends_with(['.', '!', '?'])) { + key + } else { + spelling.replace('’', "'") + }; + // Match the typed apostrophe so the completion still extends the prefix. + if typed.contains('’') { + word.replace('\'', "’") + } else { + word + } + }) + .collect()) +} + +struct Past { + keys: Vec>, + values: Vec>, + len: usize, +} + +pub struct Onnx { + session: Session, + tokenizer: Tokenizer, + layers: usize, + heads: usize, + dim: usize, + bos: i64, + context: Option<(String, Past, Vec)>, +} + +fn log_softmax(logits: &[f32]) -> Vec { + let max = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max); + let norm = max + logits.iter().map(|l| (l - max).exp()).sum::().ln(); + logits.iter().map(|l| l - norm).collect() +} + +impl Onnx { + fn open(dir: &Path) -> Result { + // Typed text must never leave the worker, so ONNX Runtime telemetry stays off. + ort::init().with_telemetry(false).commit(); + let threads = std::thread::available_parallelism().map_or(1, |n| n.get().min(4)); + let builder = Session::builder().map_err(|_| ())?; + let builder = builder + .with_optimization_level(GraphOptimizationLevel::Level3) + .map_err(|_| ())?; + let mut builder = builder.with_intra_threads(threads).map_err(|_| ())?; + let session = builder + .commit_from_file(dir.join("model.onnx")) + .map_err(|_| ())?; + let tokenizer = Tokenizer::from_file(dir.join("tokenizer.json")).map_err(|_| ())?; + let config: serde_json::Value = serde_json::from_str( + &std::fs::read_to_string(dir.join("config.json")).map_err(|_| ())?, + ) + .map_err(|_| ())?; + let number = |key: &str| config[key].as_u64().map(|n| n as usize).ok_or(()); + let layers = number("num_hidden_layers")?; + let names: HashSet<&str> = session.inputs().iter().map(|i| i.name()).collect(); + let required = ["input_ids", "attention_mask", "position_ids"]; + if !required.iter().all(|n| names.contains(n)) + || !(0..layers).all(|l| names.contains(format!("past_key_values.{l}.key").as_str())) + { + return Err(()); + } + Ok(Self { + layers, + heads: number("num_key_value_heads")?, + dim: number("head_dim")?, + bos: number("bos_token_id")? as i64, + session, + tokenizer, + context: None, + }) + } + + /// Runs `batch` rows of `width` pieces after `past`; returns the last row's + /// log-probabilities for each batch item (at `ends`) and the new past. + fn run( + &mut self, + past: Option<&Past>, + ids: Vec, + batch: usize, + width: usize, + ends: &[usize], + ) -> Result<(Vec>, Option), ()> { + let before = past.map_or(0, |p| p.len); + let tensor = + |shape: [usize; 2], data: Vec| -> Result, ()> { + Ok(Tensor::from_array((shape, data)).map_err(|_| ())?.into()) + }; + let mut inputs: Vec<(Cow, SessionInputValue)> = vec![ + ("input_ids".into(), tensor([batch, width], ids)?), + ( + "attention_mask".into(), + tensor([batch, before + width], vec![1; batch * (before + width)])?, + ), + ( + "position_ids".into(), + tensor( + [batch, width], + (0..batch) + .flat_map(|_| (before..before + width).map(|p| p as i64)) + .collect(), + )?, + ), + ]; + for layer in 0..self.layers { + for (kind, cache) in [ + ("key", past.map(|p| &p.keys)), + ("value", past.map(|p| &p.values)), + ] { + let data: Vec = cache.map_or(Vec::new(), |c| { + c[layer] + .iter() + .copied() + .cycle() + .take(c[layer].len() * batch) + .collect() + }); + let value = Tensor::from_array(([batch, self.heads, before, self.dim], data)) + .map_err(|_| ())?; + inputs.push(( + format!("past_key_values.{layer}.{kind}").into(), + value.into(), + )); + } + } + let outputs = self.session.run(inputs).map_err(|_| ())?; + let (_, logits) = outputs["logits"] + .try_extract_tensor::() + .map_err(|_| ())?; + let size = logits.len() / (batch * width); + let rows = ends + .iter() + .enumerate() + .map(|(b, end)| log_softmax(&logits[(b * width + end) * size..][..size])) + .collect(); + let present = if batch == 1 { + let mut keys = Vec::new(); + let mut values = Vec::new(); + for layer in 0..self.layers { + for (kind, out) in [("key", &mut keys), ("value", &mut values)] { + let (_, data) = outputs[format!("present.{layer}.{kind}").as_str()] + .try_extract_tensor::() + .map_err(|_| ())?; + out.push(data.to_vec()); + } + } + Some(Past { + keys, + values, + len: before + width, + }) + } else { + None + }; + Ok((rows, present)) + } +} + +impl Scorer for Onnx { + fn context(&mut self, text: &str) -> Result, ()> { + if let Some((cached, _, next)) = &self.context { + if cached == text { + return Ok(next.clone()); + } + } + self.context = None; + let mut ids = vec![self.bos]; + let encoding = self.tokenizer.encode(text, false).map_err(|_| ())?; + ids.extend(encoding.get_ids().iter().map(|&i| i as i64)); + let width = ids.len(); + let (mut rows, past) = self.run(None, ids, 1, width, &[width - 1])?; + let next = rows.pop().ok_or(())?; + self.context = Some((text.to_owned(), past.ok_or(())?, next.clone())); + Ok(next) + } + fn extend(&mut self, continuations: &[Vec]) -> Result>, ()> { + let (text, past, next) = self.context.take().ok_or(())?; + let width = continuations.iter().map(Vec::len).max().ok_or(())?; + let ids = continuations + .iter() + .flat_map(|c| c.iter().copied().chain(std::iter::repeat(0)).take(width)) + .collect(); + let ends: Vec = continuations.iter().map(|c| c.len() - 1).collect(); + let result = self.run(Some(&past), ids, continuations.len(), width, &ends); + self.context = Some((text, past, next)); + result.map(|(rows, _)| rows) + } +} + +pub struct Model { + onnx: Onnx, + vocabulary: Vocabulary, +} + +impl Model { + pub fn open(dir: &Path, blocklist: &Path) -> Result { + let onnx = Onnx::open(dir)?; + let size = onnx.tokenizer.get_vocab_size(true) as u32; + let pieces = (0..size) + .map(|id| onnx.tokenizer.decode(&[id], false).unwrap_or_default()) + .collect(); + let blocked = std::fs::read_to_string(blocklist) + .map_err(|_| ())? + .lines() + .map(|l| normalize(l.trim())) + .filter(|l| !l.is_empty()) + .collect(); + Ok(Self { + onnx, + vocabulary: Vocabulary::new(pieces, blocked), + }) + } +} + +/// The enhanced predictor behind `Database`, so tests can substitute a fake. +pub trait Predict: Send { + fn predict( + &mut self, + before: &str, + prefix: &str, + known: &dyn Fn(&str) -> Option, + deadline: Instant, + ) -> Result, ()>; +} + +impl Predict for Model { + fn predict( + &mut self, + before: &str, + prefix: &str, + known: &dyn Fn(&str) -> Option, + deadline: Instant, + ) -> Result, ()> { + spell( + &mut self.onnx, + &self.vocabulary, + recent(before).ok_or(())?, + prefix, + known, + deadline, + ) + } +} + +/// At most the last `CONTEXT_CHARS` characters of `before`, starting at a +/// word boundary. This bounds the one model pass that has no deadline. +/// `None` when the kept text holds no complete word. +fn recent(before: &str) -> Option<&str> { + let Some((cut, _)) = before.char_indices().rev().nth(CONTEXT_CHARS - 1) else { + return Some(before); + }; + let tail = &before[cut..]; + let start = tail.find(char::is_whitespace)?; + let kept = tail[start..].trim_start(); + (!kept.is_empty()).then_some(kept) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Pieces: 0 "", 1 " Wa", 2 " wa", 3 "ter", 4 "ffle", 5 " .", 6 " water", + /// 7 "Wa", 8 "zz", 9 " darn", 10 " x". Everything after a word is a boundary. + struct Fake { + next: Vec, + contexts: usize, + } + fn pieces() -> Vec { + [ + "", " Wa", " wa", "ter", "ffle", " .", " water", "Wa", "zz", " darn", " x", + ] + .map(String::from) + .to_vec() + } + fn dist(pairs: &[(usize, f32)]) -> Vec { + let mut d = vec![-30.0; 11]; + for &(i, p) in pairs { + d[i] = p.ln(); + } + d + } + impl Scorer for Fake { + fn context(&mut self, _text: &str) -> Result, ()> { + self.contexts += 1; + Ok(self.next.clone()) + } + fn extend(&mut self, continuations: &[Vec]) -> Result>, ()> { + Ok(continuations + .iter() + .map(|c| match c.last() { + Some(1) | Some(2) | Some(7) => dist(&[(3, 0.6), (4, 0.3), (5, 0.1)]), + _ => dist(&[(5, 1.0)]), + }) + .collect()) + } + } + fn vocabulary() -> Vocabulary { + Vocabulary::new(pieces(), ["darn".to_owned()].into()) + } + fn fake(pairs: &[(usize, f32)]) -> Fake { + Fake { + next: dist(pairs), + contexts: 0, + } + } + + fn run( + s: &mut Fake, + before: &str, + prefix: &str, + known: impl Fn(&str) -> Option, + ) -> Result, ()> { + let deadline = Instant::now() + std::time::Duration::from_secs(60); + spell(s, &vocabulary(), before, prefix, known, deadline) + } + fn lower(w: &str) -> Option { + Some(w.to_owned()) + } + + #[test] + fn spells_multi_piece_words_merges_case_and_applies_prefix() { + let mut s = fake(&[(1, 0.2), (2, 0.3), (6, 0.4), (9, 0.1)]); + // water = 0.4 + 0.5 * 0.6; waffle = 0.5 * 0.3; bare "wa" = 0.5 * 0.1. + let words = run(&mut s, "I would like ", "wa", lower).unwrap(); + assert_eq!(words, ["water", "waffle", "wa"]); + assert_eq!( + run(&mut s, "I would like ", "waf", lower).unwrap(), + ["waffle"] + ); + } + + #[test] + fn blocked_single_letter_and_improbable_unknown_words_are_removed() { + let mut s = fake(&[(9, 0.6), (10, 0.3), (6, 0.1)]); + assert_eq!(run(&mut s, "Say ", "", |_| None).unwrap(), ["water"]); + let mut s = fake(&[(6, 1e-5)]); + assert!(run(&mut s, "Say ", "", |_| None).unwrap().is_empty()); + assert_eq!(run(&mut s, "Say ", "", lower).unwrap(), ["water"]); + } + + #[test] + fn casing_comes_from_the_lookup_or_the_model_but_not_sentence_position() { + let mut s = fake(&[(1, 0.9)]); + let lowered = ["water", "waffle", "wa"]; + assert_eq!(run(&mut s, "Send ", "", lower).unwrap(), lowered); + assert_eq!( + run(&mut s, "Send ", "", |_| None).unwrap(), + ["Water", "Waffle", "Wa"] + ); + assert_eq!(run(&mut s, "Done. ", "", |_| None).unwrap(), lowered); + } + + #[test] + fn text_start_needs_capital_and_attached_prefixes_fall_back() { + let mut s = fake(&[(7, 0.5), (8, 0.5)]); + assert_eq!( + run(&mut s, "", "", lower).unwrap(), + ["water", "waffle", "wa"] + ); + assert!(run(&mut s, "hello-", "wa", lower).is_err()); + assert_eq!(s.contexts, 1); + } + + #[test] + fn an_expired_deadline_stops_expansion_and_bad_vocabularies_fail() { + let mut s = fake(&[(6, 0.9)]); + let expired = spell(&mut s, &vocabulary(), "Say ", "", lower, Instant::now()); + assert!(expired.unwrap().is_empty()); + s.next.pop(); + assert!(run(&mut s, "Say ", "", lower).is_err()); + } + + #[test] + fn model_context_keeps_recent_whole_words_only() { + assert_eq!(recent("I would like "), Some("I would like ")); + let long = format!("{} tea please ", "x".repeat(300)); + assert_eq!(recent(&long), Some("tea please ")); + let words = "word ".repeat(100); + let kept = recent(&words).unwrap(); + assert!(kept.chars().count() < CONTEXT_CHARS && kept.starts_with("word")); + assert_eq!(recent(&"é".repeat(300)), None); + assert_eq!(recent(&format!("{} ", "x".repeat(300))), None); + } + + #[test] + fn apostrophes_continue_words_rather_than_ending_them() { + let pieces = [" don", "’t", "'s", "’", " .", ","] + .map(String::from) + .to_vec(); + let v = Vocabulary::new(pieces, HashSet::new()); + assert_eq!(v.boundary, [true, false, false, false, true, true]); + assert_eq!(v.continues, [false, true, true, true, false, false]); + } + + #[test] + fn bundled_model_suggests_current_words() { + let root = Path::new(env!("CARGO_MANIFEST_DIR")).join("resources"); + let start = Instant::now(); + let mut model = Model::open( + &root.join("prediction-model"), + &root.join("prediction-blocklist.txt"), + ) + .expect("run `npm run prediction-model` to fetch the model"); + println!("model startup_ms={}", start.elapsed().as_millis()); + let lookup = + super::super::lookup::Lookup::open(&root.join("word-predictions.lookup")).unwrap(); + let known = |w: &str| lookup.spelling(w).map(str::to_owned); + let mut times = Vec::new(); + for (before, prefix, expected) in [ + ("I would like a cup of ", "", Some("tea")), + ("Please put the ", "ket", Some("kettle")), + ("Can you send a ", "wh", None), + ("I want to ", "fa", None), + ] { + let t = Instant::now(); + let deadline = t + std::time::Duration::from_secs(60); + let words = model.predict(before, prefix, &known, deadline).unwrap(); + times.push(t.elapsed().as_secs_f64() * 1000.0); + println!("{before:?}+{prefix:?} -> {words:?}"); + assert!(words.iter().all(|w| normalize(w).starts_with(prefix))); + if let Some(expected) = expected { + assert!(words.iter().any(|w| w == expected), "{words:?}"); + } + } + println!("model predict_ms={times:.1?}"); + } +} diff --git a/src-tauri/src/prediction/worker.rs b/src-tauri/src/prediction/worker.rs index d6dba19d..db815b5d 100644 --- a/src-tauri/src/prediction/worker.rs +++ b/src-tauri/src/prediction/worker.rs @@ -302,11 +302,12 @@ pub fn run_from_args() -> bool { return false; } let work = || -> Result<(), ()> { - if args.len() != 4 { + if args.len() != 5 { return Err(()); } let path = PathBuf::from(&args[2]); let ignored: Vec = serde_json::from_str(&args[3].to_string_lossy()).map_err(|_| ())?; + let enhanced: bool = serde_json::from_str(&args[4].to_string_lossy()).map_err(|_| ())?; if ignored.len() > 129 { return Err(()); } @@ -324,7 +325,7 @@ pub fn run_from_args() -> bool { let tracked = activity::start(); #[cfg(any(target_os = "windows", target_os = "macos"))] { - let database = Database::open(&path)?; + let database = Database::open(&path, enhanced)?; let mut engine = Engine::new(database, tracked); let mut input = std::io::stdin().lock(); let mut output = std::io::stdout().lock(); @@ -481,9 +482,10 @@ mod tests { #[ignore = "Manual database performance measurement; no desktop input"] fn bundled_database_benchmark() { let start = std::time::Instant::now(); - let db = Database::open( + let mut db = Database::open( &std::path::Path::new(env!("CARGO_MANIFEST_DIR")) .join("resources/word-predictions.lookup"), + false, ) .unwrap(); let startup = start.elapsed().as_secs_f64() * 1000.0; diff --git a/src-tauri/tauri.conf.json b/src-tauri/tauri.conf.json index 16ca554f..309011d4 100644 --- a/src-tauri/tauri.conf.json +++ b/src-tauri/tauri.conf.json @@ -4,9 +4,9 @@ "version": "1.0.0-rc.15", "identifier": "com.enaboapps.switchify.pc", "build": { - "beforeDevCommand": "npm run dev", + "beforeDevCommand": "npm run prediction-model && npm run dev", "devUrl": "http://localhost:1420", - "beforeBuildCommand": "npm run build", + "beforeBuildCommand": "npm run prediction-model && npm run build", "frontendDist": "../dist" }, "app": { @@ -35,7 +35,13 @@ } }, "bundle": { - "resources": ["resources/word-predictions.lookup"], + "resources": [ + "resources/word-predictions.lookup", + "resources/prediction-blocklist.txt", + "resources/prediction-model/model.onnx", + "resources/prediction-model/tokenizer.json", + "resources/prediction-model/config.json" + ], "active": true, "targets": "all", "icon": [ diff --git a/src/PointScan.test.tsx b/src/PointScan.test.tsx index 37d497e2..06c9ccc7 100644 --- a/src/PointScan.test.tsx +++ b/src/PointScan.test.tsx @@ -277,6 +277,16 @@ it("customises individual settings, restores defaults and returns focus without expect(screen.getByRole("heading", { name: "Keyboard" })).toHaveFocus(); expect(screen.queryByRole("checkbox", { name: "Auto select" })).not.toBeInTheDocument(); expect(screen.getByRole("checkbox", { name: "Word prediction" })).toBeEnabled(); + expect(screen.getByRole("checkbox", { name: "Enhanced word prediction" })).not.toBeChecked(); + expect(screen.getByText(/Your typing never leaves this computer/)).toBeInTheDocument(); + fireEvent.click(screen.getByRole("checkbox", { name: "Enhanced word prediction" })); + await waitFor(() => expect(mocks.invoke).toHaveBeenLastCalledWith("configure_point_scan", { + config: expect.objectContaining({ enhancedWordPrediction: true }) + })); + fireEvent.click(screen.getByRole("checkbox", { name: "Word prediction" })); + expect(screen.getByRole("checkbox", { name: "Enhanced word prediction" })).toBeDisabled(); + fireEvent.click(screen.getByRole("checkbox", { name: "Word prediction" })); + mocks.invoke.mockClear(); expect(screen.getByText(/The scanning keyboard spaces punctuation for you/)).toBeInTheDocument(); fireEvent.click(screen.getByRole("button", { name: "Reverse" })); await waitFor(() => expect(mocks.invoke).toHaveBeenLastCalledWith("configure_point_scan", { diff --git a/src/scanning/useScanning.ts b/src/scanning/useScanning.ts index 0c6bab34..a49261cc 100644 --- a/src/scanning/useScanning.ts +++ b/src/scanning/useScanning.ts @@ -6,6 +6,7 @@ export type PointScanConfig = { controlMode: "point" | "mouse"; scanPreferences?: import("./preferences").ScanPreferences; wordPrediction: boolean; + enhancedWordPrediction: boolean; keyboardWaitAfterTyping: boolean; panelAvoidsPointer: boolean; scannerColor: ScannerColor; @@ -33,6 +34,7 @@ export type PointScanState = { export const defaultPointScanConfig: PointScanConfig = { controlMode: "point", wordPrediction: true, + enhancedWordPrediction: false, keyboardWaitAfterTyping: false, panelAvoidsPointer: false, scannerColor: "blue", diff --git a/src/settings/ScannerPreferences.tsx b/src/settings/ScannerPreferences.tsx index 936b9c8c..45cb12af 100644 --- a/src/settings/ScannerPreferences.tsx +++ b/src/settings/ScannerPreferences.tsx @@ -136,6 +136,8 @@ export function ScannerPreferences({ controller }: { controller: ScanningControl } {area === 'keyboard' && update('wordPrediction', value)} /> + update('enhancedWordPrediction', value)} /> +

Uses a language model on this computer to suggest more current words. Your typing never leaves this computer and is not used for learning. Uses about 400 MB of extra memory while the keyboard is open.

} {field('automatic', 'Automatic scanning', locked => change('automatic', value)} />)}