Primeros pasos con LiteRT-LM en Android

La API de Kotlin de LiteRT-LM para Android y JVM (Linux, macOS, Windows) con funciones como aceleración de GPU y NPU, multimodalidad y uso de herramientas.

Introducción

Aquí tienes una app de chat de terminal de muestra creada con la API de Kotlin:

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) }
      }
    }
  }
}

Demostración del código de muestra de Kotlin

Para probar la muestra anterior, clona el repositorio y ejecuta con example/Main.kt:

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

Los modelos .litertlm disponibles se encuentran en la comunidad de HuggingFace LiteRT. La animación anterior usaba el Gemma3-1B-IT.

Para obtener una muestra de Android, consulta la app de Google AI Edge Gallery (disponible en Google Play).

Cómo comenzar a usar Gradle

Si bien LiteRT-LM se desarrolla con Bazel, proporcionamos los paquetes de Maven para los usuarios de Gradle o Maven.

Agrega la dependencia de 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")
}

Puedes encontrar las versiones disponibles en Google Maven en litertlm-android y litertlm-jvm.

Se puede usar latest.release para obtener la versión más reciente.

Inicializa el motor

Engine es el punto de entrada a la API. Inicialízalo con la ruta de acceso y la configuración del modelo. Recuerda cerrar el motor para liberar recursos.

Nota: El método engine.initialize() puede tardar una cantidad significativa de tiempo (p.ej., hasta 10 segundos) en cargar el modelo. Se recomienda encarecidamente llamar a este método en un subproceso o una corrutina en segundo plano para evitar bloquear el subproceso de la IU.

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()

En Android, para usar el backend de GPU, la app debe solicitar las bibliotecas nativas dependientes de forma explícita agregando lo siguiente a tu AndroidManifest.xml dentro de la etiqueta <application> tag:

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

Para usar el backend de NPU, es posible que debas especificar el directorio que contiene las bibliotecas de NPU. En Android, si las bibliotecas están empaquetadas con tu app, configúralo como context.applicationInfo.nativeLibraryDir. Consulta NPU de LiteRT-LM para obtener más detalles sobre las bibliotecas nativas de NPU.

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

Crea una conversación

Una vez que se inicializa el motor, crea una instancia de Conversation. Puedes proporcionar un ConversationConfig para personalizar su comportamiento.

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()

Conversation implementa AutoCloseable, por lo que puedes usar el bloque use para la administración automática de recursos para conversaciones únicas o de corta duración:

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

Envío de mensajes

Hay tres maneras de enviar mensajes:

  • sendMessage(contents): Message: Es una llamada síncrona que se bloquea hasta que el modelo muestra una respuesta completa. Es más simple para las interacciones básicas de solicitud y respuesta.
  • sendMessageAsync(contents, callback): Es una llamada asíncrona para transmitir respuestas. Es mejor para solicitudes de larga duración o cuando deseas mostrar la respuesta a medida que se genera.
  • sendMessageAsync(contents): Flow<Message>: Es una llamada asíncrona que muestra un flujo de Kotlin para transmitir respuestas. Este es el enfoque recomendado para los usuarios de Coroutine.

Ejemplo síncrono:

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

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

Ejemplo asíncrono con devolución de llamada:

Usa sendMessageAsync para enviar un mensaje al modelo y recibir respuestas a través de una devolución de llamada.

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)

Ejemplo asíncrono con flujo:

Usa sendMessageAsync (sin el argumento de devolución de llamada) para enviar un mensaje al modelo y recibir respuestas a través de un flujo de Kotlin.

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()) }

🔴 Nuevo: Predicción de varios tokens (MTP)

La predicción de varios tokens (MTP) es una optimización del rendimiento que acelera significativamente las velocidades de decodificación. Se recomienda MTP de forma universal para todas las tareas en backends de GPU.

Para usar MTP, habilita la decodificación especulativa con ExperimentalFlags antes de inicializar el motor.

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

// Enable MTP via speculative decoding
@OptIn(ExperimentalApi::class)
ExperimentalFlags.enableSpeculativeDecoding = true

val engineConfig = EngineConfig(
    modelPath = "/path/to/your/model.litertlm",
    backend = Backend.GPU(),
)

val engine = Engine(engineConfig)
engine.initialize()

// The same steps to create Conversation and send messages as below...

Multimodalidad

Los objetos Message pueden contener diferentes tipos de Content, incluidos Text, ImageBytes, ImageFile, AudioBytes y AudioFile.

// Initialize the `visionBackend`, `audioBackend`, or both
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."),
))

Definición y uso de herramientas

Hay dos maneras de definir herramientas:

  1. Con funciones de Kotlin (recomendado para la mayoría de los casos)
  2. Con la especificación de Open API (control total de la especificación y la ejecución de la herramienta)

Definición de herramientas con funciones de Kotlin

Puedes definir funciones personalizadas de Kotlin como herramientas que el modelo puede llamar para realizar acciones o recuperar información.

Crea una clase que implemente ToolSet y anota los métodos con @Tool y los parámetros con @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()
    }
}

Tras bambalinas, la API inspecciona estas anotaciones y la firma de la función para generar un esquema de estilo OpenAPI. Este esquema describe la funcionalidad de la herramienta, los parámetros (incluidos sus tipos y descripciones de @ToolParam) y el tipo de datos que se devuelve al modelo de lenguaje.

Tipos de parámetros

Los tipos de parámetros anotados con @ToolParam pueden ser String, Int, Boolean, Float, Double o una List de estos tipos (p.ej., List<String>). Usa tipos que acepten valores nulos (p.ej., String?) para indicar parámetros que acepten valores nulos. Establece un valor predeterminado para indicar que el parámetro es opcional y menciona el valor predeterminado en la descripción de @ToolParam.

Tipo de datos que se muestra

El tipo de datos que se devuelve de la función de la herramienta puede ser cualquier tipo de Kotlin. El resultado se convertirá en un elemento JSON antes de enviarse de vuelta al modelo.

  • Los tipos List se convierten en arrays JSON.
  • Los tipos Map se convierten en objetos JSON.
  • Los tipos primitivos (String, Number, Boolean) se convierten en el JSON primitivo correspondiente.
  • Otros tipos se convierten en cadenas con el método toString().

Para los datos estructurados, se recomienda mostrar Map o una clase de datos que se convertirá en un objeto JSON.

Definición de herramientas con la especificación de OpenAPI

Como alternativa, puedes definir una herramienta implementando la clase OpenApiTool y proporcionando la descripción de la herramienta como una cadena JSON que cumpla con la especificación de Open API. Este método es útil si ya tienes un esquema de OpenAPI para tu herramienta o si necesitas un control detallado sobre la definición de la herramienta.

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 or deserializer and
        // execute the tool.

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

Registro de herramientas

Incluye instancias de tus herramientas en 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)

El modelo decidirá cuándo llamar a la herramienta según la conversación. Los resultados de la ejecución de la herramienta se envían automáticamente al modelo para generar la respuesta final.

Llamadas manuales a herramientas

De forma predeterminada, LiteRT-LM ejecuta automáticamente las llamadas a herramientas generadas por el modelo, y los resultados de la ejecución de la herramienta se envían automáticamente al modelo para generar la siguiente respuesta.

Si deseas ejecutar herramientas de forma manual y enviar los resultados al modelo, puedes configurar automaticToolCalling en ConversationConfig como false.

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

Si inhabilitas las llamadas automáticas a herramientas, deberás ejecutar las herramientas de forma manual y enviar los resultados al modelo en el código de tu aplicación. El método execute de OpenApiTool no se llamará automáticamente cuando automaticToolCalling se establezca en false.

// 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."
}

Ejemplo

Para probar el uso de herramientas, clona el repositorio y ejecuta con example/ToolMain.kt:

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

Manejo de errores

Los métodos de la API pueden arrojar LiteRtLmJniException para errores de la capa nativa o excepciones estándar de Kotlin, como IllegalStateException para problemas del ciclo de vida. Siempre incluye las llamadas a la API en bloques try-catch. La devolución de llamada onError en MessageCallback también informará errores durante las operaciones asíncronas.