Google’s AI Just Proofread My Writing Better Than I Ever Could

[Jose Garcia]

This is Part 2 of a series of articles where I explain how to implement GenAI on Android. [Click here to view the full series.]

Let’s test out the proofreading feature on MLKit

After building the summarisation feature in SmartWriter, integrating proofreading was incredibly fast — I reused almost all the same logic and had a working feature in under an hour.

This time, instead of focusing on UI, I’ll walk you through how the entire ViewModel works — including model download handling, inference, and graceful error fallback — using the new Proofreading API from ML Kit GenAI.

Full app available here:

https://github.com/josegbel/smartwriter

What we’re building​

We’re using ML Kit’s on-device GenAI model to take raw user input and return grammatically improved alternatives. For example:

Input:

i think this idea could works, but not sure if make sense

Output:

I think this idea could work, but I’m not sure if it makes sense.

Let’s look at how the ViewModel makes this possible — no Compose needed to understand this one

The ViewModel in full​

Here’s the complete code for the ProofreadingViewModel. I’ll break it down below:

package com.example.smartwriter.viewmodel

import android.content.Context
import android.util.Log
import androidx.lifecycle.ViewModel
import androidx.lifecycle.viewModelScope
import com.example.smartwriter.ui.model.ProofreadingUiEvent
import com.example.smartwriter.ui.model.ProofreadingUiState
import com.google.mlkit.genai.common.DownloadCallback
import com.google.mlkit.genai.common.FeatureStatus
import com.google.mlkit.genai.common.GenAiException
import com.google.mlkit.genai.proofreading.Proofreader
import com.google.mlkit.genai.proofreading.ProofreaderOptions
import com.google.mlkit.genai.proofreading.Proofreading
import com.google.mlkit.genai.proofreading.ProofreadingRequest
import kotlinx.coroutines.flow.*
import kotlinx.coroutines.guava.await
import kotlinx.coroutines.launch
import javax.inject.Inject

class ProofreadingViewModel
@Inject
constructor() : ViewModel() {

companion object {
private val TAG = ProofreadingViewModel::class.java.simpleName
}

private val _uiState = MutableStateFlow(ProofreadingUiState())
val uiState: StateFlow<ProofreadingUiState> = _uiState.asStateFlow()

private val _uiEvent = MutableSharedFlow<ProofreadingUiEvent>()
val uiEvent: SharedFlow<ProofreadingUiEvent> = _uiEvent.asSharedFlow()

private var proofreader: Proofreader? = null

override fun onCleared() {
proofreader?.close()
super.onCleared()
}

fun onInputTextChanged(newText: String) {
_uiState.update { it.copy(inputText = newText) }
}

fun onProofreadClicked(context: Context) {
viewModelScope.launch {
try {
val options =
ProofreaderOptions
.builder(context)
.setInputType(ProofreaderOptions.InputType.KEYBOARD)
.setLanguage(ProofreaderOptions.Language.ENGLISH)
.build()

proofreader = Proofreading.getClient(options)

prepareAndStartProofreading()
} catch (e: Exception) {
Log.e(TAG, "Error in onProofreadClicked: ${e.message}", e)
_uiEvent.emit(ProofreadingUiEvent.Error(message = "Error: ${e.message}"))
}
}
}

Breakdown — Initialisation​

• We use ProofreaderOptions to configure the input type (keyboard text vs voice) and language (English).
• Proofreading.getClient(options) creates the proofreader instance.
• Then we delegate to prepareAndStartProofreading().

suspend fun prepareAndStartProofreading() {
val featureStatus = proofreader?.checkFeatureStatus()?.await()
Log.d(TAG, "Feature status: $featureStatus")

when (featureStatus) {
FeatureStatus.DOWNLOADABLE -> {
Log.d(TAG, "Feature DOWNLOADABLE – starting download")
downloadFeature()
}

FeatureStatus.DOWNLOADING -> {
Log.d(TAG, "Feature DOWNLOADING – will start once ready")
proofreader?.let { startProofreadingRequest(uiState.value.inputText, it) }
}

FeatureStatus.AVAILABLE -> {
Log.d(TAG, "Feature AVAILABLE – running inference")
_uiState.update { it.copy(isLoading = true) }
proofreader?.let {
Log.d(TAG, "starting proofreading request")
startProofreadingRequest(uiState.value.inputText, it)
}
}

FeatureStatus.UNAVAILABLE, null -> {
Log.e(TAG, "Feature UNAVAILABLE")
_uiEvent.emit(ProofreadingUiEvent.Error(message = "Your device does not support this feature."))
}
}
}

Breakdown — Feature availability​

• checkFeatureStatus() tells us whether the Gemini Nano model is available, downloading, or needs to be downloaded.
• If it’s downloadable, we call downloadFeature().
• If it’s already available, we immediately run the inference.
• We also handle DOWNLOADING and UNAVAILABLE.

This flexibility ensures the user gets feedback even when the model isn’t ready yet.

private fun downloadFeature() {
proofreader?.downloadFeature(
object : DownloadCallback {
override fun onDownloadStarted(bytesToDownload: Long) {
_uiState.update { it.copy(isLoading = true) }
Log.d(TAG, "Download started – bytesToDownload=$bytesToDownload")
}

override fun onDownloadProgress(totalBytesDownloaded: Long) {
_uiState.update { it.copy(isLoading = true) }
Log.d(TAG, "Download progress – totalBytesDownloaded=$totalBytesDownloaded")
}

override fun onDownloadCompleted() {
_uiState.update { it.copy(isLoading = false) }
Log.d(TAG, "Download completed – starting inference")
proofreader?.let { startProofreadingRequest(uiState.value.inputText, it) }
}

override fun onDownloadFailed(e: GenAiException) {
_uiState.update { it.copy(isLoading = false) }
Log.e(TAG, "Download failed: ${e.message}", e)
_uiEvent.tryEmit(
ProofreadingUiEvent.Error(
message = "Download failed: ${e.message}",
),
)
}
},
)
}

Breakdown — Download logic​

• This handles model download events and updates the UI accordingly.
• On successful download, we trigger the inference again — using the same input.

fun startProofreadingRequest(
text: String,
proofreader: Proofreader,
) {
val proofreadingRequest = ProofreadingRequest.builder(text).build()
_uiState.update { it.copy(isLoading = true) }

viewModelScope.launch {
try {
val results = proofreader.runInference(proofreadingRequest).await().results
_uiState.update { state ->
state.copy(correctionSuggestions = results.map { it.text })
}
} catch (e: Exception) {
_uiEvent.emit(
ProofreadingUiEvent.Error(
message = "Error during proofreading: ${e.message}",
),
)
} finally {
_uiState.update { it.copy(isLoading = false) }
}
}
}
}

Breakdown — Running the model​

• The input string is wrapped in a ProofreadingRequest.
• runInference(…).await() gives us the corrected suggestions.
• We extract and display result.text for each suggestion in the UI.

Key takeaways​

• If you’ve implemented summarisation already, adding proofreading is fast.
• The Gemini Nano model gives usable, human-like corrections.
• You get multiple suggestions, not diffs — you can choose how to present them.
• Emulators and unsupported phones won’t work. I tested on a Galaxy S25 Ultra.
• Currently only supports English.

Try it yourself​

This is Part 2 of my SmartWriter blog series. You can find the full app (including Compose UI and previews) here:

https://github.com/josegbel/smartwriter

Next up: Text Rewriting — where we’ll teach the app to rephrase user input in different tones and lengths.

Follow along and let me know what you’d improve.

---

Google’s AI Just Proofread My Writing Better Than I Ever Could was originally published in ProAndroidDev on Medium, where people are continuing the conversation by highlighting and responding to this story.

Continue reading...

 

Join 𝕋𝕄𝕋 on Telegram
Channel PREVIEW: