Task delegate 11: Expose BaseOptions in Java BertQuestionAnswer

Major changes:
(1) Created the `BertQuestionAnswererOptions` class and integrated it with BaseOptions.
(2) Added factory create methods that expose `BertQuestionAnswererOptions`.
(3) Unified the JNI method to create the C++ task using `initJniWithFileDescriptor`. Removed `initJniWithModelWithMetadataByteBuffers`.

PiperOrigin-RevId: 400357315

Author: lu-wang-g

tensorflow_lite_support/java/src/java/org/tensorflow/lite/task/text/qa/BUILD

@@ -33,6 +33,7 @@ android_library(
     javacopts = ["-source 7 -target 7"],
     deps = [
         "//tensorflow_lite_support/java/src/java/org/tensorflow/lite/task/core:base_task_api",
+        "@com_google_auto_value",
         "@org_tensorflow//tensorflow/lite/java:tensorflowlite_java",
     ],
 )

tensorflow_lite_support/java/src/java/org/tensorflow/lite/task/text/qa/BertQuestionAnswerer.java

@@ -17,13 +17,16 @@
 
 import android.content.Context;
 import android.os.ParcelFileDescriptor;
+import com.google.auto.value.AutoValue;
 import java.io.File;
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.util.List;
+import org.tensorflow.lite.task.core.BaseOptions;
 import org.tensorflow.lite.task.core.BaseTaskApi;
 import org.tensorflow.lite.task.core.TaskJniUtils;
 import org.tensorflow.lite.task.core.TaskJniUtils.EmptyHandleProvider;
+import org.tensorflow.lite.task.core.TaskJniUtils.FdAndOptionsHandleProvider;
 import org.tensorflow.lite.task.core.TaskJniUtils.MultipleBuffersHandleProvider;
 
 /**
@@ -60,22 +63,62 @@ public class BertQuestionAnswerer extends BaseTaskApi implements QuestionAnswere
    */
   public static BertQuestionAnswerer createFromFile(Context context, String modelPath)
       throws IOException {
+    return createFromFileAndOptions(
+        context, modelPath, BertQuestionAnswererOptions.builder().build());
+  }
+
+  /**
+   * Creates a {@link BertQuestionAnswerer} instance from the default {@link
+   * BertQuestionAnswererOptions}.
+   *
+   * @param modelFile a {@link File} object of the model
+   * @return a {@link BertQuestionAnswerer} instance
+   * @throws IOException if model file fails to load
+   * @throws IllegalArgumentException if an argument is invalid
+   * @throws IllegalStateException if there is an internal error
+   * @throws RuntimeException if there is an otherwise unspecified error
+   */
+  public static BertQuestionAnswerer createFromFile(File modelFile) throws IOException {
+    return createFromFileAndOptions(modelFile, BertQuestionAnswererOptions.builder().build());
+  }
+
+  /**
+   * Creates a {@link BertQuestionAnswerer} instance from {@link BertQuestionAnswererOptions}.
+   *
+   * @param context android context
+   * @param modelPath file path to the model with metadata. Note: The model should not be compressed
+   * @return a {@link BertQuestionAnswerer} instance
+   * @throws IOException if model file fails to load
+   * @throws IllegalArgumentException if an argument is invalid
+   * @throws IllegalStateException if there is an internal error
+   * @throws RuntimeException if there is an otherwise unspecified error
+   */
+  public static BertQuestionAnswerer createFromFileAndOptions(
+      Context context, String modelPath, BertQuestionAnswererOptions options) throws IOException {
     return new BertQuestionAnswerer(
-        TaskJniUtils.createHandleWithMultipleAssetFilesFromLibrary(
+        TaskJniUtils.createHandleFromFdAndOptions(
             context,
-            new MultipleBuffersHandleProvider() {
+            new FdAndOptionsHandleProvider<BertQuestionAnswererOptions>() {
               @Override
-              public long createHandle(ByteBuffer... buffers) {
-                return BertQuestionAnswerer.initJniWithModelWithMetadataByteBuffers(buffers);
+              public long createHandle(
+                  int fileDescriptor,
+                  long fileDescriptorLength,
+                  long fileDescriptorOffset,
+                  BertQuestionAnswererOptions options) {
+                return initJniWithFileDescriptor(
+                    fileDescriptor,
+                    fileDescriptorLength,
+                    fileDescriptorOffset,
+                    TaskJniUtils.createProtoBaseOptionsHandle(options.getBaseOptions()));
               }
             },
             BERT_QUESTION_ANSWERER_NATIVE_LIBNAME,
-            modelPath));
+            modelPath,
+            options));
   }
 
   /**
-   * Creates a {@link BertQuestionAnswerer} instance from the default {@link
-   * BertQuestionAnswererOptions}.
+   * Creates a {@link BertQuestionAnswerer} instance from {@link BertQuestionAnswererOptions}.
    *
    * @param modelFile a {@link File} object of the model
    * @return a {@link BertQuestionAnswerer} instance
@@ -84,7 +127,8 @@ public long createHandle(ByteBuffer... buffers) {
    * @throws IllegalStateException if there is an internal error
    * @throws RuntimeException if there is an otherwise unspecified error
    */
-  public static BertQuestionAnswerer createFromFile(File modelFile) throws IOException {
+  public static BertQuestionAnswerer createFromFileAndOptions(
+      File modelFile, final BertQuestionAnswererOptions options) throws IOException {
     try (ParcelFileDescriptor descriptor =
         ParcelFileDescriptor.open(modelFile, ParcelFileDescriptor.MODE_READ_ONLY)) {
       return new BertQuestionAnswerer(
@@ -95,7 +139,8 @@ public long createHandle() {
                   return initJniWithFileDescriptor(
                       /*fileDescriptor=*/ descriptor.getFd(),
                       /*fileDescriptorLength=*/ OPTIONAL_FD_LENGTH,
-                      /*fileDescriptorOffset=*/ OPTIONAL_FD_OFFSET);
+                      /*fileDescriptorOffset=*/ OPTIONAL_FD_OFFSET,
+                      TaskJniUtils.createProtoBaseOptionsHandle(options.getBaseOptions()));
                 }
               },
               BERT_QUESTION_ANSWERER_NATIVE_LIBNAME));
@@ -164,6 +209,26 @@ public long createHandle(ByteBuffer... buffers) {
             sentencePieceModelPath));
   }
 
+  /** Options for setting up a {@link BertQuestionAnswerer}. */
+  @AutoValue
+  public abstract static class BertQuestionAnswererOptions {
+    abstract BaseOptions getBaseOptions();
+
+    public static Builder builder() {
+      return new AutoValue_BertQuestionAnswerer_BertQuestionAnswererOptions.Builder()
+          .setBaseOptions(BaseOptions.builder().build());
+    }
+
+    /** Builder for {@link BertQuestionAnswererOptions}. */
+    @AutoValue.Builder
+    public abstract static class Builder {
+      /** Sets the general options to configure Task APIs, such as accelerators. */
+      public abstract Builder setBaseOptions(BaseOptions baseOptions);
+
+      public abstract BertQuestionAnswererOptions build();
+    }
+  }
+
   @Override
   public List<QaAnswer> answer(String context, String question) {
     checkNotClosed();
@@ -181,11 +246,11 @@ private BertQuestionAnswerer(long nativeHandle) {
   // buffer.
   private static native long initJniWithAlbertByteBuffers(ByteBuffer... modelBuffers);
 
-  // modelBuffers[0] is tflite model file buffer with metadata to specify which tokenizer to use.
-  private static native long initJniWithModelWithMetadataByteBuffers(ByteBuffer... modelBuffers);
-
   private static native long initJniWithFileDescriptor(
-      int fileDescriptor, long fileDescriptorLength, long fileDescriptorOffset);
+      int fileDescriptor,
+      long fileDescriptorLength,
+      long fileDescriptorOffset,
+      long baseOptionsHandle);
 
   private static native List<QaAnswer> answerNative(
       long nativeHandle, String context, String question);

tensorflow_lite_support/java/src/native/task/text/qa/BUILD

@@ -13,13 +13,15 @@ jni_binary_with_tflite(
     name = "libtask_text_jni.so",
     srcs = [
         "bert_question_answerer_jni.cc",
+        "//tensorflow_lite_support/java/src/native/task/core:task_jni_utils.cc",
     ],
     linkscript = "//tensorflow_lite_support/java:default_version_script.lds",
     tflite_deps = [
         "//tensorflow_lite_support/cc/task/text:bert_question_answerer",
         "//tensorflow_lite_support/cc/utils:jni_utils",
     ],
     deps = [
+        "//tensorflow_lite_support/cc/task/core/proto:base_options_proto_inc",
         "//tensorflow_lite_support/java/jni",
     ],
 )

tensorflow_lite_support/java/src/native/task/text/qa/bert_question_answerer_jni.cc

@@ -15,6 +15,7 @@ limitations under the License.
 
 #include <jni.h>
 
+#include "tensorflow_lite_support/cc/task/core/proto/base_options_proto_inc.h"
 #include "tensorflow_lite_support/cc/task/text/bert_question_answerer.h"
 #include "tensorflow_lite_support/cc/utils/jni_utils.h"
 
@@ -26,13 +27,27 @@ using ::tflite::support::utils::GetExceptionClassNameForStatusCode;
 using ::tflite::support::utils::GetMappedFileBuffer;
 using ::tflite::support::utils::JStringToString;
 using ::tflite::support::utils::ThrowException;
+using ::tflite::task::core::BaseOptions;
 using ::tflite::task::text::BertQuestionAnswerer;
 using ::tflite::task::text::BertQuestionAnswererOptions;
 using ::tflite::task::text::QaAnswer;
 using ::tflite::task::text::QuestionAnswerer;
 
 constexpr int kInvalidPointer = 0;
 
+// Creates a BertQuestionAnswererOptions proto based on the Java class.
+BertQuestionAnswererOptions ConvertToProtoOptions(jlong base_options_handle) {
+  BertQuestionAnswererOptions proto_options;
+
+  if (base_options_handle != kInvalidPointer) {
+    // proto_options will free the previous base_options and set the new one.
+    proto_options.set_allocated_base_options(
+        reinterpret_cast<BaseOptions*>(base_options_handle));
+  }
+
+  return proto_options;
+}
+
 }  // namespace
 
 extern "C" JNIEXPORT void JNICALL
@@ -41,30 +56,13 @@ Java_org_tensorflow_lite_task_text_qa_BertQuestionAnswerer_deinitJni(
   delete reinterpret_cast<QuestionAnswerer*>(native_handle);
 }
 
-extern "C" JNIEXPORT jlong JNICALL
-Java_org_tensorflow_lite_task_text_qa_BertQuestionAnswerer_initJniWithModelWithMetadataByteBuffers(
-    JNIEnv* env, jclass thiz, jobjectArray model_buffers) {
-  absl::string_view model_with_metadata =
-      GetMappedFileBuffer(env, env->GetObjectArrayElement(model_buffers, 0));
-
-  tflite::support::StatusOr<std::unique_ptr<QuestionAnswerer>> qa_status =
-      BertQuestionAnswerer::CreateFromBuffer(model_with_metadata.data(),
-                                             model_with_metadata.size());
-  if (qa_status.ok()) {
-    return reinterpret_cast<jlong>(qa_status->release());
-  } else {
-    ThrowException(
-        env, GetExceptionClassNameForStatusCode(qa_status.status().code()),
-        "Error occurred when initializing BertQuestionAnswerer: %s",
-        qa_status.status().message().data());
-    return kInvalidPointer;
-  }
-}
 extern "C" JNIEXPORT jlong JNICALL
 Java_org_tensorflow_lite_task_text_qa_BertQuestionAnswerer_initJniWithFileDescriptor(
     JNIEnv* env, jclass thiz, jint file_descriptor,
-    jlong file_descriptor_length, jlong file_descriptor_offset) {
-  BertQuestionAnswererOptions proto_options;
+    jlong file_descriptor_length, jlong file_descriptor_offset,
+    jlong base_options_handle) {
+  BertQuestionAnswererOptions proto_options =
+      ConvertToProtoOptions(base_options_handle);
   auto file_descriptor_meta = proto_options.mutable_base_options()
                                   ->mutable_model_file()
                                   ->mutable_file_descriptor_meta();