[automerge] Return StatusOr from Classify() methods 2p: 9fc1fcbd22

Original change: https://googleplex-android-review.googlesource.com/c/platform/external/tflite-support/+/19678847

Bug: 242926638
Change-Id: Ib77b4dc3da1d92570d61a73738771bc4f6b9b194

5 files changed, 55 insertions, 24 deletions

diff --git a/tensorflow_lite_support/cc/task/text/nlclassifier/bert_nl_classifier.cc b/tensorflow_lite_support/cc/task/text/nlclassifier/bert_nl_classifier.cc
index a246066b..4c8b3fd7 100644
--- a/tensorflow_lite_support/cc/task/text/nlclassifier/bert_nl_classifier.cc
+++ b/tensorflow_lite_support/cc/task/text/nlclassifier/bert_nl_classifier.cc
@@ -245,13 +245,17 @@ absl::Status BertNLClassifier::InitializeFromMetadata() {
                         segment_ids_tensor.dims->data[1]),
         TfLiteSupportStatus::kInvalidInputTensorSizeError);
   }
-  if (ids_tensor.dims_signature->data[1] == -1 &&
-      mask_tensor.dims_signature->data[1] == -1 &&
-      segment_ids_tensor.dims_signature->data[1] == -1) {
+
+  bool has_valid_dims_signature = ids_tensor.dims_signature->size == 2 &&
+                                  mask_tensor.dims_signature->size == 2 &&
+                                  segment_ids_tensor.dims_signature->size == 2;
+  int num_dynamic_tensors = ids_tensor.dims_signature->data[1] == -1 +
+                            mask_tensor.dims_signature->data[1] == -1 +
+                            segment_ids_tensor.dims_signature->data[1] == -1;
+
+  if (has_valid_dims_signature && num_dynamic_tensors == 3) {
     input_tensors_are_dynamic_ = true;
-  } else if (ids_tensor.dims_signature->data[1] == -1 ||
-             mask_tensor.dims_signature->data[1] == -1 ||
-             segment_ids_tensor.dims_signature->data[1] == -1) {
+  } else if (has_valid_dims_signature && num_dynamic_tensors > 0 && num_dynamic_tensors < 3) {
     return CreateStatusWithPayload(
         absl::StatusCode::kInternal,
         "Input tensors contain a mix of static and dynamic tensors",
diff --git a/tensorflow_lite_support/cc/task/text/nlclassifier/nl_classifier.cc b/tensorflow_lite_support/cc/task/text/nlclassifier/nl_classifier.cc
index a75fe0ff..d322a151 100644
--- a/tensorflow_lite_support/cc/task/text/nlclassifier/nl_classifier.cc
+++ b/tensorflow_lite_support/cc/task/text/nlclassifier/nl_classifier.cc
@@ -185,9 +185,16 @@ absl::Status NLClassifier::TrySetLabelFromMetadata(
 }
 
 std::vector<Category> NLClassifier::Classify(const std::string& text) {
-  // The NLClassifier implementation for Preprocess() and Postprocess() never
-  // returns errors: just call value().
-  return Infer(text).value();
+  StatusOr<std::vector<Category>> infer_result = ClassifyText(text);
+  if (!infer_result.ok()) {
+    return {};
+  }
+  return infer_result.value();
+}
+
+StatusOr<std::vector<Category>> NLClassifier::ClassifyText(
+    const std::string& text) {
+  return Infer(text);
 }
 
 std::string NLClassifier::GetModelVersion() const {
diff --git a/tensorflow_lite_support/cc/task/text/nlclassifier/nl_classifier.h b/tensorflow_lite_support/cc/task/text/nlclassifier/nl_classifier.h
index 013b7d53..8e70ba4c 100644
--- a/tensorflow_lite_support/cc/task/text/nlclassifier/nl_classifier.h
+++ b/tensorflow_lite_support/cc/task/text/nlclassifier/nl_classifier.h
@@ -109,9 +109,14 @@ class NLClassifier : public core::BaseTaskApi<std::vector<core::Category>,
       std::unique_ptr<tflite::OpResolver> resolver =
           absl::make_unique<tflite::ops::builtin::BuiltinOpResolver>());
 
-  // Performs classification on a string input, returns classified results.
+  // DEPRECATED (unannotated for backward compatibility).  Prefer using `ClassifyText`.
   std::vector<core::Category> Classify(const std::string& text);
 
+  // Performs classification on a string input, returns classified results or an
+  // error.
+  tflite::support::StatusOr<std::vector<core::Category>> ClassifyText(
+      const std::string& text);
+
   // Gets the model version, or "NO_VERSION_INFO" in case there is no version.
   std::string GetModelVersion() const;
 
diff --git a/tensorflow_lite_support/cc/test/task/text/bert_nl_classifier_test.cc b/tensorflow_lite_support/cc/test/task/text/bert_nl_classifier_test.cc
index 85ebc505..23a3ddca 100644
--- a/tensorflow_lite_support/cc/test/task/text/bert_nl_classifier_test.cc
+++ b/tensorflow_lite_support/cc/test/task/text/bert_nl_classifier_test.cc
@@ -85,15 +85,19 @@ Category* GetCategoryWithClassName(const std::string& class_name,
 void verify_classifier(std::unique_ptr<BertNLClassifier> classifier,
                        bool verify_positive) {
   if (verify_positive) {
-    std::vector<core::Category> results =
-        classifier->Classify("unflinchingly bleak and desperate");
-    EXPECT_GT(GetCategoryWithClassName("negative", results)->score,
-              GetCategoryWithClassName("positive", results)->score);
+    tflite::support::StatusOr<std::vector<core::Category>> results =
+        classifier->ClassifyText("unflinchingly bleak and desperate");
+
+    EXPECT_TRUE(results.ok());
+    EXPECT_GT(GetCategoryWithClassName("negative", results.value())->score,
+              GetCategoryWithClassName("positive", results.value())->score);
   } else {
-    std::vector<Category> results =
-        classifier->Classify("it's a charming and often affecting journey");
-    EXPECT_GT(GetCategoryWithClassName("positive", results)->score,
-              GetCategoryWithClassName("negative", results)->score);
+    tflite::support::StatusOr<std::vector<core::Category>> results =
+        classifier->ClassifyText("it's a charming and often affecting journey");
+
+    EXPECT_TRUE(results.ok());
+    EXPECT_GT(GetCategoryWithClassName("positive", results.value())->score,
+              GetCategoryWithClassName("negative", results.value())->score);
   }
 }
 
@@ -151,11 +155,12 @@ TEST(BertNLClassifierTest, TestNLClassifier_ClassifyLongPositive_notOOB) {
       BertNLClassifier::CreateFromBuffer(model_buffer.data(), model_buffer.size());
   EXPECT_TRUE(classifier.ok());
 
-  std::vector<core::Category> results =
-      classifier.value()->Classify(ss_for_positive_review.str());
+  tflite::support::StatusOr<std::vector<core::Category>> results =
+      classifier.value()->ClassifyText(ss_for_positive_review.str());
 
-  EXPECT_GT(GetCategoryWithClassName("positive", results)->score,
-            GetCategoryWithClassName("negative", results)->score);
+  EXPECT_TRUE(results.ok());
+  EXPECT_GT(GetCategoryWithClassName("positive", results.value())->score,
+            GetCategoryWithClassName("negative", results.value())->score);
 }
 
 }  // namespace
diff --git a/tensorflow_lite_support/java/src/native/task/text/nlclassifier/nl_classifier_jni_utils.cc b/tensorflow_lite_support/java/src/native/task/text/nlclassifier/nl_classifier_jni_utils.cc
index 6b002f71..d2f6e7ca 100644
--- a/tensorflow_lite_support/java/src/native/task/text/nlclassifier/nl_classifier_jni_utils.cc
+++ b/tensorflow_lite_support/java/src/native/task/text/nlclassifier/nl_classifier_jni_utils.cc
@@ -23,6 +23,9 @@ namespace task {
 namespace text {
 namespace nlclassifier {
 
+using ::tflite::support::utils::kAssertionError;
+using ::tflite::support::utils::kInvalidPointer;
+using ::tflite::support::utils::ThrowException;
 using ::tflite::support::utils::ConvertVectorToArrayList;
 using ::tflite::support::utils::JStringToString;
 using ::tflite::task::core::Category;
@@ -31,14 +34,21 @@ using ::tflite::task::text::nlclassifier::NLClassifier;
 jobject RunClassifier(JNIEnv* env, jlong native_handle, jstring text) {
   auto* nl_classifier = reinterpret_cast<NLClassifier*>(native_handle);
 
-  auto results = nl_classifier->Classify(JStringToString(env, text));
+  auto results = nl_classifier->ClassifyText(JStringToString(env, text));
+  if (!results.ok()) {
+        ThrowException(env, kAssertionError,
+                       "Error occurred when running classifier: %s",
+                       results.status().message().data());
+        return env->ExceptionOccurred();
+  }
+
   jclass category_class =
       env->FindClass("org/tensorflow/lite/support/label/Category");
   jmethodID category_init =
       env->GetMethodID(category_class, "<init>", "(Ljava/lang/String;F)V");
 
   return ConvertVectorToArrayList<Category>(
-      env, results,
+      env, results.value(),
       [env, category_class, category_init](const Category& category) {
         jstring class_name = env->NewStringUTF(category.class_name.data());
         // Convert double to float as Java interface exposes float as scores.