Skip to content

Validate row_splits bounds and monotonicity in TF.Text tokenization and utility kernels. #1562

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Validate row_splits bounds and monotonicity in TF.Text tokenization and utility kernels. #1562

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions tensorflow_text/core/kernels/BUILD
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
],
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
],
Expand Down
114 changes: 91 additions & 23 deletions tensorflow_text/core/kernels/boise_offset_converter_kernel_template.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
#include <vector>

#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"
Expand Down Expand Up @@ -265,16 +266,18 @@ absl::Status OffsetsToBoiseTagsOp<Rt>::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),
Expand All @@ -284,7 +287,7 @@ absl::Status OffsetsToBoiseTagsOp<Rt>::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(
Expand All @@ -293,7 +296,7 @@ absl::Status OffsetsToBoiseTagsOp<Rt>::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) !=
Expand All @@ -304,27 +307,67 @@ absl::Status OffsetsToBoiseTagsOp<Rt>::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<std::string> boise_tags;
boise_tags.reserve(input_token_begin_row_splits_vec(num_rows) -
input_token_begin_row_splits_vec(0));
std::vector<int32_t> input_token_begin_offsets_vec_i;
std::vector<int32_t> input_token_end_offsets_vec_i;
std::vector<int32_t> input_span_begin_offsets_vec_i;
std::vector<int32_t> input_span_end_offsets_vec_i;
std::vector<std::string> 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();
input_span_begin_offsets_vec_i.clear();
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));
Expand Down Expand Up @@ -537,63 +580,88 @@ absl::Status BoiseTagsToOffsetsOp<Rt>::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),
" BOISE tags row splits shape: ",
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<int32_t> span_begin_offsets;
std::vector<int32_t> span_end_offsets;
std::vector<std::string> span_type;
std::vector<int64_t> row_splits;
row_splits.reserve(input_token_begin_row_splits_vec.Dim(0));

row_splits.push_back(0);

// Iterate through all the input values and split them.
std::vector<int32_t> input_token_begin_offsets_vec_i;
std::vector<int32_t> input_token_end_offsets_vec_i;
std::vector<std::string> 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));
input_token_end_offsets_vec_i.push_back(input_token_end_offsets_vec(j));
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);
Expand Down
39 changes: 34 additions & 5 deletions tensorflow_text/core/kernels/byte_splitter_kernel_template.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 <cstdint>
#include <iostream>
#include <vector>

#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"
Expand Down Expand Up @@ -277,22 +279,49 @@ template <tflite::shim::Runtime Rt>
context->GetInput(kInputRowSplits));
const auto in_splits = in_splits_view->template As<int64_t, 1>();

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<absl::string_view> output_values;
std::vector<int32_t> 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());
}
Expand Down
Loading
Loading