From d6d6a70910c9fa020400b9ce60e16b856897b715 Mon Sep 17 00:00:00 2001 From: "TF.Text Team" Date: Mon, 5 Oct 2026 14:30:58 -0700 Subject: [PATCH] Validate `row_splits` bounds and monotonicity in TF.Text tokenization and utility kernels. PiperOrigin-RevId: 993905355 --- tensorflow_text/core/kernels/BUILD | 5 + .../boise_offset_converter_kernel_template.h | 114 +++++++++++--- .../kernels/byte_splitter_kernel_template.h | 39 ++++- ...fast_wordpiece_tokenizer_kernel_template.h | 41 +++-- .../core/kernels/ngrams_kernel_template.h | 68 +++++++-- .../phrase_tokenizer_kernel_template.h | 37 ++++- .../core/kernels/round_robin_trimmer.h | 37 ++++- .../round_robin_trimmer_kernel_template.h | 53 ++++++- .../kernels/split_merge_tokenize_kernel.cc | 39 ++++- .../python/ops/boise_offset_converter_test.py | 143 ++++++++++++++++++ .../python/ops/byte_splitter_test.py | 55 +++++++ .../ops/fast_wordpiece_tokenizer_test.py | 54 +++++++ tensorflow_text/python/ops/ngrams_op_test.py | 61 ++++++++ .../python/ops/phrase_tokenizer_test.py | 58 +++++++ .../python/ops/split_merge_tokenizer_test.py | 58 +++++++ .../python/ops/trimmer_ops_test.py | 99 ++++++++++++ 16 files changed, 894 insertions(+), 67 deletions(-) diff --git a/tensorflow_text/core/kernels/BUILD b/tensorflow_text/core/kernels/BUILD index 408c6e16e..1770549e0 100644 --- a/tensorflow_text/core/kernels/BUILD +++ b/tensorflow_text/core/kernels/BUILD @@ -63,6 +63,7 @@ tf_cc_library( deps = [ ":boise_offset_converter", "@com_google_absl//absl/status", + "@com_google_absl//absl/strings", # lite/kernels/shim:op_kernel tensorflow dep, # lite/kernels/shim:shape tensorflow dep, # lite/kernels/shim:status_macros tensorflow dep, @@ -113,6 +114,7 @@ tf_cc_library( deps = [ ":byte_splitter", "@com_google_absl//absl/status", + "@com_google_absl//absl/strings", # lite/kernels/shim:op_kernel tensorflow dep, # lite/kernels/shim:shape tensorflow dep, # lite/kernels/shim:status_macros tensorflow dep, @@ -498,6 +500,7 @@ cc_library( ":fast_wordpiece_tokenizer", "@com_google_absl//absl/status", "@com_google_absl//absl/strings", + "@com_google_absl//absl/types:span", # lite/kernels/shim:op_kernel tensorflow dep, # lite/kernels/shim:status_macros tensorflow dep, ], @@ -612,6 +615,7 @@ tf_cc_library( "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", + "@com_google_absl//absl/types:span", # lite/kernels/shim:op_kernel tensorflow dep, # lite/kernels/shim:status_macros tensorflow dep, # lite/kernels/shim:tensor_view tensorflow dep, @@ -1298,6 +1302,7 @@ cc_library( ":phrase_tokenizer", "@com_google_absl//absl/status", "@com_google_absl//absl/strings", + "@com_google_absl//absl/types:span", # lite/kernels/shim:op_kernel tensorflow dep, # lite/kernels/shim:status_macros tensorflow dep, ], diff --git a/tensorflow_text/core/kernels/boise_offset_converter_kernel_template.h b/tensorflow_text/core/kernels/boise_offset_converter_kernel_template.h index f49f059aa..33141793e 100644 --- a/tensorflow_text/core/kernels/boise_offset_converter_kernel_template.h +++ b/tensorflow_text/core/kernels/boise_offset_converter_kernel_template.h @@ -21,6 +21,7 @@ #include #include "absl/status/status.h" +#include "absl/strings/str_cat.h" #include "tensorflow/core/platform/tstring.h" #include "tensorflow/lite/kernels/shim/op_kernel.h" #include "tensorflow/lite/kernels/shim/shape.h" @@ -265,16 +266,18 @@ absl::Status OffsetsToBoiseTagsOp::Invoke(InvokeContext* context) { } // Check row splits are the same for token begin, end offsets. - if (input_token_begin_row_splits_vec.Dim(0) != + if (input_token_begin_row_splits_vec.Dim(0) == 0 || + input_token_begin_row_splits_vec.Dim(0) != input_token_end_row_splits_vec.Dim(0) || - input_span_begin_row_splits_vec.Dim(0) != + input_token_begin_row_splits_vec.Dim(0) != input_span_begin_row_splits_vec.Dim(0) || input_span_begin_row_splits_vec.Dim(0) != input_span_end_row_splits_vec.Dim(0) || input_span_begin_row_splits_vec.Dim(0) != input_span_type_row_splits_vec.Dim(0)) { return absl::InvalidArgumentError(absl::StrCat( - "Row splits must have the same size for token and span. ", + "Row splits must be non-empty and have the same size for token and " + "span. ", " Token begin row splits shape: ", input_token_begin_row_splits_vec.Dim(0), " Token end row splits shape: ", input_token_end_row_splits_vec.Dim(0), @@ -284,7 +287,7 @@ absl::Status OffsetsToBoiseTagsOp::Invoke(InvokeContext* context) { input_span_type_row_splits_vec.Dim(0))); } - for (int i = 0; i < input_token_begin_row_splits_vec.Dim(0) - 1; ++i) { + for (int i = 0; i < input_token_begin_row_splits_vec.Dim(0); ++i) { if (input_token_begin_row_splits_vec(i) != input_token_end_row_splits_vec(i)) { return absl::InvalidArgumentError( @@ -293,7 +296,7 @@ absl::Status OffsetsToBoiseTagsOp::Invoke(InvokeContext* context) { } // Check row splits are the same for span begin, end offsets and span type. - for (int i = 0; i < input_span_begin_row_splits_vec.Dim(0) - 1; ++i) { + for (int i = 0; i < input_span_begin_row_splits_vec.Dim(0); ++i) { if (input_span_begin_row_splits_vec(i) != input_span_end_row_splits_vec(i) || input_span_begin_row_splits_vec(i) != @@ -304,8 +307,40 @@ absl::Status OffsetsToBoiseTagsOp::Invoke(InvokeContext* context) { } } + const int num_rows = input_token_begin_row_splits_vec.Dim(0) - 1; + if (input_token_begin_row_splits_vec(0) < 0 || + input_token_begin_row_splits_vec(num_rows) > + input_token_begin_offsets_vec.Dim(0)) { + return absl::InvalidArgumentError( + "Invalid token row splits: splits must be non-negative, " + "monotonically increasing, and within token offsets bounds."); + } + if (input_span_begin_row_splits_vec(0) < 0 || + input_span_begin_row_splits_vec(num_rows) > + input_span_begin_offsets_vec.Dim(0)) { + return absl::InvalidArgumentError( + "Invalid span row splits: splits must be non-negative, " + "monotonically increasing, and within span offsets bounds."); + } + for (int i = 0; i < num_rows; ++i) { + if (input_token_begin_row_splits_vec(i + 1) < + input_token_begin_row_splits_vec(i)) { + return absl::InvalidArgumentError( + "Invalid token row splits: splits must be non-negative, " + "monotonically increasing, and within token offsets bounds."); + } + if (input_span_begin_row_splits_vec(i + 1) < + input_span_begin_row_splits_vec(i)) { + return absl::InvalidArgumentError( + "Invalid span row splits: splits must be non-negative, " + "monotonically increasing, and within span offsets bounds."); + } + } + // Outputs std::vector boise_tags; + boise_tags.reserve(input_token_begin_row_splits_vec(num_rows) - + input_token_begin_row_splits_vec(0)); std::vector input_token_begin_offsets_vec_i; std::vector input_token_end_offsets_vec_i; std::vector input_span_begin_offsets_vec_i; @@ -313,11 +348,11 @@ absl::Status OffsetsToBoiseTagsOp::Invoke(InvokeContext* context) { std::vector input_span_type_vec_i; // Iterate through all the input values and split them. - for (int i = 0; i < input_token_begin_row_splits_vec.Dim(0) - 1; ++i) { - int token_start_index = input_token_begin_row_splits_vec(i); - int token_end_index = input_token_begin_row_splits_vec(i + 1); - int span_start_index = input_span_begin_row_splits_vec(i); - int span_end_index = input_span_begin_row_splits_vec(i + 1); + for (int i = 0; i < num_rows; ++i) { + const int64_t token_start_index = input_token_begin_row_splits_vec(i); + const int64_t token_end_index = input_token_begin_row_splits_vec(i + 1); + const int64_t span_start_index = input_span_begin_row_splits_vec(i); + const int64_t span_end_index = input_span_begin_row_splits_vec(i + 1); input_token_begin_offsets_vec_i.clear(); input_token_end_offsets_vec_i.clear(); @@ -325,6 +360,14 @@ absl::Status OffsetsToBoiseTagsOp::Invoke(InvokeContext* context) { input_span_end_offsets_vec_i.clear(); input_span_type_vec_i.clear(); + const int64_t num_tokens_i = token_end_index - token_start_index; + input_token_begin_offsets_vec_i.reserve(num_tokens_i); + input_token_end_offsets_vec_i.reserve(num_tokens_i); + const int64_t num_spans_i = span_end_index - span_start_index; + input_span_begin_offsets_vec_i.reserve(num_spans_i); + input_span_end_offsets_vec_i.reserve(num_spans_i); + input_span_type_vec_i.reserve(num_spans_i); + for (int j = token_start_index; j < token_end_index; ++j) { input_token_begin_offsets_vec_i.push_back( input_token_begin_offsets_vec(j)); @@ -537,13 +580,14 @@ absl::Status BoiseTagsToOffsetsOp::Invoke(InvokeContext* context) { // Check row splits are the same for token begin, end offsets and boise tags. // First, check dimensions are the same. - if (input_token_begin_row_splits_vec.Dim(0) != + if (input_token_begin_row_splits_vec.Dim(0) == 0 || + input_token_begin_row_splits_vec.Dim(0) != input_token_end_row_splits_vec.Dim(0) || input_token_begin_row_splits_vec.Dim(0) != input_boise_tags_row_splits_vec.Dim(0)) { return absl::InvalidArgumentError(absl::StrCat( - "Row splits must have the same size for token begin/end offsets and " - "BOISE tags. ", + "Row splits must be non-empty and have the same size for token " + "begin/end offsets and BOISE tags. ", " Token begin row splits shape: ", input_token_begin_row_splits_vec.Dim(0), " Token end row splits shape: ", input_token_end_row_splits_vec.Dim(0), @@ -551,22 +595,40 @@ absl::Status BoiseTagsToOffsetsOp::Invoke(InvokeContext* context) { input_boise_tags_row_splits_vec.Dim(0))); } // Second, check values are the same. - for (int i = 0; i < input_token_begin_row_splits_vec.Dim(0) - 1; ++i) { + for (int i = 0; i < input_token_begin_row_splits_vec.Dim(0); ++i) { if (input_token_begin_row_splits_vec(i) != input_token_end_row_splits_vec(i) || input_token_begin_row_splits_vec(i) != input_boise_tags_row_splits_vec(i)) { return absl::InvalidArgumentError( - "Row splits must be the same for token begin/end offsets ad BOISE " + "Row splits must be the same for token begin/end offsets and BOISE " "tags."); } } + const int num_rows = input_token_begin_row_splits_vec.Dim(0) - 1; + if (input_token_begin_row_splits_vec(0) < 0 || + input_token_begin_row_splits_vec(num_rows) > + input_token_begin_offsets_vec.Dim(0)) { + return absl::InvalidArgumentError( + "Invalid row splits: splits must be non-negative, monotonically " + "increasing, and within token offsets bounds."); + } + for (int i = 0; i < num_rows; ++i) { + if (input_token_begin_row_splits_vec(i + 1) < + input_token_begin_row_splits_vec(i)) { + return absl::InvalidArgumentError( + "Invalid row splits: splits must be non-negative, monotonically " + "increasing, and within token offsets bounds."); + } + } + // Outputs std::vector span_begin_offsets; std::vector span_end_offsets; std::vector span_type; std::vector row_splits; + row_splits.reserve(input_token_begin_row_splits_vec.Dim(0)); row_splits.push_back(0); @@ -574,14 +636,19 @@ absl::Status BoiseTagsToOffsetsOp::Invoke(InvokeContext* context) { std::vector input_token_begin_offsets_vec_i; std::vector input_token_end_offsets_vec_i; std::vector input_boise_tags_vec_i; - for (int i = 0; i < input_token_begin_row_splits_vec.Dim(0) - 1; ++i) { - int token_start_index = input_token_begin_row_splits_vec(i); - int token_end_index = input_token_begin_row_splits_vec(i + 1); + for (int i = 0; i < num_rows; ++i) { + const int64_t token_start_index = input_token_begin_row_splits_vec(i); + const int64_t token_end_index = input_token_begin_row_splits_vec(i + 1); input_token_begin_offsets_vec_i.clear(); input_token_end_offsets_vec_i.clear(); input_boise_tags_vec_i.clear(); + const int64_t num_tokens_i = token_end_index - token_start_index; + input_token_begin_offsets_vec_i.reserve(num_tokens_i); + input_token_end_offsets_vec_i.reserve(num_tokens_i); + input_boise_tags_vec_i.reserve(num_tokens_i); + for (int j = token_start_index; j < token_end_index; ++j) { input_token_begin_offsets_vec_i.push_back( input_token_begin_offsets_vec(j)); @@ -589,11 +656,12 @@ absl::Status BoiseTagsToOffsetsOp::Invoke(InvokeContext* context) { input_boise_tags_vec_i.push_back(input_boise_tags_vec(j)); } - auto [span_begin_offsets_i, span_end_offsets_i, span_type_i] = - BoiseTagsToOffsets(input_token_begin_offsets_vec_i, - input_token_end_offsets_vec_i, - input_boise_tags_vec_i) - .value(); + SH_ASSIGN_OR_RETURN(auto span_info_i, + BoiseTagsToOffsets(input_token_begin_offsets_vec_i, + input_token_end_offsets_vec_i, + input_boise_tags_vec_i)); + const auto& [span_begin_offsets_i, span_end_offsets_i, span_type_i] = + span_info_i; const int num_span_i = span_type_i.size(); row_splits.push_back(row_splits.back() + num_span_i); diff --git a/tensorflow_text/core/kernels/byte_splitter_kernel_template.h b/tensorflow_text/core/kernels/byte_splitter_kernel_template.h index 77ab2b1ba..8064ce2f9 100644 --- a/tensorflow_text/core/kernels/byte_splitter_kernel_template.h +++ b/tensorflow_text/core/kernels/byte_splitter_kernel_template.h @@ -15,10 +15,12 @@ #ifndef THIRD_PARTY_TENSORFLOW_TEXT_CORE_KERNELS_BYTE_SPLITTER_KERNEL_TEMPLATE_H_ #define THIRD_PARTY_TENSORFLOW_TEXT_CORE_KERNELS_BYTE_SPLITTER_KERNEL_TEMPLATE_H_ +#include #include #include #include "absl/status/status.h" +#include "absl/strings/str_cat.h" #include "tensorflow/core/platform/tstring.h" #include "tensorflow/lite/kernels/shim/op_kernel.h" #include "tensorflow/lite/kernels/shim/shape.h" @@ -277,22 +279,49 @@ template context->GetInput(kInputRowSplits)); const auto in_splits = in_splits_view->template As(); + if (in_splits.Dim(0) != input_values.Dim(0) + 1) { + return absl::InvalidArgumentError(absl::StrCat( + "Expecting input_row_splits to have ", input_values.Dim(0) + 1, + " elements, got ", in_splits.Dim(0))); + } + if (starts.Dim(0) != ends.Dim(0)) { + return absl::InvalidArgumentError(absl::StrCat( + "input_start_offsets and input_end_offsets must have the same size, " + "got ", + starts.Dim(0), " and ", ends.Dim(0))); + } + + if (in_splits(0) < 0 || in_splits(in_splits.Dim(0) - 1) > starts.Dim(0)) { + return absl::InvalidArgumentError( + "Invalid input_row_splits: splits must be non-negative, " + "monotonically increasing, and within offsets bounds."); + } + for (int i = 0; i < input_values.Dim(0); ++i) { + if (in_splits(i + 1) < in_splits(i)) { + return absl::InvalidArgumentError( + "Invalid input_row_splits: splits must be non-negative, " + "monotonically increasing, and within offsets bounds."); + } + } + ByteSplitter splitter; // Outputs std::vector output_values; std::vector out_splits; + out_splits.reserve(in_splits.Dim(0)); // Iterate through all the string values and split them. out_splits.push_back(0); for (int i = 0; i < input_values.Dim(0); ++i) { - SH_ASSIGN_OR_RETURN(auto batch, + const int64_t row_start = in_splits(i); + const int64_t row_end = in_splits(i + 1); + SH_ASSIGN_OR_RETURN( + auto batch, splitter.SplitByOffsets( input_values(i), - absl::MakeSpan(starts.Ptr() + in_splits(i), - in_splits(i+1) - in_splits(i)), - absl::MakeSpan(ends.Ptr() + in_splits(i), - in_splits(i+1) - in_splits(i)))); + absl::MakeSpan(starts.Ptr() + row_start, row_end - row_start), + absl::MakeSpan(ends.Ptr() + row_start, row_end - row_start))); output_values.insert(output_values.end(), batch.begin(), batch.end()); out_splits.push_back(batch.size() + out_splits.back()); } diff --git a/tensorflow_text/core/kernels/fast_wordpiece_tokenizer_kernel_template.h b/tensorflow_text/core/kernels/fast_wordpiece_tokenizer_kernel_template.h index efc26197a..0f3831281 100644 --- a/tensorflow_text/core/kernels/fast_wordpiece_tokenizer_kernel_template.h +++ b/tensorflow_text/core/kernels/fast_wordpiece_tokenizer_kernel_template.h @@ -15,8 +15,12 @@ #ifndef THIRD_PARTY_TENSORFLOW_TEXT_CORE_KERNELS_FAST_WORDPIECE_TOKENIZER_KERNEL_TEMPLATE_H_ #define THIRD_PARTY_TENSORFLOW_TEXT_CORE_KERNELS_FAST_WORDPIECE_TOKENIZER_KERNEL_TEMPLATE_H_ +#include +#include + #include "absl/status/status.h" #include "absl/strings/str_cat.h" +#include "absl/types/span.h" #include "tensorflow/lite/kernels/shim/op_kernel.h" #include "tensorflow/lite/kernels/shim/status_macros.h" #include "tensorflow_text/core/kernels/fast_wordpiece_tokenizer.h" @@ -180,13 +184,13 @@ absl::Status FastWordpieceTokenizeWithOffsetsOp::Invoke( SH_RETURN_IF_ERROR(this->template FillOutputTensor( subwords, kOutputSubwords, context)); - SH_RETURN_IF_ERROR(this->template FillOutputTensor( + SH_RETURN_IF_ERROR(this->template FillOutputTensor( subword_ids, kOutputIds, context)); - SH_RETURN_IF_ERROR(this->template FillOutputTensor( + SH_RETURN_IF_ERROR(this->template FillOutputTensor( row_splits, kOutputRowSplits, context)); - SH_RETURN_IF_ERROR(this->template FillOutputTensor( + SH_RETURN_IF_ERROR(this->template FillOutputTensor( begin_offset, kStartValues, context)); - SH_RETURN_IF_ERROR(this->template FillOutputTensor( + SH_RETURN_IF_ERROR(this->template FillOutputTensor( end_offset, kEndValues, context)); return absl::OkStatus(); @@ -311,7 +315,24 @@ absl::Status FastWordpieceDetokenizeOp::Invoke(InvokeContext* context) { SH_ASSIGN_OR_RETURN(const auto input_row_splits, context->GetInput(kInputRowSplits)); - const auto& row_splits_vec = input_row_splits->template As(); + const auto& row_splits_vec = input_row_splits->template As(); + + if (row_splits_vec.Dim(0) == 0) { + return absl::InvalidArgumentError("input_row_splits must not be empty."); + } + if (row_splits_vec(0) < 0 || + row_splits_vec(row_splits_vec.Dim(0) - 1) > values_vec.Dim(0)) { + return absl::InvalidArgumentError( + "Invalid input_row_splits: splits must be non-negative, " + "monotonically increasing, and within input_values bounds."); + } + for (int i = 0; i < row_splits_vec.Dim(0) - 1; ++i) { + if (row_splits_vec(i + 1) < row_splits_vec(i)) { + return absl::InvalidArgumentError( + "Invalid input_row_splits: splits must be non-negative, " + "monotonically increasing, and within input_values bounds."); + } + } SH_ASSIGN_OR_RETURN(const auto wp_model, context->GetInput(kWpModel)); // OK to create on every call because FastWordpieceTokenizer is a @@ -323,15 +344,17 @@ absl::Status FastWordpieceDetokenizeOp::Invoke(InvokeContext* context) { SH_RETURN_IF_ERROR(fast_wordpiece_tokenizer.status()); std::vector sentences; + sentences.reserve(row_splits_vec.Dim(0) - 1); // Iterate through row_splits to split input_values. for (int i = 0; i < row_splits_vec.Dim(0) - 1; ++i) { - auto single_input = - absl::Span(values_vec.Ptr() + row_splits_vec(i), - row_splits_vec(i + 1) - row_splits_vec(i)); + const int64_t row_start = row_splits_vec(i); + const int64_t row_end = row_splits_vec(i + 1); + auto single_input = absl::Span(values_vec.Ptr() + row_start, + row_end - row_start); SH_ASSIGN_OR_RETURN(auto sentence, fast_wordpiece_tokenizer->Detokenize(single_input)); - sentences.push_back(sentence); + sentences.push_back(std::move(sentence)); } const int words_size = sentences.size(); diff --git a/tensorflow_text/core/kernels/ngrams_kernel_template.h b/tensorflow_text/core/kernels/ngrams_kernel_template.h index 1a8a3fc8f..7db0abc73 100644 --- a/tensorflow_text/core/kernels/ngrams_kernel_template.h +++ b/tensorflow_text/core/kernels/ngrams_kernel_template.h @@ -29,10 +29,13 @@ limitations under the License. #ifndef TENSORFLOW_TEXT_CORE_KERNELS_NGRAMS_KERNEL_TEMPLATE_H_ #define TENSORFLOW_TEXT_CORE_KERNELS_NGRAMS_KERNEL_TEMPLATE_H_ +#include + #include "absl/status/status.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_join.h" #include "absl/strings/string_view.h" +#include "absl/types/span.h" #include "tensorflow/core/platform/tstring.h" #include "tensorflow/lite/kernels/shim/op_kernel.h" #include "tensorflow/lite/kernels/shim/status_macros.h" @@ -162,12 +165,20 @@ class NgramsStringJoin : public tflite::shim::OpKernelShimGetInput(kValues)); const Shape input_values_shape(input_values->Shape()); + size_t num_input_values = 1; + for (const int dim : *input_values_shape) { + num_input_values *= dim; + } + // TfTensorView::Data() uses TotalBytes() / sizeof(T), which over-counts + // for tstring tensors because TotalBytes() includes string heap payloads. + const auto input_values_data = + input_values->template Data().subspan( + 0, num_input_values); // Tensor output if (ctx->NumOutputs() == 1) { // Generate mock input and output innermost row_splits. - int64_t total_tokens = - input_values->template Data().size(); + int64_t total_tokens = input_values_data.size(); int64_t tokens_per_element = input_values_shape->at(input_values_shape->size() - 1); tensor_output_row_splits.resize(total_tokens / tokens_per_element + 1); @@ -181,16 +192,48 @@ class NgramsStringJoin : public tflite::shim::OpKernelShimNumInputs() - kRowSplitsStart; + auto validate_splits = [](absl::Span splits, + size_t max_val) -> absl::Status { + if (splits.empty()) { + return absl::InvalidArgumentError( + "input_row_splits must not be empty."); + } + if (splits[0] < 0) { + return absl::InvalidArgumentError( + "input_row_splits must be non-negative."); + } + for (size_t i = 0; i + 1 < splits.size(); ++i) { + if (splits[i + 1] < splits[i]) { + return absl::InvalidArgumentError( + "input_row_splits must be monotonically increasing."); + } + } + if (static_cast(splits.back()) > max_val) { + return absl::InvalidArgumentError( + "input_row_splits exceeds input_values size."); + } + return absl::OkStatus(); + }; // Copy all input splits except for innermost into output splits. while (index < num_row_splits - 1) { SH_ASSIGN_OR_RETURN(const auto input_tensor_row_splits, ctx->GetInput(kRowSplitsStart + index)); + SH_ASSIGN_OR_RETURN(const auto next_tensor_row_splits, + ctx->GetInput(kRowSplitsStart + index + 1)); + const auto input_buffer = + input_tensor_row_splits->template Data(); + const auto next_buffer = + next_tensor_row_splits->template Data(); + if (next_buffer.empty()) { + return absl::InvalidArgumentError( + "input_row_splits must not be empty."); + } + SH_RETURN_IF_ERROR( + validate_splits(input_buffer, next_buffer.size() - 1)); SH_ASSIGN_OR_RETURN( const auto output_tensor_row_splits, ctx->GetOutput(kRowSplitsStart + index, Shape(input_tensor_row_splits->Shape()))); - const auto input_buffer = - input_tensor_row_splits->template Data(); const auto output_buffer = output_tensor_row_splits->template Data(); std::memcpy(output_buffer.data(), input_buffer.data(), @@ -200,28 +243,29 @@ class NgramsStringJoin : public tflite::shim::OpKernelShimGetInput(kRowSplitsStart + index)); + const auto input_splits_data = + input_tensor_row_splits->template Data(); + SH_RETURN_IF_ERROR( + validate_splits(input_splits_data, input_values_data.size())); + input_row_splits = input_splits_data.data(); + n_row_splits = input_splits_data.size(); + SH_ASSIGN_OR_RETURN( const auto output_tensor_row_splits, ctx->GetOutput(kRowSplitsStart + index, Shape(input_tensor_row_splits->Shape()))); - input_row_splits = - input_tensor_row_splits->template Data().data(); output_row_splits = output_tensor_row_splits->template Data().data(); - n_row_splits = input_tensor_row_splits->Shape().at(0); } - const auto input_values_data = - input_values->template Data(); - // Create ngrams by looping through the innermost input splits. std::vector buffer; for (int i = 0; i < n_row_splits - 1; ++i) { // Set output splits using current number of created output values. output_row_splits[i] = buffer.size(); std::vector tokens; - for (int j = input_row_splits[i]; j < input_row_splits[i + 1]; ++j) { - tokens.emplace_back(input_values_data.at(j)); + for (Tsplits j = input_row_splits[i]; j < input_row_splits[i + 1]; ++j) { + tokens.emplace_back(input_values_data[j]); if (tokens.size() < width_) continue; tokens.erase(tokens.begin(), tokens.begin() + tokens.size() - width_); buffer.push_back(absl::StrJoin(tokens, string_separator_)); diff --git a/tensorflow_text/core/kernels/phrase_tokenizer_kernel_template.h b/tensorflow_text/core/kernels/phrase_tokenizer_kernel_template.h index 67807a768..22e171975 100644 --- a/tensorflow_text/core/kernels/phrase_tokenizer_kernel_template.h +++ b/tensorflow_text/core/kernels/phrase_tokenizer_kernel_template.h @@ -15,8 +15,12 @@ #ifndef THIRD_PARTY_TENSORFLOW_TEXT_CORE_KERNELS_PHRASE_TOKENIZER_KERNEL_TEMPLATE_H_ #define THIRD_PARTY_TENSORFLOW_TEXT_CORE_KERNELS_PHRASE_TOKENIZER_KERNEL_TEMPLATE_H_ +#include +#include + #include "absl/status/status.h" #include "absl/strings/str_cat.h" +#include "absl/types/span.h" #include "tensorflow/lite/kernels/shim/op_kernel.h" #include "tensorflow/lite/kernels/shim/status_macros.h" #include "tensorflow_text/core/kernels/phrase_tokenizer.h" @@ -159,13 +163,13 @@ absl::Status PhraseTokenizeOp::Invoke(InvokeContext* context) { kOutputIds, Shape({static_cast( subword_ids.size())}))); /* same shape as `output_subwords` */ - auto output_ids_vec = output_ids->template As(); + auto output_ids_vec = output_ids->template As(); SH_ASSIGN_OR_RETURN( auto output_row_splits, context->GetOutput(kOutputRowSplits, Shape({static_cast(row_splits.size())}))); - auto output_row_splits_vec = output_row_splits->template As(); + auto output_row_splits_vec = output_row_splits->template As(); for (int i = 0; i < subwords.size(); ++i) { output_subwords_vec(i) = subwords[i]; @@ -299,7 +303,24 @@ absl::Status PhraseDetokenizeOp::Invoke(InvokeContext* context) { SH_ASSIGN_OR_RETURN(const auto input_row_splits, context->GetInput(kInputRowSplits)); - const auto& row_splits_vec = input_row_splits->template As(); + const auto& row_splits_vec = input_row_splits->template As(); + + if (row_splits_vec.Dim(0) == 0) { + return absl::InvalidArgumentError("input_row_splits must not be empty."); + } + if (row_splits_vec(0) < 0 || + row_splits_vec(row_splits_vec.Dim(0) - 1) > values_vec.Dim(0)) { + return absl::InvalidArgumentError( + "Invalid input_row_splits: splits must be non-negative, " + "monotonically increasing, and within input_values bounds."); + } + for (int i = 0; i < row_splits_vec.Dim(0) - 1; ++i) { + if (row_splits_vec(i + 1) < row_splits_vec(i)) { + return absl::InvalidArgumentError( + "Invalid input_row_splits: splits must be non-negative, " + "monotonically increasing, and within input_values bounds."); + } + } SH_ASSIGN_OR_RETURN(const auto phrase_model, context->GetInput(kPhraseModel)); // OK to create on every call because PhraseTokenizer is a @@ -310,15 +331,17 @@ absl::Status PhraseDetokenizeOp::Invoke(InvokeContext* context) { SH_RETURN_IF_ERROR(phrase_tokenizer.status()); std::vector sentences; + sentences.reserve(row_splits_vec.Dim(0) - 1); // Iterate through row_splits to split input_values. for (int i = 0; i < row_splits_vec.Dim(0) - 1; ++i) { - auto single_input = - absl::Span(values_vec.Ptr() + row_splits_vec(i), - row_splits_vec(i + 1) - row_splits_vec(i)); + const int64_t row_start = row_splits_vec(i); + const int64_t row_end = row_splits_vec(i + 1); + auto single_input = absl::Span(values_vec.Ptr() + row_start, + row_end - row_start); SH_ASSIGN_OR_RETURN(auto sentence, phrase_tokenizer->Detokenize(single_input)); - sentences.push_back(sentence); + sentences.push_back(std::move(sentence)); } SH_RETURN_IF_ERROR(this->template FillOutputTensor +#include #include #include #include -#include "tensorflow_text/core/kernels/trimmer.h" +#include "tensorflow_text/core/kernels/trimmer.h" namespace tensorflow { namespace text { @@ -153,7 +154,9 @@ std::vector RoundRobinTrimmer::GenerateMasksInternal( std::vector masks(end - begin); auto m = masks.begin(); for (auto it = begin; it != end; ++it, ++m) { - m->reserve(it->back()); + if (!it->empty() && it->back() > 0) { + m->reserve(it->back()); + } } // Process all batches, updating the masks a batch at a time. ProcessSplitsByBatch(begin, end, [&masks](std::vector* rows) { @@ -161,7 +164,8 @@ std::vector RoundRobinTrimmer::GenerateMasksInternal( const Row& row = (*rows)[s]; masks[s].reserve(row.size); masks[s].insert(masks[s].end(), row.used, true); - masks[s].insert(masks[s].end(), row.size - row.used, false); + masks[s].insert(masks[s].end(), std::max(0, row.size - row.used), + false); } }); return masks; @@ -220,9 +224,20 @@ RoundRobinTrimmer::TrimInternal( for (int s = 0; s < values_row->size(); ++s, ++values_it, ++splits_it) { Values_* vals = &trimmed.first[s]; RowSplits_* splits = &trimmed.second[s]; - auto start = values_it->begin() + (*splits_it)[splits->size()-1]; - vals->insert(vals->end(), start, start + (*values_row)[s].used); - splits->insert(splits->end(), splits->back() + (*values_row)[s].used); + const size_t split_idx = splits->size() - 1; + int used = (*values_row)[s].used; + if (split_idx < splits_it->size() && (*splits_it)[split_idx] >= 0 && + used >= 0 && + static_cast((*splits_it)[split_idx]) + + static_cast(used) <= + values_it->size()) { + Tsplits split_offset = (*splits_it)[split_idx]; + auto start = values_it->begin() + split_offset; + vals->insert(vals->end(), start, start + used); + splits->insert(splits->end(), splits->back() + used); + } else { + splits->insert(splits->end(), splits->back()); + } } }); return trimmed; @@ -296,6 +311,9 @@ template void RoundRobinTrimmer::ProcessSplitsByBatch( Iterator begin, Iterator end, std::function*)> callback) const { + if (begin == end || begin->empty()) { + return; + } int num_in_batch = begin->size() - 1; int num_values = end - begin; // Process one batch at a time. @@ -305,7 +323,12 @@ void RoundRobinTrimmer::ProcessSplitsByBatch( int idx = 0; for (auto i = begin; i < end; ++i, ++idx) { value_row_sizes[idx].idx = idx; - value_row_sizes[idx].size = (*i)[batch_idx + 1] - (*i)[batch_idx]; + Tsplits row_size = 0; + if (batch_idx + 1 < i->size() && (*i)[batch_idx] >= 0 && + (*i)[batch_idx + 1] >= (*i)[batch_idx]) { + row_size = (*i)[batch_idx + 1] - (*i)[batch_idx]; + } + value_row_sizes[idx].size = row_size; } // Perform the main processing of the batch ProcessBatch(&value_row_sizes, callback); diff --git a/tensorflow_text/core/kernels/round_robin_trimmer_kernel_template.h b/tensorflow_text/core/kernels/round_robin_trimmer_kernel_template.h index 51f17da43..48afec081 100644 --- a/tensorflow_text/core/kernels/round_robin_trimmer_kernel_template.h +++ b/tensorflow_text/core/kernels/round_robin_trimmer_kernel_template.h @@ -15,6 +15,7 @@ #ifndef THIRD_PARTY_TENSORFLOW_TEXT_CORE_KERNELS_ROUND_ROBIN_TRIMMER_KERNEL_TEMPLATE_H_ #define THIRD_PARTY_TENSORFLOW_TEXT_CORE_KERNELS_ROUND_ROBIN_TRIMMER_KERNEL_TEMPLATE_H_ +#include #include #include #include @@ -29,6 +30,38 @@ namespace tensorflow { namespace text { +template +inline absl::Status ValidateRowSplits(absl::Span splits, + size_t expected_splits_size, + size_t values_size) { + // Check for valid row splits + if (splits.empty()) { + return absl::InvalidArgumentError("input_row_splits must not be empty."); + } + // Check for consistent row splits size + if (splits.size() != expected_splits_size) { + return absl::InvalidArgumentError( + "All input_row_splits must have the same size."); + } + // Check for negative values + if (splits[0] < 0) { + return absl::InvalidArgumentError("input_row_splits must be non-negative."); + } + // Check for monotonically increasing values + for (size_t j = 0; j + 1 < splits.size(); ++j) { + if (splits[j + 1] < splits[j]) { + return absl::InvalidArgumentError( + "input_row_splits must be monotonically increasing."); + } + } + // Check if the last element exceeds the values size + if (static_cast(splits.back()) > values_size) { + return absl::InvalidArgumentError( + "input_row_splits exceeds input_values size."); + } + return absl::OkStatus(); +} + template class RoundRobinTrimOp : public tflite::shim::OpKernelShim { @@ -153,11 +186,20 @@ absl::Status RoundRobinTrimOp::Invoke(InvokeContext* context) { std::vector> list_of_splits(number_of_segments_); for (int i = 0; i < number_of_segments_; ++i) { SH_ASSIGN_OR_RETURN(const auto fv, context->GetInput(kInputValues + i)); - list_of_values[i] = fv->template Data(); + size_t values_size = 1; + for (const int dim : fv->Shape()) { + values_size *= dim; + } + // TfTensorView::Data() uses TotalBytes() / sizeof(T), which over-counts + // for tstring tensors because TotalBytes() includes string heap payloads. + list_of_values[i] = fv->template Data().subspan(0, values_size); int row_split_idx = kInputRowSplits + number_of_segments_ - 1 + i; SH_ASSIGN_OR_RETURN(const auto rs, context->GetInput(row_split_idx)); list_of_splits[i] = rs->template Data(); + + SH_RETURN_IF_ERROR(ValidateRowSplits( + list_of_splits[i], list_of_splits[0].size(), list_of_values[i].size())); } // Compute @@ -299,9 +341,18 @@ absl::Status RoundRobinGenerateMasksOp::Invoke( std::vector> list_of_splits(number_of_segments_); for (int i = 0; i < number_of_segments_; ++i) { + SH_ASSIGN_OR_RETURN(const auto fv, context->GetInput(kInputValues + i)); + size_t values_size = 1; + for (const int dim : fv->Shape()) { + values_size *= dim; + } + int row_split_idx = kInputRowSplits + number_of_segments_ - 1 + i; SH_ASSIGN_OR_RETURN(const auto rs, context->GetInput(row_split_idx)); list_of_splits[i] = rs->template Data(); + + SH_RETURN_IF_ERROR(ValidateRowSplits( + list_of_splits[i], list_of_splits[0].size(), values_size)); } // Compute diff --git a/tensorflow_text/core/kernels/split_merge_tokenize_kernel.cc b/tensorflow_text/core/kernels/split_merge_tokenize_kernel.cc index 5491fab4d..a41175ae6 100644 --- a/tensorflow_text/core/kernels/split_merge_tokenize_kernel.cc +++ b/tensorflow_text/core/kernels/split_merge_tokenize_kernel.cc @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include #include #include #include @@ -23,9 +24,12 @@ #include "icu4c/source/common/unicode/umachine.h" #include "icu4c/source/common/unicode/utf8.h" #include "tensorflow/core/framework/op_kernel.h" +#include "tensorflow/core/framework/op_requires.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor_shape.h" +#include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/platform/tstring.h" namespace tensorflow { namespace text { @@ -127,25 +131,54 @@ class SplitMergeTokenizeWithOffsetsOp : public OpKernel { void Compute(OpKernelContext* ctx) override { const Tensor* input_values; OP_REQUIRES_OK(ctx, ctx->input("input_values", &input_values)); + OP_REQUIRES( + ctx, TensorShapeUtils::IsVector(input_values->shape()), + errors::InvalidArgument("input_values must be a vector, got shape ", + input_values->shape().DebugString())); const Tensor* labels; OP_REQUIRES_OK(ctx, ctx->input("labels", &labels)); + OP_REQUIRES(ctx, TensorShapeUtils::IsVector(labels->shape()), + errors::InvalidArgument("labels must be a vector, got shape ", + labels->shape().DebugString())); const Tensor* row_splits; OP_REQUIRES_OK(ctx, ctx->input("row_splits", &row_splits)); + OP_REQUIRES( + ctx, TensorShapeUtils::IsVector(row_splits->shape()), + errors::InvalidArgument("row_splits must be a vector, got shape ", + row_splits->shape().DebugString())); OP_REQUIRES(ctx, input_values->dim_size(0) == row_splits->dim_size(0) - 1, errors::InvalidArgument("Expecting row_splits have ", input_values->dim_size(0) + 1, " elements, got ", row_splits->dim_size(0))); + const auto& values_vec = input_values->flat(); + const auto& row_splits_vec = row_splits->flat(); + OP_REQUIRES( + ctx, row_splits_vec(0) >= 0, + errors::InvalidArgument("row_splits must be non-negative, got ", + row_splits_vec(0))); + for (int i = 0; i < values_vec.size(); ++i) { + OP_REQUIRES( + ctx, row_splits_vec(i + 1) >= row_splits_vec(i), + errors::InvalidArgument( + "row_splits must be monotonically increasing, got ", + row_splits_vec(i), " followed by ", row_splits_vec(i + 1))); + } + OP_REQUIRES(ctx, row_splits_vec(values_vec.size()) <= labels->dim_size(0), + errors::InvalidArgument( + "row_splits value ", row_splits_vec(values_vec.size()), + " exceeds labels size ", labels->dim_size(0))); + std::vector tokens; std::vector begin_offset; std::vector end_offset; - std::vector output_row_splits(1, 0); + std::vector output_row_splits; + output_row_splits.reserve(row_splits->dim_size(0)); + output_row_splits.push_back(0); // Iterate through all the values and tokenize them. - const auto& values_vec = input_values->flat(); - const auto& row_splits_vec = row_splits->flat(); for (int i = 0; i < values_vec.size(); ++i) { // Tokenize into tokens and record the offset locations. int num_tokens = 0; diff --git a/tensorflow_text/python/ops/boise_offset_converter_test.py b/tensorflow_text/python/ops/boise_offset_converter_test.py index c8576e501..59ec6f33a 100644 --- a/tensorflow_text/python/ops/boise_offset_converter_test.py +++ b/tensorflow_text/python/ops/boise_offset_converter_test.py @@ -22,8 +22,12 @@ import tensorflow as tf import tensorflow_text as tf_text +from tensorflow.python.framework import errors from tensorflow.python.framework import test_util from tensorflow.python.platform import test +from tensorflow.python.framework import load_library +from tensorflow.python.platform import resource_loader +gen_boise_offset_converter = load_library.load_op_library(resource_loader.get_path_to_datafile('_boise_offset_converter.so')) @test_util.run_all_in_graph_and_eager_modes @@ -456,6 +460,145 @@ def test_ragged_input_tensor_type_error(self): tf_text.boise_tags_to_offsets(token_begin_offsets, token_end_offsets, boise_tags) + def test_invalid_row_splits(self): + # Non-monotonic row splits in OffsetsToBoiseTags. + with self.assertRaisesRegex( + (errors.InvalidArgumentError, ValueError), + r'Invalid (token |span )?row splits', + ): + self.evaluate( + gen_boise_offset_converter.tf_text_offsets_to_boise_tags( + input_token_begin_offsets=[0, 4, 8], + input_token_end_offsets=[3, 7, 11], + input_span_begin_offsets=[4], + input_span_end_offsets=[11], + input_span_type=[b'animal'], + input_token_begin_row_splits=[0, 3, 1], + input_token_end_row_splits=[0, 3, 1], + input_span_begin_row_splits=[0, 1, 1], + input_span_end_row_splits=[0, 1, 1], + input_span_type_row_splits=[0, 1, 1], + input_use_strict_boundary_mode=False, + ) + ) + # Out-of-bounds row splits in OffsetsToBoiseTags. + with self.assertRaisesRegex( + (errors.InvalidArgumentError, ValueError), + r'Invalid (token |span )?row splits', + ): + self.evaluate( + gen_boise_offset_converter.tf_text_offsets_to_boise_tags( + input_token_begin_offsets=[0, 4, 8], + input_token_end_offsets=[3, 7, 11], + input_span_begin_offsets=[4], + input_span_end_offsets=[11], + input_span_type=[b'animal'], + input_token_begin_row_splits=[0, 100], + input_token_end_row_splits=[0, 100], + input_span_begin_row_splits=[0, 1], + input_span_end_row_splits=[0, 1], + input_span_type_row_splits=[0, 1], + input_use_strict_boundary_mode=False, + ) + ) + # Negative initial row split in OffsetsToBoiseTags. + with self.assertRaisesRegex( + (errors.InvalidArgumentError, ValueError), + r'Invalid (token |span )?row splits', + ): + self.evaluate( + gen_boise_offset_converter.tf_text_offsets_to_boise_tags( + input_token_begin_offsets=[0, 4, 8], + input_token_end_offsets=[3, 7, 11], + input_span_begin_offsets=[4], + input_span_end_offsets=[11], + input_span_type=[b'animal'], + input_token_begin_row_splits=[-1, 1], + input_token_end_row_splits=[-1, 1], + input_span_begin_row_splits=[0, 1], + input_span_end_row_splits=[0, 1], + input_span_type_row_splits=[0, 1], + input_use_strict_boundary_mode=False, + ) + ) + # Empty row splits in OffsetsToBoiseTags. + with self.assertRaisesRegex( + (errors.InvalidArgumentError, ValueError), + r'Row splits must be non-empty', + ): + self.evaluate( + gen_boise_offset_converter.tf_text_offsets_to_boise_tags( + input_token_begin_offsets=[0, 4, 8], + input_token_end_offsets=[3, 7, 11], + input_span_begin_offsets=[4], + input_span_end_offsets=[11], + input_span_type=[b'animal'], + input_token_begin_row_splits=[], + input_token_end_row_splits=[], + input_span_begin_row_splits=[], + input_span_end_row_splits=[], + input_span_type_row_splits=[], + input_use_strict_boundary_mode=False, + ) + ) + # Non-monotonic row splits in BoiseTagsToOffsets. + with self.assertRaisesRegex( + (errors.InvalidArgumentError, ValueError), r'Invalid row splits' + ): + self.evaluate( + gen_boise_offset_converter.tf_text_boise_tags_to_offsets( + input_token_begin_offsets=[0, 4, 8], + input_token_end_offsets=[3, 7, 11], + input_boise_tags=[b'O', b'B-animal', b'E-animal'], + input_token_begin_row_splits=[0, 3, 1], + input_token_end_row_splits=[0, 3, 1], + input_boise_tags_row_splits=[0, 3, 1], + ) + ) + # Out-of-bounds row splits in BoiseTagsToOffsets. + with self.assertRaisesRegex( + (errors.InvalidArgumentError, ValueError), r'Invalid row splits' + ): + self.evaluate( + gen_boise_offset_converter.tf_text_boise_tags_to_offsets( + input_token_begin_offsets=[0, 4, 8], + input_token_end_offsets=[3, 7, 11], + input_boise_tags=[b'O', b'B-animal', b'E-animal'], + input_token_begin_row_splits=[0, 100], + input_token_end_row_splits=[0, 100], + input_boise_tags_row_splits=[0, 100], + ) + ) + # Negative initial row split in BoiseTagsToOffsets. + with self.assertRaisesRegex( + (errors.InvalidArgumentError, ValueError), r'Invalid row splits' + ): + self.evaluate( + gen_boise_offset_converter.tf_text_boise_tags_to_offsets( + input_token_begin_offsets=[0, 4, 8], + input_token_end_offsets=[3, 7, 11], + input_boise_tags=[b'O', b'B-animal', b'E-animal'], + input_token_begin_row_splits=[-1, 1], + input_token_end_row_splits=[-1, 1], + input_boise_tags_row_splits=[-1, 1], + ) + ) + # Empty row splits in BoiseTagsToOffsets. + with self.assertRaisesRegex( + (errors.InvalidArgumentError, ValueError), + r'Row splits must be non-empty', + ): + self.evaluate( + gen_boise_offset_converter.tf_text_boise_tags_to_offsets( + input_token_begin_offsets=[0, 4, 8], + input_token_end_offsets=[3, 7, 11], + input_boise_tags=[b'O', b'B-animal', b'E-animal'], + input_token_begin_row_splits=[], + input_token_end_row_splits=[], + input_boise_tags_row_splits=[], + ) + ) + if __name__ == '__main__': test.main() diff --git a/tensorflow_text/python/ops/byte_splitter_test.py b/tensorflow_text/python/ops/byte_splitter_test.py index d21bab47e..708aad2ac 100644 --- a/tensorflow_text/python/ops/byte_splitter_test.py +++ b/tensorflow_text/python/ops/byte_splitter_test.py @@ -24,8 +24,12 @@ import tensorflow_text as tf_text from tensorflow.lite.python import interpreter +from tensorflow.python.framework import errors from tensorflow.python.framework import test_util from tensorflow.python.platform import test +from tensorflow.python.framework import load_library +from tensorflow.python.platform import resource_loader +gen_byte_splitter = load_library.load_op_library(resource_loader.get_path_to_datafile('_byte_splitter.so')) def _split(s): @@ -428,6 +432,57 @@ def call(self, input_tensor): tflite_result = output['output_1'] self.assertAllEqual(tflite_result, tf_result) + def testByteSplitByOffsetsInvalidRowSplits(self): + # Non-monotonic row splits. + with self.assertRaisesRegex( + (errors.InvalidArgumentError, ValueError), r'Invalid input_row_splits' + ): + self.evaluate( + gen_byte_splitter.tf_text_byte_split_by_offsets( + input_values=[b'hello', b'world'], + input_start_offsets=[0, 1, 0], + input_end_offsets=[1, 2, 1], + input_row_splits=[0, 3, 1], + ) + ) + # Out-of-bounds row splits. + with self.assertRaisesRegex( + (errors.InvalidArgumentError, ValueError), r'Invalid input_row_splits' + ): + self.evaluate( + gen_byte_splitter.tf_text_byte_split_by_offsets( + input_values=[b'hello'], + input_start_offsets=[0, 1], + input_end_offsets=[1, 2], + input_row_splits=[0, 100], + ) + ) + # Negative initial row split. + with self.assertRaisesRegex( + (errors.InvalidArgumentError, ValueError), r'Invalid input_row_splits' + ): + self.evaluate( + gen_byte_splitter.tf_text_byte_split_by_offsets( + input_values=[b'hello'], + input_start_offsets=[0, 1], + input_end_offsets=[1, 2], + input_row_splits=[-1, 1], + ) + ) + # Empty row splits. + with self.assertRaisesRegex( + (errors.InvalidArgumentError, ValueError), + r'Expecting input_row_splits to have', + ): + self.evaluate( + gen_byte_splitter.tf_text_byte_split_by_offsets( + input_values=[b'hello'], + input_start_offsets=[0, 1], + input_end_offsets=[1, 2], + input_row_splits=[], + ) + ) + if __name__ == '__main__': test.main() diff --git a/tensorflow_text/python/ops/fast_wordpiece_tokenizer_test.py b/tensorflow_text/python/ops/fast_wordpiece_tokenizer_test.py index d69d2c366..9a7f86660 100644 --- a/tensorflow_text/python/ops/fast_wordpiece_tokenizer_test.py +++ b/tensorflow_text/python/ops/fast_wordpiece_tokenizer_test.py @@ -31,10 +31,14 @@ from tensorflow.python.data.ops import dataset_ops from tensorflow.python.eager import def_function from tensorflow.python.framework import dtypes +from tensorflow.python.framework import errors from tensorflow.python.framework import test_util from tensorflow.python.ops.ragged import ragged_factory_ops from tensorflow.python.platform import gfile from tensorflow.python.platform import test +from tensorflow.python.framework import load_library +from tensorflow.python.platform import resource_loader +gen_fast_wordpiece_tokenizer = load_library.load_op_library(resource_loader.get_path_to_datafile('_fast_wordpiece_tokenizer.so')) from tensorflow_text.python.ops.fast_wordpiece_tokenizer import FastWordpieceTokenizer FLAGS = flags.FLAGS @@ -487,6 +491,56 @@ def testTensors(self, subwords = tokenizer.tokenize(tokens) self.assertAllEqual(subwords, expected_subwords) + def testDetokenizeInvalidRowSplits(self): + tokenizer = FastWordpieceTokenizer( + vocab=_ENGLISH_VOCAB, support_detokenization=True + ) + # Non-monotonic row splits. + with self.assertRaisesRegex( + (errors.InvalidArgumentError, ValueError), r"Invalid input_row_splits" + ): + self.evaluate( + gen_fast_wordpiece_tokenizer.tf_text_fast_wordpiece_detokenize( + input_values=[0, 1, 2], + input_row_splits=[0, 3, 1], + wp_model=tokenizer._model, + ) + ) + # Out-of-bounds row splits. + with self.assertRaisesRegex( + (errors.InvalidArgumentError, ValueError), r"Invalid input_row_splits" + ): + self.evaluate( + gen_fast_wordpiece_tokenizer.tf_text_fast_wordpiece_detokenize( + input_values=[0, 1, 2], + input_row_splits=[0, 100], + wp_model=tokenizer._model, + ) + ) + # Negative initial row split. + with self.assertRaisesRegex( + (errors.InvalidArgumentError, ValueError), r"Invalid input_row_splits" + ): + self.evaluate( + gen_fast_wordpiece_tokenizer.tf_text_fast_wordpiece_detokenize( + input_values=[0, 1, 2], + input_row_splits=[-1, 1], + wp_model=tokenizer._model, + ) + ) + # Empty row splits. + with self.assertRaisesRegex( + (errors.InvalidArgumentError, ValueError), + r"input_row_splits must not be empty", + ): + self.evaluate( + gen_fast_wordpiece_tokenizer.tf_text_fast_wordpiece_detokenize( + input_values=[0, 1, 2], + input_row_splits=[], + wp_model=tokenizer._model, + ) + ) + # The following WordPiece setup is used in `FastWordpieceOpAdditionalTest` and # `EndToEndFastWordpieceOpTest`. diff --git a/tensorflow_text/python/ops/ngrams_op_test.py b/tensorflow_text/python/ops/ngrams_op_test.py index 6fd31b521..663b76133 100644 --- a/tensorflow_text/python/ops/ngrams_op_test.py +++ b/tensorflow_text/python/ops/ngrams_op_test.py @@ -30,6 +30,9 @@ from tensorflow.python.framework import test_util from tensorflow.python.ops.ragged import ragged_factory_ops from tensorflow.python.platform import test +from tensorflow.python.framework import load_library +from tensorflow.python.platform import resource_loader +gen_ngrams_op = load_library.load_op_library(resource_loader.get_path_to_datafile('_ngrams_op.so')) from tensorflow_text.core.pybinds import tflite_registrar from tensorflow_text.python.ops import ngrams_op @@ -482,6 +485,64 @@ def call(self, input_tensor): # Assert the results are identical. self.assertAllEqual(tflite_result, tf_result) + def testInvalidRowSplits(self): + # Non-monotonic row splits. + with self.assertRaisesRegex( + (errors.InvalidArgumentError, ValueError), + r"input_row_splits must be monotonically increasing", + ): + self.evaluate( + gen_ngrams_op.tf_text_ngrams_string_join( + input_values=[b"a", b"b", b"c"], + input_row_splits=[[0, 3, 1]], + width=2, + axis=-1, + string_separator=" ", + ) + ) + # Out-of-bounds row splits. + with self.assertRaisesRegex( + (errors.InvalidArgumentError, ValueError), + r"input_row_splits exceeds input_values size", + ): + self.evaluate( + gen_ngrams_op.tf_text_ngrams_string_join( + input_values=[b"a", b"b", b"c"], + input_row_splits=[[0, 100]], + width=2, + axis=-1, + string_separator=" ", + ) + ) + # Negative initial row split. + with self.assertRaisesRegex( + (errors.InvalidArgumentError, ValueError), + r"input_row_splits must be non-negative", + ): + self.evaluate( + gen_ngrams_op.tf_text_ngrams_string_join( + input_values=[b"a", b"b", b"c"], + input_row_splits=[[-1, 1]], + width=2, + axis=-1, + string_separator=" ", + ) + ) + # Empty row splits. + with self.assertRaisesRegex( + (errors.InvalidArgumentError, ValueError), + r"input_row_splits must not be empty", + ): + self.evaluate( + gen_ngrams_op.tf_text_ngrams_string_join( + input_values=[b"a", b"b", b"c"], + input_row_splits=[constant_op.constant([], dtype=tf.int64)], + width=2, + axis=-1, + string_separator=" ", + ) + ) + if __name__ == "__main__": test.main() diff --git a/tensorflow_text/python/ops/phrase_tokenizer_test.py b/tensorflow_text/python/ops/phrase_tokenizer_test.py index e12856c99..e87541782 100644 --- a/tensorflow_text/python/ops/phrase_tokenizer_test.py +++ b/tensorflow_text/python/ops/phrase_tokenizer_test.py @@ -26,9 +26,13 @@ from tensorflow.python.data.kernel_tests import test_base from tensorflow.python.framework import dtypes +from tensorflow.python.framework import errors from tensorflow.python.framework import test_util from tensorflow.python.ops.ragged import ragged_factory_ops from tensorflow.python.platform import test +from tensorflow.python.framework import load_library +from tensorflow.python.platform import resource_loader +gen_phrase_tokenizer = load_library.load_op_library(resource_loader.get_path_to_datafile('_phrase_tokenizer.so')) from tensorflow_text.python.ops.phrase_tokenizer import PhraseTokenizer FLAGS = flags.FLAGS @@ -91,6 +95,60 @@ def testPhraseOp(self, subwords_t = tokenizer.tokenize(tokens_t) self.assertAllEqual(subwords_t, expected_subwords) + def testDetokenizeInvalidRowSplits(self): + tokenizer = PhraseTokenizer( + vocab=_ENGLISH_VOCAB, + unknown_token="", + support_detokenization=True, + prob=0, + split_end_punctuation=True, + ) + # Non-monotonic row splits. + with self.assertRaisesRegex( + (errors.InvalidArgumentError, ValueError), r"Invalid input_row_splits" + ): + self.evaluate( + gen_phrase_tokenizer.tf_text_phrase_detokenize( + input_values=[1, 2, 3], + input_row_splits=[0, 3, 1], + phrase_model=tokenizer._model, + ) + ) + # Out-of-bounds row splits. + with self.assertRaisesRegex( + (errors.InvalidArgumentError, ValueError), r"Invalid input_row_splits" + ): + self.evaluate( + gen_phrase_tokenizer.tf_text_phrase_detokenize( + input_values=[1, 2, 3], + input_row_splits=[0, 100], + phrase_model=tokenizer._model, + ) + ) + # Negative initial row split. + with self.assertRaisesRegex( + (errors.InvalidArgumentError, ValueError), r"Invalid input_row_splits" + ): + self.evaluate( + gen_phrase_tokenizer.tf_text_phrase_detokenize( + input_values=[1, 2, 3], + input_row_splits=[-1, 1], + phrase_model=tokenizer._model, + ) + ) + # Empty row splits. + with self.assertRaisesRegex( + (errors.InvalidArgumentError, ValueError), + r"input_row_splits must not be empty", + ): + self.evaluate( + gen_phrase_tokenizer.tf_text_phrase_detokenize( + input_values=[1, 2, 3], + input_row_splits=[], + phrase_model=tokenizer._model, + ) + ) + @parameterized.parameters([ dict( diff --git a/tensorflow_text/python/ops/split_merge_tokenizer_test.py b/tensorflow_text/python/ops/split_merge_tokenizer_test.py index 60b0bd3d9..346de9f4a 100644 --- a/tensorflow_text/python/ops/split_merge_tokenizer_test.py +++ b/tensorflow_text/python/ops/split_merge_tokenizer_test.py @@ -21,6 +21,7 @@ from __future__ import print_function from tensorflow.python.framework import constant_op +from tensorflow.python.framework import errors from tensorflow.python.framework import ops from tensorflow.python.framework import test_util from tensorflow.python.ops import array_ops @@ -29,6 +30,9 @@ from tensorflow.python.ops.ragged import ragged_factory_ops from tensorflow.python.ops.ragged import ragged_tensor from tensorflow.python.platform import test +from tensorflow.python.framework import load_library +from tensorflow.python.platform import resource_loader +gen_split_merge_tokenizer = load_library.load_op_library(resource_loader.get_path_to_datafile('_split_merge_tokenizer.so')) from tensorflow_text.python.ops.split_merge_tokenizer import SplitMergeTokenizer @@ -402,6 +406,60 @@ def testRaggedInputHigherRank(self): tokens = self.tokenizer.tokenize(test_value, test_label) self.assertAllEqual(tokens, expected_tokens) + def testInvalidRowSplits(self): + # Non-monotonic row splits. + with self.assertRaisesRegex( + (errors.InvalidArgumentError, ValueError), + r'row_splits must be monotonically increasing', + ): + self.evaluate( + gen_split_merge_tokenizer.split_merge_tokenize_with_offsets( + input_values=[b'hello', b'world'], + labels=[0, 1, 1, 1, 1, 0, 1, 1, 1, 1], + row_splits=[0, 8, 5], + force_split_at_break_character=True, + ) + ) + # Out-of-bounds row splits. + with self.assertRaisesRegex( + (errors.InvalidArgumentError, ValueError), + r'row_splits value .* exceeds labels size', + ): + self.evaluate( + gen_split_merge_tokenizer.split_merge_tokenize_with_offsets( + input_values=[b'hello'], + labels=[0, 1, 1, 1, 1], + row_splits=[0, 100], + force_split_at_break_character=True, + ) + ) + # Negative initial row split. + with self.assertRaisesRegex( + (errors.InvalidArgumentError, ValueError), + r'row_splits must be non-negative', + ): + self.evaluate( + gen_split_merge_tokenizer.split_merge_tokenize_with_offsets( + input_values=[b'hello'], + labels=[0, 1, 1, 1, 1], + row_splits=[-1, 1], + force_split_at_break_character=True, + ) + ) + # Empty row splits. + with self.assertRaisesRegex( + (errors.InvalidArgumentError, ValueError), + r'Expecting row_splits have', + ): + self.evaluate( + gen_split_merge_tokenizer.split_merge_tokenize_with_offsets( + input_values=[b'hello'], + labels=[0, 1, 1, 1, 1], + row_splits=[], + force_split_at_break_character=True, + ) + ) + if __name__ == '__main__': test.main() diff --git a/tensorflow_text/python/ops/trimmer_ops_test.py b/tensorflow_text/python/ops/trimmer_ops_test.py index 798ebfa31..e44d39e2f 100644 --- a/tensorflow_text/python/ops/trimmer_ops_test.py +++ b/tensorflow_text/python/ops/trimmer_ops_test.py @@ -18,9 +18,14 @@ from absl.testing import parameterized from tensorflow.python.framework import constant_op +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import errors from tensorflow.python.framework import test_util from tensorflow.python.ops.ragged import ragged_factory_ops from tensorflow.python.platform import test +from tensorflow.python.framework import load_library +from tensorflow.python.platform import resource_loader +gen_trimmer_ops = load_library.load_op_library(resource_loader.get_path_to_datafile('_trimmer_ops.so')) from tensorflow_text.python.ops import trimmer_ops @@ -539,6 +544,100 @@ def testPerBatchBudgetTrimmer(self, for expected_seg, actual_seg in zip(expected, actual): self.assertAllEqual(expected_seg, actual_seg) + def testInvalidRowSplits(self): + # Non-monotonic row splits. + with self.assertRaisesRegex( + (errors.InvalidArgumentError, ValueError), + r"input_row_splits must be monotonically increasing", + ): + self.evaluate( + gen_trimmer_ops.tf_text_round_robin_generate_masks( + max_sequence_length=3, + input_values=[[b"a", b"b", b"c"]], + input_row_splits=[[0, 2, 1]], + ) + ) + with self.assertRaisesRegex( + (errors.InvalidArgumentError, ValueError), + r"input_row_splits must be monotonically increasing", + ): + self.evaluate( + gen_trimmer_ops.tf_text_round_robin_trim( + max_sequence_length=3, + input_values=[[b"a", b"b", b"c"]], + input_row_splits=[[0, 2, 1]], + ) + ) + # Out-of-bounds row splits. + with self.assertRaisesRegex( + (errors.InvalidArgumentError, ValueError), + r"input_row_splits exceeds input_values size", + ): + self.evaluate( + gen_trimmer_ops.tf_text_round_robin_generate_masks( + max_sequence_length=3, + input_values=[[b"a", b"b", b"c"]], + input_row_splits=[[0, 100]], + ) + ) + with self.assertRaisesRegex( + (errors.InvalidArgumentError, ValueError), + r"input_row_splits exceeds input_values size", + ): + self.evaluate( + gen_trimmer_ops.tf_text_round_robin_trim( + max_sequence_length=3, + input_values=[[b"a", b"b", b"c"]], + input_row_splits=[[0, 100]], + ) + ) + # Negative initial row split. + with self.assertRaisesRegex( + (errors.InvalidArgumentError, ValueError), + r"input_row_splits must be non-negative", + ): + self.evaluate( + gen_trimmer_ops.tf_text_round_robin_generate_masks( + max_sequence_length=3, + input_values=[[b"a", b"b", b"c"]], + input_row_splits=[[-1, 1]], + ) + ) + with self.assertRaisesRegex( + (errors.InvalidArgumentError, ValueError), + r"input_row_splits must be non-negative", + ): + self.evaluate( + gen_trimmer_ops.tf_text_round_robin_trim( + max_sequence_length=3, + input_values=[[b"a", b"b", b"c"]], + input_row_splits=[[-1, 1]], + ) + ) + # Empty row splits. + with self.assertRaisesRegex( + (errors.InvalidArgumentError, ValueError), + r"input_row_splits must not be empty", + ): + self.evaluate( + gen_trimmer_ops.tf_text_round_robin_generate_masks( + max_sequence_length=3, + input_values=[[b"a", b"b", b"c"]], + input_row_splits=[constant_op.constant([], dtype=dtypes.int64)], + ) + ) + with self.assertRaisesRegex( + (errors.InvalidArgumentError, ValueError), + r"input_row_splits must not be empty", + ): + self.evaluate( + gen_trimmer_ops.tf_text_round_robin_trim( + max_sequence_length=3, + input_values=[[b"a", b"b", b"c"]], + input_row_splits=[constant_op.constant([], dtype=dtypes.int64)], + ) + ) + # These two tests started to segfault after new TF integration, # Investigate after brunch cut tf.text #