diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 534fa320..74112b1e 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -32,7 +32,6 @@ jobs: cargo audit --file src-tauri/Cargo.lock cargo audit --file src-tauri/startup-launcher/Cargo.lock cargo audit --file tools/updater-signature-verifier/Cargo.lock - cargo audit --file tools/prediction-data/Cargo.lock frontend: runs-on: ubuntu-latest @@ -118,10 +117,6 @@ jobs: - 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 - - run: cargo fmt --manifest-path tools/prediction-data/Cargo.toml --check - - run: cargo clippy --locked --manifest-path tools/prediction-data/Cargo.toml --all-targets -- -D warnings - - run: cargo test --locked --manifest-path tools/prediction-data/Cargo.toml - - run: cargo run --release --locked --manifest-path tools/prediction-data/Cargo.toml -- --check src-tauri/resources/WordData2017051601.db src-tauri/resources/word-predictions.lookup - run: cargo fmt --manifest-path tools/updater-signature-verifier/Cargo.toml --check - run: cargo clippy --locked --manifest-path tools/updater-signature-verifier/Cargo.toml --all-targets -- -D warnings - run: cargo test --locked --manifest-path tools/updater-signature-verifier/Cargo.toml diff --git a/docs/word-prediction.md b/docs/word-prediction.md index 3962c480..50d57357 100644 --- a/docs/word-prediction.md +++ b/docs/word-prediction.md @@ -8,42 +8,25 @@ The buffer holds at most 512 characters, retaining complete Unicode graphemes at A passive observer records only an activity counter and timestamp, never external text. Prediction is unavailable if this observer cannot start or loses access. Edits made before the observer is ready are discarded because intervening activity cannot be verified. Each queued edit is scoped to its foreground target and the time before injection, so edits preceding an observed external change are discarded. Changes within an application that produce no observed input cannot be detected without inspecting its fields. -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. +## On-device model -## Enhanced word prediction +One Word prediction setting controls an offline SmolLM2-135M int8 ONNX model. The saved `enhancedWordPrediction` field is retained for compatibility but does not choose an engine. The model spells candidates from subword pieces and reads only the last 256 characters of the temporary buffer, beginning at a word boundary. It never learns from typing. -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. +For a first typed prefix, a fixed local context keeps the model from favoring website names at the start of a document; it adds no user text. The model loads in the worker while keyboard input remains available. Suggestions are blank until loading finishes. A passive badge beside the scan prompt distinguishes loading, a ready keyboard awaiting typed context, no matching suggestions, available suggestions, paused activity tracking, and prediction failure. It never shows typed text and is not a scan target. If loading or inference fails, or a call exceeds one second, the keyboard continues accepting input and offers **Retry predictions** in its toolbar. Retry restarts only the prediction worker, clears its private text context, and leaves the keyboard open; type a new prefix afterward. A 400 ms search budget bounds candidate exploration. A clipped buffer with no complete earlier word yields no suggestions. -- **Loading and fallback:** the model loads in the background when the worker starts, taking about 0.7 s. Until it is ready, and whenever it fails, the lookup answers as before. Turning the setting on or off restarts the worker. -- **Filters:** - - words contain only ASCII letters and apostrophes; - - a word at the start of the text must be capitalised, which rules out subword fragments; - - single letters must be in the lookup; - - blocklisted words are never suggested; - - words the lookup does not know must reach a minimum probability. -- **Spelling:** words the lookup knows take its usual spelling. Unknown words keep the model's spelling, except that a capital coming only from sentence position is removed. -- **Merging:** lookup suggestions fill any remaining positions. -- **Time limits:** the model stops exploring new words after 400 ms. If one call ever takes longer than 1 s, the worker uses the lookup for the rest of its life, so replies stay well inside the two-second deadline. -- **Clipped buffers:** when clipping leaves no complete earlier word, the lookup answers. -- **Cost:** predictions take roughly 40–80 ms on a recent laptop, and the worker needs about 300–450 MB more memory. +Candidates contain ASCII letters and apostrophes, must have sufficient model probability, and cannot appear in the blocklist. Single-letter candidates are limited to “a” and “I”. Up to five distinct suggestions are shown. Ordinary words use lowercase; mixed-case names retain model casing. Shift, Caps and sentence capitalization are applied when the keyboard displays a candidate. 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. +A separate process owns the buffer, activity observer and model. 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. ## Validation -Automated tests use fixture databases, fake foreground/activity sources, fake input adapters and sleeping subprocesses. They never inject desktop input. They cover first-letter completion, Unicode Backspace, bounded context, stale replies, foreground/external changes, failure cleanup and switch action compatibility. The database benchmark prints counts and timings only: - -```powershell -cargo test --manifest-path src-tauri/Cargo.toml bundled_database_benchmark -- --ignored --nocapture -``` - -`cargo test` also runs the language model against the fetched files with synthetic text. For release-build suggestions and timings: +Automated tests use fake predictors, foreground/activity sources, input adapters and sleeping subprocesses. They never inject desktop input. They cover first-letter completion, Unicode Backspace, bounded context, stale replies, foreground/external changes, failure cleanup and switch action compatibility. `cargo test` also runs the language model against the fetched files with synthetic text. For release-build suggestions and timings: ```powershell cargo test --release --manifest-path src-tauri/Cargo.toml bundled_model_suggests_current_words -- --nocapture diff --git a/scripts/fetch-prediction-model.mjs b/scripts/fetch-prediction-model.mjs index 68d6434e..fac824c0 100644 --- a/scripts/fetch-prediction-model.mjs +++ b/scripts/fetch-prediction-model.mjs @@ -1,4 +1,4 @@ -// Downloads the enhanced word prediction model into src-tauri/resources/prediction-model. +// Downloads the 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"; diff --git a/src-tauri/Cargo.lock b/src-tauri/Cargo.lock index 26030140..a2d7bd19 100644 --- a/src-tauri/Cargo.lock +++ b/src-tauri/Cargo.lock @@ -1332,18 +1332,6 @@ dependencies = [ "pin-project-lite", ] -[[package]] -name = "fallible-iterator" -version = "0.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2acce4a10f12dc2fb14a218589d4f1f62ef011b2d0cc4b3cb1bba8e94da14649" - -[[package]] -name = "fallible-streaming-iterator" -version = "0.1.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7360491ce676a36bf9bb3c56c1aa791658183a54d2744120f27285738d90465a" - [[package]] name = "fancy-regex" version = "0.14.0" @@ -1856,32 +1844,11 @@ version = "0.12.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888" -[[package]] -name = "hashbrown" -version = "0.16.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100" -dependencies = [ - "foldhash", -] - [[package]] name = "hashbrown" version = "0.17.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" -dependencies = [ - "foldhash", -] - -[[package]] -name = "hashlink" -version = "0.12.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a596f1b20ed2cc5ecac41a164aaebc7258057060f06c0cf7a2ba3991ee7990fb" -dependencies = [ - "hashbrown 0.17.1", -] [[package]] name = "heck" @@ -2487,17 +2454,6 @@ dependencies = [ "libc", ] -[[package]] -name = "libsqlite3-sys" -version = "0.38.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f1d20bef17f513b9b3004532233187769cd072d790971f4e4da0e346eb6401e8" -dependencies = [ - "cc", - "pkg-config", - "vcpkg", -] - [[package]] name = "linux-raw-sys" version = "0.12.1" @@ -3814,31 +3770,6 @@ dependencies = [ "windows-sys 0.52.0", ] -[[package]] -name = "rsqlite-vfs" -version = "0.1.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c51c9ae4df8a7fba42103df5c621fa3c37eccf3a3c650879e90fc48b11cc192c" -dependencies = [ - "hashbrown 0.16.1", - "thiserror 2.0.19", -] - -[[package]] -name = "rusqlite" -version = "0.40.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "23f2a97da3e3873c73cb2a2e71b35c40ff95e0b1eefa8d72d8499a6928c3b5b3" -dependencies = [ - "bitflags 2.13.1", - "fallible-iterator", - "fallible-streaming-iterator", - "hashlink", - "libsqlite3-sys", - "smallvec", - "sqlite-wasm-rs", -] - [[package]] name = "rustc-hash" version = "2.1.3" @@ -4408,18 +4339,6 @@ dependencies = [ "unicode-segmentation", ] -[[package]] -name = "sqlite-wasm-rs" -version = "0.5.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dc3efc0da82635d7e1ced0053bbbfa8c7ab9645d0bf36ceb4f7127bb85315d75" -dependencies = [ - "cc", - "js-sys", - "rsqlite-vfs", - "wasm-bindgen", -] - [[package]] name = "stable_deref_trait" version = "1.2.1" @@ -4508,7 +4427,6 @@ dependencies = [ "ort", "rand 0.8.7", "reqwest", - "rusqlite", "serde", "serde_json", "sha2", @@ -5577,12 +5495,6 @@ dependencies = [ "wasm-bindgen", ] -[[package]] -name = "vcpkg" -version = "0.2.15" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "accd4ea62f7bb7a82fe23066fb0957d48ef677f6eeb8215f372f52e48bb32426" - [[package]] name = "version-compare" version = "0.2.1" diff --git a/src-tauri/Cargo.toml b/src-tauri/Cargo.toml index 01ccb6df..7a72fea5 100644 --- a/src-tauri/Cargo.toml +++ b/src-tauri/Cargo.toml @@ -75,6 +75,3 @@ winreg = "0.55" [patch.crates-io] corebluetooth-rs = { path = "../vendor/corebluetooth-rs" } - -[dev-dependencies] -rusqlite = { version = "0.40.2", features = ["bundled"] } diff --git a/src-tauri/resources/PROVENANCE.md b/src-tauri/resources/PROVENANCE.md index 4e7d5f14..6c0e855d 100644 --- a/src-tauri/resources/PROVENANCE.md +++ b/src-tauri/resources/PROVENANCE.md @@ -1,27 +1,4 @@ -# English prediction database - -`WordData2017051601.db` was recovered from the user's private `enaboapps/sayit-ios` history, path `Assets/Databases/WordData2017051601.db`, commit `95bda9265758e50787c5da2779b16f36b53c319d`, dated May 20, 2017. Recovery took place September 18, 2026. The user authorized bundling this database with Switchify. - -- Size: 102,690,816 bytes. -- SHA-256: `dedd65d263bde8315e7e5ed7d2c8e04f17c33598a68506c9d70661a6d1f57318`. -- SQLite integrity check during recovery: `ok`. -- Rows: WORDS 165,420; BIGRAMS 976,804; TRIGRAMS 996,846; QUADGRAMS 953,630. - -This is the unchanged original SQLite file. The later Realm archive is not used. It is retained for conversion and parity tests, not shipped in the app. The converter ignores personal frequency columns. No captured text or user vocabulary is written into this resource. - -## Packaged dedicated lookup - -`word-predictions.lookup` is generated from the unchanged database. It contains frequency-ranked words and n-grams, a sorted prefix index, a versioned header and a SHA-256 payload checksum. The runtime validates counts, bounds, references and ordering before serving predictions. The app packages only this lookup and does not write to it. - -Regenerate with Rust 1.97.1 from the repository root: - -```sh -cargo run --release --locked --manifest-path tools/prediction-data/Cargo.toml -- src-tauri/resources/WordData2017051601.db src-tauri/resources/word-predictions.lookup -``` - -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 +# Word 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: @@ -33,8 +10,8 @@ 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. +`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. Word 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. +Word prediction statically links ONNX Runtime 1.28.0 through the `ort` crate (`=2.0.0-rc.13`, MIT/Apache-2.0). At build time, `ort-sys` downloads the prebuilt archive for the target from `cdn.pyke.io` and verifies it against a SHA-256 embedded in the crate. So every build, including signed releases, depends on that download. Only the CPU provider is used, and ONNX Runtime telemetry is disabled when the worker starts it. diff --git a/src-tauri/resources/WordData2017051601.db b/src-tauri/resources/WordData2017051601.db deleted file mode 100644 index eb93d5d5..00000000 Binary files a/src-tauri/resources/WordData2017051601.db and /dev/null differ diff --git a/src-tauri/resources/word-predictions.lookup b/src-tauri/resources/word-predictions.lookup deleted file mode 100644 index d1ba4250..00000000 Binary files a/src-tauri/resources/word-predictions.lookup and /dev/null differ diff --git a/src-tauri/src/point_scan.rs b/src-tauri/src/point_scan.rs index 06d98cdf..87f97782 100644 --- a/src-tauri/src/point_scan.rs +++ b/src-tauri/src/point_scan.rs @@ -26,8 +26,8 @@ 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. + /// Retained for saved-settings compatibility; the model is always used + /// when `word_prediction` is enabled. pub enhanced_word_prediction: bool, pub keyboard_wait_after_typing: bool, /// Moves the keyboard or mouse panel away while the pointer is over it. @@ -113,7 +113,6 @@ 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, @@ -140,7 +139,6 @@ 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, diff --git a/src-tauri/src/point_scan_runtime.rs b/src-tauri/src/point_scan_runtime.rs index 6b920790..caf69a2a 100644 --- a/src-tauri/src/point_scan_runtime.rs +++ b/src-tauri/src/point_scan_runtime.rs @@ -124,6 +124,10 @@ impl Adapter for PointScan { Request::Prediction { token, index } => { return crate::prediction::select(token, index).map(|()| None) } + Request::PredictionRetry => { + crate::prediction::retry(); + return Ok(None); + } Request::Keyboard(stroke) => { let scope = crate::prediction::InputScope::capture(); let result = crate::scan_executor::activate(request); @@ -268,6 +272,7 @@ impl Adapter for PointScan { request, Request::Keyboard(_) | Request::Prediction { .. } + | Request::PredictionRetry | Request::MouseMove { .. } | Request::MouseMoveAbsolute { .. } | Request::MouseScroll { .. } @@ -282,14 +287,7 @@ impl Adapter for PointScan { crate::prediction::start_keyboard_activity(captured_keys); } let enabled = technique.prediction_enabled(); - let enhanced = technique.prediction_enhanced(); - crate::prediction::poll( - app, - technique.prediction_keyboard(), - enabled, - enhanced, - captured_keys, - ); + crate::prediction::poll(app, technique.prediction_keyboard(), enabled, captured_keys); } fn cleanup(_app: &AppHandle) -> Result<(), String> { crate::prediction::stop(); @@ -425,6 +423,7 @@ mod tests { token: 1, index: 0 })); + assert!(PointScan::preserve_visuals(&Request::PredictionRetry)); assert!(PointScan::preserve_visuals(&Request::MouseScroll { dy: 5 })); assert!(!PointScan::preserve_visuals( &crate::point_workflow::default_click((10, 20)) diff --git a/src-tauri/src/point_workflow.rs b/src-tauri/src/point_workflow.rs index 550daefa..cc93a061 100644 --- a/src-tauri/src/point_workflow.rs +++ b/src-tauri/src/point_workflow.rs @@ -12,6 +12,7 @@ pub enum Request { token: u64, index: usize, }, + PredictionRetry, OpenKeyboard, OpenMouse, OpenPoint, @@ -308,9 +309,7 @@ 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; } @@ -664,6 +663,9 @@ impl Technique for Workflow { Some(crate::scan_keyboard::Output::Prediction { token, index }) => { return Some(Request::Prediction { token, index }) } + Some(crate::scan_keyboard::Output::RetryPrediction) => { + return Some(Request::PredictionRetry) + } Some(crate::scan_keyboard::Output::Close) => self.keyboard_closed(), None => {} } diff --git a/src-tauri/src/prediction/context.rs b/src-tauri/src/prediction/context.rs index 2b5418f3..e7b968e7 100644 --- a/src-tauri/src/prediction/context.rs +++ b/src-tauri/src/prediction/context.rs @@ -1,7 +1,6 @@ use unicode_segmentation::UnicodeSegmentation; pub struct Context { - pub words: Vec, pub prefix: String, pub sentence_start: bool, /// Buffered text before `prefix`, without a clipped leading fragment. @@ -42,13 +41,6 @@ fn extract_words(text: &str, clipped: bool) -> Context { }; Context { sentence_start: tokens.is_empty() && (!clipped || boundary.is_some()), - words: tokens - .iter() - .rev() - .take(3) - .rev() - .map(|(_, w)| (*w).to_owned()) - .collect(), prefix, before: String::new(), partial: false, @@ -60,9 +52,7 @@ mod tests { #[test] fn unicode_boundaries_and_clipped_words() { let c = extract("Discard this. I can’t find cafe\u{301}", false); - assert_eq!(c.words, ["I", "can’t", "find"]); 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, diff --git a/src-tauri/src/prediction/database.rs b/src-tauri/src/prediction/database.rs index 4747c7ad..c6d6785b 100644 --- a/src-tauri/src/prediction/database.rs +++ b/src-tauri/src/prediction/database.rs @@ -1,6 +1,5 @@ use super::{ context::Context, - lookup::Lookup, model::{Model, Predict}, }; use std::{ @@ -9,100 +8,100 @@ use std::{ 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 { +#[derive(Clone, Copy, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +pub enum Status { + Loading, + Ready, + Unavailable, +} + +enum State { Loading(mpsc::Receiver>), Ready(Box), Unavailable, } -// The lookup answers until the optional model has loaded in the background, -// and whenever the model fails or runs slowly, so it never delays typing. +/// Loading is asynchronous so every keyboard edit remains usable. There is no +/// secondary suggestion source when the model is unavailable. pub struct Database { - lookup: Lookup, - enhanced: Enhanced, + state: State, slow_call: Duration, } impl Database { - 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 open(resources: &Path) -> Self { + 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)); + }); + Self { + state: State::Loading(rx), slow_call: SLOW_CALL, - }) + } } - 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, + pub fn status(&mut self) -> Status { + if let State::Loading(rx) = &self.state { + self.state = match rx.try_recv() { + Ok(Ok(model)) => State::Ready(Box::new(model)), + Err(mpsc::TryRecvError::Empty) => return Status::Loading, + _ => State::Unavailable, }; } - let Enhanced::Ready(model) = &mut self.enhanced else { - return Ok(base); - }; - if context.partial { - return Ok(base); + match self.state { + State::Loading(_) => Status::Loading, + State::Ready(_) => Status::Ready, + State::Unavailable => Status::Unavailable, } - let lookup = &self.lookup; - let known = |w: &str| lookup.spelling(w).map(str::to_owned); - let start = Instant::now(); - let result = model.predict( - &context.before, - &context.prefix, - &known, - start + SEARCH_BUDGET, - ); - if start.elapsed() > self.slow_call { - self.enhanced = Enhanced::Unavailable; + } + pub fn predict(&mut self, context: &Context) -> (Status, Vec) { + let status = self.status(); + if status != Status::Ready || context.partial { + return (status, Vec::new()); } - match result { - Ok(mut words) => { - for word in base { - if words.len() < 5 - && !words - .iter() - .any(|w| w.to_lowercase() == word.to_lowercase()) - { - words.push(word); - } - } - words.truncate(5); - Ok(words) - } - Err(()) => Ok(base), + let State::Ready(model) = &mut self.state else { + unreachable!() + }; + let start = Instant::now(); + let result = model.predict(&context.before, &context.prefix, start + SEARCH_BUDGET); + if start.elapsed() > self.slow_call || result.is_err() { + self.state = State::Unavailable; + return (Status::Unavailable, Vec::new()); } + let mut words = result.unwrap(); + words.truncate(5); + (Status::Ready, words) } #[cfg(test)] pub fn fixture() -> Self { Self { - lookup: Lookup::fixture(), - enhanced: Enhanced::Unavailable, + state: State::Ready(Box::new(FakeModel)), slow_call: SLOW_CALL, } } } +#[cfg(test)] +struct FakeModel; +#[cfg(test)] +impl Predict for FakeModel { + fn predict( + &mut self, + _before: &str, + prefix: &str, + _deadline: Instant, + ) -> Result, ()> { + Ok(["water", "waffle", "walk"] + .into_iter() + .filter(|w| w.starts_with(&prefix.to_lowercase())) + .map(str::to_owned) + .collect()) + } +} + #[cfg(test)] mod tests { use super::*; @@ -112,7 +111,6 @@ mod tests { &mut self, _before: &str, _prefix: &str, - _known: &dyn Fn(&str) -> Option, _deadline: Instant, ) -> Result, ()> { std::thread::sleep(self.1); @@ -121,167 +119,41 @@ mod tests { .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))); + fn context() -> Context { + super::super::context::extract("I would like wa", false) } #[test] - fn model_words_lead_and_the_lookup_fills_without_duplicates() { - let ctx = super::super::context::extract("I would like wa", false); - let mut db = Database::fixture(); - let lookup = db.predict(&ctx).unwrap(); - assert!(lookup.len() >= 2); - ready(&mut db, Ok(vec!["WhatsApp", "WATER"]), Duration::ZERO); - let words = db.predict(&ctx).unwrap(); - assert_eq!(words[..2], ["WhatsApp", "WATER"]); - assert!(!words[2..].iter().any(|w| w.eq_ignore_ascii_case("water"))); - assert!(words.len() <= 5); - ready( - &mut db, - Ok(vec!["a", "b", "c", "d", "e", "f"]), - Duration::ZERO, - ); - assert_eq!(db.predict(&ctx).unwrap(), ["a", "b", "c", "d", "e"]); - ready(&mut db, Err(()), Duration::ZERO); - assert_eq!(db.predict(&ctx).unwrap(), lookup); - assert!(matches!(db.enhanced, Enhanced::Ready(_))); - } - #[test] - fn slow_models_and_clipped_fragments_use_the_lookup() { - let mut db = Database::fixture(); - let clipped = super::super::context::extract("agmentwa", true); - let lookup = db.predict(&clipped).unwrap(); - ready(&mut db, Ok(vec!["model"]), Duration::ZERO); - assert_eq!(db.predict(&clipped).unwrap(), lookup); - let ctx = super::super::context::extract("I would like wa", false); - db.slow_call = Duration::from_millis(20); - ready(&mut db, Ok(vec!["model"]), Duration::from_millis(40)); - assert_eq!(db.predict(&ctx).unwrap()[0], "model"); - assert!(matches!(db.enhanced, Enhanced::Unavailable)); - assert_ne!(db.predict(&ctx).unwrap()[0], "model"); - } - #[test] - fn lookup_answers_while_the_model_loads_or_after_it_fails() { - let ctx = super::super::context::extract("I would like wa", false); - let mut db = Database::fixture(); - let expected = db.predict(&ctx).unwrap(); + fn loading_failure_and_missing_model_have_no_fallback() { 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(_))); + let mut db = Database { + state: State::Loading(rx), + slow_call: SLOW_CALL, + }; + assert_eq!(db.predict(&context()), (Status::Loading, vec![])); 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_eq!(db.predict(&context()), (Status::Unavailable, vec![])); + let missing = + std::env::temp_dir().join(format!("switchify-missing-model-{}", std::process::id())); + let mut db = Database::open(&missing); + let end = Instant::now() + Duration::from_secs(5); + while db.status() == Status::Loading && Instant::now() < end { + std::thread::sleep(Duration::from_millis(10)); } - assert!(matches!(db.enhanced, Enhanced::Unavailable)); - assert!(!db.predict(&ctx).unwrap().is_empty()); - std::fs::remove_dir_all(dir).unwrap(); + assert_eq!(db.predict(&context()), (Status::Unavailable, vec![])); } #[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 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, - ) - .unwrap(); - let mut texts: Vec = [ - "", - "a", - "w", - "wa", - "water", - "zzzzzz", - "I ", - "I would ", - "I would like ", - "I would like wa", - "the ", - "the a", - "unknownword wa", - "café", - "cafe\u{301}", - "İ", - "你好", - "I can’t find ", - "Done! ", - "hello\nwa", - "W", - "WATER", - "🙂 ", - ] - .iter() - .map(|s| s.to_string()) - .collect(); - let ids: std::collections::HashMap = c - .prepare("SELECT ID,WORD FROM WORDS") - .unwrap() - .query_map([], |r| Ok((r.get(0)?, r.get(1)?))) - .unwrap() - .map(Result::unwrap) - .collect(); - for offset in [0, 499] { - for (n, table) in [(2, "BIGRAMS"), (3, "TRIGRAMS"), (4, "QUADGRAMS")] { - let fields = (1..=n) - .map(|i| format!("ID{i}")) - .collect::>() - .join(","); - let mut q = c - .prepare(&format!( - "SELECT {fields} FROM {table} WHERE ID%997={offset} ORDER BY ID LIMIT 80" - )) - .unwrap(); - let rows = q - .query_map([], |r| { - (0..n) - .map(|i| r.get::<_, u32>(i)) - .collect::, _>>() - }) - .unwrap(); - for row in rows { - let row = row.unwrap(); - let context = row[..n - 1] - .iter() - .map(|id| ids[id].as_str()) - .collect::>() - .join(" "); - for size in [0, 1, 3] { - texts.push(format!( - "{context} {}", - ids[&row[n - 1]].chars().take(size).collect::() - )); - } - } - } - } - assert_eq!(texts.len(), 1463); - for (i, text) in texts.iter().enumerate() { - let context = super::super::context::extract(text, i % 17 == 0); - assert_eq!( - lookup.predict(&context), - reference.predict(&context), - "case {i}" - ); - } + fn model_only_limit_failure_and_slow_call() { + let mut db = Database::fixture(); + db.state = State::Ready(Box::new(Fake( + Ok(vec!["one", "two", "three", "four", "five", "six"]), + Duration::ZERO, + ))); + assert_eq!(db.predict(&context()).1.len(), 5); + db.state = State::Ready(Box::new(Fake(Err(()), Duration::ZERO))); + assert_eq!(db.predict(&context()), (Status::Unavailable, vec![])); + db.state = State::Ready(Box::new(Fake(Ok(vec!["late"]), Duration::from_millis(20)))); + db.slow_call = Duration::from_millis(10); + assert_eq!(db.predict(&context()), (Status::Unavailable, vec![])); + assert_eq!(db.predict(&context()), (Status::Unavailable, vec![])); } } diff --git a/src-tauri/src/prediction/database_reference.rs b/src-tauri/src/prediction/database_reference.rs deleted file mode 100644 index 78117738..00000000 --- a/src-tauri/src/prediction/database_reference.rs +++ /dev/null @@ -1,123 +0,0 @@ -use super::context::Context; -use rusqlite::{Connection, OpenFlags}; -use std::{ - collections::{HashMap, HashSet}, - path::Path, -}; - -pub struct Database { - connection: Connection, - words: Vec<(i64, String, String)>, - ids: HashMap, -} -impl Database { - pub fn open(path: &Path) -> Result { - let connection = Connection::open_with_flags( - path, - OpenFlags::SQLITE_OPEN_READ_ONLY | OpenFlags::SQLITE_OPEN_NO_MUTEX, - ) - .map_err(|_| ())?; - Self::from_connection(connection) - } - fn from_connection(connection: Connection) -> Result { - connection - .execute_batch("PRAGMA query_only=ON; PRAGMA cache_size=-8192;") - .map_err(|_| ())?; - let words = { - let mut q = connection - .prepare("SELECT ID, WORD FROM WORDS ORDER BY BASE_FREQUENCY DESC, ID ASC") - .map_err(|_| ())?; - let result = q - .query_map([], |r| Ok((r.get::<_, i64>(0)?, r.get::<_, String>(1)?))) - .map_err(|_| ())? - .collect::, _>>() - .map_err(|_| ())?; - result - }; - let words: Vec<_> = words - .into_iter() - .map(|(id, word)| { - let folded = word.to_lowercase(); - (id, word, folded) - }) - .collect(); - let mut ids = HashMap::new(); - for (id, _, folded) in &words { - ids.entry(folded.clone()).or_insert(*id); - } - Ok(Self { - connection, - words, - ids, - }) - } - pub fn predict(&self, context: &Context) -> Result, ()> { - let prefix = context.prefix.to_lowercase(); - let mut result = Vec::new(); - let mut seen = HashSet::new(); - let mut push = |word: String| { - let folded = word.to_lowercase(); - if result.len() < 5 - && word.chars().count() <= 48 - && !word.chars().any(char::is_whitespace) - && folded.starts_with(&prefix) - && seen.insert(folded) - { - result.push(word); - } - result.len() == 5 - }; - for count in (1..=context.words.len().min(3)).rev() { - let ids: Option> = context.words[context.words.len() - count..] - .iter() - .map(|w| self.ids.get(&w.to_lowercase()).copied()) - .collect(); - let Some(ids) = ids else { - continue; - }; - let table = ["", "BIGRAMS", "TRIGRAMS", "QUADGRAMS"][count]; - let conditions = (1..=count) - .map(|i| format!("g.ID{i}=?{i}")) - .collect::>() - .join(" AND "); - let sql = format!("SELECT w.WORD FROM {table} g JOIN WORDS w ON w.ID=g.ID{} WHERE {conditions} ORDER BY g.BASE_FREQUENCY DESC, w.ID ASC", count+1); - let mut statement = self.connection.prepare_cached(&sql).map_err(|_| ())?; - let mut rows = statement - .query(rusqlite::params_from_iter(ids)) - .map_err(|_| ())?; - // Context indexes bound this candidate set; the entire query is also - // contained by the parent process's two-second deadline. - while let Some(row) = rows.next().map_err(|_| ())? { - if push(row.get(0).map_err(|_| ())?) { - break; - } - } - } - for (_, word, _) in self - .words - .iter() - .filter(|(_, _, folded)| folded.starts_with(&prefix)) - { - if push(word.clone()) { - break; - } - } - Ok(result) - } -} - -#[cfg(test)] -mod tests { - use super::*; - #[test] - fn longest_context_wins_and_backoff_deduplicates() { - let c = Connection::open_in_memory().unwrap(); - c.execute_batch("CREATE TABLE WORDS(ID INTEGER,WORD TEXT,BASE_FREQUENCY INTEGER); INSERT INTO WORDS VALUES(1,'I',9),(2,'would',8),(3,'like',7),(4,'water',1),(5,'wine',20),(6,'walk',30); CREATE TABLE BIGRAMS(ID1 INTEGER,ID2 INTEGER,BASE_FREQUENCY INTEGER); INSERT INTO BIGRAMS VALUES(3,5,9); CREATE TABLE TRIGRAMS(ID1 INTEGER,ID2 INTEGER,ID3 INTEGER,BASE_FREQUENCY INTEGER); CREATE TABLE QUADGRAMS(ID1 INTEGER,ID2 INTEGER,ID3 INTEGER,ID4 INTEGER,BASE_FREQUENCY INTEGER); INSERT INTO QUADGRAMS VALUES(1,2,3,4,1);").unwrap(); - let db = Database::from_connection(c).unwrap(); - let result = db - .predict(&super::super::context::extract("I would like w", false)) - .unwrap(); - assert_eq!(&result[..3], ["water", "wine", "walk"]); - assert!(db.connection.execute("DELETE FROM WORDS", []).is_err()); - } -} diff --git a/src-tauri/src/prediction/lookup.rs b/src-tauri/src/prediction/lookup.rs deleted file mode 100644 index 13283ea0..00000000 --- a/src-tauri/src/prediction/lookup.rs +++ /dev/null @@ -1,398 +0,0 @@ -use sha2::{Digest, Sha256}; -use std::{ - collections::{HashMap, HashSet}, - fs::File, - io::{BufReader, Read}, - path::Path, -}; - -pub const MAGIC: &[u8; 8] = b"SWPRED02"; -const MAX_FILE: u64 = 128 * 1024 * 1024; -const MAX_WORDS: usize = 250_000; -const MAX_ROWS: usize = 1_500_000; -const MAX_WORD_BYTES: usize = 512; -const MAX_TEXT_BYTES: usize = 32 * 1024 * 1024; -const NONE: u32 = u32::MAX; - -struct Reader { - input: R, - hash: Sha256, -} -impl Reader { - fn bytes(&mut self, bytes: &mut [u8]) -> Result<(), ()> { - self.input.read_exact(bytes).map_err(|_| ())?; - self.hash.update(bytes); - Ok(()) - } - fn number(&mut self) -> Result { - let mut bytes = [0; 4]; - self.bytes(&mut bytes)?; - Ok(u32::from_le_bytes(bytes)) - } -} -struct Word { - id: u32, - text: String, - folded: String, -} -pub struct Lookup { - words: Vec, - ids: HashMap, - groups: HashMap<[u32; 3], Vec>, - prefixes: Vec, -} -impl Lookup { - pub fn open(path: &Path) -> Result { - let file = File::open(path).map_err(|_| ())?; - let metadata = file.metadata().map_err(|_| ())?; - if !metadata.is_file() || !(60..=MAX_FILE).contains(&metadata.len()) { - return Err(()); - } - Self::read(BufReader::new(file.take(MAX_FILE + 1))) - } - fn read(mut input: impl Read) -> Result { - let mut magic = [0; 8]; - input.read_exact(&mut magic).map_err(|_| ())?; - if &magic != MAGIC { - return Err(()); - } - let mut expected = [0; 32]; - input.read_exact(&mut expected).map_err(|_| ())?; - let mut r = Reader { - input, - hash: Sha256::new(), - }; - let count = r.number()? as usize; - if count == 0 || count > MAX_WORDS { - return Err(()); - } - let mut words: Vec = Vec::new(); - words.try_reserve_exact(count).map_err(|_| ())?; - let mut ids = HashMap::new(); - ids.try_reserve(count).map_err(|_| ())?; - let mut unique = HashSet::new(); - unique.try_reserve(count).map_err(|_| ())?; - let mut previous = None; - let mut total = 0usize; - for rank in 0..count { - let id = r.number()?; - let frequency = r.number()?; - if id == 0 - || !unique.insert(id) - || previous.is_some_and(|(f, i)| frequency > f || (frequency == f && id <= i)) - { - return Err(()); - } - previous = Some((frequency, id)); - let size = r.number()? as usize; - total = total.checked_add(size).ok_or(())?; - if size > MAX_WORD_BYTES || total > MAX_TEXT_BYTES { - return Err(()); - } - let mut bytes = Vec::new(); - bytes.try_reserve_exact(size).map_err(|_| ())?; - bytes.resize(size, 0); - r.bytes(&mut bytes)?; - let text = String::from_utf8(bytes).map_err(|_| ())?; - let folded = text.to_lowercase(); - ids.entry(folded.clone()).or_insert(rank as u32); - words.push(Word { id, text, folded }); - } - drop(unique); - let mut groups: HashMap<[u32; 3], Vec> = HashMap::new(); - for length in 1..=3 { - let rows = r.number()? as usize; - if rows > MAX_ROWS { - return Err(()); - } - let mut previous = None; - for _ in 0..rows { - let mut key = [NONE; 3]; - for slot in &mut key[..length] { - *slot = r.number()?; - if *slot as usize >= count { - return Err(()); - } - } - let rank = r.number()?; - let frequency = r.number()?; - let word = words.get(rank as usize).ok_or(())?; - if previous.is_some_and(|(f, i)| frequency > f || (frequency == f && word.id < i)) { - return Err(()); - } - previous = Some((frequency, word.id)); - if !groups.contains_key(&key) { - groups.try_reserve(1).map_err(|_| ())?; - } - let list = groups.entry(key).or_default(); - list.try_reserve(1).map_err(|_| ())?; - list.push(rank); - } - } - if r.number()? as usize != count { - return Err(()); - } - let mut prefixes = Vec::new(); - prefixes.try_reserve_exact(count).map_err(|_| ())?; - let mut seen = Vec::new(); - seen.try_reserve_exact(count).map_err(|_| ())?; - seen.resize(count, false); - for _ in 0..count { - let rank = r.number()?; - let word = words.get(rank as usize).ok_or(())?; - if seen[rank as usize] { - return Err(()); - } - seen[rank as usize] = true; - if let Some(&p) = prefixes.last() { - let previous: &Word = &words[p as usize]; - if (&previous.folded, p) >= (&word.folded, rank) { - return Err(()); - } - } - prefixes.push(rank); - } - let mut trailing = [0]; - if r.input.read(&mut trailing).map_err(|_| ())? != 0 - || r.hash.finalize().as_slice() != expected - { - return Err(()); - } - Ok(Self { - words, - ids, - groups, - prefixes, - }) - } - pub fn predict(&self, context: &[String], prefix: &str) -> Vec { - let prefix = prefix.to_lowercase(); - let mut result = Vec::new(); - let mut seen = HashSet::new(); - for count in (1..=context.len().min(3)).rev() { - let mut key = [NONE; 3]; - let mut known = true; - for (slot, word) in key.iter_mut().zip(&context[context.len() - count..]) { - match self.ids.get(&word.to_lowercase()) { - Some(rank) => *slot = *rank, - None => { - known = false; - break; - } - } - } - if let Some(list) = known.then(|| self.groups.get(&key)).flatten() { - for &rank in list { - self.push(rank, &prefix, &mut result, &mut seen); - if result.len() == 5 { - return result; - } - } - } - } - if prefix.is_empty() { - for rank in 0..self.words.len() { - self.push(rank as u32, &prefix, &mut result, &mut seen); - if result.len() == 5 { - break; - } - } - } else { - let start = self - .prefixes - .partition_point(|&i| self.words[i as usize].folded.as_str() < prefix.as_str()); - let mut best: Vec = Vec::new(); - for &rank in self.prefixes[start..] - .iter() - .take_while(|&&i| self.words[i as usize].folded.starts_with(&prefix)) - { - let word = &self.words[rank as usize]; - if usable(&word.text) && !seen.contains(&word.folded) { - if let Some(i) = best - .iter() - .position(|&r| self.words[r as usize].folded == word.folded) - { - best[i] = best[i].min(rank); - } else { - best.push(rank); - } - best.sort_unstable(); - best.truncate(5 - result.len()); - } - } - for rank in best { - self.push(rank, &prefix, &mut result, &mut seen); - } - } - 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 - && usable(&word.text) - && word.folded.starts_with(prefix) - && seen.insert(word.folded.clone()) - { - result.push(word.text.clone()); - } - } -} -fn usable(word: &str) -> bool { - word.chars().count() <= 48 && !word.chars().any(char::is_whitespace) -} - -#[cfg(test)] -impl Lookup { - pub fn fixture() -> Self { - Self::read(&tests::fixture_bytes()[..]).unwrap() - } - pub fn fixture_bytes() -> Vec { - tests::fixture_bytes() - } -} -#[cfg(test)] -mod tests { - use super::*; - fn n(out: &mut Vec, value: u32) { - out.extend_from_slice(&value.to_le_bytes()); - } - fn seal(payload: Vec) -> Vec { - let mut out = MAGIC.to_vec(); - out.extend_from_slice(&Sha256::digest(&payload)); - out.extend(payload); - out - } - pub(super) fn fixture_bytes() -> Vec { - let mut out = Vec::new(); - n(&mut out, 5); - for (id, text, freq) in [ - (1, "water", 9), - (2, "walk", 8), - (3, "want", 7), - (4, "I", 6), - (5, "like", 5), - ] { - n(&mut out, id); - n(&mut out, freq); - n(&mut out, text.len() as u32); - out.extend_from_slice(text.as_bytes()); - } - for _ in 0..3 { - n(&mut out, 0); - } - n(&mut out, 5); - for rank in [3, 4, 1, 2, 0] { - n(&mut out, rank); - } - seal(out) - } - fn change(bytes: &[u8], offset: usize, value: u32) -> Vec { - let mut payload = bytes[40..].to_vec(); - payload[offset..offset + 4].copy_from_slice(&value.to_le_bytes()); - seal(payload) - } - #[test] - fn malformed_artifacts_fail_closed() { - let bytes = fixture_bytes(); - assert!(Lookup::read(&bytes[..]).is_ok()); - for size in 0..bytes.len() { - assert!(Lookup::read(&bytes[..size]).is_err()); - } - for i in 0..bytes.len() { - let mut bad = bytes.clone(); - bad[i] ^= 0x80; - assert!(Lookup::read(&bad[..]).is_err()); - } - let mut trailing = bytes.clone(); - trailing.push(0); - assert!(Lookup::read(&trailing[..]).is_err()); - for (offset, value) in [ - (0, 0), - (0, u32::MAX), - (4, 0), - (8, 0), - (12, u32::MAX), - (21, 1), - ] { - assert!(Lookup::read(&change(&bytes, offset, value)[..]).is_err()); - } - let prefix_count = bytes.len() - 40 - 24; - assert!(Lookup::read(&change(&bytes, prefix_count, 4)[..]).is_err()); - assert!(Lookup::read(&change(&bytes, prefix_count + 4, 5)[..]).is_err()); - assert!(Lookup::read(&change(&bytes, prefix_count + 8, 3)[..]).is_err()); - assert!(Lookup::read(&change(&bytes, prefix_count + 4, 0)[..]).is_err()); - assert!(Lookup::read(&change(&bytes, prefix_count - 12, u32::MAX)[..]).is_err()); - } - #[test] - fn ranking_unicode_deduplication_and_context_backoff() { - let entries = [ - (1, "I", 100), - (2, "like", 90), - (3, "water", 20), - (4, "walk", 20), - (5, "Water", 19), - (6, "café", 15), - (7, "CAFÉ", 14), - (8, "café", 13), - (9, "two words", 12), - (10, "İstanbul", 11), - (11, "你好", 10), - ]; - let mut out = Vec::new(); - n(&mut out, entries.len() as u32); - for (id, text, freq) in entries { - n(&mut out, id); - n(&mut out, freq); - n(&mut out, text.len() as u32); - out.extend_from_slice(text.as_bytes()); - } - n(&mut out, 3); - for (rank, freq) in [(4, 40), (2, 30), (3, 30)] { - n(&mut out, 1); - n(&mut out, rank); - n(&mut out, freq); - } - n(&mut out, 1); - for value in [0, 1, 2, 1] { - n(&mut out, value); - } - n(&mut out, 1); - for value in [0, 0, 1, 3, 1] { - n(&mut out, value); - } - let mut prefixes: Vec<_> = (0..entries.len()).collect(); - prefixes.sort_by_key(|&i| (entries[i].1.to_lowercase(), i)); - n(&mut out, prefixes.len() as u32); - for i in prefixes { - n(&mut out, i as u32); - } - let bytes = seal(out); - let lookup = Lookup::read(&bytes[..]).unwrap(); - assert_eq!( - lookup.predict(&["I".into(), "I".into(), "like".into()], "w"), - ["walk", "water"] - ); - assert_eq!(lookup.predict(&["like".into()], "w"), ["Water", "walk"]); - assert_eq!(lookup.predict(&[], "CAF"), ["café", "café"]); - assert_eq!(lookup.predict(&[], "İ"), ["İstanbul"]); - assert_eq!(lookup.predict(&[], "你"), ["你好"]); - assert_eq!(lookup.predict(&["unknown".into()], "w"), ["water", "walk"]); - assert!(!lookup.predict(&[], "").contains(&"two words".into())); - assert!(lookup.predict(&[], "unknown").is_empty()); - } - #[test] - fn missing_and_oversized_files_fail_closed() { - let path = std::env::temp_dir().join(format!("switchify-lookup-{}", std::process::id())); - assert!(Lookup::open(&path).is_err()); - let file = File::create(&path).unwrap(); - file.set_len(MAX_FILE + 1).unwrap(); - drop(file); - assert!(Lookup::open(&path).is_err()); - std::fs::remove_file(path).unwrap(); - } -} diff --git a/src-tauri/src/prediction/mod.rs b/src-tauri/src/prediction/mod.rs index e7bcf086..a26a7a2f 100644 --- a/src-tauri/src/prediction/mod.rs +++ b/src-tauri/src/prediction/mod.rs @@ -1,9 +1,6 @@ mod activity; mod context; mod database; -#[cfg(test)] -mod database_reference; -mod lookup; mod model; pub mod worker; @@ -66,13 +63,12 @@ struct Client { _job: Job, } impl Client { - fn start(path: &Path, ignored: Vec, enhanced: bool) -> Result { + fn start(path: &Path, ignored: Vec) -> 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()); @@ -143,7 +139,6 @@ impl Drop for Client { struct Service { client: Option, failed: bool, - received_reply: bool, generation: u64, outstanding: Option, last: Option, @@ -155,8 +150,6 @@ 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 { @@ -197,13 +190,8 @@ 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) - }; + // Model loading is asynchronous; only inference uses the query deadline. + let deadline = Duration::from_secs(2); self.outstanding .is_some_and(|sent| now.duration_since(sent) >= deadline) } @@ -232,6 +220,7 @@ impl Service { revision: u64, batch: Option, tracking: bool, + status: database::Status, ) { if revision < self.acknowledged_revision || revision > self.edit_revision { self.fail(keyboard); @@ -239,7 +228,7 @@ impl Service { } self.acknowledged_revision = revision; if generation == self.generation && revision == self.edit_revision { - self.suggestions(keyboard, batch, tracking); + self.suggestions(keyboard, batch, tracking, status); } else { self.tracking = tracking; if !tracking { @@ -253,6 +242,7 @@ impl Service { keyboard: &mut Keyboard, batch: Option, tracking: bool, + status: database::Status, ) { self.tracking = tracking; if !tracking { @@ -260,7 +250,12 @@ impl Service { } let visible = keyboard.page == Page::Letters && !keyboard.modifiers[1..].iter().any(|m| *m != Modifier::Off); - keyboard.predictions(if visible { batch } else { None }, false); + keyboard.predictions( + if visible { batch } else { None }, + status == database::Status::Unavailable, + ); + keyboard.prediction_loading = status == database::Status::Loading; + keyboard.prediction_tracking(tracking); } fn fail(&mut self, keyboard: &mut Keyboard) { self.client = None; @@ -323,6 +318,11 @@ pub fn keyboard_input_context() -> Option { pub fn stop() { SERVICE.with(|s| *s.borrow_mut() = Service::default()); } +/// Restart only prediction. The keyboard and its scan session remain open, +/// while the old worker and its private text context are discarded. +pub fn retry() { + stop(); +} pub fn record(stroke: Stroke, success: bool, scope: InputScope) { SERVICE.with(|s| { let mut s = s.borrow_mut(); @@ -375,16 +375,16 @@ fn resource(app: &AppHandle) -> Result { let bundled = app .path() .resolve( - "resources/word-predictions.lookup", + "resources/prediction-blocklist.txt", tauri::path::BaseDirectory::Resource, ) .map_err(|_| ()); let development = cfg!(debug_assertions) - .then(|| Path::new(env!("CARGO_MANIFEST_DIR")).join("resources/word-predictions.lookup")); - database_resource(bundled, development) + .then(|| Path::new(env!("CARGO_MANIFEST_DIR")).join("resources/prediction-blocklist.txt")); + model_resource(bundled, development) } -fn database_resource( +fn model_resource( bundled: Result, development: Option, ) -> Result { @@ -395,13 +395,7 @@ fn database_resource( .ok_or(()) } -pub fn poll( - app: &AppHandle, - keyboard: Option<&mut Keyboard>, - enabled: bool, - enhanced: bool, - ignored: &[String], -) { +pub fn poll(app: &AppHandle, keyboard: Option<&mut Keyboard>, enabled: bool, ignored: &[String]) { let Some(keyboard) = keyboard else { stop(); return; @@ -413,17 +407,6 @@ pub fn poll( } SERVICE.with(|slot| { let mut s = slot.borrow_mut(); - if s.enhanced != enhanced { - // Changing the engine restarts the worker, like reopening the keyboard. - if s.accepting || s.accept.is_some() { - keyboard.failed(); - } - *s = Service { - enhanced, - ..Service::default() - }; - keyboard.predictions(None, false); - } let case = ( keyboard.prediction_shift(), keyboard.caps, @@ -433,7 +416,7 @@ pub fn poll( s.case = Some(case); s.generation = s.generation.wrapping_add(1); s.last = None; - keyboard.predictions(None, false); + keyboard.predictions(None, s.failed); } if s.failed { keyboard.predictions(None, true); @@ -445,13 +428,14 @@ pub fn poll( .filter_map(|name| crate::switch_input::prediction_key_code(name)) .collect(); s.client = resource(app) - .and_then(|path| Client::start(&path, ignored, enhanced)) + .and_then(|path| Client::start(path.parent().ok_or(())?, ignored)) .ok(); if s.client.is_none() { s.failed = true; keyboard.predictions(None, true); return; } + keyboard.prediction_loading = true; } if s.response_timed_out(Instant::now()) { s.fail(keyboard); @@ -461,15 +445,17 @@ pub fn poll( match reply { Ok(Ok(response)) => { s.outstanding = None; - s.received_reply = true; match response { Response::Suggestions { generation, batch, revision, tracking, + status, } => { - s.received_suggestions(keyboard, generation, revision, batch, tracking); + s.received_suggestions( + keyboard, generation, revision, batch, tracking, status, + ); } Response::Insert { generation, @@ -560,6 +546,28 @@ pub fn poll( mod tests { use super::*; #[test] + fn retry_discards_failed_worker_state_and_private_edits() { + stop(); + SERVICE.with(|slot| { + let mut service = slot.borrow_mut(); + service.failed = true; + service.generation = 17; + service + .edit + .push(InputScope::capture().record(Edit::Append("test".into()))); + service.accept = Some((3, 0)); + }); + retry(); + SERVICE.with(|slot| { + let service = slot.borrow(); + assert!(!service.failed); + assert_eq!(service.generation, 0); + assert!(service.edit.is_empty()); + assert!(service.accept.is_none()); + assert!(service.client.is_none()); + }); + } + #[test] fn activity_observer_retries_after_failure_or_lost_hook() { let state = AtomicU8::new(0); assert!(claim_keyboard_activity_start(&state, false)); @@ -605,7 +613,6 @@ 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()); @@ -614,16 +621,15 @@ mod tests { } } #[test] - fn first_worker_reply_has_a_startup_deadline_then_queries_use_two_seconds() { + fn worker_replies_have_two_second_deadline() { let now = Instant::now(); let mut service = Service { outstanding: Some(now - Duration::from_secs(7)), ..Default::default() }; - assert!(!service.response_timed_out(now)); + 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)); @@ -671,7 +677,7 @@ mod tests { ..Default::default() }; let mut keyboard = Keyboard::new(false); - service.received_suggestions(&mut keyboard, 0, 4, None, true); + service.received_suggestions(&mut keyboard, 0, 4, None, true, database::Status::Ready); assert!(service.failed); } @@ -684,30 +690,30 @@ mod tests { std::fs::write(&bundled, []).unwrap(); std::fs::write(&source, []).unwrap(); assert_eq!( - database_resource(Ok(bundled.clone()), Some(source.clone())), + model_resource(Ok(bundled.clone()), Some(source.clone())), Ok(bundled.clone()) ); assert_eq!( - database_resource(Ok(bundled.clone()), None), + model_resource(Ok(bundled.clone()), None), Ok(bundled.clone()) ); std::fs::remove_file(&source).unwrap(); assert_eq!( - database_resource(Ok(bundled.clone()), Some(source.clone())), + model_resource(Ok(bundled.clone()), Some(source.clone())), Ok(bundled.clone()) ); std::fs::write(&source, []).unwrap(); std::fs::remove_file(&bundled).unwrap(); assert_eq!( - database_resource(Ok(bundled.clone()), Some(source.clone())), + model_resource(Ok(bundled.clone()), Some(source.clone())), Ok(source.clone()) ); assert_eq!( - database_resource(Err(()), Some(source.clone())), + model_resource(Err(()), Some(source.clone())), Ok(source.clone()) ); - assert_eq!(database_resource(Ok(bundled), None), Err(())); - assert_eq!(database_resource(Ok(root.clone()), None), Err(())); + assert_eq!(model_resource(Ok(bundled), None), Err(())); + assert_eq!(model_resource(Ok(root.clone()), None), Err(())); std::fs::remove_file(source).unwrap(); std::fs::remove_dir(root).unwrap(); } diff --git a/src-tauri/src/prediction/model.rs b/src-tauri/src/prediction/model.rs index 04cb779f..fe562030 100644 --- a/src-tauri/src/prediction/model.rs +++ b/src-tauri/src/prediction/model.rs @@ -1,10 +1,9 @@ -//! Optional on-device language model for word prediction. +//! 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. +//! typed prefix. //! 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 +//! case variants merged. Low probability and blocked words never //! appear. Runs only inside the prediction worker; text is never logged. use ort::{ session::{builder::GraphOptimizationLevel, Session, SessionInputValue}, @@ -18,12 +17,12 @@ use std::{ }; use tokenizers::Tokenizer; -const BEAM: usize = 8; +const BEAM: usize = 16; 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; +const PRUNE: f32 = -18.0; +/// Minimum log-probability for a candidate word. +const UNKNOWN_MIN: f32 = -16.0; /// Most recent buffered characters the model reads. const CONTEXT_CHARS: usize = 256; @@ -71,14 +70,12 @@ 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; @@ -88,7 +85,7 @@ pub fn spell( // 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(()); + return Ok(Vec::new()); } let compatible = |word: &str| { let w = normalize(word); @@ -130,7 +127,7 @@ pub fn spell( // 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()) + .filter(|(_, f)| f.0.ln() >= UNKNOWN_MIN) .map(|(_, f)| f.0) .collect(); p.sort_by(|a, b| b.total_cmp(a)); @@ -157,7 +154,7 @@ pub fn spell( return Err(()); } let key = normalize(word); - let plausible = key.chars().count() > 1 || known(&key).is_some(); + let plausible = key.chars().count() > 1 || key == "a" || key == "i"; if plausible && key.starts_with(&prefix) && !vocabulary.blocked.contains(&key) { let boundary: f32 = row .iter() @@ -190,24 +187,22 @@ pub fn spell( } 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)), - }) + .filter_map(|(key, (p, _, spelling))| (p.ln() >= UNKNOWN_MIN).then_some((p, key, spelling))) .collect(); words.sort_by(|a, b| b.0.total_cmp(&a.0).then_with(|| a.1.cmp(&b.1))); Ok(words .into_iter() .take(5) .map(|(_, key, spelling)| { - // Capitals on unknown words at a sentence start come from position. - let word = if known(&key).is_none() && (start || text.ends_with(['.', '!', '?'])) { - key - } else { + // Sentence position and all-caps tokens should not force uppercase + // onto ordinary completions. Keep mixed case names from the model. + let word = if spelling.chars().skip(1).any(|c| c.is_ascii_uppercase()) + && !spelling.chars().all(|c| c.is_ascii_uppercase()) + { spelling.replace('’', "'") + } else { + key }; - // Match the typed apostrophe so the completion still extends the prefix. if typed.contains('’') { word.replace('\'', "’") } else { @@ -418,15 +413,10 @@ impl Model { } } -/// The enhanced predictor behind `Database`, so tests can substitute a fake. +/// Predictor behind `Database`, so tests can substitute a fake. pub trait Predict: Send { - fn predict( - &mut self, - before: &str, - prefix: &str, - known: &dyn Fn(&str) -> Option, - deadline: Instant, - ) -> Result, ()>; + fn predict(&mut self, before: &str, prefix: &str, deadline: Instant) + -> Result, ()>; } impl Predict for Model { @@ -434,17 +424,17 @@ impl Predict for Model { &mut self, before: &str, prefix: &str, - known: &dyn Fn(&str) -> Option, deadline: Instant, ) -> Result, ()> { - spell( - &mut self.onnx, - &self.vocabulary, - recent(before).ok_or(())?, - prefix, - known, - deadline, - ) + // The model's beginning-of-document distribution favors site names. + // A fixed, local context yields useful initial completions without + // reading any additional user text. + let context = match recent(before) { + Some("") if !prefix.is_empty() => "I ", + Some(text) => text, + None => return Ok(Vec::new()), + }; + spell(&mut self.onnx, &self.vocabulary, context, prefix, deadline) } } @@ -510,70 +500,51 @@ mod tests { } } - fn run( - s: &mut Fake, - before: &str, - prefix: &str, - known: impl Fn(&str) -> Option, - ) -> Result, ()> { + fn run(s: &mut Fake, before: &str, prefix: &str) -> 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()) + spell(s, &vocabulary(), before, prefix, deadline) } #[test] fn spells_multi_piece_words_merges_case_and_applies_prefix() { let mut s = fake(&[(1, 0.2), (2, 0.3), (6, 0.4), (9, 0.1)]); // water = 0.4 + 0.5 * 0.6; waffle = 0.5 * 0.3; bare "wa" = 0.5 * 0.1. - let words = run(&mut s, "I would like ", "wa", lower).unwrap(); + let words = run(&mut s, "I would like ", "wa").unwrap(); assert_eq!(words, ["water", "waffle", "wa"]); - assert_eq!( - run(&mut s, "I would like ", "waf", lower).unwrap(), - ["waffle"] - ); + assert_eq!(run(&mut s, "I would like ", "waf").unwrap(), ["waffle"]); } #[test] fn blocked_single_letter_and_improbable_unknown_words_are_removed() { let mut s = fake(&[(9, 0.6), (10, 0.3), (6, 0.1)]); - assert_eq!(run(&mut s, "Say ", "", |_| None).unwrap(), ["water"]); - let mut s = fake(&[(6, 1e-5)]); - assert!(run(&mut s, "Say ", "", |_| None).unwrap().is_empty()); - assert_eq!(run(&mut s, "Say ", "", lower).unwrap(), ["water"]); + assert_eq!(run(&mut s, "Say ", "").unwrap(), ["water"]); + let mut s = fake(&[(6, 1e-8)]); + assert!(run(&mut s, "Say ", "").unwrap().is_empty()); } #[test] - fn casing_comes_from_the_lookup_or_the_model_but_not_sentence_position() { + fn casing_comes_from_model_except_sentence_position() { let mut s = fake(&[(1, 0.9)]); let lowered = ["water", "waffle", "wa"]; - assert_eq!(run(&mut s, "Send ", "", lower).unwrap(), lowered); - assert_eq!( - run(&mut s, "Send ", "", |_| None).unwrap(), - ["Water", "Waffle", "Wa"] - ); - assert_eq!(run(&mut s, "Done. ", "", |_| None).unwrap(), lowered); + assert_eq!(run(&mut s, "Send ", "").unwrap(), lowered); + assert_eq!(run(&mut s, "Done. ", "").unwrap(), lowered); } #[test] fn text_start_needs_capital_and_attached_prefixes_fall_back() { let mut s = fake(&[(7, 0.5), (8, 0.5)]); - assert_eq!( - run(&mut s, "", "", lower).unwrap(), - ["water", "waffle", "wa"] - ); - assert!(run(&mut s, "hello-", "wa", lower).is_err()); + assert_eq!(run(&mut s, "", "").unwrap(), ["water", "waffle", "wa"]); + assert!(run(&mut s, "hello-", "wa").unwrap().is_empty()); assert_eq!(s.contexts, 1); } #[test] fn an_expired_deadline_stops_expansion_and_bad_vocabularies_fail() { let mut s = fake(&[(6, 0.9)]); - let expired = spell(&mut s, &vocabulary(), "Say ", "", lower, Instant::now()); + let expired = spell(&mut s, &vocabulary(), "Say ", "", Instant::now()); assert!(expired.unwrap().is_empty()); s.next.pop(); - assert!(run(&mut s, "Say ", "", lower).is_err()); + assert!(run(&mut s, "Say ", "").is_err()); } #[test] @@ -608,11 +579,10 @@ mod tests { ) .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 [ + ("", "w", Some("want")), + ("I would like ", "wa", None), ("I would like a cup of ", "", Some("tea")), ("Please put the ", "ket", Some("kettle")), ("Can you send a ", "wh", None), @@ -620,7 +590,7 @@ mod tests { ] { let t = Instant::now(); let deadline = t + std::time::Duration::from_secs(60); - let words = model.predict(before, prefix, &known, deadline).unwrap(); + let words = model.predict(before, prefix, deadline).unwrap(); times.push(t.elapsed().as_secs_f64() * 1000.0); println!("{before:?}+{prefix:?} -> {words:?}"); assert!(words.iter().all(|w| normalize(w).starts_with(prefix))); diff --git a/src-tauri/src/prediction/worker.rs b/src-tauri/src/prediction/worker.rs index db815b5d..85fd054b 100644 --- a/src-tauri/src/prediction/worker.rs +++ b/src-tauri/src/prediction/worker.rs @@ -1,4 +1,7 @@ -use super::{activity, context, database::Database}; +use super::{ + activity, context, + database::{Database, Status}, +}; use serde::{Deserialize, Serialize}; use std::{ io::{Read, Write}, @@ -53,6 +56,7 @@ pub enum Response { batch: Option, revision: u64, tracking: bool, + status: Status, }, Insert { generation: u64, @@ -84,6 +88,7 @@ pub fn receive(reader: &mut impl Read) -> Result } pub struct Engine { database: Database, + status: Status, foreground: fn() -> Result, observe: fn() -> (u64, bool), last_activity: fn() -> u64, @@ -103,6 +108,7 @@ impl Engine { pub fn new(database: Database, tracked: bool) -> Self { Self { database, + status: Status::Loading, tracked, foreground: || crate::scan_host::foreground().map_err(|_| ()), observe: activity::snapshot, @@ -200,17 +206,16 @@ impl Engine { self.snapshot = None; return None; } - if self.snapshot.as_ref() == Some(&self.buffer) && self.case == (shift, caps) { + self.status = self.database.status(); + if self.snapshot.as_ref() == Some(&self.buffer) + && self.case == (shift, caps) + && self.status != Status::Loading + { return self.batch.clone(); } let ctx = context::extract(&self.buffer, self.clipped); - let words = match self.database.predict(&ctx) { - Ok(words) => words, - Err(()) => { - self.clear(); - return None; - } - }; + let (status, words) = self.database.predict(&ctx); + self.status = status; self.token = self.token.wrapping_add(1); self.suffixes.clear(); let mut labels = Vec::new(); @@ -268,6 +273,7 @@ impl Engine { caps, } => { let batch = self.query(edits, revision, shift, caps); + self.status = self.database.status(); let tracking = self .target .is_some_and(|target| self.stable(target, self.activity)); @@ -276,6 +282,7 @@ impl Engine { batch, revision: self.revision, tracking, + status: self.status, } } Request::Accept { @@ -302,12 +309,11 @@ pub fn run_from_args() -> bool { return false; } let work = || -> Result<(), ()> { - if args.len() != 5 { + if args.len() != 4 { 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(()); } @@ -325,7 +331,7 @@ pub fn run_from_args() -> bool { let tracked = activity::start(); #[cfg(any(target_os = "windows", target_os = "macos"))] { - let database = Database::open(&path, enhanced)?; + let database = Database::open(&path); let mut engine = Engine::new(database, tracked); let mut input = std::io::stdin().lock(); let mut output = std::io::stdout().lock(); @@ -478,32 +484,4 @@ mod tests { }; assert!(send(&mut Vec::new(), &request).is_err()); } - #[test] - #[ignore = "Manual database performance measurement; no desktop input"] - fn bundled_database_benchmark() { - let start = std::time::Instant::now(); - 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; - let mut times = Vec::new(); - for i in 0..200 { - let text = [ - "I would like wa", - "the", - "hello ", - "I would like some ", - "zzzx", - ][i % 5]; - let ctx = context::extract(text, false); - let t = std::time::Instant::now(); - assert!(db.predict(&ctx).is_ok()); - times.push(t.elapsed().as_secs_f64() * 1000.0); - } - times.sort_by(f64::total_cmp); - println!("database startup_ms={startup:.2} samples=200 failures=0 median_ms={:.2} p95_ms={:.2} max_ms={:.2}",times[99],times[189],times[199]); - } } diff --git a/src-tauri/src/scan_executor.rs b/src-tauri/src/scan_executor.rs index 08a54431..5b318c0e 100644 --- a/src-tauri/src/scan_executor.rs +++ b/src-tauri/src/scan_executor.rs @@ -15,7 +15,9 @@ pub fn execute( ); } match request { - Request::Prediction { .. } => Err("Prediction requires the scan controller.".into()), + Request::Prediction { .. } | Request::PredictionRetry => { + Err("Prediction requires the scan controller.".into()) + } Request::OpenKeyboard | Request::OpenMouse | Request::OpenPoint => input.release_all(), Request::MouseMove { dx, dy } => input.move_pointer_pixels(dx, dy).map(|_| ()), Request::MouseMoveAbsolute { x, y } => input.move_pointer_pixels_absolute(x, y).map(|_| ()), diff --git a/src-tauri/src/scan_keyboard.rs b/src-tauri/src/scan_keyboard.rs index c1094d38..d1846317 100644 --- a/src-tauri/src/scan_keyboard.rs +++ b/src-tauri/src/scan_keyboard.rs @@ -36,6 +36,7 @@ impl Modifier { #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum Key { Prediction(usize), + RetryPrediction, Character(char, char), Named(&'static str), Modifier(usize), @@ -90,6 +91,7 @@ impl Stroke { #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum Output { Prediction { token: u64, index: usize }, + RetryPrediction, Stroke(Stroke), Punctuation(Punctuation), Close, @@ -202,6 +204,8 @@ pub struct Keyboard { activation: Option, prediction_enabled: bool, prediction_failed: bool, + pub(crate) prediction_loading: bool, + prediction_tracking: bool, predictions: Option, queued_predictions: Option, prefer_predictions: bool, @@ -234,6 +238,8 @@ impl Keyboard { activation: None, prediction_enabled: false, prediction_failed: false, + prediction_loading: false, + prediction_tracking: true, predictions: None, queued_predictions: None, prefer_predictions: true, @@ -255,6 +261,8 @@ impl Keyboard { return; } self.prediction_enabled = enabled; + self.prediction_loading = enabled; + self.prediction_failed = false; self.rebuild_rows(); } fn rebuild_rows(&mut self) { @@ -266,6 +274,9 @@ impl Keyboard { if self.prediction_enabled && self.page == Page::Letters && !self.positioning { self.rows.insert(0, (0..5).map(Key::Prediction).collect()); } + if self.prediction_enabled && self.prediction_failed && !self.positioning { + self.rows.last_mut().unwrap().push(Key::RetryPrediction); + } self.scan.replace(ItemScanner::nodes(&self.rows)); self.scan.restart(); self.activation = None; @@ -278,10 +289,24 @@ impl Keyboard { && self.scan.position(&self.rows).0 == 0 } pub fn predictions(&mut self, batch: Option, failed: bool) { + let retry_changed = self.prediction_failed != failed; self.prediction_failed = failed; + self.prediction_loading = false; + if failed { + self.predictions = None; + self.queued_predictions = None; + } + if retry_changed { + let suspended = self.scan.suspended; + self.rebuild_rows(); + self.scan.suspended = suspended; + } if batch.is_none() { self.predictions = None; self.queued_predictions = None; + if retry_changed && !self.scan.suspended { + self.skip_disabled(false); + } return; } if self.waiting_after_typing { @@ -310,6 +335,40 @@ impl Keyboard { } } } + pub fn prediction_tracking(&mut self, tracking: bool) { + self.prediction_tracking = tracking; + } + fn prediction_note(&self) -> Option { + if !self.prediction_enabled { + return None; + } + Some(if self.prediction_failed { + "Prediction unavailable".into() + } else if self.prediction_loading { + "Prediction loading".into() + } else if !self.prediction_tracking { + "Prediction paused".into() + } else if self.queued_predictions.is_some() { + "Suggestions updating".into() + } else { + match self.predictions.as_ref() { + Some(batch) if batch.words.len() == 1 => "1 suggestion".into(), + Some(batch) if !batch.words.is_empty() => { + format!("{} suggestions", batch.words.len()) + } + Some(_) => "No suggestions".into(), + None => "Type for suggestions".into(), + } + }) + } + fn retry_predictions(&mut self) { + self.prediction_failed = false; + self.prediction_loading = true; + self.prediction_tracking = true; + self.predictions = None; + self.queued_predictions = None; + self.rebuild_rows(); + } fn disabled(&self) -> bool { if !self.prediction_row_active() || self.scan.nav.escaping() { return false; @@ -434,6 +493,11 @@ impl Keyboard { fn choose_with_context(&mut self, key: Key, context: Option) -> Option { self.discard_stale_context(context); match key { + Key::RetryPrediction if self.prediction_failed => { + self.retry_predictions(); + return Some(Output::RetryPrediction); + } + Key::RetryPrediction => return None, Key::Prediction(index) => { let batch = self.predictions.as_ref()?; if index >= batch.words.len() { @@ -596,7 +660,7 @@ impl Keyboard { self.auto_shift = false; self.capital_context = None; self.pending_typed = None; - self.predictions(None, false); + self.predictions(None, self.prediction_failed); self.restart(); } pub fn failed(&mut self) { @@ -620,6 +684,7 @@ impl Keyboard { .and_then(|b| b.words.get(i)) .cloned() .unwrap_or_default(), + Key::RetryPrediction => "Retry predictions".into(), Key::Character(' ', _) => "Space".into(), Key::Character(..) => self.stroke(key).character().unwrap().to_string(), Key::Named(name) => match name { @@ -681,7 +746,7 @@ impl Keyboard { Key::Character(' ', _) => 4.0, Key::Named("Backspace" | "Enter") => 1.9, Key::Named("Tab") | Key::Caps | Key::Modifier(_) => 1.6, - Key::Close | Key::Dock => 1.5, + Key::Close | Key::Dock | Key::RetryPrediction => 1.5, _ => 1.0, } } @@ -689,7 +754,9 @@ impl Keyboard { TileStyle { role: match key { Key::Character(..) => TileRole::Character, - Key::Page(_) | Key::Dock | Key::Back | Key::Close => TileRole::Toolbar, + Key::Page(_) | Key::Dock | Key::Back | Key::Close | Key::RetryPrediction => { + TileRole::Toolbar + } _ => TileRole::Utility, }, active: match key { @@ -721,8 +788,6 @@ impl Keyboard { crate::scan_panel::BACK_TO_ROWS.to_owned() } else if self.disabled() { "Suggestions updating · Select to continue".to_owned() - } else if self.prediction_failed { - "Predictions unavailable · Keyboard ready".to_owned() } else { crate::scan_panel::scanning_status( page, @@ -750,6 +815,7 @@ impl Keyboard { }) .collect(), status, + note: self.prediction_note(), dock: self.dock, selected: (!self.scan.suspended && !self.disabled() && !self.scan.nav.escaping()) .then_some((active_row, active_column)), @@ -765,6 +831,79 @@ impl Keyboard { mod tests { use super::*; + #[test] + fn prediction_badge_explains_empty_slots_without_replacing_scan_prompt() { + let mut keyboard = Keyboard::new(false); + keyboard.enable_predictions(true); + let screen = Rect { + x: 0.0, + y: 0.0, + width: 1920.0, + height: 1080.0, + }; + let badge = |keyboard: &Keyboard| { + let frame = keyboard.frame(screen, 1.0, ScannerColor::default()); + let prompt = frame.tiles.last().unwrap(); + let badge = &frame.tiles[frame.tiles.len() - 2]; + assert!(prompt.text.contains("Select")); + assert!(!badge.selected); + assert!(badge.rect.x >= prompt.rect.x + prompt.rect.width); + badge.text.clone() + }; + assert_eq!(badge(&keyboard), "Prediction loading"); + keyboard.predictions(None, false); + assert_eq!(badge(&keyboard), "Type for suggestions"); + keyboard.predictions( + Some(crate::prediction::worker::Batch { + token: 1, + words: vec![], + }), + false, + ); + assert_eq!(badge(&keyboard), "No suggestions"); + keyboard.predictions( + Some(crate::prediction::worker::Batch { + token: 2, + words: vec!["water".into()], + }), + false, + ); + assert_eq!(badge(&keyboard), "1 suggestion"); + keyboard.prediction_tracking(false); + assert_eq!(badge(&keyboard), "Prediction paused"); + } + + #[test] + fn retry_tile_restarts_prediction_without_closing_keyboard() { + let mut keyboard = Keyboard::new(false); + keyboard.enable_predictions(true); + keyboard.scan.suspended = true; + keyboard.predictions(None, true); + assert!(keyboard.suspended()); + assert_eq!( + keyboard.prediction_note().as_deref(), + Some("Prediction unavailable") + ); + assert_eq!( + keyboard.rows.last().unwrap().last(), + Some(&Key::RetryPrediction) + ); + assert_eq!( + keyboard.choose(Key::RetryPrediction), + Some(Output::RetryPrediction) + ); + assert_eq!( + keyboard.prediction_note().as_deref(), + Some("Prediction loading") + ); + assert!(!keyboard + .rows + .last() + .unwrap() + .contains(&Key::RetryPrediction)); + assert_eq!(keyboard.page, Page::Letters); + } + fn context(epoch: u64) -> Option { Some(TypingContext { foreground: 7, diff --git a/src-tauri/src/scan_mouse.rs b/src-tauri/src/scan_mouse.rs index e7cbc4e0..051ea4e9 100644 --- a/src-tauri/src/scan_mouse.rs +++ b/src-tauri/src/scan_mouse.rs @@ -250,6 +250,7 @@ impl MousePanel { }) .collect(), status, + note: None, dock: self.dock, selected: (scanning && !escaping).then_some((active_row, active_column)), status_selected: scanning && escaping, diff --git a/src-tauri/src/scan_panel.rs b/src-tauri/src/scan_panel.rs index 4c9b55dd..80ec26bd 100644 --- a/src-tauri/src/scan_panel.rs +++ b/src-tauri/src/scan_panel.rs @@ -136,6 +136,8 @@ pub fn position_label(dock: Dock, current: Dock) -> String { pub struct Panel { pub rows: Vec>, pub status: String, + /// Passive feedback beside the scanning prompt; never a scan target. + pub note: Option, pub dock: Dock, /// Scanner position to highlight, or `None` when no key may be chosen. pub selected: Option<(usize, Option)>, @@ -211,6 +213,28 @@ impl Panel { c += 1; } } + let note_width = self.note.as_ref().map_or(0.0, |_| width * 0.28); + if let Some(note) = &self.note { + frame.tiles.push(FrameTile { + thickness: self.thickness, + color, + text: note.clone(), + icon: Item::KeyboardKey, + style: Some(TileStyle { + role: TileRole::Status, + active: false, + row_scan: false, + }), + rect: Rect { + x: x + width - note_width, + y, + width: note_width, + height: (header_height - gap).max(1.0), + }, + scale: scale * 0.8, + selected: false, + }); + } frame.tiles.push(FrameTile { thickness: self.thickness, color, @@ -224,7 +248,7 @@ impl Panel { rect: Rect { x, y, - width, + width: width - note_width - if note_width > 0.0 { gap } else { 0.0 }, height: (header_height - gap).max(1.0), }, scale, @@ -263,6 +287,7 @@ mod tests { let panel = Panel { rows: vec![vec![key("a", 1.0), key("b", 2.0)], vec![key("c", 1.0)]], status: "Page · Select a row".into(), + note: None, dock: Dock::default(), selected: Some((1, None)), status_selected: false, @@ -307,6 +332,7 @@ mod tests { let panel = Panel { rows: vec![vec![key("a", 1.0)]], status: BACK_TO_ROWS.into(), + note: None, dock: Dock { column: 1, row: 0 }, selected: None, status_selected: true, @@ -399,6 +425,7 @@ mod tests { let mut frame = Panel { rows: vec![vec![key("a", 1.0)]], status: String::new(), + note: None, dock: Dock::default(), selected: None, status_selected: false, @@ -428,6 +455,7 @@ mod tests { let panel = Panel { rows: vec![vec![key("a", 1.0)]], status: String::new(), + note: None, dock, selected: None, status_selected: false, diff --git a/src-tauri/tauri.conf.json b/src-tauri/tauri.conf.json index 309011d4..9fff767f 100644 --- a/src-tauri/tauri.conf.json +++ b/src-tauri/tauri.conf.json @@ -36,7 +36,6 @@ }, "bundle": { "resources": [ - "resources/word-predictions.lookup", "resources/prediction-blocklist.txt", "resources/prediction-model/model.onnx", "resources/prediction-model/tokenizer.json", diff --git a/src/PointScan.test.tsx b/src/PointScan.test.tsx index 06c9ccc7..57daaac6 100644 --- a/src/PointScan.test.tsx +++ b/src/PointScan.test.tsx @@ -277,14 +277,9 @@ 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.queryByRole("checkbox", { name: "Enhanced word prediction" })).not.toBeInTheDocument(); 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(); diff --git a/src/settings/ScannerPreferences.tsx b/src/settings/ScannerPreferences.tsx index 45cb12af..d1f657b3 100644 --- a/src/settings/ScannerPreferences.tsx +++ b/src/settings/ScannerPreferences.tsx @@ -136,8 +136,7 @@ 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.

+

Uses an on-device language model. The keyboard shows when suggestions are loading, unavailable, or have no match. If prediction fails, choose Retry predictions on the keyboard; typing still works. 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)} />)} diff --git a/tools/prediction-data/Cargo.lock b/tools/prediction-data/Cargo.lock deleted file mode 100644 index 8ab0207a..00000000 --- a/tools/prediction-data/Cargo.lock +++ /dev/null @@ -1,361 +0,0 @@ -# This file is automatically @generated by Cargo. -# It is not intended for manual editing. -version = 4 - -[[package]] -name = "bitflags" -version = "2.13.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3ded4057c258ba199e2d26386d3af3780957ecaee6c4ef4041c6b4b8b97c0b06" - -[[package]] -name = "block-buffer" -version = "0.10.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" -dependencies = [ - "generic-array", -] - -[[package]] -name = "bumpalo" -version = "3.20.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" - -[[package]] -name = "cc" -version = "1.4.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "54413ede23c2daf518f35156dfde027feb2374004d63bd497f983c8db9c0e313" -dependencies = [ - "find-msvc-tools", - "shlex", -] - -[[package]] -name = "cfg-if" -version = "1.0.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4e7648175b45a9a48536d676f68d918270699102aa8dab5496df06904c914600" - -[[package]] -name = "cpufeatures" -version = "0.2.17" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" -dependencies = [ - "libc", -] - -[[package]] -name = "crypto-common" -version = "0.1.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" -dependencies = [ - "generic-array", - "typenum", -] - -[[package]] -name = "digest" -version = "0.10.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" -dependencies = [ - "block-buffer", - "crypto-common", -] - -[[package]] -name = "fallible-iterator" -version = "0.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2acce4a10f12dc2fb14a218589d4f1f62ef011b2d0cc4b3cb1bba8e94da14649" - -[[package]] -name = "fallible-streaming-iterator" -version = "0.1.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7360491ce676a36bf9bb3c56c1aa791658183a54d2744120f27285738d90465a" - -[[package]] -name = "find-msvc-tools" -version = "0.1.13" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ef25905e51abafe4dcea6c15fec58c57b601cdbd0ee53d22ea1d3016c587d39b" - -[[package]] -name = "foldhash" -version = "0.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb" - -[[package]] -name = "generic-array" -version = "0.14.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" -dependencies = [ - "typenum", - "version_check", -] - -[[package]] -name = "hashbrown" -version = "0.16.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100" -dependencies = [ - "foldhash", -] - -[[package]] -name = "hashbrown" -version = "0.17.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" -dependencies = [ - "foldhash", -] - -[[package]] -name = "hashlink" -version = "0.12.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a596f1b20ed2cc5ecac41a164aaebc7258057060f06c0cf7a2ba3991ee7990fb" -dependencies = [ - "hashbrown 0.17.1", -] - -[[package]] -name = "js-sys" -version = "0.3.105" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ce57d20d1ea864ce2ac172ab472d409214f4fd359f0b2a2775abdf522e2af99e" -dependencies = [ - "cfg-if", - "wasm-bindgen", -] - -[[package]] -name = "libc" -version = "0.2.189" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2" - -[[package]] -name = "libsqlite3-sys" -version = "0.38.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f1d20bef17f513b9b3004532233187769cd072d790971f4e4da0e346eb6401e8" -dependencies = [ - "cc", - "pkg-config", - "vcpkg", -] - -[[package]] -name = "once_cell" -version = "1.21.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" - -[[package]] -name = "pkg-config" -version = "0.3.34" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f6b464fbc74e149a392436b17d523f769e057cb6877f6a5c4618bc6f11800548" - -[[package]] -name = "proc-macro2" -version = "1.0.107" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" -dependencies = [ - "unicode-ident", -] - -[[package]] -name = "quote" -version = "1.0.47" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" -dependencies = [ - "proc-macro2", -] - -[[package]] -name = "rsqlite-vfs" -version = "0.1.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c51c9ae4df8a7fba42103df5c621fa3c37eccf3a3c650879e90fc48b11cc192c" -dependencies = [ - "hashbrown 0.16.1", - "thiserror", -] - -[[package]] -name = "rusqlite" -version = "0.40.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "23f2a97da3e3873c73cb2a2e71b35c40ff95e0b1eefa8d72d8499a6928c3b5b3" -dependencies = [ - "bitflags", - "fallible-iterator", - "fallible-streaming-iterator", - "hashlink", - "libsqlite3-sys", - "smallvec", - "sqlite-wasm-rs", -] - -[[package]] -name = "rustversion" -version = "1.0.23" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" - -[[package]] -name = "sha2" -version = "0.10.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" -dependencies = [ - "cfg-if", - "cpufeatures", - "digest", -] - -[[package]] -name = "shlex" -version = "2.0.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" - -[[package]] -name = "smallvec" -version = "1.16.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ba467056f1b547ed52077911161fc86985becbc60e8e1857c8a144dab0def891" - -[[package]] -name = "sqlite-wasm-rs" -version = "0.5.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dc3efc0da82635d7e1ced0053bbbfa8c7ab9645d0bf36ceb4f7127bb85315d75" -dependencies = [ - "cc", - "js-sys", - "rsqlite-vfs", - "wasm-bindgen", -] - -[[package]] -name = "switchify-prediction-data" -version = "0.1.0" -dependencies = [ - "rusqlite", - "sha2", -] - -[[package]] -name = "syn" -version = "3.0.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8593e8e72159ed2257d083c7a454a85cbf854f37a0966d8d483aff8c8a3ebcee" -dependencies = [ - "proc-macro2", - "quote", - "unicode-ident", -] - -[[package]] -name = "thiserror" -version = "2.0.20" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ec86235f5fcc2a73650310756d2ac5b138a5780bbbdfae3eeccec992c435ba4f" -dependencies = [ - "thiserror-impl", -] - -[[package]] -name = "thiserror-impl" -version = "2.0.20" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bc04cd3e1236dd4a98afca4569f2deb3f120e5422a4023be2cb683f8486292af" -dependencies = [ - "proc-macro2", - "quote", - "syn", -] - -[[package]] -name = "typenum" -version = "1.20.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" - -[[package]] -name = "unicode-ident" -version = "1.0.26" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d245f478577f809a851594d02313b640fb437e0bb33866753cff937863096954" - -[[package]] -name = "vcpkg" -version = "0.2.15" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "accd4ea62f7bb7a82fe23066fb0957d48ef677f6eeb8215f372f52e48bb32426" - -[[package]] -name = "version_check" -version = "0.9.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" - -[[package]] -name = "wasm-bindgen" -version = "0.2.128" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "aecb87a33d3b0c5e3b7aa46336eaf486cffafbd281b195e4c8b80d50df2351bf" -dependencies = [ - "cfg-if", - "once_cell", - "rustversion", - "wasm-bindgen-macro", - "wasm-bindgen-shared", -] - -[[package]] -name = "wasm-bindgen-macro" -version = "0.2.128" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a690d511e3c1a8b3a55e33511e3c2c00c78415cd23650f32b808627f5696b9ed" -dependencies = [ - "quote", - "wasm-bindgen-macro-support", -] - -[[package]] -name = "wasm-bindgen-macro-support" -version = "0.2.128" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "411e4887f0071ef2d2164a9d5fdf2d20efbef78fccd3a78b0c10a1dc5295e48a" -dependencies = [ - "bumpalo", - "proc-macro2", - "quote", - "syn", - "wasm-bindgen-shared", -] - -[[package]] -name = "wasm-bindgen-shared" -version = "0.2.128" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "81941cd78d0c92026c33e5e01312845a4cb1e9af3407f9134b100dd03144103e" -dependencies = [ - "unicode-ident", -] diff --git a/tools/prediction-data/Cargo.toml b/tools/prediction-data/Cargo.toml deleted file mode 100644 index 21aa3f39..00000000 --- a/tools/prediction-data/Cargo.toml +++ /dev/null @@ -1,9 +0,0 @@ -[package] -name = "switchify-prediction-data" -version = "0.1.0" -edition = "2021" -rust-version = "1.97.1" - -[dependencies] -rusqlite = { version = "0.40.2", features = ["bundled"] } -sha2 = "0.10" diff --git a/tools/prediction-data/src/main.rs b/tools/prediction-data/src/main.rs deleted file mode 100644 index 2a6a68d6..00000000 --- a/tools/prediction-data/src/main.rs +++ /dev/null @@ -1,135 +0,0 @@ -#[allow(dead_code)] -#[path = "../../../src-tauri/src/prediction/lookup.rs"] -mod lookup; -use rusqlite::{Connection, OpenFlags}; -use sha2::{Digest, Sha256}; -use std::{ - collections::HashMap, - fs::{self, File}, - io::{BufWriter, Read, Seek, SeekFrom, Write}, - path::Path, -}; -type Error = Box; -fn number(out: &mut impl Write, n: u32) -> Result<(), Error> { - out.write_all(&n.to_le_bytes())?; - Ok(()) -} -fn convert(source: &Path, output: &Path) -> Result<(), Error> { - let c = Connection::open_with_flags(source, OpenFlags::SQLITE_OPEN_READ_ONLY)?; - let mut q = - c.prepare("SELECT ID,WORD,BASE_FREQUENCY FROM WORDS ORDER BY BASE_FREQUENCY DESC,ID ASC")?; - let words = q - .query_map([], |r| { - Ok(( - r.get::<_, u32>(0)?, - r.get::<_, String>(1)?, - r.get::<_, u32>(2)?, - )) - })? - .collect::, _>>()?; - let ranks: HashMap<_, _> = words - .iter() - .enumerate() - .map(|(rank, w)| (w.0, rank as u32)) - .collect(); - let temporary = output.with_extension("pending"); - let mut out = BufWriter::new(File::create(&temporary)?); - out.write_all(lookup::MAGIC)?; - out.write_all(&[0; 32])?; - number(&mut out, u32::try_from(words.len())?)?; - for (id, text, frequency) in &words { - number(&mut out, *id)?; - number(&mut out, *frequency)?; - number(&mut out, u32::try_from(text.len())?)?; - out.write_all(text.as_bytes())?; - } - for (n, table) in [(2, "BIGRAMS"), (3, "TRIGRAMS"), (4, "QUADGRAMS")] { - let count: u32 = c.query_row(&format!("SELECT count(*) FROM {table}"), [], |r| r.get(0))?; - number(&mut out, count)?; - let fields = (1..=n) - .map(|i| format!("ID{i}")) - .collect::>() - .join(","); - let context = (1..n) - .map(|i| format!("ID{i}")) - .collect::>() - .join(","); - let mut q=c.prepare(&format!("SELECT {fields},BASE_FREQUENCY FROM {table} ORDER BY BASE_FREQUENCY DESC,ID{n} ASC,{context}"))?; - let mut rows = q.query([])?; - while let Some(row) = rows.next()? { - for i in 0..n { - let id: u32 = row.get(i)?; - number(&mut out, *ranks.get(&id).ok_or("Missing word reference")?)?; - } - number(&mut out, row.get(n)?)?; - } - } - let folded: Vec<_> = words.iter().map(|w| w.1.to_lowercase()).collect(); - let mut prefixes: Vec<_> = (0..words.len()).collect(); - prefixes.sort_unstable_by(|&a, &b| folded[a].cmp(&folded[b]).then(a.cmp(&b))); - number(&mut out, u32::try_from(prefixes.len())?)?; - for rank in prefixes { - number(&mut out, u32::try_from(rank)?)?; - } - out.flush()?; - drop(out); - let mut input = File::open(&temporary)?; - input.seek(SeekFrom::Start(40))?; - let mut hash = Sha256::new(); - let mut buffer = [0; 65536]; - loop { - let n = input.read(&mut buffer)?; - if n == 0 { - break; - } - hash.update(&buffer[..n]); - } - drop(input); - let mut file = fs::OpenOptions::new().write(true).open(&temporary)?; - file.seek(SeekFrom::Start(8))?; - file.write_all(&hash.finalize())?; - file.sync_all()?; - drop(file); - lookup::Lookup::open(&temporary).map_err(|_| "Generated lookup failed validation")?; - fs::rename(&temporary, output)?; - println!("Generated {} bytes", fs::metadata(output)?.len()); - Ok(()) -} -fn checksum(path: &Path) -> Result, Error> { - let mut input = File::open(path)?; - let mut hash = Sha256::new(); - let mut buffer = [0; 65536]; - loop { - let n = input.read(&mut buffer)?; - if n == 0 { - break; - } - hash.update(&buffer[..n]); - } - Ok(hash.finalize().to_vec()) -} -fn main() -> Result<(), Error> { - let args: Vec<_> = std::env::args_os().collect(); - if args.len() == 4 && args[1] == "--check" { - let generated = std::env::temp_dir().join(format!( - "switchify-predictions-{}.lookup", - std::process::id() - )); - let result = (|| { - convert(Path::new(&args[2]), &generated)?; - if checksum(&generated)? != checksum(Path::new(&args[3]))? { - return Err("Generated lookup differs from packaged artifact".into()); - } - Ok(()) - })(); - let _ = fs::remove_file(&generated); - let _ = fs::remove_file(generated.with_extension("pending")); - return result; - } - if args.len() != 3 { - return Err( - "Usage: switchify-prediction-data [--check] ".into(), - ); - } - convert(Path::new(&args[1]), Path::new(&args[2])) -}