From a7e56cbf8915e08ca569533b2161b05f4403930f Mon Sep 17 00:00:00 2001 From: Owen McGirr Date: Thu, 24 Sep 2026 21:21:16 +0100 Subject: [PATCH 1/4] Add optional on-device language-model word prediction Enhanced word prediction runs SmolLM2-135M (int8 ONNX) inside the existing prediction worker. The model spells words from subword pieces, so it can suggest current words the 2017 lookup lacks. Filters keep out fragments, blocked words and improbable unknown words. The lookup answers while the model loads and whenever it fails. The setting is off by default because it needs about 300-450 MB more memory. The model is fetched from a pinned revision with SHA-256 verification, because it exceeds GitHub's file size limit. Refs #878 Co-Authored-By: Claude Opus 5.5 (1M context) --- .github/workflows/ci.yml | 2 + .github/workflows/release-tauri.yml | 6 + .gitignore | 1 + docs/word-prediction.md | 25 + package.json | 1 + scripts/fetch-prediction-model.mjs | 57 ++ src-tauri/Cargo.lock | 526 +++++++++++++++++- src-tauri/Cargo.toml | 2 + src-tauri/resources/PROVENANCE.md | 14 + src-tauri/resources/prediction-blocklist.txt | 403 ++++++++++++++ src-tauri/src/point_scan.rs | 15 + src-tauri/src/point_scan_runtime.rs | 9 +- src-tauri/src/point_workflow.rs | 3 + src-tauri/src/prediction/context.rs | 22 + src-tauri/src/prediction/database.rs | 108 +++- src-tauri/src/prediction/lookup.rs | 8 + src-tauri/src/prediction/mod.rs | 24 +- src-tauri/src/prediction/model.rs | 544 +++++++++++++++++++ src-tauri/src/prediction/worker.rs | 8 +- src-tauri/tauri.conf.json | 12 +- src/PointScan.test.tsx | 10 + src/scanning/useScanning.ts | 2 + src/settings/ScannerPreferences.tsx | 2 + 23 files changed, 1779 insertions(+), 25 deletions(-) create mode 100644 scripts/fetch-prediction-model.mjs create mode 100644 src-tauri/resources/prediction-blocklist.txt create mode 100644 src-tauri/src/prediction/model.rs 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/docs/word-prediction.md b/docs/word-prediction.md index 9c262dd2..28b8f50f 100644 --- a/docs/word-prediction.md +++ b/docs/word-prediction.md @@ -10,6 +10,25 @@ 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, 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. +- **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 +41,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 ``` +The language model check uses synthetic text and prints its suggestions and timings: + +```powershell +cargo test --release --manifest-path src-tauri/Cargo.toml bundled_model_suggests_current_words -- --ignored --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..4a43a08c --- /dev/null +++ b/scripts/fetch-prediction-model.mjs @@ -0,0 +1,57 @@ +// 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 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..716cb135 100644 --- a/src-tauri/resources/PROVENANCE.md +++ b/src-tauri/resources/PROVENANCE.md @@ -20,3 +20,17 @@ 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. 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..3acbc61a 100644 --- a/src-tauri/src/prediction/context.rs +++ b/src-tauri/src/prediction/context.rs @@ -4,8 +4,23 @@ 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, } 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 { + 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 +46,7 @@ pub fn extract(text: &str, clipped: bool) -> Context { .map(|(_, w)| (*w).to_owned()) .collect(), prefix, + before: String::new(), } } #[cfg(test)] @@ -43,5 +59,11 @@ 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, ""); } } diff --git a/src-tauri/src/prediction/database.rs b/src-tauri/src/prediction/database.rs index d240ec96..60fb67a4 100644 --- a/src-tauri/src/prediction/database.rs +++ b/src-tauri/src/prediction/database.rs @@ -1,17 +1,72 @@ -use super::{context::Context, lookup::Lookup}; -use std::path::Path; +use super::{context::Context, lookup::Lookup, model::Model}; +use std::{path::Path, sync::mpsc}; -pub struct Database(Lookup); +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, so a slow or broken model never delays typing. +pub struct Database { + lookup: Lookup, + enhanced: Enhanced, +} 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 }) } - 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); + }; + let lookup = &self.lookup; + match model.predict(&context.before, &context.prefix, |w| { + lookup.spelling(w).map(str::to_owned) + }) { + Ok(mut words) => { + for word in base { + if words.len() < 5 + && !words + .iter() + .any(|w| w.to_lowercase() == word.to_lowercase()) + { + words.push(word); + } + } + Ok(words) + } + Err(()) => Ok(base), + } } #[cfg(test)] pub fn fixture() -> Self { - Self(Lookup::fixture()) + Self { + lookup: Lookup::fixture(), + enhanced: Enhanced::Unavailable, + } } } @@ -19,11 +74,46 @@ impl Database { mod tests { use super::*; #[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..8619d8d7 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()); @@ -152,6 +154,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 { @@ -378,7 +382,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 +400,14 @@ 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. + *s = Service { + enhanced, + ..Service::default() + }; + keyboard.predictions(None, false); + } let case = ( keyboard.prediction_shift(), keyboard.caps, @@ -411,7 +429,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; diff --git a/src-tauri/src/prediction/model.rs b/src-tauri/src/prediction/model.rs new file mode 100644 index 00000000..6d6274a9 --- /dev/null +++ b/src-tauri/src/prediction/model.rs @@ -0,0 +1,544 @@ +//! 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, +}; +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; + +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 == '\'') + }) + .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. +pub fn spell( + scorer: &mut impl Scorer, + vocabulary: &Vocabulary, + before: &str, + prefix: &str, + known: impl Fn(&str) -> Option, +) -> Result, ()> { + 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)?; + 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| { + let mut p: Vec = finished.values().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 { + 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. + if known(&key).is_none() && (start || text.ends_with(['.', '!', '?'])) { + key + } else { + spelling.replace('’', "'") + } + }) + .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 { + 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), + }) + } + pub fn predict( + &mut self, + before: &str, + prefix: &str, + known: impl Fn(&str) -> Option, + ) -> Result, ()> { + spell(&mut self.onnx, &self.vocabulary, before, prefix, known) + } +} + +#[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, + } + } + + #[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)]); + let words = spell(&mut s, &vocabulary(), "I would like ", "wa", |w| { + Some(w.to_owned()) + }) + .unwrap(); + // water = 0.4 + 0.5 * 0.6; waffle = 0.5 * 0.3; bare "wa" = 0.5 * 0.1. + assert_eq!(words, ["water", "waffle", "wa"]); + let words = spell(&mut s, &vocabulary(), "I would like ", "waf", |w| { + Some(w.to_owned()) + }) + .unwrap(); + assert_eq!(words, ["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)]); + let words = spell(&mut s, &vocabulary(), "Say ", "", |_| None).unwrap(); + assert_eq!(words, ["water"]); + let mut s = fake(&[(6, 1e-5)]); + assert!(spell(&mut s, &vocabulary(), "Say ", "", |_| None) + .unwrap() + .is_empty()); + assert_eq!( + spell(&mut s, &vocabulary(), "Say ", "", |w| Some(w.to_owned())).unwrap(), + ["water"] + ); + } + + #[test] + fn casing_comes_from_the_lookup_or_the_model_but_not_sentence_position() { + let mut s = fake(&[(1, 0.9)]); + let lower = |w: &str| Some(w.to_owned()); + assert_eq!( + spell(&mut s, &vocabulary(), "Send ", "", lower).unwrap(), + ["water", "waffle", "wa"] + ); + assert_eq!( + spell(&mut s, &vocabulary(), "Send ", "", |_| None).unwrap(), + ["Water", "Waffle", "Wa"] + ); + assert_eq!( + spell(&mut s, &vocabulary(), "Done. ", "", |_| None).unwrap(), + ["water", "waffle", "wa"] + ); + } + + #[test] + fn text_start_needs_capital_and_attached_prefixes_fall_back() { + let mut s = fake(&[(7, 0.5), (8, 0.5)]); + let words = spell(&mut s, &vocabulary(), "", "", |w| Some(w.to_owned())).unwrap(); + assert_eq!(words, ["water", "waffle", "wa"]); + assert!(spell(&mut s, &vocabulary(), "hello-", "wa", |w| Some( + w.to_owned() + )) + .is_err()); + assert_eq!(s.contexts, 1); + } + + #[test] + #[ignore = "Needs the fetched model: npm run prediction-model"] + fn bundled_model_suggests_current_words() { + let root = Path::new(env!("CARGO_MANIFEST_DIR")).join("resources"); + let start = std::time::Instant::now(); + let mut model = Model::open( + &root.join("prediction-model"), + &root.join("prediction-blocklist.txt"), + ) + .unwrap(); + println!("model startup_ms={}", start.elapsed().as_millis()); + let lookup = + super::super::lookup::Lookup::open(&root.join("word-predictions.lookup")).unwrap(); + let mut times = Vec::new(); + for (before, prefix, expected) in [ + ("Can you send a ", "wh", "whatsapp"), + ("I would like a cup of ", "", "tea"), + ("Please put the ", "ket", "kettle"), + ("I want to ", "fa", "facetime"), + ] { + let t = std::time::Instant::now(); + let words = model + .predict(before, prefix, |w| lookup.spelling(w).map(str::to_owned)) + .unwrap(); + times.push(t.elapsed().as_secs_f64() * 1000.0); + println!("{before:?}+{prefix:?} -> {words:?} (hoped for {expected})"); + assert!(words.iter().all(|w| normalize(w).starts_with(prefix))); + } + 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)} />)} From cf83e10d84d19fa091ad840c0d8c49e1e2e85044 Mon Sep 17 00:00:00 2001 From: Owen McGirr Date: Thu, 24 Sep 2026 21:39:30 +0100 Subject: [PATCH 2/4] Bound enhanced prediction time and address review findings - Stop the model's word search after 400 ms, and switch the worker to the lookup for good if one call exceeds 1 s, so a slow machine cannot hit the two-second deadline and lose all suggestions. - Treat curly apostrophes as part of a word, and keep a typed curly apostrophe in completions. - Use the lookup when clipping leaves no complete earlier word. - Prune only with words that can pass the final filter, and check the context vocabulary length. - Release the keyboard when the setting changes mid-acceptance. - Disable ONNX Runtime telemetry explicitly. - Run the real-model test in CI and cover the model/lookup merge path with a fake predictor. - Document the fetch step in AGENTS.md validation and the build-time ONNX Runtime download in PROVENANCE, and back off fetch retries. Refs #878 Co-Authored-By: Claude Opus 5.5 (1M context) --- AGENTS.md | 1 + docs/word-prediction.md | 6 +- scripts/fetch-prediction-model.mjs | 1 + src-tauri/resources/PROVENANCE.md | 4 + src-tauri/src/prediction/context.rs | 8 ++ src-tauri/src/prediction/database.rs | 102 ++++++++++++++-- src-tauri/src/prediction/mod.rs | 3 + src-tauri/src/prediction/model.rs | 166 ++++++++++++++++++--------- 8 files changed, 229 insertions(+), 62 deletions(-) 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 28b8f50f..8e1ebb2f 100644 --- a/docs/word-prediction.md +++ b/docs/word-prediction.md @@ -23,6 +23,8 @@ Enhanced word prediction is off by default. When it is on, a small language mode - 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. @@ -41,10 +43,10 @@ Automated tests use fixture databases, fake foreground/activity sources, fake in cargo test --manifest-path src-tauri/Cargo.toml bundled_database_benchmark -- --ignored --nocapture ``` -The language model check uses synthetic text and prints its suggestions and timings: +`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 -- --ignored --nocapture +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. diff --git a/scripts/fetch-prediction-model.mjs b/scripts/fetch-prediction-model.mjs index 4a43a08c..68d6434e 100644 --- a/scripts/fetch-prediction-model.mjs +++ b/scripts/fetch-prediction-model.mjs @@ -47,6 +47,7 @@ async function download(file) { } 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)); } } } diff --git a/src-tauri/resources/PROVENANCE.md b/src-tauri/resources/PROVENANCE.md index 716cb135..4e7d5f14 100644 --- a/src-tauri/resources/PROVENANCE.md +++ b/src-tauri/resources/PROVENANCE.md @@ -34,3 +34,7 @@ It is an int8 ONNX conversion of [HuggingFaceTB/SmolLM2-135M](https://huggingfac # 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/src/prediction/context.rs b/src-tauri/src/prediction/context.rs index 3acbc61a..2b5418f3 100644 --- a/src-tauri/src/prediction/context.rs +++ b/src-tauri/src/prediction/context.rs @@ -6,6 +6,9 @@ pub struct Context { 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); @@ -16,6 +19,7 @@ pub fn extract(text: &str, clipped: bool) -> Context { .map_or("", |i| &before[i..]); } Context { + partial: clipped && before.trim().is_empty(), before: before.to_owned(), ..context } @@ -47,6 +51,7 @@ fn extract_words(text: &str, clipped: bool) -> Context { .collect(), prefix, before: String::new(), + partial: false, } } #[cfg(test)] @@ -65,5 +70,8 @@ mod tests { ); 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 60fb67a4..4747c7ad 100644 --- a/src-tauri/src/prediction/database.rs +++ b/src-tauri/src/prediction/database.rs @@ -1,17 +1,32 @@ -use super::{context::Context, lookup::Lookup, model::Model}; -use std::{path::Path, sync::mpsc}; +use super::{ + context::Context, + lookup::Lookup, + model::{Model, Predict}, +}; +use std::{ + path::Path, + sync::mpsc, + time::{Duration, Instant}, +}; + +/// 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), + Ready(Box), Unavailable, } // The lookup answers until the optional model has loaded in the background, -// and whenever the model fails, so a slow or broken model never delays typing. +// 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, enhanced: bool) -> Result { @@ -28,7 +43,11 @@ impl Database { } None => Enhanced::Unavailable, }; - Ok(Self { lookup, enhanced }) + Ok(Self { + lookup, + enhanced, + slow_call: SLOW_CALL, + }) } pub fn predict(&mut self, context: &Context) -> Result, ()> { let base = self.lookup.predict(&context.words, &context.prefix); @@ -42,10 +61,22 @@ impl Database { let Enhanced::Ready(model) = &mut self.enhanced else { return Ok(base); }; + if context.partial { + return Ok(base); + } let lookup = &self.lookup; - match model.predict(&context.before, &context.prefix, |w| { - lookup.spelling(w).map(str::to_owned) - }) { + 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 @@ -56,6 +87,7 @@ impl Database { words.push(word); } } + words.truncate(5); Ok(words) } Err(()) => Ok(base), @@ -66,6 +98,7 @@ impl Database { Self { lookup: Lookup::fixture(), enhanced: Enhanced::Unavailable, + slow_call: SLOW_CALL, } } } @@ -73,6 +106,59 @@ impl Database { #[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); diff --git a/src-tauri/src/prediction/mod.rs b/src-tauri/src/prediction/mod.rs index 8619d8d7..02e2c0bf 100644 --- a/src-tauri/src/prediction/mod.rs +++ b/src-tauri/src/prediction/mod.rs @@ -402,6 +402,9 @@ pub fn poll( 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() diff --git a/src-tauri/src/prediction/model.rs b/src-tauri/src/prediction/model.rs index 6d6274a9..c135a365 100644 --- a/src-tauri/src/prediction/model.rs +++ b/src-tauri/src/prediction/model.rs @@ -14,6 +14,7 @@ use std::{ borrow::Cow, collections::{HashMap, HashSet}, path::Path, + time::Instant, }; use tokenizers::Tokenizer; @@ -53,7 +54,7 @@ impl Vocabulary { .map(|p| { !p.chars() .next() - .is_some_and(|c| c.is_alphanumeric() || c == '\'') + .is_some_and(|c| c.is_alphanumeric() || c == '\'' || c == '’') }) .collect(), continues: pieces.iter().map(|p| wordy(p)).collect(), @@ -69,13 +70,16 @@ fn normalize(word: &str) -> String { /// 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 @@ -89,6 +93,9 @@ pub fn spell( 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 { @@ -118,11 +125,19 @@ pub fn spell( // Probability mass per normalized word, plus its most probable spelling. let mut finished: HashMap = HashMap::new(); let fifth = |finished: &HashMap| { - let mut p: Vec = finished.values().map(|f| f.0).collect(); + // 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); @@ -185,10 +200,16 @@ pub fn spell( .take(5) .map(|(_, key, spelling)| { // Capitals on unknown words at a sentence start come from position. - if known(&key).is_none() && (start || text.ends_with(['.', '!', '?'])) { + 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()) @@ -218,6 +239,8 @@ fn log_softmax(logits: &[f32]) -> Vec { 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 @@ -391,13 +414,35 @@ impl Model { vocabulary: Vocabulary::new(pieces, blocked), }) } - pub fn predict( +} + +/// The enhanced predictor behind `Database`, so tests can substitute a fake. +pub trait Predict: Send { + fn predict( &mut self, before: &str, prefix: &str, - known: impl Fn(&str) -> Option, + 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, before, prefix, known) + spell( + &mut self.onnx, + &self.vocabulary, + before, + prefix, + known, + deadline, + ) } } @@ -450,94 +495,111 @@ mod tests { } } + 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)]); - let words = spell(&mut s, &vocabulary(), "I would like ", "wa", |w| { - Some(w.to_owned()) - }) - .unwrap(); // 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"]); - let words = spell(&mut s, &vocabulary(), "I would like ", "waf", |w| { - Some(w.to_owned()) - }) - .unwrap(); - assert_eq!(words, ["waffle"]); + 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)]); - let words = spell(&mut s, &vocabulary(), "Say ", "", |_| None).unwrap(); - assert_eq!(words, ["water"]); + assert_eq!(run(&mut s, "Say ", "", |_| None).unwrap(), ["water"]); let mut s = fake(&[(6, 1e-5)]); - assert!(spell(&mut s, &vocabulary(), "Say ", "", |_| None) - .unwrap() - .is_empty()); - assert_eq!( - spell(&mut s, &vocabulary(), "Say ", "", |w| Some(w.to_owned())).unwrap(), - ["water"] - ); + 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 lower = |w: &str| Some(w.to_owned()); - assert_eq!( - spell(&mut s, &vocabulary(), "Send ", "", lower).unwrap(), - ["water", "waffle", "wa"] - ); + let lowered = ["water", "waffle", "wa"]; + assert_eq!(run(&mut s, "Send ", "", lower).unwrap(), lowered); assert_eq!( - spell(&mut s, &vocabulary(), "Send ", "", |_| None).unwrap(), + run(&mut s, "Send ", "", |_| None).unwrap(), ["Water", "Waffle", "Wa"] ); - assert_eq!( - spell(&mut s, &vocabulary(), "Done. ", "", |_| 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)]); - let words = spell(&mut s, &vocabulary(), "", "", |w| Some(w.to_owned())).unwrap(); - assert_eq!(words, ["water", "waffle", "wa"]); - assert!(spell(&mut s, &vocabulary(), "hello-", "wa", |w| Some( - w.to_owned() - )) - .is_err()); + 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] - #[ignore = "Needs the fetched model: npm run prediction-model"] + 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 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 = std::time::Instant::now(); + let start = Instant::now(); let mut model = Model::open( &root.join("prediction-model"), &root.join("prediction-blocklist.txt"), ) - .unwrap(); + .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 [ - ("Can you send a ", "wh", "whatsapp"), - ("I would like a cup of ", "", "tea"), - ("Please put the ", "ket", "kettle"), - ("I want to ", "fa", "facetime"), + ("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 = std::time::Instant::now(); - let words = model - .predict(before, prefix, |w| lookup.spelling(w).map(str::to_owned)) - .unwrap(); + 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:?} (hoped for {expected})"); + 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?}"); } From a7e6c5908aa2a72159cb45560503670afcf4de68 Mon Sep 17 00:00:00 2001 From: Owen McGirr Date: Thu, 24 Sep 2026 21:43:32 +0100 Subject: [PATCH 3/4] Cap the model's context to the last 256 characters The context pass is the one model step without a deadline. Reading at most the last 256 buffered characters, starting at a word boundary, bounds its cost on slow machines. Buffers with no complete word in that window use the lookup. Refs #878 Co-Authored-By: Claude Opus 5.5 (1M context) --- docs/word-prediction.md | 2 +- src-tauri/src/prediction/model.rs | 29 ++++++++++++++++++++++++++++- 2 files changed, 29 insertions(+), 2 deletions(-) diff --git a/docs/word-prediction.md b/docs/word-prediction.md index 8e1ebb2f..3962c480 100644 --- a/docs/word-prediction.md +++ b/docs/word-prediction.md @@ -12,7 +12,7 @@ The bundled English lookup database works offline. Ranking backs off from the lo ## 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, including earlier sentences still in the buffer. It runs offline in the prediction worker, never learns, and nothing it reads is logged or stored. +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:** diff --git a/src-tauri/src/prediction/model.rs b/src-tauri/src/prediction/model.rs index c135a365..04cb779f 100644 --- a/src-tauri/src/prediction/model.rs +++ b/src-tauri/src/prediction/model.rs @@ -24,6 +24,8 @@ const MAX_PIECES: usize = 5; 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`. @@ -438,7 +440,7 @@ impl Predict for Model { spell( &mut self.onnx, &self.vocabulary, - before, + recent(before).ok_or(())?, prefix, known, deadline, @@ -446,6 +448,19 @@ impl Predict for Model { } } +/// 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::*; @@ -561,6 +576,18 @@ mod tests { 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", "’", " .", ","] From a2a329d622c6b3f2a22eb66d46125668526a227b Mon Sep 17 00:00:00 2001 From: Owen McGirr Date: Fri, 25 Sep 2026 09:27:33 +0100 Subject: [PATCH 4/4] Allow prediction worker time to load bundled lookup --- src-tauri/src/prediction/mod.rs | 37 ++++++++++++++++++++++++++++++--- 1 file changed, 34 insertions(+), 3 deletions(-) diff --git a/src-tauri/src/prediction/mod.rs b/src-tauri/src/prediction/mod.rs index 02e2c0bf..e7bcf086 100644 --- a/src-tauri/src/prediction/mod.rs +++ b/src-tauri/src/prediction/mod.rs @@ -143,6 +143,7 @@ impl Drop for Client { struct Service { client: Option, failed: bool, + received_reply: bool, generation: u64, outstanding: Option, last: Option, @@ -195,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 { @@ -440,9 +453,7 @@ pub fn poll( 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; } @@ -450,6 +461,7 @@ pub fn poll( match reply { Ok(Ok(response)) => { s.outstanding = None; + s.received_reply = true; match response { Response::Suggestions { generation, @@ -593,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()); @@ -601,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 {