add support for the tokenized input in NLClassifier.

PiperOrigin-RevId: 322493255

Author: flamearrow

tensorflow_lite_support/cc/task/text/nlclassifier/BUILD

@@ -19,6 +19,9 @@ cc_library(
         "//tensorflow_lite_support/cc/task/core:category",
         "//tensorflow_lite_support/cc/task/core:task_api_factory",
         "//tensorflow_lite_support/cc/task/core:task_utils",
+        "//tensorflow_lite_support/cc/text/tokenizers:regex_tokenizer",
+        "//tensorflow_lite_support/cc/text/tokenizers:tokenizer",
+        "//tensorflow_lite_support/cc/text/tokenizers:tokenizer_utils",
         "//tensorflow_lite_support/cc/utils:common_utils",
         "//tensorflow_lite_support/metadata/cc:metadata_extractor",
         "@com_google_absl//absl/algorithm:container",
@@ -27,6 +30,7 @@ cc_library(
         "@com_google_absl//absl/strings",
         "@flatbuffers",
         "@org_tensorflow//tensorflow/lite:string",
+        "@org_tensorflow//tensorflow/lite:type_to_tflitetype",
         "@org_tensorflow//tensorflow/lite/c:common",
         "@org_tensorflow//tensorflow/lite/core/api",
         "@org_tensorflow//tensorflow/lite/kernels:builtin_ops",

tensorflow_lite_support/cc/task/text/nlclassifier/bert_nl_classifier.cc

@@ -47,7 +47,7 @@ namespace nlclassifier {
 
 using ::tflite::support::task::core::FindTensorByName;
 using ::tflite::support::task::core::PopulateTensor;
-using ::tflite::support::text::tokenizer::CreateTokenizerFromMetadata;
+using ::tflite::support::text::tokenizer::CreateTokenizerFromProcessUnit;
 using ::tflite::support::text::tokenizer::TokenizerResult;
 
 namespace {
@@ -57,6 +57,7 @@ constexpr char kSegmentIdsTensorName[] = "segment_ids";
 constexpr char kScoreTensorName[] = "probability";
 constexpr char kClassificationToken[] = "[CLS]";
 constexpr char kSeparator[] = "[SEP]";
+constexpr int kTokenizerProcessUnitIndex = 0;
 }  // namespace
 
 absl::Status BertNLClassifier::Preprocess(
@@ -160,8 +161,17 @@ BertNLClassifier::CreateBertNLClassifierWithMetadataFromBinary(
 
 absl::Status BertNLClassifier::InitializeFromMetadata() {
   // Set up mandatory tokenizer.
+  const ProcessUnit* tokenizer_process_unit =
+      GetMetadataExtractor()->GetInputProcessUnit(kTokenizerProcessUnitIndex);
+  if (tokenizer_process_unit == nullptr) {
+    return CreateStatusWithPayload(
+        absl::StatusCode::kInvalidArgument,
+        "No input process unit found from metadata.",
+        TfLiteSupportStatus::kMetadataInvalidTokenizerError);
+  }
   ASSIGN_OR_RETURN(tokenizer_,
-                   CreateTokenizerFromMetadata(*GetMetadataExtractor()));
+                   CreateTokenizerFromProcessUnit(tokenizer_process_unit,
+                                                  GetMetadataExtractor()));
 
   // Set up optional label vector.
   TrySetLabelFromMetadata(

tensorflow_lite_support/cc/task/text/nlclassifier/nl_classifier.cc

@@ -29,12 +29,16 @@ limitations under the License.
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/core/api/op_resolver.h"
 #include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/type_to_tflitetype.h"
 #include "tensorflow_lite_support/cc/common.h"
 #include "tensorflow_lite_support/cc/port/status_macros.h"
 #include "tensorflow_lite_support/cc/port/statusor.h"
 #include "tensorflow_lite_support/cc/task/core/category.h"
 #include "tensorflow_lite_support/cc/task/core/task_api_factory.h"
 #include "tensorflow_lite_support/cc/task/core/task_utils.h"
+#include "tensorflow_lite_support/cc/text/tokenizers/regex_tokenizer.h"
+#include "tensorflow_lite_support/cc/text/tokenizers/tokenizer.h"
+#include "tensorflow_lite_support/cc/text/tokenizers/tokenizer_utils.h"
 #include "tensorflow_lite_support/cc/utils/common_utils.h"
 
 namespace tflite {
@@ -52,8 +56,16 @@ using ::tflite::support::StatusOr;
 using ::tflite::support::task::core::Dequantize;
 using ::tflite::support::task::core::GetStringAtIndex;
 using ::tflite::support::task::core::PopulateTensor;
+using ::tflite::support::text::tokenizer::CreateTokenizerFromProcessUnit;
+using ::tflite::support::text::tokenizer::RegexTokenizer;
+using ::tflite::support::text::tokenizer::TokenizerResult;
 using ::tflite::support::utils::LoadVocabFromBuffer;
 
+namespace {
+constexpr int kRegexTokenizerInputTensorIndex = 0;
+constexpr int kRegexTokenizerProcessUnitIndex = 0;
+}  // namespace
+
 const NLClassifierOptions& NLClassifier::GetOptions() const { return options_; }
 
 absl::Status NLClassifier::TrySetLabelFromMetadata(
@@ -102,11 +114,59 @@ std::vector<core::Category> NLClassifier::Classify(const std::string& text) {
 
 absl::Status NLClassifier::Preprocess(
     const std::vector<TfLiteTensor*>& input_tensors, const std::string& input) {
-  PopulateTensor(
-      input,
-      FindTensorWithNameOrIndex(
-          input_tensors, GetMetadataExtractor()->GetInputTensorMetadata(),
-          options_.input_tensor_name, options_.input_tensor_index));
+  TfLiteTensor* input_tensor = FindTensorWithNameOrIndex(
+      input_tensors, GetMetadataExtractor()->GetInputTensorMetadata(),
+      options_.input_tensor_name, options_.input_tensor_index);
+  if (input_tensor == nullptr) {
+    return CreateStatusWithPayload(
+        absl::StatusCode::kInvalidArgument,
+        "No input tensor found from NLClassifierOptions.",
+        TfLiteSupportStatus::kInputTensorNotFoundError);
+  }
+
+  if (HasRegexTokenizerMetadata()) {
+    RETURN_IF_ERROR(SetupRegexTokenizer());
+
+    //                              |<-------sentence_length-------->|
+    // input_tensor                 <START>, t1, t2... <PAD>, <PAD>...
+    // <START> is optional, t1, t2... will be replaced by <UNKNOWN> if it's not
+    // found in tokenizer vocab.
+    TokenizerResult result = tokenizer_->Tokenize(input);
+
+    size_t max_sentence_length = input_tensor->dims->size == 2
+                                     ? input_tensor->dims->data[1]
+                                     : input_tensor->dims->data[0];
+
+    int unknown_token_id = 0;
+    tokenizer_->LookupId(RegexTokenizer::kUnknown, &unknown_token_id);
+
+    int pad_token_id = 0;
+    tokenizer_->LookupId(RegexTokenizer::kPad, &pad_token_id);
+
+    std::vector<float> input_tokens(max_sentence_length, pad_token_id);
+    int start_token_id = 0;
+    size_t input_token_index = 0;
+    if (tokenizer_->LookupId(RegexTokenizer::kStart, &start_token_id)) {
+      input_tokens[0] = start_token_id;
+      input_token_index = 1;
+    }
+
+    for (size_t i = 0; (i < result.subwords.size()) &&
+                       (input_token_index < max_sentence_length);
+         ++i) {
+      const std::string& token = result.subwords[i];
+      int token_id = 0;
+      if (tokenizer_->LookupId(token, &token_id)) {
+        input_tokens[input_token_index] = token_id;
+      } else {
+        input_tokens[input_token_index] = unknown_token_id;
+      }
+    }
+
+    PopulateTensor(input_tokens, input_tensor);
+  } else {
+    PopulateTensor(input, input_tensor);
+  }
   return absl::OkStatus();
 }
 
@@ -172,7 +232,7 @@ absl::Status NLClassifier::Initialize(const NLClassifierOptions& options) {
                      options.input_tensor_index),
         TfLiteSupportStatus::kInputTensorNotFoundError);
   }
-  if (input_tensor->type != kTfLiteString) {
+  if (!HasRegexTokenizerMetadata() && input_tensor->type != kTfLiteString) {
     return CreateStatusWithPayload(
         StatusCode::kInvalidArgument,
         absl::StrCat("Type mismatch for input tensor ", input_tensor->name,
@@ -278,6 +338,38 @@ StatusOr<std::unique_ptr<NLClassifier>> NLClassifier::CreateNLClassifier(
   return std::move(nl_classifier);
 }
 
+bool NLClassifier::HasRegexTokenizerMetadata() {
+  if (GetMetadataExtractor()->GetInputTensorMetadata(
+          kRegexTokenizerInputTensorIndex) == nullptr ||
+      GetMetadataExtractor()
+              ->GetInputTensorMetadata(kRegexTokenizerInputTensorIndex)
+              ->process_units() == nullptr ||
+      GetMetadataExtractor()
+              ->GetInputTensorMetadata(kRegexTokenizerInputTensorIndex)
+              ->process_units()
+              ->Get(kRegexTokenizerProcessUnitIndex) == nullptr) {
+    return false;
+  }
+  return GetMetadataExtractor()
+             ->GetInputTensorMetadata(kRegexTokenizerInputTensorIndex)
+             ->process_units()
+             ->Get(kRegexTokenizerProcessUnitIndex)
+             ->options_type() == ProcessUnitOptions_RegexTokenizerOptions;
+}
+
+absl::Status NLClassifier::SetupRegexTokenizer() {
+  ASSIGN_OR_RETURN(
+      tokenizer_,
+      CreateTokenizerFromProcessUnit(
+          GetMetadataExtractor()
+              ->GetInputTensorMetadata(kRegexTokenizerInputTensorIndex)
+              ->process_units()
+              ->Get(kRegexTokenizerProcessUnitIndex),
+          GetMetadataExtractor()));
+
+  return absl::OkStatus();
+}
+
 }  // namespace nlclassifier
 }  // namespace text
 }  // namespace task

tensorflow_lite_support/cc/task/text/nlclassifier/nl_classifier.h

@@ -34,6 +34,7 @@ limitations under the License.
 #include "tensorflow_lite_support/cc/port/statusor.h"
 #include "tensorflow_lite_support/cc/task/core/base_task_api.h"
 #include "tensorflow_lite_support/cc/task/core/category.h"
+#include "tensorflow_lite_support/cc/text/tokenizers/tokenizer.h"
 
 namespace tflite {
 namespace support {
@@ -59,6 +60,8 @@ struct NLClassifierOptions {
 // The API expects a TFLite model with the following input/output tensor:
 // Input tensor:
 //   (kTfLiteString) - input of the model, accepts a string.
+//      or
+//   (kTfLiteFloat32) - input of the model, accepts a tokenized input of a string
 // Output score tensor:
 //   (kTfLiteUInt8/kTfLiteInt8/kTfLiteInt16/kTfLiteFloat32/kTfLiteFloat64)
 //    - output scores for each class, if type is one of the Int types,
@@ -155,10 +158,14 @@ class NLClassifier : public core::BaseTaskApi<std::vector<core::Category>,
   }
 
  private:
+  bool HasRegexTokenizerMetadata();
+  absl::Status SetupRegexTokenizer();
+
   NLClassifierOptions options_;
   // labels vector initialized from output tensor's associated file, if one
   // exists.
   std::unique_ptr<std::vector<std::string>> labels_vector_;
+  std::unique_ptr<tflite::support::text::tokenizer::Tokenizer> tokenizer_;
 };
 
 }  // namespace nlclassifier

tensorflow_lite_support/cc/task/text/qa/bert_question_answerer.cc

@@ -34,10 +34,14 @@ using ::tflite::support::task::core::PopulateTensor;
 using ::tflite::support::task::core::PopulateVector;
 using ::tflite::support::task::core::ReverseSortIndices;
 using ::tflite::support::text::tokenizer::BertTokenizer;
-using ::tflite::support::text::tokenizer::CreateTokenizerFromMetadata;
+using ::tflite::support::text::tokenizer::CreateTokenizerFromProcessUnit;
 using ::tflite::support::text::tokenizer::SentencePieceTokenizer;
 using ::tflite::support::text::tokenizer::TokenizerResult;
 
+namespace {
+constexpr int kTokenizerProcessUnitIndex = 0;
+}
+
 StatusOr<std::unique_ptr<QuestionAnswerer>>
 BertQuestionAnswerer::CreateQuestionAnswererWithMetadata(
     const std::string& path_to_model_with_metadata) {
@@ -327,8 +331,17 @@ std::string BertQuestionAnswerer::ConvertIndexToString(int start, int end) {
 }
 
 absl::Status BertQuestionAnswerer::InitializeFromMetadata() {
+  const ProcessUnit* tokenizer_process_unit =
+      GetMetadataExtractor()->GetInputProcessUnit(kTokenizerProcessUnitIndex);
+  if (tokenizer_process_unit == nullptr) {
+    return CreateStatusWithPayload(
+        absl::StatusCode::kInvalidArgument,
+        "No input process unit found from metadata.",
+        TfLiteSupportStatus::kMetadataInvalidTokenizerError);
+  }
   ASSIGN_OR_RETURN(tokenizer_,
-                   CreateTokenizerFromMetadata(*GetMetadataExtractor()));
+                   CreateTokenizerFromProcessUnit(tokenizer_process_unit,
+                                                  GetMetadataExtractor()));
   return absl::OkStatus();
 }