Android应用中使用TF-Lite实现句子嵌入的Java与C++集成问题

Java API推理可行性 #

直接用TensorFlow Lite Java API无法完成该模型的推理，因为模型依赖TFSentencepieceTokenizeOp自定义算子，而官方默认的Java版TF Lite库并未包含该算子的实现，所以加载模型时会出现 unresolved custom op 错误。目前官方没有提供该算子的Java扩展支持，建议采用C++ NDK路线实现。

C++ NDK高效实现方案 #

不需要引入完整的TF-Lite仓库，最简洁的方式是利用TensorFlow Lite Support Library的预编译组件，步骤如下：

1. 配置CMakeLists.txt #

在你的CMake配置中，引入TF Lite和Universal Sentence Encoder QA的依赖：

cmake_minimum_required(VERSION 3.18.1)

project("usencoderqa")

# 引入TF Lite及Support Library预编译库
find_package(tensorflow-lite REQUIRED)
find_package(tflite-support REQUIRED)

add_library(
             native-lib
             SHARED
             native-lib.cpp )

# 链接依赖库
target_link_libraries(
                       native-lib
                       android
                       log
                       tensorflow-lite::tensorflow-lite
                       tflite-support::tflite-support )

2. 修改native-lib.cpp实现推理逻辑 #

替换现有代码，实现模型加载、输入处理和推理：

#include <jni.h>
#include <string>
#include "universal_sentence_encoder_qa.h"

using namespace tflite::task::text::qa;

// 全局模型实例
std::unique_ptr<BertQuestionAnswerer> qa_model = nullptr;

extern "C"
JNIEXPORT jboolean JNICALL
Java_com_dummy_Caller_initModel(JNIEnv *env, jobject thiz, jstring model_path) {
    const char* model_cpath = env->GetStringUTFChars(model_path, nullptr);
    std::string model_path_str(model_cpath);
    env->ReleaseStringUTFChars(model_path, model_cpath);

    // 初始化模型
    BertQuestionAnswererOptions options;
    options.mutable_model_file()->set_file_name(model_path_str);
    auto model_result = BertQuestionAnswerer::CreateFromOptions(options);
    if (model_result.ok()) {
        qa_model = std::move(model_result.value());
        return JNI_TRUE;
    }
    return JNI_FALSE;
}

extern "C"
JNIEXPORT jstring JNICALL
Java_com_dummy_Caller_answerQuestion(JNIEnv *env, jobject thiz, jstring question, jstring context) {
    if (!qa_model) {
        return env->NewStringUTF("Model not initialized");
    }

    const char* q_cstr = env->GetStringUTFChars(question, nullptr);
    const char* c_cstr = env->GetStringUTFChars(context, nullptr);
    std::string question_str(q_cstr);
    std::string context_str(c_cstr);
    env->ReleaseStringUTFChars(question, q_cstr);
    env->ReleaseStringUTFChars(context, c_cstr);

    // 执行推理
    std::vector<QaAnswer> results = qa_model->Answer(context_str, question_str);
    if (results.empty()) {
        return env->NewStringUTF("No answer found");
    }

    // 拼接返回结果（示例取置信度最高的第一个结果）
    std::string result = "Answer: " + results[0].text + ", Score: " + std::to_string(results[0].pos.score);
    return env->NewStringUTF(result.c_str());
}

3. Java层调用 #

在Java代码中添加对应的native方法，并调用初始化和推理：

public class Caller {
    static {
        System.loadLibrary("native-lib");
    }

    public native boolean initModel(String modelPath);
    public native String answerQuestion(String question, String context);
}

补充说明 #

• 将模型文件放置在assets目录，初始化时传递正确的模型路径（可通过getAssets().openFd("model.tflite").getFileDescriptor()获取路径，或直接传递assets相对路径）。
• 在app/build.gradle中引入依赖库：

dependencies {
    implementation 'org.tensorflow:tensorflow-lite-support:0.4.4'
    implementation 'org.tensorflow:tensorflow-lite:2.15.0'
}

内容的提问来源于stack exchange，提问作者paganpasta