Running LLMs locally on an Android smartphone

import kotlinx.coroutines.flow.Flow

interface LocalChat : AutoCloseable {
    /**
     * Emits the answer accumulated so far and completes when the answer ends.
     * Cancelling the collection stops the generation. One call at a time.
     */
    fun send(message: String): Flow<String>
}

import androidx.lifecycle.ViewModel
import androidx.lifecycle.viewModelScope
import kotlinx.coroutines.Job
import kotlinx.coroutines.flow.MutableStateFlow
import kotlinx.coroutines.flow.catch
import kotlinx.coroutines.flow.launchIn
import kotlinx.coroutines.flow.onCompletion
import kotlinx.coroutines.flow.onEach

class ChatViewModel(private val chat: LocalChat) : ViewModel() {

    val answer = MutableStateFlow("")
    val error = MutableStateFlow<Throwable?>(null)
    val isGenerating = MutableStateFlow(false)

    private var generation: Job? = null

    fun send(message: String) {
        isGenerating.value = true
        generation = chat.send(message)
            .onEach { answer.value = it }
            .catch { error.value = it }
            .onCompletion { isGenerating.value = false }
            .launchIn(viewModelScope)
    }

    fun stop() {
        generation?.cancel()
    }
}

implementation "com.google.ai.edge.litertlm:litertlm-android:0.19.0"

<uses-native-library android:name="libvndksupport.so" android:required="false" />
<uses-native-library android:name="libOpenCL.so" android:required="false" />

import com.google.ai.edge.litertlm.Backend
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.CancellationException
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.conflate
import kotlinx.coroutines.flow.flow
import kotlinx.coroutines.flow.flowOn
import kotlinx.coroutines.flow.onCompletion

class LiteRtLmChat(modelPath: String, cacheDir: String) : LocalChat {

    private val engine = Engine(
        EngineConfig(
            modelPath = modelPath,
            backend = Backend.GPU(),
            cacheDir = cacheDir, // speeds up every load after the first one
        )
    ).apply { initialize() }

    // The conversation keeps the dialogue history itself.
    private val conversation = engine.createConversation(
        ConversationConfig(samplerConfig = SamplerConfig(topK = 40, topP = 0.95, temperature = 0.8))
    )

    override fun send(message: String): Flow<String> = flow {
        val answer = StringBuilder()
        conversation.sendMessageAsync(message).collect { chunk ->
            answer.append(chunk.toString())
            emit(answer.toString())
        }
    }.onCompletion { cause ->
        // Cancelling the collection does not reach the engine.
        if (cause is CancellationException) conversation.cancelProcess()
    }.conflate().flowOn(Dispatchers.IO)

    override fun close() {
        conversation.close()
        engine.close()
    }
}

add_subdirectory(llama_cpp build-llama)

implementation "com.google.mediapipe:tasks-genai:0.10.35"

import android.content.Context
import com.google.mediapipe.tasks.genai.llminference.LlmInference
import com.google.mediapipe.tasks.genai.llminference.LlmInferenceSession
import com.google.mediapipe.tasks.genai.llminference.ProgressListener
import java.util.concurrent.Future
import java.util.concurrent.TimeUnit
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.channels.awaitClose
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.callbackFlow
import kotlinx.coroutines.flow.conflate
import kotlinx.coroutines.flow.flowOn

class MediaPipeChat(
    context: Context,
    modelPath: String,
    backend: LlmInference.Backend = LlmInference.Backend.GPU,
) : LocalChat {

    private val inference = LlmInference.createFromOptions(
        context,
        LlmInference.LlmInferenceOptions.builder()
            .setModelPath(modelPath)
            .setMaxTokens(1024)
            .setMaxTopK(64)
            .setPreferredBackend(backend)
            .build(),
    )

    private val session = LlmInferenceSession.createFromOptions(
        inference,
        LlmInferenceSession.LlmInferenceSessionOptions.builder()
            .setTopK(40)
            .setTopP(1.0f)
            .setTemperature(0.8f)
            .build(),
    )

    private var generation: Future<String>? = null

    override fun send(message: String): Flow<String> = callbackFlow {
        awaitPreviousGeneration()
        val answer = StringBuilder()
        session.addQueryChunk(message)
        val future = session.generateResponseAsync(ProgressListener { partial, done ->
            answer.append(partial)
            trySend(answer.toString())
            if (done) close()
        })
        generation = future
        awaitClose {
            if (!future.isDone) session.cancelGenerateResponseAsync()
        }
    }.conflate().flowOn(Dispatchers.IO)

    override fun close() {
        awaitPreviousGeneration()
        session.close()
        inference.close()
    }

    // After a cancel the engine stays busy until its final done=true.
    private fun awaitPreviousGeneration() {
        runCatching { generation?.get(30, TimeUnit.SECONDS) }
    }
}

{
  "device": "android",
  "model_list": [
    {
      "model": "HF://mlc-ai/Phi-3.5-mini-instruct-q4f16_0-MLC",
      "model_id": "Phi-3.5-mini-instruct-q4f16_0-MLC",
      "estimated_vram_bytes": 4250586449
    }
  ]
}

{
  "device": "android",
  "model_list": [
    {
      "model": "../Recompiled",
      "model_id": "gemma-3-1b-it-q4f16_0-MLC",
      "estimated_vram_bytes": 3000000000,
      "bundle_weight": true
    }
  ]
}

include ':mlc4j'
project(':mlc4j').projectDir = file('dist/lib/mlc4j')

<uses-native-library android:name="libOpenCL.so" android:required="false" />
<uses-native-library android:name="libOpenCL-pixel.so" android:required="false" />

import ai.mlc.mlcllm.MLCEngine
import ai.mlc.mlcllm.OpenAIProtocol.ChatCompletionMessage
import ai.mlc.mlcllm.OpenAIProtocol.ChatCompletionRole
import ai.mlc.mlcllm.OpenAIProtocol.StreamOptions
import android.util.Log
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.conflate
import kotlinx.coroutines.flow.flow
import kotlinx.coroutines.flow.flowOn

class MlcChat : LocalChat {

    // One per process: the constructor starts two background threads that cannot be stopped.
    private val engine = MLCEngine()
    private val history = mutableListOf<ChatCompletionMessage>()

    // Also switches models.
    fun load(modelPath: String, modelLib: String) {
        engine.unload()
        engine.reload(modelPath, modelLib)
        history.clear()
    }

    override fun send(message: String): Flow<String> = flow {
        history += ChatCompletionMessage(role = ChatCompletionRole.user, content = message)
        val answer = StringBuilder()
        var finished = false
        try {
            val responses = engine.chat.completions.create(
                messages = history,
                stream_options = StreamOptions(include_usage = true),
            )
            for (response in responses) {
                response.choices.firstOrNull()?.delta?.content?.asText()?.let { chunk ->
                    answer.append(chunk)
                    emit(answer.toString())
                }
                // The last message carries the statistics: this is where to read the real speed.
                response.usage?.extra?.let { stats ->
                    Log.d("MlcChat", "prefill ${stats.prefill_tokens_per_s}, decode ${stats.decode_tokens_per_s} tok/s")
                }
            }
            finished = true
        } finally {
            if (!finished) engine.reset()
            if (answer.isNotEmpty()) {
                history += ChatCompletionMessage(role = ChatCompletionRole.assistant, content = answer.toString())
            } else {
                history.removeAt(history.lastIndex)
            }
        }
    }.conflate().flowOn(Dispatchers.IO)

    override fun close() = engine.unload()
}

implementation "com.microsoft.onnxruntime:onnxruntime-android:latest.release"
implementation files("libs/onnxruntime-genai-android-0.19.0.aar")

import ai.onnxruntime.genai.SimpleGenAI
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.ensureActive
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.channelFlow
import kotlinx.coroutines.flow.conflate
import kotlinx.coroutines.flow.flowOn

class OnnxChat(modelDir: String) : LocalChat {

    private val genAi = SimpleGenAI(modelDir)
    private val dialogue = StringBuilder()

    override fun send(message: String): Flow<String> = channelFlow {
        // The markup is Phi-3's; another model needs its own.
        dialogue.append("<|user|>\n").append(message).append("<|end|>\n<|assistant|>\n")
        val answer = StringBuilder()
        try {
            genAi.createGeneratorParams().use { params ->
                params.setSearchOption("max_length", 2048.0)
                params.setSearchOption("do_sample", true)
                params.setSearchOption("temperature", 0.7)
                // The answer is assembled from the chunks: what generate() returns
                // is the whole sequence decoded, with the prompt at the start of it.
                genAi.generate(params, dialogue.toString()) { chunk ->
                    ensureActive() // an exception is the only way to stop generate() from inside
                    answer.append(chunk)
                    trySend(answer.toString())
                }
            }
        } finally {
            dialogue.append(answer).append("<|end|>\n")
        }
    }.conflate().flowOn(Dispatchers.IO)

    override fun close() = genAi.close()
}

import android.content.Context
import android.content.res.AssetManager
import java.io.File

fun copyModelsFromAssets(context: Context, modelsVersion: Int, assetsDir: String = "models"): File {
    val target = File(context.filesDir, assetsDir)
    // The marker is written last, so an interrupted copy starts over on the next launch.
    val marker = File(target, ".copied-$modelsVersion")
    if (marker.exists()) return target
    target.deleteRecursively()
    copyAssets(context.assets, assetsDir, target)
    marker.createNewFile()
    return target
}

private fun copyAssets(assets: AssetManager, path: String, target: File) {
    val children = assets.list(path).orEmpty()
    if (children.isEmpty()) {
        target.parentFile?.mkdirs()
        assets.open(path).use { input -> target.outputStream().use { input.copyTo(it) } }
    } else {
        children.forEach { copyAssets(assets, "$path/$it", File(target, it)) }
    }
}