Initial commit: HeliBoard + Gemma 4 on-device AI correction

Integrates LiteRT-LM (Google) with Gemma 4 E2B-it for on-device
spell/grammar correction in HeliBoard. After sentence-ending punctuation,
the model checks the sentence and shows a correction in the suggestion strip.
Tapping it replaces the original text — fully offline, no cloud API.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
nova
2026-04-10 14:57:10 +02:00
commit aa8fff2c0d
17 changed files with 1038 additions and 0 deletions
+38
View File
@@ -0,0 +1,38 @@
From: AIKeyboard AI Integration <ai@aikeyboard.local>
Date: Wed, 09 Apr 2025 00:00:00 +0000
Subject: [PATCH 1/3] Add AI spell correction trigger hook in InputLogic
After a sentence-ending punctuation character (. ! ? ...) is committed to the
InputConnection, notify AiTriggerHook with the text before the cursor so the
AI correction engine can run in the background.
Field name confirmed: mLatinIME (type LatinIME, first constructor parameter).
The hook is debounced (500ms) and async — the IME main thread is never blocked.
---
.../keyboard/latin/inputlogic/InputLogic.java | 13 +++++++++++++
1 file changed, 13 insertions(+)
diff --git a/app/src/main/java/helium314/keyboard/latin/inputlogic/InputLogic.java b/app/src/main/java/helium314/keyboard/latin/inputlogic/InputLogic.java
--- a/app/src/main/java/helium314/keyboard/latin/inputlogic/InputLogic.java
+++ b/app/src/main/java/helium314/keyboard/latin/inputlogic/InputLogic.java
@@ -1228,6 +1228,19 @@ final class InputLogic {
mConnection.commitCodePoint(codePoint);
+ // === AI SPELL CORRECTION HOOK ===
+ // Runs async on Dispatchers.IO via coroutine — does not block the IME thread.
+ // mLatinIME is confirmed non-null here (it's set in the constructor).
+ if (helium314.keyboard.latin.ai.AiTriggerHook.isSentenceEnder(codePoint)
+ && mLatinIME.getAiTriggerHook() != null) {
+ final CharSequence beforeCursor =
+ mConnection.getTextBeforeCursor(500, 0);
+ if (beforeCursor != null) {
+ mLatinIME.getAiTriggerHook()
+ .onSentenceEndDetected(beforeCursor.toString());
+ }
+ }
+ // === END AI SPELL CORRECTION HOOK ===
+
if (isInlineEmojiSearchAction()) {
inputTransaction.setRequiresUpdateSuggestions();
} else {
+99
View File
@@ -0,0 +1,99 @@
From: AIKeyboard AI Integration <ai@aikeyboard.local>
Date: Wed, 09 Apr 2025 00:00:00 +0000
Subject: [PATCH 2/3] Wire AI engine lifecycle into LatinIME
- Declare AI component fields after existing field declarations
- Initialize AiCorrectionEngine, AiSuggestionManager, AiTriggerHook in onCreate()
- Clean up (cancel coroutines, close LlmInference) in onDestroy()
- Expose getAiTriggerHook() accessor used by InputLogic patch
---
.../keyboard/latin/LatinIME.java | 54 ++++++++++++++++++
1 file changed, 54 insertions(+)
diff --git a/app/src/main/java/helium314/keyboard/latin/LatinIME.java b/app/src/main/java/helium314/keyboard/latin/LatinIME.java
index xxxxxxx..xxxxxxx 100644
--- a/app/src/main/java/helium314/keyboard/latin/LatinIME.java
+++ b/app/src/main/java/helium314/keyboard/latin/LatinIME.java
@@ -133,6 +133,15 @@ public class LatinIME extends InputMethodService implements
private RichInputMethodManager mRichImm;
final KeyboardSwitcher mKeyboardSwitcher;
+ // ── AI spell correction components ──────────────────────────────────────
+ // Lazily initialised in onCreate(); destroyed in onDestroy().
+ private helium314.keyboard.latin.ai.AiCorrectionEngine mAiCorrectionEngine;
+ private helium314.keyboard.latin.ai.AiSuggestionManager mAiSuggestionManager;
+ private helium314.keyboard.latin.ai.AiTriggerHook mAiTriggerHook;
+ // CoroutineScope type requires kotlin.coroutines on classpath (satisfied by
+ // kotlinx-coroutines-android dependency in app/build.gradle.kts).
+ private kotlinx.coroutines.CoroutineScope mAiScope;
+ // ────────────────────────────────────────────────────────────────────────
+
private final SubtypeState mSubtypeState = new SubtypeState(
(InputMethodSubtype subtype) -> { switchToSubtype(subtype); return Unit.INSTANCE; });
private final StatsUtilsManager mStatsUtilsManager;
@@ -572,6 +581,21 @@ public class LatinIME extends InputMethodService implements
registerReceiver(mRestartAfterDeviceUnlockReceiver, restartAfterUnlockFilter);
StatsUtils.onCreate(mSettings.getCurrent(), mRichImm);
+
+ // ── Initialise AI components ─────────────────────────────────────────
+ // AiSuggestionManager must be attached before AiTriggerHook is created
+ // so that the first correction result has somewhere to go.
+ mAiSuggestionManager = new helium314.keyboard.latin.ai.AiSuggestionManager();
+ mAiSuggestionManager.attachAccessor(this);
+ mAiCorrectionEngine = new helium314.keyboard.latin.ai.AiCorrectionEngine(this);
+ // createImeScope() returns a SupervisorJob-backed CoroutineScope.
+ // Using @JvmStatic companion factory avoids complex Java↔Kotlin interop.
+ mAiScope = helium314.keyboard.latin.ai.AiSuggestionManager.createImeScope();
+ mAiTriggerHook = new helium314.keyboard.latin.ai.AiTriggerHook(
+ mAiCorrectionEngine,
+ mAiSuggestionManager,
+ mAiScope
+ );
+ // ────────────────────────────────────────────────────────────────────
}
@@ -683,6 +697,22 @@ public class LatinIME extends InputMethodService implements
@Override
public void onDestroy() {
mClipboardHistoryManager.onDestroy();
+
+ // ── Tear down AI components before releasing dictionaries ────────────
+ if (mAiTriggerHook != null) {
+ mAiTriggerHook.cancel();
+ }
+ if (mAiScope != null) {
+ // Cancel the coroutine scope to stop any in-flight correction jobs.
+ kotlinx.coroutines.JobKt.cancel(
+ (kotlinx.coroutines.Job) mAiScope.getCoroutineContext()
+ .get(kotlinx.coroutines.Job.Key),
+ null
+ );
+ }
+ if (mAiCorrectionEngine != null) {
+ mAiCorrectionEngine.close();
+ }
+ if (mAiSuggestionManager != null) {
+ mAiSuggestionManager.detachAccessor();
+ }
+ // ────────────────────────────────────────────────────────────────────
+
mDictionaryFacilitator.closeDictionaries();
mSettings.onDestroy();
unregisterReceiver(mRingerModeChangeReceiver);
@@ -695,4 +725,14 @@ public class LatinIME extends InputMethodService implements
mHandler.removeCallbacksAndMessages(null);
deallocateMemory();
}
+
+ // ── AI accessor ──────────────────────────────────────────────────────────
+ /**
+ * Returns the AiTriggerHook instance for use by InputLogic.
+ * Returns null if AI initialisation has not yet completed (race-safe).
+ */
+ @androidx.annotation.Nullable
+ public helium314.keyboard.latin.ai.AiTriggerHook getAiTriggerHook() {
+ return mAiTriggerHook;
+ }
+ // ────────────────────────────────────────────────────────────────────────
}
@@ -0,0 +1,74 @@
From: AIKeyboard AI Integration <ai@aikeyboard.local>
Date: Wed, 09 Apr 2025 00:00:00 +0000
Subject: [PATCH 3/3] Add italic colour styling for AI correction suggestions
Modifies two methods in SuggestionStripLayoutHelper:
1. getSuggestionTextColor(): Returns mColorAutoCorrect for AI suggestions
(KIND_APP_DEFINED / KIND_AI_CORRECTION = 6), giving them the same accent
colour as auto-corrections.
2. getStyledSuggestedWord(): Adds an italic StyleSpan for AI suggestions so
users can visually distinguish them from regular word predictions.
Both method signatures and bodies verified against HeliBoard main (Apr 2025).
---
.../suggestions/SuggestionStripLayoutHelper.java | 31 ++++++++++++++++++
1 file changed, 31 insertions(+)
diff --git a/app/src/main/java/helium314/keyboard/latin/suggestions/SuggestionStripLayoutHelper.java b/app/src/main/java/helium314/keyboard/latin/suggestions/SuggestionStripLayoutHelper.java
--- a/app/src/main/java/helium314/keyboard/latin/suggestions/SuggestionStripLayoutHelper.java
+++ b/app/src/main/java/helium314/keyboard/latin/suggestions/SuggestionStripLayoutHelper.java
@@ -1,3 +1,5 @@
+import android.graphics.Typeface;
+import android.text.style.StyleSpan;
/*
* Copyright (C) 2013 The Android Open Source Project
*
@@ -1,0 +1,0 @@ private int getSuggestionTextColor(final SuggestedWords suggestedWords,
private int getSuggestionTextColor(final SuggestedWords suggestedWords,
final int indexInSuggestedWords) {
- final boolean isTypedWord = suggestedWords.getInfo(indexInSuggestedWords).isKindOf(SuggestedWordInfo.KIND_TYPED);
+ final SuggestedWordInfo wordInfo = suggestedWords.getInfo(indexInSuggestedWords);
+ final boolean isTypedWord = wordInfo.isKindOf(SuggestedWordInfo.KIND_TYPED);
+
+ // AI correction: use auto-correct accent colour so it stands out
+ if (wordInfo.isKindOf(SuggestedWordInfo.KIND_APP_DEFINED)) {
+ return mColorAutoCorrect;
+ }
final int color;
if (indexInSuggestedWords == SuggestedWords.INDEX_OF_AUTO_CORRECTION && suggestedWords.mWillAutoCorrect) {
@@ -1,0 +1,0 @@ private CharSequence getStyledSuggestedWord(final SuggestedWords suggestedWords,
private CharSequence getStyledSuggestedWord(final SuggestedWords suggestedWords,
final int indexInSuggestedWords) {
if (indexInSuggestedWords >= suggestedWords.size()) {
return null;
}
final String word = suggestedWords.getLabel(indexInSuggestedWords);
final boolean isAutoCorrection = suggestedWords.mWillAutoCorrect
&& indexInSuggestedWords == SuggestedWords.INDEX_OF_AUTO_CORRECTION;
final boolean isTypedWordValid = suggestedWords.mTypedWordValid
&& indexInSuggestedWords == SuggestedWords.INDEX_OF_TYPED_WORD;
- if (!isAutoCorrection && !isTypedWordValid) {
+ final boolean isAiCorrection = suggestedWords.getInfo(indexInSuggestedWords)
+ .isKindOf(SuggestedWordInfo.KIND_APP_DEFINED);
+ if (!isAutoCorrection && !isTypedWordValid && !isAiCorrection) {
return word;
}
final Spannable spannedWord = new SpannableString(word);
final int options = mSuggestionStripOptions;
if ((isAutoCorrection && (options & AUTO_CORRECT_BOLD) != 0)
|| (isTypedWordValid && (options & VALID_TYPED_WORD_BOLD) != 0)) {
addStyleSpan(spannedWord, BOLD_SPAN);
}
if (isAutoCorrection && (options & AUTO_CORRECT_UNDERLINE) != 0) {
addStyleSpan(spannedWord, UNDERLINE_SPAN);
}
+ // AI correction: italic to signal model-generated suggestion
+ if (isAiCorrection) {
+ addStyleSpan(spannedWord, new StyleSpan(Typeface.ITALIC));
+ }
return spannedWord;
}