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>

.gitignore

@@ -0,0 +1,26 @@
+*.iml
+.gradle/
+local.properties
+.idea/
+.DS_Store
+build/
+/captures/
+.externalNativeBuild/
+.cxx/
+*.jks
+*.keystore
+
+# Gradle wrapper jar (commit only the .properties file)
+gradle/wrapper/gradle-wrapper.jar
+gradlew
+gradlew.bat
+
+# HeliBoard sources are cloned by setup_heliboard.sh — not tracked here
+heliboard/
+
+# AI model files are too large for git — push to device via ADB
+*.litertlm
+*.task
+
+# Claude internal files
+.claude/

README.md

@@ -0,0 +1,139 @@
+# 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 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.
+
+---
+
+## How it works
+
+```
+User types "guten rag."
+          ↓
+InputLogic detects sentence-ending punctuation
+          ↓
+AiTriggerHook (500 ms debounce) extracts last sentence
+          ↓
+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
+          ↓
+User taps → original text replaced with correction
+```
+
+### 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 |
+| `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/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.
+
+---
+
+## 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()`.
+
+**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).
+
+**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.
+
+---
+
+## 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`).
+
+---
+
+## 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/`)
+is original work added to this project.
+
+HeliBoard itself is licensed under **GPL-3.0-only**. See
+`heliboard/LICENSE` after running `setup_heliboard.sh`.

app/build.gradle.kts

@@ -0,0 +1,128 @@
+plugins {
+    id("com.android.application")
+    id("org.jetbrains.kotlin.android")
+    id("org.jetbrains.kotlin.plugin.serialization")
+    id("org.jetbrains.kotlin.plugin.compose")
+}
+
+android {
+    // HeliBoard code expects this namespace to find its R class and BuildConfig
+    namespace = "helium314.keyboard.latin"
+    compileSdk = 35
+
+    defaultConfig {
+        applicationId = "helium314.keyboard.ai"
+        minSdk = 31        // MediaPipe LlmInference practical minimum
+        targetSdk = 35
+        versionCode = 1
+        versionName = "1.0-ai"
+
+        ndk {
+            abiFilters += listOf("arm64-v8a", "armeabi-v7a")
+        }
+    }
+
+    sourceSets {
+        getByName("main") {
+            java.srcDirs(
+                "src/main/java",                            // AI layer
+                "../heliboard/app/src/main/java"            // HeliBoard sources
+            )
+            res.srcDirs(
+                "src/main/res",
+                "../heliboard/app/src/main/res"
+            )
+            assets.srcDirs(
+                "../heliboard/app/src/main/assets"
+            )
+            jniLibs.srcDirs(
+                "../heliboard/app/src/main/jniLibs"
+            )
+            manifest.srcFile("../heliboard/app/src/main/AndroidManifest.xml")
+        }
+    }
+
+    externalNativeBuild {
+        ndkBuild {
+            path = file("../heliboard/app/src/main/jni/Android.mk")
+        }
+    }
+
+    ndkVersion = "28.0.13004108"
+
+    compileOptions {
+        sourceCompatibility = JavaVersion.VERSION_17
+        targetCompatibility = JavaVersion.VERSION_17
+        isCoreLibraryDesugaringEnabled = true
+    }
+
+    kotlinOptions {
+        jvmTarget = "17"
+    }
+
+    buildFeatures {
+        viewBinding = true
+        buildConfig = true
+        compose = true
+    }
+
+    buildTypes {
+        release {
+            isMinifyEnabled = true
+            proguardFiles(
+                getDefaultProguardFile("proguard-android-optimize.txt"),
+                "proguard-rules.pro"
+            )
+            signingConfig = signingConfigs.getByName("debug")
+        }
+        debug {
+            isMinifyEnabled = false
+            applicationIdSuffix = ".debug"
+            versionNameSuffix = "-debug"
+        }
+    }
+
+    packaging {
+        jniLibs {
+            pickFirsts += setOf(
+                "**/libjni_sentencepiece.so",
+                "**/libc++_shared.so"
+            )
+        }
+        resources {
+            excludes += setOf(
+                "META-INF/DEPENDENCIES",
+                "META-INF/LICENSE*",
+                "META-INF/NOTICE*",
+                "META-INF/*.kotlin_module"
+            )
+        }
+    }
+}
+
+dependencies {
+    coreLibraryDesugaring("com.android.tools:desugar_jdk_libs:2.1.5")
+
+    implementation("androidx.core:core-ktx:1.16.0")
+    implementation("androidx.recyclerview:recyclerview:1.4.0")
+    implementation("androidx.autofill:autofill:1.3.0")
+    implementation("androidx.viewpager2:viewpager2:1.1.0")
+    implementation("androidx.preference:preference-ktx:1.2.1")
+    implementation("androidx.constraintlayout:constraintlayout:2.2.1")
+    implementation("com.google.android.material:material:1.12.0")
+
+    implementation(platform("androidx.compose:compose-bom:2025.05.00"))
+    implementation("androidx.compose.material3:material3")
+    implementation("androidx.compose.ui:ui-tooling-preview")
+    implementation("androidx.navigation:navigation-compose:2.8.9")
+    debugImplementation("androidx.compose.ui:ui-tooling")
+
+    implementation("org.jetbrains.kotlinx:kotlinx-serialization-json:1.8.1")
+    implementation("sh.calvin.reorderable:reorderable:2.4.3")
+    implementation("com.github.skydoves:colorpicker-compose:1.1.3")
+
+    implementation("com.google.ai.edge.litertlm:litertlm-android:latest.release")
+
+    implementation("org.jetbrains.kotlinx:kotlinx-coroutines-android:1.8.1")
+    implementation("org.jetbrains.kotlinx:kotlinx-coroutines-core:1.8.1")
+}

app/proguard-rules.pro

@@ -0,0 +1,61 @@
+# ============================================================
+# MediaPipe tasks-genai / LlmInference
+# ============================================================
+
+-keep class com.google.mediapipe.** { *; }
+-dontwarn com.google.mediapipe.**
+
+# Protobuf (used internally by MediaPipe)
+-keep class com.google.mediapipe.proto.** { *; }
+-keepclassmembers class * extends com.google.protobuf.GeneratedMessageLite {
+    <fields>;
+}
+-keep class com.google.protobuf.** { *; }
+-dontwarn com.google.protobuf.**
+
+# Google Flogger (used by MediaPipe internals)
+-keep class com.google.common.flogger.** { *; }
+-dontwarn com.google.common.flogger.**
+
+# JNI methods must not be renamed
+-keepclasseswithmembernames class * {
+    native <methods>;
+}
+
+# MediaPipe framework interop
+-keep public interface com.google.mediapipe.framework.* { public *; }
+-keepclassmembers class com.google.mediapipe.framework.Packet {
+    public static *** create(...);
+    public long getNativeHandle();
+    public void release();
+}
+-keepclassmembers class com.google.mediapipe.framework.PacketCreator {
+    public void releaseWithSyncToken(...);
+}
+-keep class com.google.mediapipe.framework.MediaPipeException {
+    <init>(int, byte[]);
+}
+
+# ============================================================
+# AI engine — never obfuscate so Java↔Kotlin interop is stable
+# ============================================================
+-keep class helium314.keyboard.latin.ai.** { *; }
+-keepclassmembers class helium314.keyboard.latin.ai.** { *; }
+
+# ============================================================
+# Kotlin coroutines
+# ============================================================
+-keep class kotlinx.coroutines.** { *; }
+-dontwarn kotlinx.coroutines.**
+
+# Keep coroutine debug metadata (optional, remove in prod to shrink further)
+-keepattributes SourceFile,LineNumberTable
+
+# ============================================================
+# HeliBoard — keep JNI-accessible classes and reflection targets
+# ============================================================
+-keep class helium314.keyboard.latin.BinaryDictionary { *; }
+-keep class helium314.keyboard.latin.DicTraverseSession { *; }
+-keep class helium314.keyboard.latin.NgramContext { *; }
+-keep class helium314.keyboard.latin.utils.BinaryDictionaryUtils { *; }
+-dontwarn helium314.keyboard.**

app/src/main/java/helium314/keyboard/latin/ai/AiCorrectionEngine.kt

@@ -0,0 +1,121 @@
+package helium314.keyboard.latin.ai
+
+import android.content.Context
+import android.util.Log
+import com.google.ai.edge.litertlm.Backend
+import com.google.ai.edge.litertlm.Content
+import com.google.ai.edge.litertlm.Contents
+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.Dispatchers
+import kotlinx.coroutines.withContext
+import java.io.File
+
+/**
+ * Wraps the LiteRT-LM [Engine] to provide on-device AI spell/grammar correction
+ * using the Gemma 4 E2B-it model in .litertlm format.
+ *
+ * Model placement:
+ *   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
+ */
+class AiCorrectionEngine(private val context: Context) {
+
+    companion object {
+        private const val TAG = "AiCorrectionEngine"
+        const val MODEL_PATH = "/data/local/tmp/gemma-4-E2B-it.litertlm"
+
+        private const val SYSTEM_INSTRUCTION =
+            "You are a helpful assistant. Correct only spelling and grammar errors. " +
+            "Return ONLY the corrected sentence, nothing else."
+    }
+
+    @Volatile private var engine: Engine? = null
+    // Set to true after a GPU inference failure so we permanently switch to CPU.
+    @Volatile private var forceCpu: Boolean = false
+
+    private fun getOrCreate(): Engine {
+        engine?.let { return it }
+        synchronized(this) {
+            engine?.let { return it }
+
+            if (!File(MODEL_PATH).exists()) {
+                throw IllegalStateException("Gemma 4 model not found at $MODEL_PATH")
+            }
+
+            Log.i(TAG, "Loading Gemma 4 E2B-it (.litertlm) model...")
+            val newEngine = createEngine()
+            engine = newEngine
+            Log.i(TAG, "Gemma 4 loaded successfully")
+            return newEngine
+        }
+    }
+
+    // Try GPU first for performance, fall back to CPU on init failure.
+    private fun createEngine(): Engine {
+        val backends = if (forceCpu) listOf(Backend.CPU()) else listOf(Backend.GPU(), Backend.CPU())
+        for (backend in backends) {
+            try {
+                val e = Engine(EngineConfig(modelPath = MODEL_PATH, backend = backend))
+                e.initialize()
+                Log.i(TAG, "Using backend: ${backend::class.simpleName}")
+                return e
+            } catch (ex: Exception) {
+                Log.w(TAG, "Backend ${backend::class.simpleName} failed at init: ${ex.message}")
+            }
+        }
+        throw IllegalStateException("All backends failed to load model at $MODEL_PATH")
+    }
+
+    // Switches to CPU engine and returns it.
+    private fun switchToCpu(): Engine {
+        synchronized(this) {
+            forceCpu = true
+            engine?.close()
+            engine = null
+        }
+        return getOrCreate()
+    }
+
+    suspend fun correctSentence(sentence: String): String = withContext(Dispatchers.IO) {
+        if (sentence.isBlank()) return@withContext sentence
+        return@withContext try {
+            runInference(sentence, getOrCreate())
+        } catch (ex: Exception) {
+            if (!forceCpu && ex.message?.let { it.contains("OpenCL") || it.contains("GPU") } == true) {
+                Log.w(TAG, "GPU inference failed, switching to CPU permanently")
+                try {
+                    runInference(sentence, switchToCpu())
+                } catch (e2: Exception) {
+                    Log.e(TAG, "CPU inference failed: ${e2.message}")
+                    sentence
+                }
+            } else {
+                Log.e(TAG, "Gemma 4 correction failed: ${ex.message}")
+                sentence
+            }
+        }
+    }
+
+    private fun runInference(sentence: String, e: Engine): String {
+        val config = ConversationConfig(
+            systemInstruction = Contents.of(SYSTEM_INSTRUCTION),
+            samplerConfig = SamplerConfig(topK = 64, topP = 0.95, temperature = 1.0)
+        )
+        return e.createConversation(config).use { conversation ->
+            val message = conversation.sendMessage("Sentence: $sentence")
+            val text = message.contents.contents
+                .filterIsInstance<Content.Text>()
+                .joinToString("") { it.text }
+                .trim()
+            text.ifBlank { sentence }
+        }
+    }
+
+    fun close() {
+        engine?.close()
+        engine = null
+    }
+}

app/src/main/java/helium314/keyboard/latin/ai/AiSuggestionManager.kt

@@ -0,0 +1,139 @@
+package helium314.keyboard.latin.ai
+
+import android.util.Log
+import helium314.keyboard.latin.SuggestedWords
+import helium314.keyboard.latin.SuggestedWords.SuggestedWordInfo
+import helium314.keyboard.latin.suggestions.SuggestionStripViewAccessor
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.SupervisorJob
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.asStateFlow
+
+/**
+ * Bridges AI correction results into HeliBoard's suggestion strip.
+ *
+ * Holds the current AI suggestion as a [StateFlow] and pushes it to the
+ * attached [SuggestionStripViewAccessor] whenever a new correction arrives.
+ *
+ * Lifecycle
+ * ---------
+ * Call [attachAccessor] in LatinIME.onCreate() and [detachAccessor] in
+ * LatinIME.onDestroy(). The [createImeScope] factory (@JvmStatic) creates
+ * a SupervisorJob-based scope that can be cancelled from Java.
+ *
+ * When the user taps the AI suggestion in the strip, HeliBoard's normal
+ * suggestion-pick flow will call pickSuggestionManually(). The corrected
+ * sentence (entire sentence, not just a word) replaces the typed text.
+ * This works because we set the suggestion score to Int.MAX_VALUE and
+ * mark it as the typed-word so HeliBoard uses commitText for the whole thing.
+ *
+ * Constructor signatures verified against HeliBoard main branch (Apr 2025).
+ */
+class AiSuggestionManager {
+
+    companion object {
+        private const val TAG = "AiSuggestionManager"
+
+        /**
+         * Custom flag bit added to KIND_CORRECTION to mark AI suggestions.
+         * Using 0x10000 keeps the base kind as KIND_CORRECTION (= 1) so HeliBoard's
+         * normal suggestion-pick path handles it as a text commit, not commitCompletion.
+         * SuggestionStripLayoutHelper checks this bit for italic/colour styling.
+         */
+        const val KIND_AI_FLAG = 0x10000
+        const val KIND_AI_CORRECTION = SuggestedWordInfo.KIND_CORRECTION or KIND_AI_FLAG
+
+        /**
+         * Factory method accessible from Java (LatinIME.java patch).
+         * Returns a [CoroutineScope] tied to a SupervisorJob so individual
+         * coroutine failures do not cancel the whole scope.
+         */
+        @JvmStatic
+        fun createImeScope(): CoroutineScope =
+            CoroutineScope(SupervisorJob() + Dispatchers.Main.immediate)
+    }
+
+    // The current AI-corrected suggestion, or null if none is pending.
+    private val _currentSuggestion = MutableStateFlow<String?>(null)
+    val currentSuggestion: StateFlow<String?> = _currentSuggestion.asStateFlow()
+
+    // Original (uncorrected) sentence paired with the correction — used when
+    // the user accepts the suggestion to know how much text to replace.
+    @Volatile private var originalSentence: String? = null
+
+    @Volatile private var accessor: SuggestionStripViewAccessor? = null
+
+    /** Attach the IME's SuggestionStripViewAccessor. Called in LatinIME.onCreate(). */
+    fun attachAccessor(a: SuggestionStripViewAccessor) {
+        accessor = a
+    }
+
+    /** Called in LatinIME.onDestroy(). */
+    fun detachAccessor() {
+        accessor = null
+    }
+
+    /**
+     * Posts an AI-corrected suggestion to the suggestion strip.
+     *
+     * [correctedSentence] — the model's output
+     * [original]          — the original sentence (used for length calculation)
+     *
+     * Thread-safe: can be called from any coroutine context.
+     * HeliBoard's [SuggestionStripViewAccessor.setSuggestions] is main-thread-safe.
+     */
+    fun postSuggestion(correctedSentence: String, original: String) {
+        _currentSuggestion.value = correctedSentence
+        originalSentence = original
+        Log.d(TAG, "Posting AI suggestion: \"$correctedSentence\"")
+
+        // Constructor signature verified against HeliBoard main (Apr 2025):
+        // SuggestedWordInfo(word, prevWordsContext, score, kindAndFlags,
+        //                   sourceDict, indexOfTouchPointOfSecondWord, autoCommitFirstWordConfidence)
+        val wordInfo = SuggestedWordInfo(
+            /* word = */                          correctedSentence,
+            /* prevWordsContext = */              "",
+            /* score = */                         Int.MAX_VALUE,  // rank first in strip
+            /* kindAndFlags = */                  KIND_AI_CORRECTION,
+            /* sourceDict = */                    null,
+            /* indexOfTouchPointOfSecondWord = */ SuggestedWordInfo.NOT_AN_INDEX,
+            /* autoCommitFirstWordConfidence = */ SuggestedWordInfo.NOT_A_CONFIDENCE
+        )
+
+        val words = SuggestedWords(
+            /* suggestions = */        arrayListOf(wordInfo),
+            /* rawSuggestions = */     null,
+            /* typedWordInfo = */      wordInfo,
+            /* typedWordValid = */     false,
+            /* willAutoCorrect = */    false,
+            /* isObsoleteSuggestions = */ false,
+            /* inputStyle = */         SuggestedWords.INPUT_STYLE_PREDICTION,
+            /* sequenceNumber = */     SuggestedWords.NOT_A_SEQUENCE_NUMBER
+        )
+
+        val a = accessor
+        if (a == null) {
+            Log.w(TAG, "accessor is null — suggestion not shown")
+            return
+        }
+        // showAiSuggestion bypasses isSuggestionsEnabledPerUserSettings() so the AI
+        // correction appears regardless of toolbar mode or per-field suggestion settings.
+        a.showAiSuggestion(words)
+    }
+
+    /** Returns the original (pre-correction) sentence, or null if none is pending. */
+    fun getOriginalSentence(): String? = originalSentence
+
+    /**
+     * Clears the AI suggestion from the strip.
+     * Called when the correction matches the original (no change needed)
+     * or when the user starts typing a new sentence.
+     */
+    fun clearSuggestion() {
+        _currentSuggestion.value = null
+        originalSentence = null
+        accessor?.setNeutralSuggestionStrip()
+    }
+}

app/src/main/java/helium314/keyboard/latin/ai/AiTriggerHook.kt

@@ -0,0 +1,106 @@
+package helium314.keyboard.latin.ai
+
+import android.util.Log
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Job
+import kotlinx.coroutines.delay
+import kotlinx.coroutines.launch
+
+/**
+ * Detects sentence-ending punctuation in the input stream and triggers
+ * AI spell/grammar correction with a debounce to avoid rapid-fire calls.
+ *
+ * Integration
+ * -----------
+ * Called from InputLogic.java (patched via 0001-InputLogic-ai-hook.patch).
+ * The static [isSentenceEnder] method is used from Java to check individual
+ * code points before invoking [onSentenceEndDetected].
+ *
+ * Threading
+ * ---------
+ * [onSentenceEndDetected] is called on the IME main thread (from InputLogic).
+ * Coroutines launched here inherit the [scope] provided by LatinIME and
+ * switch to IO automatically inside [AiCorrectionEngine.correctSentence].
+ */
+class AiTriggerHook(
+    private val engine: AiCorrectionEngine,
+    private val manager: AiSuggestionManager,
+    private val scope: CoroutineScope
+) {
+    companion object {
+        private const val TAG = "AiTriggerHook"
+        private const val DEBOUNCE_MS = 500L
+        private const val MIN_SENTENCE_LENGTH = 4  // ignore trivially short fragments
+        private const val MAX_LOOKBACK_CHARS = 500  // chars to request from InputConnection
+
+        private val SENTENCE_ENDERS = setOf('.', '!', '?', '\u2026' /* … */)
+
+        /**
+         * Called from InputLogic.java to check whether a code point ends a sentence.
+         * Must be @JvmStatic so Java can call it without an instance.
+         */
+        @JvmStatic
+        fun isSentenceEnder(codePoint: Int): Boolean =
+            codePoint.toChar() in SENTENCE_ENDERS
+    }
+
+    private var debounceJob: Job? = null
+
+    /**
+     * Called after a sentence-ending punctuation character is committed to the
+     * input connection. [textBeforeCursor] is the raw text from
+     * InputConnection.getTextBeforeCursor(500, 0).
+     *
+     * Debounced at [DEBOUNCE_MS] ms to handle burst punctuation (e.g. "…").
+     */
+    fun onSentenceEndDetected(textBeforeCursor: String) {
+        val sentence = extractLastSentence(textBeforeCursor)
+        if (sentence.length < MIN_SENTENCE_LENGTH) {
+            Log.v(TAG, "Sentence too short, skipping: \"$sentence\"")
+            return
+        }
+
+        debounceJob?.cancel()
+        debounceJob = scope.launch {
+            delay(DEBOUNCE_MS)
+            Log.d(TAG, "Triggering AI correction for: \"$sentence\"")
+            val corrected = engine.correctSentence(sentence)
+            if (corrected != sentence) {
+                manager.postSuggestion(corrected, sentence)
+                Log.d(TAG, "Correction: \"$sentence\" → \"$corrected\"")
+            } else {
+                manager.clearSuggestion()
+                Log.d(TAG, "No correction needed for: \"$sentence\"")
+            }
+        }
+    }
+
+    /**
+     * Extracts the last complete sentence from [text].
+     *
+     * Walks backwards to find the sentence boundary just before the one the
+     * user just typed, so we get the full sentence including its trailing
+     * punctuation. Example:
+     *   "Hello world. This is a test." → "This is a test."
+     */
+    private fun extractLastSentence(text: String): String {
+        if (text.isBlank()) return ""
+        val trimmed = text.trimEnd()
+
+        // Find the index of the last sentence-ender (the one just typed)
+        val lastEnd = trimmed.indexOfLast { it in SENTENCE_ENDERS }
+        if (lastEnd < 0) return trimmed  // no boundary found, return all text
+
+        // Find the sentence-ender before that (start of our target sentence)
+        val prevEnd = trimmed.substring(0, lastEnd).indexOfLast { it in SENTENCE_ENDERS }
+
+        val start = if (prevEnd >= 0) prevEnd + 1 else 0
+        return trimmed.substring(start, lastEnd + 1).trim()
+    }
+
+    /** Cancels any pending debounce job. Call when the IME is hidden or destroyed. */
+    fun cancel() {
+        debounceJob?.cancel()
+        debounceJob = null
+    }
+}

build.gradle.kts

@@ -0,0 +1,6 @@
+plugins {
+    id("com.android.application") version "9.0.1" apply false
+    id("org.jetbrains.kotlin.android") version "2.2.10" apply false
+    id("org.jetbrains.kotlin.plugin.serialization") version "2.0.21" apply false
+    id("org.jetbrains.kotlin.plugin.compose") version "2.2.10" apply false
+}

gradle.properties

@@ -0,0 +1,16 @@
+org.gradle.jvmargs=-Xmx4096m -XX:MaxMetaspaceSize=1024m -XX:+HeapDumpOnOutOfMemoryError -Dfile.encoding=UTF-8
+android.useAndroidX=true
+kotlin.code.style=official
+android.enableJetifier=false
+org.gradle.parallel=true
+org.gradle.caching=true
+android.defaults.buildfeatures.resvalues=true
+android.sdk.defaultTargetSdkToCompileSdkIfUnset=false
+android.enableAppCompileTimeRClass=false
+android.usesSdkInManifest.disallowed=false
+android.uniquePackageNames=false
+android.dependency.useConstraints=true
+android.r8.strictFullModeForKeepRules=false
+android.r8.optimizedResourceShrinking=false
+android.builtInKotlin=false
+android.newDsl=false

gradle/gradle-daemon-jvm.properties

@@ -0,0 +1,13 @@
+#This file is generated by updateDaemonJvm
+toolchainUrl.FREE_BSD.AARCH64=https\://api.foojay.io/disco/v3.0/ids/56a19bc915b9ba2eb62ba7554c61b919/redirect
+toolchainUrl.FREE_BSD.X86_64=https\://api.foojay.io/disco/v3.0/ids/398ffe3949748bfb1d5636f023d228fd/redirect
+toolchainUrl.LINUX.AARCH64=https\://api.foojay.io/disco/v3.0/ids/56a19bc915b9ba2eb62ba7554c61b919/redirect
+toolchainUrl.LINUX.X86_64=https\://api.foojay.io/disco/v3.0/ids/398ffe3949748bfb1d5636f023d228fd/redirect
+toolchainUrl.MAC_OS.AARCH64=https\://api.foojay.io/disco/v3.0/ids/e99bae143b75f9a10ead10248f02055e/redirect
+toolchainUrl.MAC_OS.X86_64=https\://api.foojay.io/disco/v3.0/ids/04e088f8677de3b384108493cc9481d0/redirect
+toolchainUrl.UNIX.AARCH64=https\://api.foojay.io/disco/v3.0/ids/56a19bc915b9ba2eb62ba7554c61b919/redirect
+toolchainUrl.UNIX.X86_64=https\://api.foojay.io/disco/v3.0/ids/398ffe3949748bfb1d5636f023d228fd/redirect
+toolchainUrl.WINDOWS.AARCH64=https\://api.foojay.io/disco/v3.0/ids/e55dccbfe27cb97945148c61a39c89c5/redirect
+toolchainUrl.WINDOWS.X86_64=https\://api.foojay.io/disco/v3.0/ids/dbd05c4936d573642f94cd149e1356c8/redirect
+toolchainVendor=JETBRAINS
+toolchainVersion=21

gradle/wrapper/gradle-wrapper.properties

@@ -0,0 +1,7 @@
+distributionBase=GRADLE_USER_HOME
+distributionPath=wrapper/dists
+distributionUrl=https\://services.gradle.org/distributions/gradle-9.2.1-bin.zip
+networkTimeout=10000
+validateDistributionUrl=true
+zipStoreBase=GRADLE_USER_HOME
+zipStorePath=wrapper/dists

local.properties.template

@@ -0,0 +1,5 @@
+# Copy this file to local.properties and fill in your SDK/NDK paths.
+# local.properties is NOT committed to git.
+
+sdk.dir=C\:\\Users\\thors\\AppData\\Local\\Android\\Sdk
+ndk.dir=C\:\\Users\\thors\\AppData\\Local\\Android\\Sdk\\ndk\\28.0.13004108

patches/0001-InputLogic-ai-hook.patch

@@ -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 {

patches/0002-LatinIME-ai-lifecycle.patch

@@ -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;
++    }
++    // ────────────────────────────────────────────────────────────────────────
+ }

patches/0003-SuggestionStrip-ai-style.patch

@@ -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;
+     }

settings.gradle.kts

@@ -0,0 +1,19 @@
+pluginManagement {
+    repositories {
+        google()
+        mavenCentral()
+        gradlePluginPortal()
+    }
+}
+plugins {
+    id("org.gradle.toolchains.foojay-resolver-convention") version "1.0.0"
+}
+dependencyResolutionManagement {
+    repositoriesMode.set(RepositoriesMode.FAIL_ON_PROJECT_REPOS)
+    repositories {
+        google()
+        mavenCentral()
+    }
+}
+rootProject.name = "AIKeyboard"
+include(":app")

setup_heliboard.sh

@@ -0,0 +1,41 @@
+#!/usr/bin/env bash
+# setup_heliboard.sh
+# Run once after cloning this repo to pull HeliBoard and apply AI patches.
+#
+# Model: https://huggingface.co/litert-community/gemma-4-E2B-it-litert-lm
+
+set -euo pipefail
+
+HELIBOARD_REPO="https://github.com/Helium314/HeliBoard.git"
+HELIBOARD_COMMIT="main"
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+HELIBOARD_DIR="$SCRIPT_DIR/heliboard"
+PATCHES_DIR="$SCRIPT_DIR/patches"
+
+echo "==> AIKeyboard Gemma 4 E2B Setup"
+
+if [ ! -d "$HELIBOARD_DIR/.git" ]; then
+    git clone "$HELIBOARD_REPO" "$HELIBOARD_DIR"
+fi
+
+cd "$HELIBOARD_DIR"
+git fetch origin
+git checkout "$HELIBOARD_COMMIT"
+
+echo "--> Applying Gemma 4 integration patches ..."
+for patch in "$PATCHES_DIR"/*.patch; do
+    if [ -f "$patch" ]; then
+        git apply --check "$patch" 2>/dev/null && git apply "$patch" \
+            || echo "    SKIP: $(basename "$patch") already applied or conflicts."
+    fi
+done
+
+cd "$SCRIPT_DIR"
+
+echo ""
+echo "Next steps:"
+echo "  1. Push Gemma 4 LiteRT model file to your device:"
+echo "     adb push gemma-4-E2B-it.litertlm /data/local/tmp/gemma-4-E2B-it.litertlm"
+echo "     adb shell chmod 644 /data/local/tmp/gemma-4-E2B-it.litertlm"
+echo "  2. Build & Install: ./gradlew installDebug"