Android で LiteRT-LM を使ってみる

AndroidJVM(Linux、MacOS、Windows)用の LiteRT-LM の Kotlin API。GPU と NPU のアクセラレーションマルチモーダルツール使用などの機能があります。

はじめに

Kotlin API で構築されたターミナル チャットアプリのサンプルを次に示します。

import com.google.ai.edge.litertlm.*

suspend fun main() {
  Engine.setNativeMinLogSeverity(LogSeverity.ERROR) // Hide log for TUI app

  val engineConfig = EngineConfig(modelPath = "/path/to/model.litertlm")
  Engine(engineConfig).use { engine ->
    engine.initialize()

    engine.createConversation().use { conversation ->
      while (true) {
        print("\n>>> ")
        conversation.sendMessageAsync(readln()).collect { print(it) }
      }
    }
  }
}

Kotlin サンプルコードのデモ

上記のサンプルを試すには、リポジトリのクローンを作成し、example/Main.kt で実行します。

bazel run -c opt //kotlin/java/com/google/ai/edge/litertlm/example:main -- <abs_model_path>

使用可能な .litertlm モデルは、HuggingFace LiteRT コミュニティにあります。上記のアニメーションでは Gemma3-1B-IT を使用しています。

Android のサンプルについては、Google AI Edge ギャラリー アプリをご覧ください。

Gradle のスタートガイド

LiteRT-LM は Bazel で開発されていますが、Gradle/Maven ユーザー向けに Maven パッケージを提供しています。

1. Gradle 依存関係を追加する

dependencies {
    // For Android
    implementation("com.google.ai.edge.litertlm:litertlm-android:latest.release")

    // For JVM (Linux, MacOS, Windows)
    implementation("com.google.ai.edge.litertlm:litertlm-jvm:latest.release")
}

利用可能なバージョンは、Google Maven の litertlm-androidlitertlm-jvm で確認できます。

latest.release を使用して最新のリリースを取得できます。

2. エンジンを初期化する

Engine は API のエントリ ポイントです。モデルのパスと構成で初期化します。リソースを解放するために、エンジンを閉じる必要があります。

注: engine.initialize() メソッドは、モデルの読み込みにかなりの時間(最大 10 秒など)を要することがあります。UI スレッドのブロックを回避するため、バックグラウンド スレッドまたはコルーチンでこのメソッドを呼び出すことを強くおすすめします。

import com.google.ai.edge.litertlm.Backend
import com.google.ai.edge.litertlm.Engine
import com.google.ai.edge.litertlm.EngineConfig

val engineConfig = EngineConfig(
    modelPath = "/path/to/your/model.litertlm", // Replace with your model path
    backend = Backend.GPU(), // Or Backend.NPU(nativeLibraryDir = "...")
    // Optional: Pick a writable dir. This can improve 2nd load time.
    // cacheDir = "/tmp/" or context.cacheDir.path (for Android)
)

val engine = Engine(engineConfig)
engine.initialize()
// ... Use the engine to create a conversation ...

// Close the engine when done
engine.close()

Android で GPU バックエンドを使用するには、<application> タグ内の AndroidManifest.xml に次の行を追加して、依存するネイティブ ライブラリを明示的にリクエストする必要があります。

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

NPU バックエンドを使用するには、NPU ライブラリを含むディレクトリを指定する必要がある場合があります。Android では、ライブラリがアプリにバンドルされている場合は、context.applicationInfo.nativeLibraryDir に設定します。NPU ネイティブ ライブラリの詳細については、LiteRT-LM NPU をご覧ください。

val engineConfig = EngineConfig(
    modelPath = modelPath,
    backend = Backend.NPU(nativeLibraryDir = context.applicationInfo.nativeLibraryDir)
)

3. 会話を作成する

エンジンを初期化したら、Conversation インスタンスを作成します。ConversationConfig を指定して、その動作をカスタマイズできます。

import com.google.ai.edge.litertlm.ConversationConfig
import com.google.ai.edge.litertlm.Message
import com.google.ai.edge.litertlm.SamplerConfig

// Optional: Configure the system instruction, initial messages, sampling
// parameters, etc.
val conversationConfig = ConversationConfig(
    systemInstruction = Contents.of("You are a helpful assistant."),
    initialMessages = listOf(
        Message.user("What is the capital city of the United States?"),
        Message.model("Washington, D.C."),
    ),
    samplerConfig = SamplerConfig(topK = 10, topP = 0.95, temperature = 0.8),
)

val conversation = engine.createConversation(conversationConfig)
// Or with default config:
// val conversation = engine.createConversation()

// ... Use the conversation ...

// Close the conversation when done
conversation.close()

ConversationAutoCloseable を実装するため、1 回限りの会話や短期間の会話のリソース管理を自動化するために use ブロックを使用できます。

engine.createConversation(conversationConfig).use { conversation ->
    // Interact with the conversation
}

4. メッセージの送信

メッセージを送信するには、次の 3 つの方法があります。

  • sendMessage(contents): Message: モデルが完全なレスポンスを返すまでブロックする同期呼び出し。これは、基本的なリクエスト/レスポンスのやり取りに適しています。
  • sendMessageAsync(contents, callback): レスポンスをストリーミングするための非同期呼び出し。これは、長時間実行されるリクエストや、レスポンスの生成中にレスポンスを表示する場合に適しています。
  • sendMessageAsync(contents): Flow<Message>: レスポンスをストリーミングするための Kotlin Flow を返す非同期呼び出し。これは、コルーチン ユーザーにおすすめの方法です。

同期の例:

import com.google.ai.edge.litertlm.Content
import com.google.ai.edge.litertlm.Message

print(conversation.sendMessage("What is the capital of France?"))

コールバックを使用した非同期の例:

sendMessageAsync を使用してモデルにメッセージを送信し、コールバックを通じてレスポンスを受け取ります。

import com.google.ai.edge.litertlm.Content
import com.google.ai.edge.litertlm.Message
import com.google.ai.edge.litertlm.MessageCallback
import java.util.concurrent.CountDownLatch
import java.util.concurrent.TimeUnit

val callback = object : MessageCallback {
    override fun onMessage(message: Message) {
        print(message)
    }

    override fun onDone() {
        // Streaming completed
    }

    override fun onError(throwable: Throwable) {
        // Error during streaming
    }
}

conversation.sendMessageAsync("What is the capital of France?", callback)

フローを使用した非同期の例:

sendMessageAsync(コールバック引数なし)を使用してモデルにメッセージを送信し、Kotlin Flow を介してレスポンスを受け取ります。

import com.google.ai.edge.litertlm.Content
import com.google.ai.edge.litertlm.Message
import kotlinx.coroutines.flow.catch
import kotlinx.coroutines.launch

// Within a coroutine scope
conversation.sendMessageAsync("What is the capital of France?")
    .catch { ... } // Error during streaming
    .collect { print(it.toString()) }

5. マルチモダリティ

Message オブジェクトには、TextImageBytesImageFileAudioBytesAudioFile など、さまざまなタイプの Content を含めることができます。

// Initialize the `visionBackend` and/or the `audioBackend`
val engineConfig = EngineConfig(
    modelPath = "/path/to/your/model.litertlm", // Replace with your model path
    backend = Backend.CPU(), // Or Backend.GPU() or Backend.NPU(...)
    visionBackend = Backend.GPU(), // Or Backend.NPU(...)
    audioBackend = Backend.CPU(), // Or Backend.NPU(...)
)

// Sends a message with multi-modality.
// See the Content class for other variants.
conversation.sendMessage(Contents.of(
    Content.ImageFile("/path/to/image"),
    Content.AudioBytes(audioBytes), // ByteArray of the audio
    Content.Text("Describe this image and audio."),
))

6. ツールの定義と使用

ツールを定義する方法は次の 2 つです。

  1. Kotlin 関数を使用する(ほとんどの場合におすすめ)
  2. Open API 仕様(ツール仕様と実行を完全に制御)

Kotlin 関数でツールを定義する

モデルが呼び出してアクションを実行したり、情報を取得したりできるツールとして、カスタム Kotlin 関数を定義できます。

ToolSet を実装するクラスを作成し、メソッドに @Tool、パラメータに @ToolParam のアノテーションを付けます。

import com.google.ai.edge.litertlm.Tool
import com.google.ai.edge.litertlm.ToolParam

class SampleToolSet: ToolSet {
    @Tool(description = "Get the current weather for a city")
    fun getCurrentWeather(
        @ToolParam(description = "The city name, e.g., San Francisco") city: String,
        @ToolParam(description = "Optional country code, e.g., US") country: String? = null,
        @ToolParam(description = "Temperature unit (celsius or fahrenheit). Default: celsius") unit: String = "celsius"
    ): Map<String, Any> {
        // In a real application, you would call a weather API here
        return mapOf("temperature" to 25, "unit" to  unit, "condition" to "Sunny")
    }

    @Tool(description = "Get the sum of a list of numbers.")
    fun sum(
        @ToolParam(description = "The numbers, could be floating point.") numbers: List<Double>,
    ): Double {
        return numbers.sum()
    }
}

API は、これらのアノテーションと関数シグネチャを検査して、OpenAPI スタイルのスキーマを生成します。このスキーマは、ツールの機能、パラメータ(@ToolParam の型と説明を含む)、戻り値の型を言語モデルに記述します。

パラメータ タイプ

@ToolParam でアノテーションが付けられたパラメータの型は、StringIntBooleanFloatDouble、またはこれらの型の List(例: List<String>)を使用します。null 許容型(String?)を使用して、null 許容パラメータを示します。パラメータが省略可能であることを示すデフォルト値を設定し、@ToolParam の説明でデフォルト値について言及します。

戻り値の型

ツール関数の戻り値の型は、任意の Kotlin 型にできます。結果は JSON 要素に変換されてから、モデルに送り返されます。

  • List 型は JSON 配列に変換されます。
  • Map 型は JSON オブジェクトに変換されます。
  • プリミティブ型(StringNumberBoolean)は、対応する JSON プリミティブに変換されます。
  • 他の型は toString() メソッドで文字列に変換されます。

構造化データの場合は、Map または JSON オブジェクトに変換されるデータクラスを返すことをおすすめします。

OpenAPI 仕様を使用したツールの定義

また、OpenApiTool クラスを実装し、Open API 仕様に準拠した JSON 文字列としてツールの説明を提供することで、ツールを定義することもできます。この方法は、ツールの OpenAPI スキーマがすでに存在する場合や、ツールの定義をきめ細かく制御する必要がある場合に便利です。

import com.google.ai.edge.litertlm.OpenApiTool

class SampleOpenApiTool : OpenApiTool {

    override fun getToolDescriptionJsonString(): String {
        return """
        {
          "name": "addition",
          "description": "Add all numbers.",
          "parameters": {
            "type": "object",
            "properties": {
              "numbers": {
                "type": "array",
                "items": {
                  "type": "number"
                }
              },
              "description": "The list of numbers to sum."
            },
            "required": [
              "numbers"
            ]
          }
        }
        """.trimIndent() // Tip: trim to save tokens
    }

    override fun execute(paramsJsonString: String): String {
        // Parse paramsJsonString with your choice of parser/deserializer and
        // execute the tool.

        // Return the result as a JSON string
        return """{"result": 1.4142}"""
    }
}

登録ツール

ConversationConfig にツールのインスタンスを含めます。

val conversation = engine.createConversation(
    ConversationConfig(
        tools = listOf(
            tool(SampleToolSet()),
            tool(SampleOpenApiTool()),
        ),
        // ... other configs
    )
)

// Send messages that might trigger the tool
conversation.sendMessageAsync("What's the weather like in London?", callback)

モデルは、会話に基づいてツールを呼び出すタイミングを決定します。ツール実行の結果は、最終的なレスポンスを生成するためにモデルに自動的に送り返されます。

手動ツール呼び出し

デフォルトでは、モデルによって生成されたツール呼び出しは LiteRT-LM によって自動的に実行され、ツール実行の結果はモデルに自動的に送り返されて、次のレスポンスが生成されます。

ツールを手動で実行して結果をモデルに送信する場合は、ConversationConfigautomaticToolCallingfalse に設定します。

val conversation = engine.createConversation(
    ConversationConfig(
        tools = listOf(
            tool(SampleOpenApiTool()),
        ),
        automaticToolCalling = false,
    )
)

ツールの自動呼び出しを無効にした場合は、アプリケーション コードでツールを手動で実行し、結果をモデルに送り返す必要があります。automaticToolCallingfalse に設定されている場合、OpenApiToolexecute メソッドは自動的に呼び出されません。

// Send a message that triggers a tool call.
val responseMessage = conversation.sendMessage("What's the weather like in London?")

// The model returns a Message with `toolCalls` populated.
if (responseMessage.toolCalls.isNotEmpty()) {
    val toolResponses = mutableListOf<Content.ToolResponse>()
    // There can be multiple tool calls in a single response.
    for (toolCall in responseMessage.toolCalls) {
        println("Model wants to call: ${toolCall.name} with arguments: ${toolCall.arguments}")

        // Execute the tool manually with your own logic. `executeTool` is just an example here.
        val toolResponseJson = executeTool(toolCall.name, toolCall.arguments)

        // Collect tool responses.
        toolResponses.add(Content.ToolResponse(toolCall.name, toolResponseJson))
    }

    // Use Message.tool to create the tool response message.
    val toolResponseMessage = Message.tool(Contents.of(toolResponses))

    // Send the tool response message to the model.
    val finalMessage = conversation.sendMessage(toolResponseMessage)
    println("Final answer: ${finalMessage.text}") // e.g., "The weather in London is 25c."
}

ツール使用を試すには、リポジトリのクローンを作成し、example/ToolMain.kt で実行します。

bazel run -c opt //kotlin/java/com/google/ai/edge/litertlm/example:tool -- <abs_model_path>

エラー処理

API メソッドは、ネイティブ レイヤのエラーに対して LiteRtLmJniException をスローしたり、ライフサイクルに関する問題に対して IllegalStateException などの標準の Kotlin 例外をスローしたりする可能性があります。API 呼び出しは常に try-catch ブロックでラップします。MessageCallbackonError コールバックは、非同期オペレーション中のエラーも報告します。