移植Mediapipe LLM Demo到Kotlin Multiplatform

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

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

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

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

KMP 兩種常見的調(diào)用 iOS SDK 的方式:

Kotlin 直接調(diào)用 Cocoapods 引入的第三方庫。

Kotlin 通過 iOS 工程調(diào)用第三方庫。

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

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

On-Device Model 本地模型

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

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

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

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

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

LLM Inference

Android Sample

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

打開 Android 倉庫會發(fā)現(xiàn)幾個特點:

純 Kotlin 實現(xiàn)。

UI 是純 Jetpack Compose 實現(xiàn)。

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

再查看 iOS 的版本:

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

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

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

Kotlin Multiplatform

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

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

移植流程

初始化

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

封裝和調(diào)用 LLM Inference

我們在 commonMain 中，根據(jù) 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 實現(xiàn)的，所以除了初始化加載模型文件，其余的部分基本就是代理原有的 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，我們先嘗試第一種調(diào)用方式:直接調(diào)用 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")
    }
}

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

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

之后，我們在 iosMain 中便可直接 import 相關(guān)的庫代碼，如法炮制 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 初始化結(jié)束的一瞬間有 NPE 拋出。最可能的問題是因為 Kotlin 現(xiàn)在 interop 的目標是 Objective-C，MPPLLMInference 的構(gòu)造器比 Swift 版本多一個 error 參數(shù)，而我們傳入的是 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 工程調(diào)用第三方庫。

// 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 工程里實現(xiàn)這個接口
// 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 上的實現(xiàn)細節(jié)直接代理給該對象（通過構(gòu)造器注入）
class LLMOperatorIOSImpl(
   private val delegate: LLMOperatorSwift) : LLMOperator {   
   ...
}

細心的朋友可能已經(jīng)發(fā)現(xiàn)，兩端的 Impl 實例需要不同的構(gòu)造器參數(shù)，這個需求一般使用 KMP 的 expect 與 actual 關(guān)鍵字解決。下面的代碼中:

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

利用了 Koin 實現(xiàn)各自平臺所需參數(shù)的注入，再統(tǒng)一把創(chuàng)建的接口實例注入到 Common 層所需的地方。

// Common
expect class LLMOperatorFactory {
    fun create(): LLMOperator
}
val sharedModule = module {
   // 從不同的 LLMOperatorFactory 創(chuàng)建出 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 函數(shù)傳入的 delegate
    single { LLMOperatorFactory(llmInferenceDelegate) }
}

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

移植 UI 和 ViewModel

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

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

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

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


// KMP 版本，轉(zhuǎn)換 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
) {


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

對應的 ViewModel 內(nèi)部的 LLM 功能調(diào)用接口也要進行替換:

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


// KMP 版本，把 Flow 的返回前置了，兼容了兩個平臺的 SDK 設(shè)計
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))

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

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

LoadingScreen 我們?nèi)绶ㄅ谥苽魅?LLMOperator 進行初始化 (替換原有 InferenceModel)。

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

ChatUiState 則完全不用動。

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

小結(jié): 倘若略去 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 模型，更適合聊天用途，因為它們經(jīng)過微調(diào)能更好地理解指令，并生成更準確的回答。

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

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

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

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

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

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

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

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

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

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

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

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

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