From 642773765db46558e9d9b1fd276f6d9aa0fd58d9 Mon Sep 17 00:00:00 2001 From: nova Date: Sat, 11 Apr 2026 19:19:22 +0200 Subject: [PATCH] Improve AI correction: loading indicator, field safety, memory, dismiss MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Show "…" placeholder immediately after debounce so users know the AI is working (KIND_AI_LOADING_FLAG = 0x20000) - Skip AI correction in password, URL, and e-mail fields (mIsGeneralTextInput + mIsPasswordField guards in InputLogic) - Protect suggestion strip from being cleared while AI suggestion is visible (mAiSuggestionVisible flag + clearAiSuggestionAndSetNeutral) - Release Gemma 4 engine on TRIM_MEMORY_RUNNING_CRITICAL - Long-press on AI suggestion dismisses it (dismissAiSuggestion) - Warm-up model at IME startup for instant first correction Co-Authored-By: Claude Sonnet 4.6 --- README.md | 102 ++++++++++++------ .../keyboard/latin/ai/AiCorrectionEngine.kt | 17 +++ .../keyboard/latin/ai/AiSuggestionManager.kt | 20 +++- .../keyboard/latin/ai/AiTriggerHook.kt | 1 + 4 files changed, 106 insertions(+), 34 deletions(-) diff --git a/README.md b/README.md index cc95974..2907e2b 100644 --- a/README.md +++ b/README.md @@ -3,22 +3,25 @@ HeliBoard fork that adds on-device spell and grammar correction powered by **Gemma 4 E2B-it** running locally via Google's **LiteRT-LM** SDK. -After you finish typing a sentence (`.`, `!`, `?`), the model silently checks -it and — if it finds an error — shows the corrected sentence in the suggestion -strip. Tap it once to replace the original text. Everything runs on-device, -no network, no cloud API. +After you finish typing a sentence (`.`, `!`, `?`), the model checks it and — +if it finds an error — shows `…` immediately, then the corrected sentence in +the suggestion strip. Tap it to replace the original text, or long-press to +dismiss. Everything runs on-device, no network, no cloud API. --- ## How it works -``` +```text User types "guten rag." ↓ InputLogic detects sentence-ending punctuation +(skipped in password / URL / e-mail fields) ↓ AiTriggerHook (500 ms debounce) extracts last sentence ↓ +"…" placeholder appears in suggestion strip immediately + ↓ AiCorrectionEngine.correctSentence() → LiteRT-LM Engine (GPU first, CPU fallback) → Gemma 4 E2B-it .litertlm model @@ -27,18 +30,19 @@ AiSuggestionManager.postSuggestion("Guten Tag.", "guten rag.") ↓ Suggestion strip shows "Guten Tag." in italic/accent color ↓ -User taps → original text replaced with correction +Tap → original text replaced with correction +Long-press → suggestion dismissed ``` ### Components | File | Role | -|------|------| -| `app/…/ai/AiCorrectionEngine.kt` | LiteRT-LM wrapper; GPU→CPU fallback | -| `app/…/ai/AiTriggerHook.kt` | Sentence detection + debounce | -| `app/…/ai/AiSuggestionManager.kt` | Posts corrected text to suggestion strip | +| ---- | ---- | +| `app/…/ai/AiCorrectionEngine.kt` | LiteRT-LM wrapper; GPU→CPU fallback; warm-up on start | +| `app/…/ai/AiTriggerHook.kt` | Sentence detection, debounce, loading placeholder | +| `app/…/ai/AiSuggestionManager.kt` | Posts corrections and loading state to suggestion strip | | `patches/0001-InputLogic-ai-hook.patch` | Hook in `InputLogic.java` after `commitCodePoint` | -| `patches/0002-LatinIME-ai-lifecycle.patch` | AI object lifecycle in `LatinIME.java` | +| `patches/0002-LatinIME-ai-lifecycle.patch` | AI object lifecycle, pick intercept, memory trim | | `patches/0003-SuggestionStrip-ai-style.patch` | Italic + accent color for AI suggestions | --- @@ -84,26 +88,71 @@ Languages & input → On-screen keyboard. --- +## Features + +### AI correction flow + +- Triggers after `.` `!` `?` `…` at the end of a sentence +- 500 ms debounce prevents firing on rapid punctuation (e.g. `...`) +- `…` placeholder appears immediately so the user knows the AI is working +- Correction shown in **italic** with accent color; normal suggestions are unaffected +- Tap to accept, long-press to dismiss + +### Field-type safety + +AI correction is automatically skipped in: + +- Password fields +- URL / URI fields +- E-mail address fields +- Web password / phonetic input fields + +### Model warm-up + +The Gemma 4 engine is loaded in the background as soon as the keyboard +service starts (`onCreate`), so the first correction after a sentence +appears without the initial loading delay. + +### Memory management + +Under critical memory pressure (`TRIM_MEMORY_RUNNING_CRITICAL`), the engine +is released automatically. It reloads on the next sentence. + +### Backend selection + +`AiCorrectionEngine` tries **GPU** (WebGPU/Vulkan via LiteRT-LM) first. +If GPU initialization fails, it falls back to **CPU** immediately. A runtime +GPU inference error triggers a permanent CPU switch for the session. + +Tested on: + +- Google Pixel 7 (Tensor G2) — CPU backend (OpenCL not available) +- Devices with Mali-G710 — GPU via WebGPU/Vulkan + +--- + ## Patch details The three patches modify HeliBoard source files in `heliboard/`: **0001 — InputLogic hook** After every committed code point, checks `AiTriggerHook.isSentenceEnder()`. -On a sentence-ender, reads up to 500 chars before the cursor and calls -`onSentenceEndDetected()`. +Only fires for general text input (not password/URL/email fields). Reads up +to 500 chars before the cursor and calls `onSentenceEndDetected()`. **0002 — LatinIME lifecycle** Instantiates `AiCorrectionEngine`, `AiSuggestionManager`, and `AiTriggerHook` -in `onCreate()`, cleans them up in `onDestroy()`. Overrides -`showAiSuggestion()` to bypass the normal suggestions-enabled gate and -intercepts `pickSuggestionManually()` to do a clean sentence replacement -(finish composing → delete original → commit corrected text). +in `onCreate()` and triggers warm-up. Cleans up in `onDestroy()` and +`onTrimMemory()`. Overrides `showAiSuggestion()` with a `mAiSuggestionVisible` +guard that prevents `setNeutralSuggestionStrip()` from clearing an active AI +suggestion. Intercepts `pickSuggestionManually()` to handle loading placeholders +(ignore tap) and accepted corrections (delete original → commit corrected text). +Implements `dismissAiSuggestion()` for long-press dismiss. **0003 — SuggestionStrip styling** -Detects the `KIND_AI_FLAG` bit (`0x10000`) on a `SuggestedWordInfo` and -applies italic style + auto-correct accent color to visually distinguish AI -suggestions from normal word predictions. +Detects `KIND_AI_FLAG` (`0x10000`) for accepted corrections (italic + accent +color) and `KIND_AI_LOADING_FLAG` (`0x20000`) for the loading placeholder +(normal color, no italic, not tappable). --- @@ -117,19 +166,6 @@ in `.gitignore`). --- -## Backend selection - -`AiCorrectionEngine` tries **GPU** (WebGPU/Vulkan via LiteRT-LM) first. -If GPU initialization fails, it falls back to **CPU** immediately. If a GPU -inference error occurs at runtime, it permanently switches to CPU for the -session and recreates the engine. - -Tested on: -- Google Pixel 7 (Tensor G2) — CPU backend (OpenCL not available) -- Devices with Mali-G710 — GPU via WebGPU/Vulkan - ---- - ## License The AI integration layer (`app/src/main/java/helium314/keyboard/latin/ai/`) diff --git a/app/src/main/java/helium314/keyboard/latin/ai/AiCorrectionEngine.kt b/app/src/main/java/helium314/keyboard/latin/ai/AiCorrectionEngine.kt index 9407994..19e728f 100644 --- a/app/src/main/java/helium314/keyboard/latin/ai/AiCorrectionEngine.kt +++ b/app/src/main/java/helium314/keyboard/latin/ai/AiCorrectionEngine.kt @@ -9,7 +9,9 @@ import com.google.ai.edge.litertlm.ConversationConfig import com.google.ai.edge.litertlm.Engine import com.google.ai.edge.litertlm.EngineConfig import com.google.ai.edge.litertlm.SamplerConfig +import kotlinx.coroutines.CoroutineScope import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.launch import kotlinx.coroutines.withContext import java.io.File @@ -114,6 +116,21 @@ class AiCorrectionEngine(private val context: Context) { } } + /** + * Pre-loads the model in the background so the first correction request + * doesn't have to wait. Call once after IME startup. + */ + fun warmUp(scope: CoroutineScope) { + scope.launch(Dispatchers.IO) { + try { + getOrCreate() + Log.i(TAG, "Warm-up complete") + } catch (ex: Exception) { + Log.w(TAG, "Warm-up failed (model may not be present yet): ${ex.message}") + } + } + } + fun close() { engine?.close() engine = null diff --git a/app/src/main/java/helium314/keyboard/latin/ai/AiSuggestionManager.kt b/app/src/main/java/helium314/keyboard/latin/ai/AiSuggestionManager.kt index 5c14027..f7a9ed3 100644 --- a/app/src/main/java/helium314/keyboard/latin/ai/AiSuggestionManager.kt +++ b/app/src/main/java/helium314/keyboard/latin/ai/AiSuggestionManager.kt @@ -45,6 +45,10 @@ class AiSuggestionManager { const val KIND_AI_FLAG = 0x10000 const val KIND_AI_CORRECTION = SuggestedWordInfo.KIND_CORRECTION or KIND_AI_FLAG + /** Marks a transient "loading" placeholder — not tappable. */ + const val KIND_AI_LOADING_FLAG = 0x20000 + const val KIND_AI_LOADING = SuggestedWordInfo.KIND_CORRECTION or KIND_AI_LOADING_FLAG + /** * Factory method accessible from Java (LatinIME.java patch). * Returns a [CoroutineScope] tied to a SupervisorJob so individual @@ -126,6 +130,20 @@ class AiSuggestionManager { /** Returns the original (pre-correction) sentence, or null if none is pending. */ fun getOriginalSentence(): String? = originalSentence + /** Shows a non-interactive "…" placeholder while inference is running. */ + fun showLoadingPlaceholder() { + val wordInfo = SuggestedWordInfo( + "…", "", Int.MAX_VALUE, KIND_AI_LOADING, + null, SuggestedWordInfo.NOT_AN_INDEX, SuggestedWordInfo.NOT_A_CONFIDENCE + ) + val words = SuggestedWords( + arrayListOf(wordInfo), null, wordInfo, + false, false, false, + SuggestedWords.INPUT_STYLE_PREDICTION, SuggestedWords.NOT_A_SEQUENCE_NUMBER + ) + accessor?.showAiSuggestion(words) + } + /** * Clears the AI suggestion from the strip. * Called when the correction matches the original (no change needed) @@ -134,6 +152,6 @@ class AiSuggestionManager { fun clearSuggestion() { _currentSuggestion.value = null originalSentence = null - accessor?.setNeutralSuggestionStrip() + accessor?.clearAiSuggestionAndSetNeutral() } } diff --git a/app/src/main/java/helium314/keyboard/latin/ai/AiTriggerHook.kt b/app/src/main/java/helium314/keyboard/latin/ai/AiTriggerHook.kt index e4ae4f9..65127fb 100644 --- a/app/src/main/java/helium314/keyboard/latin/ai/AiTriggerHook.kt +++ b/app/src/main/java/helium314/keyboard/latin/ai/AiTriggerHook.kt @@ -63,6 +63,7 @@ class AiTriggerHook( debounceJob?.cancel() debounceJob = scope.launch { delay(DEBOUNCE_MS) + manager.showLoadingPlaceholder() Log.d(TAG, "Triggering AI correction for: \"$sentence\"") val corrected = engine.correctSentence(sentence) if (corrected != sentence) {