以下文章來源于Android高效開發，作者2BAB

作者 / Android 谷歌開發者專家 El Zhang (2BAB)

在今年的廈門和廣州 Google I/O Extended 上，我分享了《On-Device Model 集成 (KMP) 與用例》。本文是當時 Demo 的深入細節分析，同時也是后面幾篇同類型文章的開頭。通過本文你將了解到:

移植 Mediapipe 的 LLM Inference Android 官方 Demo 到 KMP，支持在 iOS 上運行。

KMP 兩種常見的調用 iOS SDK 的方式:

Kotlin 直接調用 Cocoapods 引入的第三方庫。

Kotlin 通過 iOS 工程調用第三方庫。

KMP 與多平臺依賴注入時的小技巧 (基于 Koin)。

On-Device Model 與 LLM 模型 Gemma 1.1 2B 的簡單背景。

On-Device Model 本地模型

大語言模型 (LLM) 持續火熱了很長一段時間，而今年開始這股風正式吹到了移動端，包括 Google 在內的最新手機與系統均深度集成了此類 On-Device Model 的相關功能。對于 Google 目前的公開戰略中，On-Device Model 這塊的大語言模型主要分為兩個:

Gemini Nano: 非開源，支持機型較少 (某些機型支持特定芯片加速如 Tensor G4)，具有強勁的表現。目前可以在桌面平臺 (Chrome) 和部分 Android 手機上使用 (Pixel 8/9 Samsung 和小米部分機型)。據報道晚些時候會公開給更多的開發者進行使用和測試。

Gemma: 開源，支持所有滿足最低要求的機型，同樣有不俗的性能表現，與 Nano 使用類似的技術路線進行訓練。目前可以在多平臺上體驗 (Android/iOS/Desktop)。

目前多數移動端開發者尚無法直接基于 Gemini Nano 開發，所以今天的主角便是 Gemma 1 的 2B 版本。想在移動平臺上直接使用 Gemma，Google 已給我們提供一個開箱即用的工具: Mediapipe。MediaPipe 是一個跨平臺的框架，它封裝了一系列預構建的 On-Device 機器學習模型和工具，支持實時的手勢識別、面部檢測、姿態估計等任務，還可應用于生成圖片、聊天機器人等各種應用場景。感興趣的朋友可以試玩它的 Web 版 Demo，以及相關文檔。

而其中的 LLM Inference API (上表第一行)，用于運行大語言模型推理的組件，支持 Gemma 2B/7B，Phi-2，Falcon-RW-1B，StableLM-3B 等模型。針對 Gemma 的預轉換模型 (基于 TensorFlow Lite) 可在 Kaggle 下載，并在稍后直接放入 Mediapipe 中加載。

LLM Inference

Android Sample

Mediapipe 官方的 LLM Inference Demo 包含了 Android/iOS/Web 前端等平臺。

打開 Android 倉庫會發現幾個特點:

純 Kotlin 實現。

UI 是純 Jetpack Compose 實現。

依賴的 LLM Task SDK 已經高度封裝，暴露出來的方法僅 3 個。

再查看 iOS 的版本:

UI 是 SwiftUI 實現，做的事情和 Compose 一模一樣，稍微再簡化掉一些元素 (例如 Topbar 和發送按鈕)。

依賴的 LLM Task SDK 已經高度封裝，暴露出來的方法一樣為 3 個。

所以，一個好玩的想法出現了:Android 版本的這個 Demo 具備移植到 iOS 上的基礎；移植可使兩邊的代碼高度高度一致，大幅縮減維護成本，而核心要實現的僅僅是橋接下 iOS 上的 LLM Inference SDK。

Kotlin Multiplatform

移植工程所使用的技術叫做 Kotlin Multiplatform (縮寫為 KMP)，它是 Kotlin 團隊開發的一種支持跨平臺開發的技術，允許開發者使用相同的代碼庫來構建 Android、iOS、Web 等多個平臺的應用程序。通過共享業務邏輯代碼，KMP 能顯著減少開發時間和維護成本，同時盡量保留每個平臺的原生性能和體驗。Google 在今年的 I/O 大會上也宣布對 KMP 提供一等的支持，把一些 Android 平臺上的庫和工具遷移到了多平臺，KMP 的開發者可以方便的使用它到 iOS 等其他平臺。

盡管 Mediapipe 也支持多個平臺，但我們這次主要聚焦在 Android 和 iOS。一方面更貼近現實，各行各業使用 KMP 的公司的用例更多在移動端上；另外一方面也更方便對標其他移動端開發技術棧。

移植流程

初始化

使用 IDEA 或 Android Studio 創建一個 KMP 的基礎工程，你可以借助 KMP Wizard 或者第三方 KMP App 的模版。如果你沒有 KMP 的相關經驗，可以看到它其實就是一個非常類似 Android 工程的結構，只不過這一次我們把 iOS 的殼工程也放到根目錄，并且在 app 模塊的 build.gradle.kts 內同時配置了 iOS 的相關依賴。

封裝和調用 LLM Inference

我們在 commonMain 中，根據 Mediapipe LLM Task SDK 的特征抽象一個簡單的接口，使用 Kotlin 編寫，用以滿足 Android 和 iOS 兩端的需要。該接口取代了原有倉庫里的 InferenceModel.kt 類。

// app/src/commonMain/.../llm/LLMOperator
interface LLMOperator {


    /**
     * To load the model into current context.
     * @return 1. null if it went well 2. an error message in string
     */
    suspend fun initModel(): String?


    fun sizeInTokens(text: String): Int


    suspend fun generateResponse(inputText: String): String


    suspend fun generateResponseAsync(inputText: String): Flow>


}

在 Android 上面，因為 LLM Task SDK 原先就是 Kotlin 實現的，所以除了初始化加載模型文件，其余的部分基本就是代理原有的 SDK 功能。

class LLMInferenceAndroidImpl(private val ctx: Context): LLMOperator {


    private lateinit var llmInference: LlmInference
    private val initialized = AtomicBoolean(false)
    private val partialResultsFlow = MutableSharedFlow>(...)


    override suspend fun initModel(): String? {
        if (initialized.get()) {
            return null
        }
        return try {
            val modelPath = ...
            if (File(modelPath).exists().not()) {
                return "Model not found at path: $modelPath"
            }
            loadModel(modelPath)
            initialized.set(true)
            null
        } catch (e: Exception) {
            e.message
        }
    }
    private fun loadModel(modelPath: String) {
        val options = LlmInference.LlmInferenceOptions.builder()
            .setModelPath(modelPath)
            .setMaxTokens(1024)
            .setResultListener { partialResult, done ->
                // Transforming the listener to flow,
                // making it easy on UI integration.
                partialResultsFlow.tryEmit(partialResult to done)
            }
            .build()


        llmInference = LlmInference.createFromOptions(ctx, options)
    }


    override fun sizeInTokens(text: String): Int = llmInference.sizeInTokens(text)


    override suspend fun generateResponse(inputText: String): String {
        ...
        return llmInference.generateResponse(inputText)
    }


    override suspend fun generateResponseAsync(inputText: String): Flow> {
        ...
        llmInference.generateResponseAsync(inputText)
        return partialResultsFlow.asSharedFlow()
    }


}

而針對 iOS，我們先嘗試第一種調用方式:直接調用 Cocoapods 引入的庫。在 app 模塊引入 cocoapods 的插件，同時添加 Mediapipe 的 LLM Task 庫:

// app/build.gradle.kts
plugins {
    ...
    alias(libs.plugins.cocoapods)
}
cocoapods {
    ...
    ios.deploymentTarget = "15"


    pod("MediaPipeTasksGenAIC") {
        version = "0.10.14"
        extraOpts += listOf("-compiler-option", "-fmodules")
    }
    pod("MediaPipeTasksGenAI") {
        version = "0.10.14"
        extraOpts += listOf("-compiler-option", "-fmodules")
    }
}

注意上面的引入配置中要添加一個編譯參數為 -fmodules 才可正常生成 Kotlin 的引用 (參考鏈接)。

一些 Objective-C 庫，尤其是那些作為 Swift 庫包裝器的庫，在它們的頭文件中使用了 @import 指令。默認情況下，cinterop 不支持這些指令。要啟用對 @import 指令的支持，可以在 pod() 函數的配置塊中指定 -fmodules 選項。

之后，我們在 iosMain 中便可直接 import 相關的庫代碼，如法炮制 Android 端的代理思路:

// 注意這些 import 是 cocoapods 開頭的
import cocoapods.MediaPipeTasksGenAI.MPPLLMInference
import cocoapods.MediaPipeTasksGenAI.MPPLLMInferenceOptions
import platform.Foundation.NSBundle
...
class LLMOperatorIOSImpl: LLMOperator {


    private val inference: MPPLLMInference


        init {
        val modelPath = NSBundle.mainBundle.pathForResource(..., "bin")


        val options = MPPLLMInferenceOptions(modelPath!!)
        options.setModelPath(modelPath!!)
        options.setMaxTokens(2048)
        options.setTopk(40)
        options.setTemperature(0.8f)
        options.setRandomSeed(102)


        // NPE was thrown here right after it printed the success initialization message internally.
        inference = MPPLLMInference(options, null) 
    }


    override fun generateResponse(inputText: String): String {...}
    override fun generateResponseAsync(inputText: String, ...) :... {
        ...
    }
    ...
}

但這回我們沒那么幸運，MPPLLMInference 初始化結束的一瞬間有 NPE 拋出。最可能的問題是因為 Kotlin 現在 interop 的目標是 Objective-C，MPPLLMInference 的構造器比 Swift 版本多一個 error 參數，而我們傳入的是 null。

constructor(
  options: cocoapods.MediaPipeTasksGenAI.MPPLLMInferenceOptions, 
error:CPointer>?)

但幾番測試各種指針傳入，也并未解決這個問題:

// 其中一種嘗試
memScoped {
    val pp: CPointerVar> = allocPointerTo()
    val inference = MPPLLMInference(options, pp.value)
    Napier.i(pp.value.toString())
}

于是只能另辟蹊徑采用第二種方案: 通過 iOS 工程調用第三方庫。

// 1. 聲明一個類似 LLMOperator 的接口但更簡單，方便適配 iOS 的 SDK。
// app/src/iosMain/.../llm/LLMOperator.kt
interface LLMOperatorSwift {
    suspend fun loadModel(modelName: String)
    fun sizeInTokens(text: String): Int
    suspend fun generateResponse(inputText: String): String
    suspend fun generateResponseAsync(
        inputText: String,
        progress: (partialResponse: String) -> Unit,
        completion: (completeResponse: String) -> Unit
    )
}


// 2. 在 iOS 工程里實現這個接口
// iosApp/iosApp/LLMInferenceDelegate.swift
class LLMOperatorSwiftImpl: LLMOperatorSwift {
    ...
    var llmInference: LlmInference?


    func loadModel(modelName: String) async throws {
        let path = Bundle.main.path(forResource: modelName, ofType: "bin")!
        let llmOptions =  LlmInference.Options(modelPath: path)
        llmOptions.maxTokens = 4096
        llmOptions.temperature = 0.9


        llmInference = try LlmInference(options: llmOptions)
    }


    func generateResponse(inputText: String) async throws -> String {
        return try llmInference!.generateResponse(inputText: inputText)
    }


    func generateResponseAsync(inputText: String, progress: @escaping (String) -> Void, completion: @escaping (String) -> Void) async throws {
        try llmInference!.generateResponseAsync(inputText: inputText) { partialResponse, error in
            // progress
            if let e = error {
                print("(self.errorTag) (e)")
                completion(e.localizedDescription)
                return
            }
            if let partial = partialResponse {
                progress(partial)
            }
        } completion: {
            completion("")
        }
    }
    ...    
}


// 3. iOS 再把代理好的（重點是初始化）類傳回給 Kotlin
// iosApp/iosApp/iosApp.swift
class AppDelegate: UIResponder, UIApplicationDelegate {
    ...
    func application(）{
        ...
        let delegate = try LLMOperatorSwiftImpl()
        MainKt.onStartup(llmInferenceDelegate: delegate)        
    }
}


// 4. 最初 iOS 在 KMP 上的實現細節直接代理給該對象（通過構造器注入）
class LLMOperatorIOSImpl(
   private val delegate: LLMOperatorSwift) : LLMOperator {   
   ...
}

細心的朋友可能已經發現，兩端的 Impl 實例需要不同的構造器參數，這個需求一般使用 KMP 的 expect 與 actual 關鍵字解決。下面的代碼中:

利用了 expect class 不需要構造器參數聲明的特點加了層封裝 (類似接口)。

利用了 Koin 實現各自平臺所需參數的注入，再統一把創建的接口實例注入到 Common 層所需的地方。

// Common
expect class LLMOperatorFactory {
    fun create(): LLMOperator
}
val sharedModule = module {
   // 從不同的 LLMOperatorFactory 創建出 Common 層所需的 LLMOperator
  single { get().create() }
}


// Android
actual class LLMOperatorFactory(private val context: Context){
    actual fun create(): LLMOperator = LLMInferenceAndroidImpl(context)
}
val androidModule = module {
    // Android 注入 App 的 Context
    single { LLMOperatorFactory(androidContext()) }
}


// iOS
actual class LLMOperatorFactory(private val llmInferenceDelegate: LLMOperatorSwift) {
    actual fun create(): LLMOperator = LLMOperatorIOSImpl(llmInferenceDelegate)
}


module {
    // iOS 注入 onStartup 函數傳入的 delegate
    single { LLMOperatorFactory(llmInferenceDelegate) }
}

小結: 我們通過一個小小的案例，領略到了 Kotlin 和 Swift 的深度交互。還借助 expect/actual 關鍵字與 Koin 的依賴注入，讓整體方案更流暢和自動化，達到了在 KMP 的 Common 模塊調用 Android 和 iOS Native SDK 的目標。

移植 UI 和 ViewModel

原項目里的 InferenceMode 已經被上一節的 LLMOperator 所取代，因此我們拷貝除 Activity 的剩下 5 個類:

下面我們修改幾處代碼使 Jetpack Compose 的代碼可以方便的遷移到 Compose Multiplatform。

首先是外圍的 ViewModel，KMP 版本我在這里使用了 Voyage，因此替換為 ScreenModel。不過官方 ViewModel 的方案也在實驗中了，請參考這個文檔。

// Android 版本
class ChatViewModel(
    private val inferenceModel: InferenceModel
) : ViewModel() {...}


// KMP 版本，轉換 ViewModel 為 ScreenModel，并修改傳入對象
class ChatViewModel(
    private val llmOperator: LLMOperator
):ScreenModel{...}

Voyage https://github.com/adrielcafe/voyager

文檔 https://www.jetbrains.com/help/kotlin-multiplatform-dev/compose-viewmodel.html

相應的 ViewModel 初始化方式也更改成 ScreenModel 的方法:

// Android 版本
@Composable
internal fun ChatRoute(
    chatViewModel: ChatViewModel = viewModel(
        factory = ChatViewModel.getFactory(LocalContext.current.applicationContext)
    )
) {
    ...
    ChatScreen(...) {...}
}


// KMP 版本，改成外部初始化后傳入
@Composable
internal fun ChatRoute(
    chatViewModel: ChatViewModel
) {


// 此處采用了默認參數注入的方案，便于解耦。
// koinInject() 是 Koin 官方提供的針對 Compose 
// 的 @Composable 函數注入的一個方法。
@Composable
fun AiScreen(llmOperator:LLMOperator = koinInject()) {
    // 使用 ScreenModel 的 remember 方法
    val chatViewModel = rememberScreenModel { ChatViewModel(llmOperator) }
    ...
    Column {
        ...
        Box(...) {
            if (showLoading) {
                ...
            } else {
                ChatRoute(chatViewModel)
            }
        }
    }
}

對應的 ViewModel 內部的 LLM 功能調用接口也要進行替換:

// Android 版本
inferenceModel.generateResponseAsync(fullPrompt)
inferenceModel.partialResults
    .collectIndexed { index, (partialResult, done) ->
        ...
    }


// KMP 版本，把 Flow 的返回前置了，兼容了兩個平臺的 SDK 設計
llmOperator.generateResponseAsync(fullPrompt)
    .collectIndexed { index, (partialResult, done) ->
        ...
    }

然后是 Compose Multiplatform 特定的資源加載方式，把 R 文件替換為 Res:

// Android 版本
Text(stringResource(R.string.chat_label))


// KMP 版本，該引用是使用插件從 xml 映射而來
// (commonMain/composeResources/values/strings.xml)
import mediapiper.app.generated.resources.chat_label
...
Text(stringResource(Res.string.chat_label))

至此我們已經完成了 ChatScreen ChatViewModel 的主頁面功能遷移。

最后是其他的幾個輕微改動:

LoadingScreen 我們如法炮制傳入 LLMOperator 進行初始化 (替換原有 InferenceModel)。

ChatMessage 只需修改了 UUID 調用的一行 API 到原生實現 (Kotlin 2.0.20 后就不需要了)。

ChatUiState 則完全不用動。

剩下的就只有整體修改下 Log 庫的引用等小細節。

小結: 倘若略去 Log、R 文件的引用替換以及 import 替換等，核心的修改其實僅十幾行，便能把整個 UI 部分也跑起來了。

簡單測試

那 Gemma 2B 的性能如何，我們看幾個簡單的例子。此處主要使用三個版本的模型進行測試，模型的定義在 me.xx2bab.mediapiper.llm.LLMOperator (模型在兩端部署請參考項目 README)。

gemma-2b-it-gpu-int4

gemma-2b-it-cpu-int4

gemma-2b-it-cpu-int8

其中:

it 指代一種變體，即 Instruction Tuned 模型，更適合聊天用途，因為它們經過微調能更好地理解指令，并生成更準確的回答。

int4/8 指代模型量化，即將模型中的浮點數轉換為低精度整數，從而減小模型的大小和計算量以適配小型的本地設備例如手機。當然，模型的精度和回答準確度也會有一些下降。

CPU 和 GPU指針對的硬件平臺，這方便了設備 GPU 較弱甚至沒有時可選擇 CPU 執行。從下面的測試結果你會發現當前移動設備上 CPU 版本也常常會占優，因為模型規模小、簡單對話計算操作也不大，并且 Int 量化也有利于 CPU 的指令執行。

首先我們測試一個簡單的邏輯: "蘆筍是不是一種動物"？可以看到下圖的 CPU 版本答案比兩個 GPU (iOS 和 Android) 更合理。而下一個測試是翻譯答案為中文，則是三個嘗試都不太行。

接著我們提高了測試問題的難度，讓它執行區分動植物的單詞分類: 不管是 GPU 或者 CPU 的版本都不錯。

再次升級上個問題，讓它用 JSON 的方式輸出答案，就出現明顯的問題:

圖 1 沒有輸出完整的代碼片段，缺少了結尾的三個點 ```。

圖二分類錯誤，把山竹放到動物，植物出現了兩次向日葵。

圖三同二的錯誤，但這三次都沒有純輸出一個 JSON，實際上還是不夠嚴格執行作為 JSON Responder 的角色。

最后，這其實不是極限，如果我們使用 cpu-int8 的版本，則可以高準確率地解答上面問題。以及，如果把本 Demo 的 iOS 入口代碼發送給它分析，也能答的不錯。

Gemma 1 的 2B 版本測試至此，我們發覺其推理效果還有不少進步空間，勝在回復速度不錯。而事實上 Gemma 2 的 2B 版本前不久已推出，并且據官方測試其綜合水平已超過 GPT 3.5。這意味著在一臺小小的手機里，本地的推理已經可以達到一年半前的主流模型效果。總結實現這個本地聊天 Demo 的遷移和測試，給了我們些一手的經驗:

LLM 的 On-Device Model 發展非常迅速，而借助 Google 的一系列基礎設施可以讓第三方 Mobile App 開發者也迅速地集成相關的功能，并跨越 Android 與 iOS 雙平臺。

觀望目前情況綜合判斷，LLM 的 On-Device Model 有望在今年達到初步可用狀態，推理速度已經不錯，準確度還有待進一步測試 (例如 Gemma 2 的 2B 版本 + Mediapipe)。

遵循 Android 團隊目前的策略 "Kotlin First"并大膽使用 Compose，是頗具前景的——在基礎設施完備的情況下，一個聊天的小模塊僅寥寥數行修改即可遷移到 iOS。