![](https://www.droidcon.com/wp-content/uploads/2024/07/0_REasIJSMVUyTNMFf-600x401.webp)
A while ago, I developed an Android app, Android-Doc-QA which is an instance of on-device RAG for PDF/DOCX documents. It used ObjectBox as a vector database, Gemini Cloud API as a LLM and Mediapipe’s Text Embedder as an embedding provider.
For RAG applications to perform well, having a good embedding model that performs a text to vector/embedding transformation is crucial. The embeddings must capture the semantics of the text as they are used to find the most similar chunks (text subsequences from the documents given) from the database through a vector search. Chunks that are most similar to the query serve as context to the LLM and go in the prompt directly.
From my experience in testing the app, Mediapipe’s Text Embedder that uses Google’s Universal Sentence Encoder was not good enough to understand the semantics of the text/chunk or the query and felt more like a text search matching word occurrences rather than understanding the true meaning of the sentences. Moreover, the size of the embeddings was 100 that felt less when compared to models from HuggingFace sentence-transformers.
The all-MiniLM-L6-V2
model from sentence-transformers seemed powerful from the RAG applications that I built in Python. It had an embedding size of 384, which can theoretically capture more information than an embedding of size 100. I wanted to use this model in my Android, but I couldn’t find any clear implementation doing so.
In this blog, I’ll share my thoughts on how I solved the problem on bringing the all-MiniLM-L6-V2
to an Android app and build an Android library and demo app for the same,
![](https://miro.medium.com/v2/resize:fit:640/format:webp/1*jJwVYaUyL3yTsSE36MDFlQ.gif)
Here’s the GitHub repository for the library and demo app:
Approach
In order to execute any ML model on Android, it is good to have the model converted to the ONNX or TensorFlow Lite format. For ONNX, Microsoft provides onnxruntime
the TensorFlow team provides an Android library to execute TFLite models. Both of these libraries are well-documented and have numerous examples for models operating on different modalities.
Solving Model Execution — ONNX
In our case, the sentence-transformers
team has already provided an ONNX version of theall-MiniLM-L6-V2
model in their HuggingFace repository. With onnxruntime
‘s Python package, we can load the ONNX model and check its input/output shapes, data-types and tensor-names,
use tokenizers::tokenizer::Tokenizer;
fn main() {
let sent_bert_tokenizer = Tokenizer::from_file("hf-tokenizer/tokenizer.json").unwrap();
let encoding = sent_bert_tokenizer.encode("what is the population of london", true).unwrap();
let ids: Vec<i64> = encoding.get_ids().iter().map(|id| *id as i64).collect();
let attention_mask: Vec<i64> = encoding
.get_attention_mask()
.iter()
.map(|x| *x as i64)
.collect();
println!( "{:?}" , ids ) ;
println!( "{:?}" , attention_mask ) ;
}
[101, 2054, 2003, 1996, 2313, 1997, 2414, 102, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
[1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
Rust supports compilation to Android-based targets and can build a dynamic/static C library also. Similar to C/C++, we can make use of JNI (Java Native Interface) in Rust and expose functions that interact with the Tokenizer
in Rust. Java/Kotlin classes can declare these functions as native
/ external
and search for their implementations in dynamic libraries (shared objects, .so
) we’ll build from the Rust source code.
This Rust-Java interface will help us access the Rust-based Tokenizer
from Android and we will be able to tokenize a String
and obtain input_ids
and attention_mask
for the ONNX model.
JNI and Rust Library
To start building the Rust source code, we add tokenizers
, jni
and serde_json
as dependencies in our crate and modify the crate_type
to cdylib
(C dynamic library),
[dependencies]
tokenizers = "0.19.1"
jni = { version = "0.21.1" }
serde = { version = "1.0.192" , features = ["derive"] }
serde_json = "1.0.108"
onig = "6.4.0"
[lib]
crate-type = ["cdylib"]
name = "hftokenizer"
Next in src/lib.rs
, with JNI we define functions that create and destroy the tokenizer instance and which use the instance to tokenize the given text. By instance, we mean a pointer/address that references the Tokenizer
object in memory. The pointer will be represented with a Box
in Rust and as a Long
in Kotlin. The first function, createTokenizer
, takes the file-bytes of tokenizer.json
and returns the instance pointer,
use jni::objects::{JByteArray, JClass, JString, ReleaseMode};
use jni::sys::{jbyteArray, jlong};
use jni::JNIEnv;
use serde::Serialize;
use serde_json;
use tokenizers::Tokenizer;
#[no_mangle]
pub extern "C" fn Java_com_ml_shubham0204_sentence_1embeddings_HFTokenizer_createTokenizer<'a>(
mut env: JNIEnv<'a>,
_: JClass<'a>,
tokenizer_bytes: JByteArray<'a>,
) -> jlong {
let tokenizer_bytes_rs: Vec<u8> = env
.get_array_elements(&tokenizer_bytes, ReleaseMode::CopyBack)
.expect("Could not read tokenizer_bytes")
.into();
Box::into_raw(Box::new(Tokenizer::from_bytes(&tokenizer_bytes_rs))) as jlong
}
Next, the tokenize
function accepts an instance pointer and a text, returning a JSON string contains the ids
and attention_mask
,
#[no_mangle]
pub extern "C" fn Java_com_ml_shubham0204_sentence_1embeddings_HFTokenizer_tokenize<'a>(
mut env: JNIEnv<'a>,
_: JClass<'a>,
tokenizer_ptr: jlong,
text: JString<'a>,
) -> JString<'a> {
let tokenizer = unsafe { &mut *(tokenizer_ptr as *mut Tokenizer) };
let text: String = env
.get_string(&text)
.expect("Could not convert text to Rust String")
.into();
let encoding = tokenizer.encode(text, true).expect("Could not encode text");
let result = TokenizationResult {
ids: encoding.get_ids().to_vec(),
attention_mask: encoding.get_attention_mask().to_vec(),
};
let result_json_str =
serde_json::to_string(&result).expect("Could not convert tokenization result to JSON");
env.new_string(result_json_str)
.expect("Could not convert result_json_str to jstring")
}
Last, we define a deleteTokenizer
function which allows Rust to deallocate the instance of the Tokenizer
,
#[no_mangle]
pub extern "C" fn Java_com_ml_shubham0204_sentence_1embeddings_HFTokenizer_deleteTokenizer(
_: JNIEnv,
_: JClass,
tokenizer_ptr: jlong
) {
let _ptr = unsafe { Box::from_raw(tokenizer_ptr as *mut Tokenizer) };
// _ptr is not referred further, hence it will be deallocated
// as its reference-count becomes zero
}
To compile a dynamic library for Android architectures, we add the necessary toolchains with rustup
,
$> rustup target add aarch64-linux-android
$> rustup target add armv7-linux-androideabi
$> rustup target add i686-linux-android
$> rustup target add x86_64-linux-android
We also need to add paths for clang
and clang++
compilers from Android NDK in .cargo/config.toml
. Crates like onig
require these compilers to build C/C++ code on which they depend. Paths to the linker for each architecture also have to be included.
[target.aarch64-linux-android]
linker = "~/android-ndk-r26d/toolchains/llvm/prebuilt/linux-x86_64/bin/aarch64-linux-android33-clang"
[target.armv7-linux-androideabi]
linker = "~/android-ndk-r26d/toolchains/llvm/prebuilt/linux-x86_64/bin/armv7a-linux-androideabi33-clang"
[target.i686-linux-android]
linker = "~/android-ndk-r26d/toolchains/llvm/prebuilt/linux-x86_64/bin/i686-linux-android33-clang"
[target.x86_64-linux-android]
linker = "~/android-ndk-r26d/toolchains/llvm/prebuilt/linux-x86_64/bin/x86_64-linux-android33-clang"
[env]
AR_aarch64-linux-android="~/android-ndk-r26d/toolchains/llvm/prebuilt/linux-x86_64/bin/llvm-ar"
CC_aarch64-linux-android="~/android-ndk-r26d/toolchains/llvm/prebuilt/linux-x86_64/bin/aarch64-linux-android33-clang"
CXX_aarch64-linux-android="~/android-ndk-r26d/toolchains/llvm/prebuilt/linux-x86_64/bin/aarch64-linux-android33-clang++"
AR_armv7-linux-androideabi="~/android-ndk-r26d/toolchains/llvm/prebuilt/linux-x86_64/bin/llvm-ar"
CC_armv7-linux-androideabi="~/android-ndk-r26d/toolchains/llvm/prebuilt/linux-x86_64/bin/armv7a-linux-androideabi33-clang"
CXX_armv7-linux-androideabi="~/android-ndk-r26d/toolchains/llvm/prebuilt/linux-x86_64/bin/armv7a-linux-androideabi33-clang++"
AR_i686-linux-android="~/android-ndk-r26d/toolchains/llvm/prebuilt/linux-x86_64/bin/llvm-ar"
CC_i686-linux-android="~/android-ndk-r26d/toolchains/llvm/prebuilt/linux-x86_64/bin/i686-linux-android33-clang"
CXX_i686-linux-android="~/android-ndk-r26d/toolchains/llvm/prebuilt/linux-x86_64/bin/i686-linux-android33-clang++"
AR_x86_64-linux-android="~/android-ndk-r26d/toolchains/llvm/prebuilt/linux-x86_64/bin/llvm-ar"
CC_x86_64-linux-android="~/android-ndk-r26d/toolchains/llvm/prebuilt/linux-x86_64/bin/x86_64-linux-android33-clang"
CXX_x86_64-linux-android="~/android-ndk-r26d/toolchains/llvm/prebuilt/linux-x86_64/bin/x86_64-linux-android33-clang++"
We have completed our build setup for Android now. We execute cargo run
now, passing --release
and --target
flags for each architecture,
$> cargo build --release --target armv7-linux-androideabi
$> cargo build --release --target aarch64-linux-android
$> cargo build --release --target i686-linux-android
$> cargo build --release --target x86_64-linux-android
The .so
dynamic libraries can be found in the path target/<architecture>/release/libhftokenizer.so
.
You can view the Rust source code here,
https://github.com/shubham0204/Sentence-Embeddings-Android/tree/libs?source=post_page—–1ecea0ce94d8——————————–
Job Offers
Android Library
Once we’ve built the dynamic libraries from our Rust code, we add them to our library module in the following paths: sentence_embeddings/src/main/jniLibs/<architecture>/libhftokenizer.so
where <architecture>
will be x86
, x86_64
, armeabi-v7a
and arm64-v8a
.
Next, we create a Kotlin class HFTokenizer
that declares the JNI functions and holds an instancePtr
, an address/pointer that references the Tokenizer
instance created by library and returned by create_tokenizer
,
import org.json.JSONObject
class HFTokenizer(tokenizerBytes: ByteArray) {
data class Result(
val ids: LongArray = longArrayOf(),
val attentionMask: LongArray = longArrayOf()
)
private val tokenizerPtr: Long = createTokenizer(tokenizerBytes)
fun tokenize(
text: String
): Result {
val output = tokenize(tokenizerPtr, text)
// Deserialize the string
// and read `ids` and `attention_mask` as LongArray
val jsonObject = JSONObject(output)
val idsArray = jsonObject.getJSONArray("ids")
val ids = LongArray(idsArray.length())
for (i in 0 until idsArray.length()) {
ids[i] = (idsArray.get(i) as Int).toLong()
}
val attentionMaskArray = jsonObject.getJSONArray("attention_mask")
val attentionMask = LongArray(attentionMaskArray.length())
for (i in 0 until attentionMaskArray.length()) {
attentionMask[i] = (attentionMaskArray.get(i) as Int).toLong()
}
return Result(ids, attentionMask)
}
fun close() {
deleteTokenizer(tokenizerPtr)
}
// Given the bytes of the file `tokenizer.json`,
// return a pointer
private external fun createTokenizer(
tokenizerBytes: ByteArray
): Long
// Given the pointer to `Tokenizer` and the text,
// return `ids` and `attention_mask` in JSON format
private external fun tokenize(
tokenizerPtr: Long,
text: String
): String
// Pass the `tokenizerPtr` which is then deallocated
// by the library
private external fun deleteTokenizer(
tokenizerPtr: Long
)
companion object {
init {
System.loadLibrary("hftokenizer")
}
}
}
The externaltokenize
method will return a String
that contains ids
and attention_mask
encoded in JSON. We deserialize the string, reading arrays with keys ids
and attention_mask
as LongArray
and pack them in a data class Result
.
Next, we create the SentenceEmbedding
class that uses onnxruntime
to execute the model and use HFTokenizer
to tokenize the given text,
package com.ml.shubham0204.sentence_embeddings
import ai.onnxruntime.OnnxTensor
import ai.onnxruntime.OrtEnvironment
import ai.onnxruntime.OrtSession
import ai.onnxruntime.providers.NNAPIFlags
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.withContext
import java.nio.LongBuffer
import java.util.EnumSet
class SentenceEmbedding {
private lateinit var hfTokenizer: HFTokenizer
private lateinit var ortEnvironment: OrtEnvironment
private lateinit var ortSession: OrtSession
suspend fun init(
modelBytes: ByteArray,
tokenizerBytes: ByteArray,
useFP16: Boolean = false,
useXNNPack: Boolean = false
) = withContext(Dispatchers.IO) {
// Initialize HFTokenizer and OrtSession
hfTokenizer = HFTokenizer(tokenizerBytes)
ortEnvironment = OrtEnvironment.getEnvironment()
val options = OrtSession.SessionOptions().apply{
if (useFP16) {
addNnapi(EnumSet.of(NNAPIFlags.USE_FP16, NNAPIFlags.CPU_DISABLED))
}
if (useXNNPack) {
addXnnpack(mapOf(
"intra_op_num_threads" to "2"
))
}
}
ortSession = ortEnvironment.createSession(modelBytes,options)
}
suspend fun encode(
sentence: String
): FloatArray = withContext(Dispatchers.IO) {
val result = hfTokenizer.tokenize(sentence)
// Create input tensors for `ids` and `attention_mask`
val idsTensor =
OnnxTensor.createTensor(
ortEnvironment,
LongBuffer.wrap(result.ids),
longArrayOf(1, result.ids.size.toLong()),
)
val attentionMaskTensor =
OnnxTensor.createTensor(
ortEnvironment,
LongBuffer.wrap(result.attentionMask),
longArrayOf(1, result.attentionMask.size.toLong()),
)
val outputs =
ortSession.run(mapOf("input_ids" to idsTensor, "attention_mask" to attentionMaskTensor))
val embeddingTensor = outputs.get("sentence_embedding").get() as OnnxTensor
return@withContext embeddingTensor.floatBuffer.array()
}
}
The library can then be distributed through Maven Central or Jitpack and sentence embeddings can be obtained by the use of the SentenceEmbedding
class.
Extended Resources
You may find additional blogs on compiling Rust for Android apps, the process being very similar to that for C/C++ with CMake,
These are some helpful resources using onnxruntime
in Android,
Conclusion
I’m an on-device ML enthusiast and developing ML apps on Android is my passion. If you would like to view on-device ML projects, ahead here,
https://github.com/shubham0204?source=post_page—–1ecea0ce94d8——————————–#on-device-ml
Do check out my website and share your thoughts on this story. Keep learning and have a nice day ahead!
This article is previously published on proandroiddev.com