Task delegate 12: Expose BaseOptions in Java BertNLClassifier

PiperOrigin-RevId: 400355143

Author: lu-wang-g

tensorflow_lite_support/java/src/java/org/tensorflow/lite/task/text/nlclassifier/BertNLClassifier.java

@@ -25,6 +25,7 @@
 import java.util.List;
 import org.tensorflow.lite.annotations.UsedByReflection;
 import org.tensorflow.lite.support.label.Category;
+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;
@@ -54,14 +55,21 @@ public abstract static class BertNLClassifierOptions {
     @UsedByReflection("bert_nl_classifier_jni.cc")
     abstract int getMaxSeqLen();
 
-    public static BertNLClassifierOptions.Builder builder() {
+    abstract BaseOptions getBaseOptions();
+
+    public static Builder builder() {
       return new AutoValue_BertNLClassifier_BertNLClassifierOptions.Builder()
-          .setMaxSeqLen(DEFAULT_MAX_SEQ_LEN);
+          .setMaxSeqLen(DEFAULT_MAX_SEQ_LEN)
+          .setBaseOptions(BaseOptions.builder().build());
     }
 
     /** Builder for {@link BertNLClassifierOptions}. */
     @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 BertNLClassifierOptions.Builder setMaxSeqLen(int value);
 
       public abstract BertNLClassifierOptions build();
@@ -139,7 +147,10 @@ public static BertNLClassifier createFromFileAndOptions(
               new EmptyHandleProvider() {
                 @Override
                 public long createHandle() {
-                  return initJniWithFileDescriptor(descriptor.getFd(), options);
+                  return initJniWithFileDescriptor(
+                      descriptor.getFd(),
+                      options,
+                      TaskJniUtils.createProtoBaseOptionsHandle(options.getBaseOptions()));
                 }
               },
               BERT_NL_CLASSIFIER_NATIVE_LIBNAME));
@@ -179,7 +190,10 @@ public static BertNLClassifier createFromBufferAndOptions(
             new EmptyHandleProvider() {
               @Override
               public long createHandle() {
-                return initJniWithByteBuffer(modelBuffer, options);
+                return initJniWithByteBuffer(
+                    modelBuffer,
+                    options,
+                    TaskJniUtils.createProtoBaseOptionsHandle(options.getBaseOptions()));
               }
             },
             BERT_NL_CLASSIFIER_NATIVE_LIBNAME));
@@ -196,9 +210,10 @@ public List<Category> classify(String text) {
   }
 
   private static native long initJniWithByteBuffer(
-      ByteBuffer modelBuffer, BertNLClassifierOptions options);
+      ByteBuffer modelBuffer, BertNLClassifierOptions options, long baseOptionsHandle);
 
-  private static native long initJniWithFileDescriptor(int fd, BertNLClassifierOptions options);
+  private static native long initJniWithFileDescriptor(
+      int fd, BertNLClassifierOptions options, long baseOptionsHandle);
 
   private static native List<Category> classifyNative(long nativeHandle, String text);
 

tensorflow_lite_support/java/src/native/task/text/nlclassifier/bert/BUILD

@@ -20,6 +20,7 @@ jni_binary_with_tflite(
     name = "libtask_text_jni.so",
     srcs = [
         "bert_nl_classifier_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 = [
@@ -28,6 +29,7 @@ jni_binary_with_tflite(
         "//tensorflow_lite_support/java/src/native/task/text/nlclassifier:nl_classifier_jni_utils",
     ],
     deps = [
+        "//tensorflow_lite_support/cc/task/core/proto:base_options_proto_inc",
         "//tensorflow_lite_support/cc/task/text/proto:bert_nl_classifier_options_proto_inc",
         "//tensorflow_lite_support/java/jni",
     ],

tensorflow_lite_support/java/src/native/task/text/nlclassifier/bert/bert_nl_classifier_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_nl_classifier.h"
 #include "tensorflow_lite_support/cc/task/text/proto/bert_nl_classifier_options_proto_inc.h"
 #include "tensorflow_lite_support/cc/utils/jni_utils.h"
@@ -25,13 +26,21 @@ namespace {
 using ::tflite::support::utils::kAssertionError;
 using ::tflite::support::utils::kInvalidPointer;
 using ::tflite::support::utils::ThrowException;
+using ::tflite::task::core::BaseOptions;
 using ::tflite::task::text::BertNLClassifier;
 using ::tflite::task::text::BertNLClassifierOptions;
 using ::tflite::task::text::nlclassifier::RunClassifier;
 
 BertNLClassifierOptions ConvertJavaBertNLClassifierOptions(
-    JNIEnv* env, jobject java_options) {
+    JNIEnv* env, jobject java_options, jlong base_options_handle) {
   BertNLClassifierOptions 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));
+  }
+
   jclass java_options_class = env->FindClass(
       "org/tensorflow/lite/task/text/nlclassifier/"
       "BertNLClassifier$BertNLClassifierOptions");
@@ -41,6 +50,7 @@ BertNLClassifierOptions ConvertJavaBertNLClassifierOptions(
       env->CallIntMethod(java_options, max_seq_len_id));
   return proto_options;
 }
+}  // namespace
 
 extern "C" JNIEXPORT void JNICALL
 Java_org_tensorflow_lite_task_text_nlclassifier_BertNLClassifier_deinitJni(
@@ -50,9 +60,10 @@ Java_org_tensorflow_lite_task_text_nlclassifier_BertNLClassifier_deinitJni(
 
 extern "C" JNIEXPORT jlong JNICALL
 Java_org_tensorflow_lite_task_text_nlclassifier_BertNLClassifier_initJniWithByteBuffer(
-    JNIEnv* env, jclass thiz, jobject model_buffer, jobject java_options) {
-  BertNLClassifierOptions proto_options =
-      ConvertJavaBertNLClassifierOptions(env, java_options);
+    JNIEnv* env, jclass thiz, jobject model_buffer, jobject java_options,
+    jlong base_options_handle) {
+  BertNLClassifierOptions proto_options = ConvertJavaBertNLClassifierOptions(
+      env, java_options, base_options_handle);
   proto_options.mutable_base_options()->mutable_model_file()->set_file_content(
       static_cast<char*>(env->GetDirectBufferAddress(model_buffer)),
       static_cast<size_t>(env->GetDirectBufferCapacity(model_buffer)));
@@ -71,9 +82,10 @@ Java_org_tensorflow_lite_task_text_nlclassifier_BertNLClassifier_initJniWithByte
 
 extern "C" JNIEXPORT jlong JNICALL
 Java_org_tensorflow_lite_task_text_nlclassifier_BertNLClassifier_initJniWithFileDescriptor(
-    JNIEnv* env, jclass thiz, jint fd, jobject java_options) {
-  BertNLClassifierOptions proto_options =
-      ConvertJavaBertNLClassifierOptions(env, java_options);
+    JNIEnv* env, jclass thiz, jint fd, jobject java_options,
+    jlong base_options_handle) {
+  BertNLClassifierOptions proto_options = ConvertJavaBertNLClassifierOptions(
+      env, java_options, base_options_handle);
   proto_options.mutable_base_options()
       ->mutable_model_file()
       ->mutable_file_descriptor_meta()
@@ -96,5 +108,3 @@ Java_org_tensorflow_lite_task_text_nlclassifier_BertNLClassifier_classifyNative(
     JNIEnv* env, jclass clazz, jlong native_handle, jstring text) {
   return RunClassifier(env, native_handle, text);
 }
-
-}  // namespace