642773765d
- 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 <noreply@anthropic.com>
176 lines
5.5 KiB
Markdown
176 lines
5.5 KiB
Markdown
# AIKeyboard — HeliBoard + Gemma 4 On-Device AI Correction
|
|
|
|
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 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
|
|
↓
|
|
AiSuggestionManager.postSuggestion("Guten Tag.", "guten rag.")
|
|
↓
|
|
Suggestion strip shows "Guten Tag." in italic/accent color
|
|
↓
|
|
Tap → original text replaced with correction
|
|
Long-press → suggestion dismissed
|
|
```
|
|
|
|
### Components
|
|
|
|
| File | Role |
|
|
| ---- | ---- |
|
|
| `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, pick intercept, memory trim |
|
|
| `patches/0003-SuggestionStrip-ai-style.patch` | Italic + accent color for AI suggestions |
|
|
|
|
---
|
|
|
|
## Requirements
|
|
|
|
- Android device with **ARM64** (arm64-v8a), Android 11+ (API 31)
|
|
- Android Studio **Ladybug** or newer / Gradle 8+
|
|
- NDK **28.0.13004108**
|
|
- **Gemma 4 E2B-it** model in `.litertlm` format (~2.6 GB)
|
|
Download: <https://huggingface.co/litert-community/gemma-4-E2B-it-litert-lm>
|
|
|
|
---
|
|
|
|
## Setup
|
|
|
|
### 1. Clone and prepare HeliBoard sources
|
|
|
|
```bash
|
|
git clone http://172.17.2.68:3001/nova/AIKeyboard.git
|
|
cd AIKeyboard
|
|
bash setup_heliboard.sh
|
|
```
|
|
|
|
`setup_heliboard.sh` clones HeliBoard from GitHub and applies the three AI
|
|
integration patches automatically.
|
|
|
|
### 2. Push the model to the device
|
|
|
|
```bash
|
|
adb push gemma-4-E2B-it.litertlm /data/local/tmp/gemma-4-E2B-it.litertlm
|
|
adb shell chmod 644 /data/local/tmp/gemma-4-E2B-it.litertlm
|
|
```
|
|
|
|
### 3. Build and install
|
|
|
|
```bash
|
|
./gradlew installDebug
|
|
```
|
|
|
|
Enable **AIKeyboard** as your input method in Android Settings → System →
|
|
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()`.
|
|
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()` 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 `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).
|
|
|
|
---
|
|
|
|
## Model placement
|
|
|
|
The model must be at `/data/local/tmp/gemma-4-E2B-it.litertlm` on the device.
|
|
This path is configured in `AiCorrectionEngine.MODEL_PATH`.
|
|
|
|
The model file is **not** tracked in this repository (2.6 GB, `.litertlm` is
|
|
in `.gitignore`).
|
|
|
|
---
|
|
|
|
## License
|
|
|
|
The AI integration layer (`app/src/main/java/helium314/keyboard/latin/ai/`)
|
|
is original work added to this project.
|
|
|
|
HeliBoard itself is licensed under **GPL-3.0-only**. See
|
|
`heliboard/LICENSE` after running `setup_heliboard.sh`.
|