AIKeyboard/README.md

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

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

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

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

./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.