Add optional on-device language-model word prediction

.github/workflows/ci.yml

@@ -113,6 +113,8 @@ jobs:
           toolchain: 1.97.1
           components: rustfmt, clippy
       - run: npm ci
+      - name: Fetch prediction model
+        run: npm run prediction-model
       - run: cargo fmt --manifest-path src-tauri/Cargo.toml --check
       - run: cargo clippy --locked --manifest-path src-tauri/Cargo.toml --all-targets -- -D warnings
       - run: cargo test --locked --manifest-path src-tauri/Cargo.toml

.github/workflows/release-tauri.yml

@@ -164,6 +164,9 @@ jobs:
       - name: Install dependencies
         run: npm ci
 
+      - name: Fetch prediction model
+        run: npm run prediction-model
+
       - name: Validate Rust at release tag
         run: |
           cargo fmt --manifest-path src-tauri/Cargo.toml --check
@@ -340,6 +343,9 @@ jobs:
       - name: Install dependencies
         run: npm ci
 
+      - name: Fetch prediction model
+        run: npm run prediction-model
+
       - name: Validate Rust at release tag
         shell: powershell
         run: |

.gitignore

@@ -16,3 +16,4 @@ src-tauri/icons/android/
 src-tauri/icons/ios/
 src-tauri/tauri.release*.generated.json
 tools/prediction-data/target/
+src-tauri/resources/prediction-model/

AGENTS.md

@@ -24,6 +24,7 @@ The frozen C# application is not maintained in this repository. Do not reintrodu
 Use Node.js 24 and Rust 1.97.1. Run checks from the repository root:
 
 ```bash
+npm run prediction-model
 npm run lint
 npm test
 npm run build

docs/word-prediction.md

@@ -10,6 +10,27 @@ A passive observer records only an activity counter and timestamp, never externa
 
 The bundled English lookup database works offline. Ranking backs off from the longest available three-word context to shorter contexts, then to word frequency. Prefix matching is case-insensitive. Duplicate words are removed. The data is fixed; typing is never used to train or update it.
 
+## Enhanced word prediction
+
+Enhanced word prediction is off by default. When it is on, a small language model (SmolLM2-135M, int8 ONNX, run with ONNX Runtime on the CPU) spells suggestions from subword pieces, so it can suggest words the lookup lacks, such as WhatsApp. It sees only the buffer described above: at most its last 256 characters, starting at a word boundary, including earlier sentences still in the buffer. It runs offline in the prediction worker, never learns, and nothing it reads is logged or stored.
+
+- **Loading and fallback:** the model loads in the background when the worker starts, taking about 0.7 s. Until it is ready, and whenever it fails, the lookup answers as before. Turning the setting on or off restarts the worker.
+- **Filters:**
+  - words contain only ASCII letters and apostrophes;
+  - a word at the start of the text must be capitalised, which rules out subword fragments;
+  - single letters must be in the lookup;
+  - blocklisted words are never suggested;
+  - words the lookup does not know must reach a minimum probability.
+- **Spelling:** words the lookup knows take its usual spelling. Unknown words keep the model's spelling, except that a capital coming only from sentence position is removed.
+- **Merging:** lookup suggestions fill any remaining positions.
+- **Time limits:** the model stops exploring new words after 400 ms. If one call ever takes longer than 1 s, the worker uses the lookup for the rest of its life, so replies stay well inside the two-second deadline.
+- **Clipped buffers:** when clipping leaves no complete earlier word, the lookup answers.
+- **Cost:** predictions take roughly 40–80 ms on a recent laptop, and the worker needs about 300–450 MB more memory.
+
+The model files are too large to commit. `npm run prediction-model` downloads them from a pinned upstream revision and verifies their SHA-256. The Tauri dev and build commands and CI run it automatically. Run it once before `cargo test` or `cargo clippy` in a fresh checkout.
+
+## Worker process
+
 A separate process owns the buffer, activity observer and database. Private bounded inherited pipes carry successful edits and results to native rendering. Text and suggestions are not sent to the React UI, diagnostic history or telemetry. One request is outstanding at a time with a two-second deadline. Timeout stops predictions until the keyboard is reopened; ordinary keyboard operation remains available. Closing the keyboard, ending scanning, or exiting kills and reaps the worker.
 
 Before accepting a suggestion, the worker checks its token, edit revision, foreground identity and external activity. The main process checks its generation and foreground again before injection. Verification and native insertion cannot be atomic across applications; the target can still change in that short interval.
@@ -22,6 +43,12 @@ Automated tests use fixture databases, fake foreground/activity sources, fake in
 cargo test --manifest-path src-tauri/Cargo.toml bundled_database_benchmark -- --ignored --nocapture
 ```
 
+`cargo test` also runs the language model against the fetched files with synthetic text. For release-build suggestions and timings:
+
+```powershell
+cargo test --release --manifest-path src-tauri/Cargo.toml bundled_model_suggests_current_words -- --nocapture
+```
+
 For native validation, use disposable synthetic text in Notepad and a browser on Windows, and TextEdit and a browser on macOS. Launch macOS with `npm run macos:run` to retain its signed Accessibility identity.
 
 1. Open the keyboard from the action menu and an assigned Open keyboard switch, including from idle, a paused scan and an active drag. Verify no click or focus change occurs and owned drag input is released.

package.json

@@ -9,6 +9,7 @@
   "scripts": {
     "dev": "vite",
     "build": "tsc --noEmit && vite build",
+    "prediction-model": "node scripts/fetch-prediction-model.mjs",
     "test": "vitest run && node --test scripts/create-update-feed.node-test.mjs",
     "lint": "tsc --noEmit",
     "macos:setup-signing": "./scripts/setup-macos-dev-signing.sh",

scripts/fetch-prediction-model.mjs

@@ -0,0 +1,58 @@
+// Downloads the enhanced word prediction model into src-tauri/resources/prediction-model.
+// The model is too large to commit, so every file is pinned to one upstream
+// revision and verified by SHA-256. Verified files are not downloaded again.
+import { createHash } from "node:crypto";
+import { createReadStream, createWriteStream } from "node:fs";
+import { mkdir, rename, rm, stat } from "node:fs/promises";
+import { dirname, join } from "node:path";
+import { Readable } from "node:stream";
+import { pipeline } from "node:stream/promises";
+import { fileURLToPath } from "node:url";
+
+const revision = "d0ae6834f1df45e0e95b5fdae95e536f9ca7cd3f";
+const base = `https://huggingface.co/onnx-community/SmolLM2-135M-ONNX/resolve/${revision}`;
+const files = [
+  { name: "model.onnx", source: "onnx/model_int8.onnx", size: 135658354, sha256: "50ba80511ce74634d232a043b6c37775cca756b826b49d0a4a8eff958c4bbcc9" },
+  { name: "tokenizer.json", source: "tokenizer.json", size: 2053526, sha256: "139d2f4b4919b90953bdd3c0c40c94c9b23074799a766508dc3bf5eb8ab73351" },
+  { name: "config.json", source: "config.json", size: 1035, sha256: "2c5f23fddabecdf9c47d0048f555899822ec87ec4a28169393840ac7e74192c4" },
+];
+const target = join(dirname(fileURLToPath(import.meta.url)), "..", "src-tauri", "resources", "prediction-model");
+
+async function digest(path) {
+  const hash = createHash("sha256");
+  await pipeline(createReadStream(path), hash);
+  return hash.digest("hex");
+}
+
+async function verified(path, file) {
+  try {
+    return (await stat(path)).size === file.size && (await digest(path)) === file.sha256;
+  } catch {
+    return false;
+  }
+}
+
+async function download(file) {
+  const path = join(target, file.name);
+  if (await verified(path, file)) return false;
+  const pending = `${path}.pending`;
+  for (let attempt = 1; ; attempt++) {
+    try {
+      const response = await fetch(`${base}/${file.source}`);
+      if (!response.ok) throw new Error(`HTTP ${response.status}`);
+      await pipeline(Readable.fromWeb(response.body), createWriteStream(pending));
+      if (!(await verified(pending, file))) throw new Error("size or SHA-256 mismatch");
+      await rename(pending, path);
+      return true;
+    } catch (error) {
+      await rm(pending, { force: true });
+      if (attempt === 3) throw new Error(`Could not fetch ${file.name}: ${error.message}`);
+      await new Promise((resolve) => setTimeout(resolve, attempt * 2000));
+    }
+  }
+}
+
+await mkdir(target, { recursive: true });
+for (const file of files) {
+  if (await download(file)) console.log(`Fetched ${file.name}`);
+}