Add a tflite status code for Tokenization

PiperOrigin-RevId: 320523560
tensorflow · Jul 10, 2020 · 10db72d · 10db72d
1 parent 30cfc18
commit 10db72d
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 14 deletions.
diff --git a/tensorflow_lite_support/cc/common.h b/tensorflow_lite_support/cc/common.h
@@ -101,6 +101,9 @@ enum class TfLiteSupportStatus {
   // none was found or it was empty.
   // E.g. current task requires labels but none were found.
   kMetadataMissingLabelsError,
+  // The ProcessingUnit for tokenizer is not correctly configured.
+  // E.g BertTokenizer doesn't have a valid vocab file associated.
+  kMetadataInvalidTokenizerError,
 
   // Input tensor(s) error codes.
 

diff --git a/tensorflow_lite_support/cc/task/text/qa/bert_question_answerer.h b/tensorflow_lite_support/cc/task/text/qa/bert_question_answerer.h
@@ -97,9 +97,6 @@ class BertQuestionAnswerer : public QuestionAnswerer {
   static constexpr int kNumLiteThreads = 4;
   static constexpr bool kUseLowerCase = true;
 
-  // Constant for model metadata
-  static constexpr int kTokenizerProcessUnitIndex = 0;
-
   static StatusOr<std::unique_ptr<QuestionAnswerer>>
   CreateQuestionAnswererWithMetadata(
       const std::string& path_to_model_with_metadata);

diff --git a/tensorflow_lite_support/cc/text/tokenizers/tokenizer_utils.cc b/tensorflow_lite_support/cc/text/tokenizers/tokenizer_utils.cc
@@ -24,7 +24,6 @@ limitations under the License.
 
 namespace tflite::support::text::tokenizer {
 
-using ::absl::StatusCode;
 using ::tflite::ProcessUnit;
 using ::tflite::SentencePieceTokenizerOptions;
 using ::tflite::support::CreateStatusWithPayload;
@@ -35,7 +34,9 @@ StatusOr<std::unique_ptr<Tokenizer>> CreateTokenizerFromMetadata(
       metadata_extractor.GetInputProcessUnit(kTokenizerProcessUnitIndex);
   if (tokenizer_process_unit == nullptr) {
     return CreateStatusWithPayload(
-        StatusCode::kNotFound, "No input process unit found from metadata.");
+        absl::StatusCode::kInvalidArgument,
+        "No input process unit found from metadata.",
+        TfLiteSupportStatus::kMetadataInvalidTokenizerError);
   }
   if (tokenizer_process_unit->options_type() ==
       ProcessUnitOptions_BertTokenizerOptions) {
@@ -45,8 +46,9 @@ StatusOr<std::unique_ptr<Tokenizer>> CreateTokenizerFromMetadata(
     if (options->vocab_file() == nullptr || options->vocab_file()->size() < 1 ||
         options->vocab_file()->Get(0)->name() == nullptr) {
       return CreateStatusWithPayload(
-          StatusCode::kInvalidArgument,
-          "Invalid vocab_file from input process unit.");
+          absl::StatusCode::kInvalidArgument,
+          "Invalid vocab_file from input process unit.",
+          TfLiteSupportStatus::kMetadataInvalidTokenizerError);
     }
     ASSIGN_OR_RETURN(vocab_buffer,
                      metadata_extractor.GetAssociatedFile(
@@ -64,8 +66,9 @@ StatusOr<std::unique_ptr<Tokenizer>> CreateTokenizerFromMetadata(
         options->sentencePiece_model()->size() < 1 ||
         options->sentencePiece_model()->Get(0)->name() == nullptr) {
       return CreateStatusWithPayload(
-          StatusCode::kInvalidArgument,
-          "Invalid sentencePiece_model from input process unit.");
+          absl::StatusCode::kInvalidArgument,
+          "Invalid sentencePiece_model from input process unit.",
+          TfLiteSupportStatus::kMetadataInvalidTokenizerError);
     }
     ASSIGN_OR_RETURN(
         model_buffer,
@@ -74,12 +77,13 @@ StatusOr<std::unique_ptr<Tokenizer>> CreateTokenizerFromMetadata(
     // TODO(b/160647204): Extract sentence piece model vocabulary
     return absl::make_unique<SentencePieceTokenizer>(model_buffer.data(),
                                                      model_buffer.size());
+  } else {
+    return CreateStatusWithPayload(
+        absl::StatusCode::kNotFound,
+        absl::StrCat("Incorrect options_type:",
+                     tokenizer_process_unit->options_type()),
+        TfLiteSupportStatus::kMetadataInvalidTokenizerError);
   }
-
-  return CreateStatusWithPayload(
-      StatusCode::kNotFound,
-      absl::StrCat("Incorrect options_type:",
-                   tokenizer_process_unit->options_type()));
 }
 
 }  // namespace tflite::support::text::tokenizer