nova/AIKeyboard
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files
app/build.gradle.kts
+128
View File
@@ -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
Vendored
+61
View File
@@ -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
+121
View File
@@ -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
+139
View File
@@ -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
+106
View File
@@ -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
}
}