Compare commits
1 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 642773765d |
@@ -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/`)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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) {
|
||||
|
||||
Reference in New Issue
Block a user