Allow specifying the BERT text classifier's model output tensor name …

…via options. PiperOrigin-RevId: 638652416
tensorflow · May 30, 2024 · 8ed4a7b · 8ed4a7b
1 parent 3247254
commit 8ed4a7b
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/tensorflow_lite_support/cc/task/text/bert_nl_classifier.cc b/tensorflow_lite_support/cc/task/text/bert_nl_classifier.cc
@@ -74,7 +74,8 @@ StatusOr<std::vector<core::Category>> BertNLClassifier::Postprocess(
   }
   const TfLiteTensor* scores = FindTensorByName(
       output_tensors, GetMetadataExtractor()->GetOutputTensorMetadata(),
-      kScoreTensorName);
+      options_->has_output_tensor_name() ? options_->output_tensor_name()
+                                         : kScoreTensorName);
 
   // optional labels extracted from metadata
   return BuildResults(scores, /*labels=*/nullptr);

diff --git a/tensorflow_lite_support/cc/task/text/proto/bert_nl_classifier_options.proto b/tensorflow_lite_support/cc/task/text/proto/bert_nl_classifier_options.proto
@@ -20,7 +20,7 @@ package tflite.task.text;
 import "tensorflow_lite_support/cc/task/core/proto/base_options.proto";
 
 // Options for setting up a BertNLClassifier.
-// Next Id: 3
+// Next Id: 4
 message BertNLClassifierOptions {
   // Base options for configuring BertNLClassifier, such as specifying the
   // TfLite model file with metadata, accelerator options, etc.
@@ -31,4 +31,9 @@ message BertNLClassifierOptions {
   // Deprecated: max_seq_len is now read from the model (i.e. input tensor size)
   // automatically.
   optional int32 max_seq_len = 2 [default = 128];
+
+  // The name of the output tensor.
+  //
+  // If not provided, defaults to "probability".
+  optional string output_tensor_name = 3;
 }