Build you own Task API

Similar to Native APIs, to build an API object, the client needs to provide the following information by extending BaseTaskApi, which provides JNI handlings for all Java Task APIs.

• Determine the API I/O - This usually mirrors the native interfaces. e.g. BertQuestionAnswerer takes (String context, String question) as input and outputs List<QaAnswer>. The implementation calls a private native function with similar signature, except it has an additional parameter long nativeHandle, which is the pointer returned from C++.

  class BertQuestionAnswerer extends BaseTaskApi {
    public List<QaAnswer> answer(String context, String question) {
      return answerNative(getNativeHandle(), context, question);
    }
  
    private static native List<QaAnswer> answerNative(
                                          long nativeHandle, // C++ pointer
                                          String context, String question // API I/O
                                         );
  
  }
  

• Create factory functions of the API - This also mirrors native factory functions, except Android factory functions also need to take Context for file access. The implementation calls one of the utilities in TaskJniUtils to build the corresponding C++ API object and pass its pointer to the BaseTaskApi constructor.

    class BertQuestionAnswerer extends BaseTaskApi {
      private static final String BERT_QUESTION_ANSWERER_NATIVE_LIBNAME =
                                                "bert_question_answerer_jni";
  
      // Extending super constructor by providing the
      // native handle(pointer of corresponding C++ API object)
      private BertQuestionAnswerer(long nativeHandle) {
        super(nativeHandle);
      }
  
      public static BertQuestionAnswerer createBertQuestionAnswerer(
                                          Context context, // Accessing Android files
                                          String pathToModel, String pathToVocab) {
        return new BertQuestionAnswerer(
            // The util first try loads the JNI module with name
            // BERT_QUESTION_ANSWERER_NATIVE_LIBNAME, then opens two files,
            // converts them into ByteBuffer, finally ::initJniWithBertByteBuffers
            // is called with the buffer for a C++ API object pointer
            TaskJniUtils.createHandleWithMultipleAssetFilesFromLibrary(
                context,
                BertQuestionAnswerer::initJniWithBertByteBuffers,
                BERT_QUESTION_ANSWERER_NATIVE_LIBNAME,
                pathToModel,
                pathToVocab));
      }
  
      // modelBuffers[0] is tflite model file buffer, and modelBuffers[1] is vocab file buffer.
      // returns C++ API object pointer casted to long
      private static native long initJniWithBertByteBuffers(ByteBuffer... modelBuffers);
  
    }
  

• Implement the JNI module for native functions - All Java native methods are implemented by calling a corresponding native function from the JNI module. The factory functions would create a native API object and return its pointer as a long type to Java. In later calls to Java API, the long type pointer is passed back to JNI and cast back to the native API object. The native API results are then converted back to Java results.

  For example, this is how bert_question_answerer_jni is implemented.

    // Implements BertQuestionAnswerer::initJniWithBertByteBuffers
    extern "C" JNIEXPORT jlong JNICALL
    Java_org_tensorflow_lite_task_text_qa_BertQuestionAnswerer_initJniWithBertByteBuffers(
        JNIEnv* env, jclass thiz, jobjectArray model_buffers) {
      // Convert Java ByteBuffer object into a buffer that can be read by native factory functions
      absl::string_view model =
          GetMappedFileBuffer(env, env->GetObjectArrayElement(model_buffers, 0));
  
      // Creates the native API object
      absl::StatusOr<std::unique_ptr<QuestionAnswerer>> status =
          BertQuestionAnswerer::CreateFromBuffer(
              model.data(), model.size());
      if (status.ok()) {
        // converts the object pointer to jlong and return to Java.
        return reinterpret_cast<jlong>(status->release());
      } else {
        return kInvalidPointer;
      }
    }
  
    // Implements BertQuestionAnswerer::answerNative
    extern "C" JNIEXPORT jobject JNICALL
    Java_org_tensorflow_lite_task_text_qa_BertQuestionAnswerer_answerNative(
    JNIEnv* env, jclass thiz, jlong native_handle, jstring context, jstring question) {
    // Convert long to native API object pointer
    QuestionAnswerer* question_answerer = reinterpret_cast<QuestionAnswerer*>(native_handle);
  
    // Calls the native API
    std::vector<QaAnswer> results = question_answerer->Answer(JStringToString(env, context),
                                           JStringToString(env, question));
  
    // Converts native result(std::vector<QaAnswer>) to Java result(List<QaAnswerer>)
    jclass qa_answer_class =
      env->FindClass("org/tensorflow/lite/task/text/qa/QaAnswer");
    jmethodID qa_answer_ctor =
      env->GetMethodID(qa_answer_class, "<init>", "(Ljava/lang/String;IIF)V");
    return ConvertVectorToArrayList<QaAnswer>(
      env, results,
      [env, qa_answer_class, qa_answer_ctor](const QaAnswer& ans) {
        jstring text = env->NewStringUTF(ans.text.data());
        jobject qa_answer =
            env->NewObject(qa_answer_class, qa_answer_ctor, text, ans.pos.start,
                           ans.pos.end, ans.pos.logit);
        env->DeleteLocalRef(text);
        return qa_answer;
      });
    }
  
    // Implements BaseTaskApi::deinitJni by delete the native object
    extern "C" JNIEXPORT void JNICALL Java_task_core_BaseTaskApi_deinitJni(
        JNIEnv* env, jobject thiz, jlong native_handle) {
      delete reinterpret_cast<QuestionAnswerer*>(native_handle);
    }
  

iOS API

Create iOS APIs by wrapping a native API object into a ObjC API object. The created API object can be used in either ObjC or Swift. iOS API requires the native API to be built first.

Sample usage

Here is an example using ObjC TFLBertQuestionAnswerer for MobileBert in Swift.

  static let mobileBertModelPath = "path/to/model.tflite";
  // Create the API from a model file and vocabulary file
  let mobileBertAnswerer = TFLBertQuestionAnswerer.mobilebertQuestionAnswerer(
      modelPath: mobileBertModelPath)

  static let context = ...; // context of a question to be answered
  static let question = ...; // question to be answered
  // ask a question
  let answers = mobileBertAnswerer.answer(
      context: TFLBertQuestionAnswererTest.context, question: TFLBertQuestionAnswererTest.question)
  // answers.[0].text is the best answer

Building the API

Figure 4. iOS Task API

iOS API is a simple ObjC wrapper on top of native API. Build the API by following the steps below:

• Define the ObjC wrapper - Define an ObjC class and delegate the implementations to the corresponding native API object. Note the native dependencies can only appear in a .mm file due to Swift's inability to interop with C++.

  • .h file

    @interface TFLBertQuestionAnswerer : NSObject
  
    // Delegate calls to the native BertQuestionAnswerer::CreateBertQuestionAnswerer
    + (instancetype)mobilebertQuestionAnswererWithModelPath:(NSString*)modelPath
                                                  vocabPath:(NSString*)vocabPath
        NS_SWIFT_NAME(mobilebertQuestionAnswerer(modelPath:vocabPath:));
  
    // Delegate calls to the native BertQuestionAnswerer::Answer
    - (NSArray<TFLQAAnswer*>*)answerWithContext:(NSString*)context
                                       question:(NSString*)question
        NS_SWIFT_NAME(answer(context:question:));
  }
  

  • .mm file

    using BertQuestionAnswererCPP = ::tflite::task::text::BertQuestionAnswerer;
  
    @implementation TFLBertQuestionAnswerer {
      // define an iVar for the native API object
      std::unique_ptr<QuestionAnswererCPP> _bertQuestionAnswerwer;
    }
  
    // Initialize the native API object
    + (instancetype)mobilebertQuestionAnswererWithModelPath:(NSString *)modelPath
                                            vocabPath:(NSString *)vocabPath {
      absl::StatusOr<std::unique_ptr<QuestionAnswererCPP>> cQuestionAnswerer =
          BertQuestionAnswererCPP::CreateBertQuestionAnswerer(MakeString(modelPath),
                                                              MakeString(vocabPath));
      _GTMDevAssert(cQuestionAnswerer.ok(), @"Failed to create BertQuestionAnswerer");
      return [[TFLBertQuestionAnswerer alloc]
          initWithQuestionAnswerer:std::move(cQuestionAnswerer.value())];
    }
  
    // Calls the native API and converts C++ results into ObjC results
    - (NSArray<TFLQAAnswer *> *)answerWithContext:(NSString *)context question:(NSString *)question {
      std::vector<QaAnswerCPP> results =
        _bertQuestionAnswerwer->Answer(MakeString(context), MakeString(question));
      return [self arrayFromVector:results];
    }
  }
  

Send feedback