/*
|
* Copyright (C) 2018 The Android Open Source Project
|
*
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
* you may not use this file except in compliance with the License.
|
* You may obtain a copy of the License at
|
*
|
* http://www.apache.org/licenses/LICENSE-2.0
|
*
|
* Unless required by applicable law or agreed to in writing, software
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
* See the License for the specific language governing permissions and
|
* limitations under the License.
|
*/
|
|
#include <memory>
|
#include <vector>
|
|
#include "utils/base/logging.h"
|
#include "utils/sentencepiece/double_array_trie.h"
|
#include "utils/sentencepiece/encoder.h"
|
#include "utils/sentencepiece/normalizer.h"
|
#include "utils/sentencepiece/sorted_strings_table.h"
|
#include "utils/strings/stringpiece.h"
|
#include "utils/tflite/encoder_common.h"
|
#include "utils/tflite/text_encoder.h"
|
#include "utils/tflite/text_encoder_config_generated.h"
|
#include "flatbuffers/flatbuffers.h"
|
#include "flatbuffers/flexbuffers.h"
|
#include "tensorflow/lite/kernels/kernel_util.h"
|
#include "tensorflow/lite/model.h"
|
#include "tensorflow/lite/string_util.h"
|
|
namespace libtextclassifier3 {
|
namespace {
|
|
struct TextEncoderOp {
|
std::unique_ptr<SentencePieceNormalizer> normalizer;
|
std::unique_ptr<Encoder> encoder;
|
std::unique_ptr<SentencePieceMatcher> matcher;
|
};
|
|
// Input parameters for the op.
|
// The conversation message as a (1, conversation length) string tensor.
|
constexpr const int kInputTexts = 0;
|
|
// The number of messages, the conversation length, int scalar.
|
constexpr const int kInputNumInputs = 1;
|
|
// Maximum output length of the encoding, int scalar.
|
constexpr const int kInputMaxLength = 2;
|
|
// Additional attributes to align to the sentence pieces, e.g. user ids per
|
// message.
|
constexpr const int kInputAttr = 3;
|
|
// Output parameters for the op.
|
// The text sentence piece encodings as ids, (1, max output length) int tensor.
|
constexpr const int kOutputEncoded = 0;
|
|
// Relative position of each sentence piece in the input text,
|
// (1, max output length) int tensor.
|
constexpr const int kOutputPosition = 1;
|
|
// Output length after trimming to the maximum output length specified.
|
// int scalar.
|
constexpr const int kOutputLengths = 2;
|
|
// Padded and sentence piece aligned provided attributes, e.g. user id per
|
// sentence piece.
|
constexpr const int kOutputAttr = 3;
|
|
const char kTextEncoderConfigAttr[] = "text_encoder_config";
|
|
// Initializes text encoder object from serialized options:
|
// The options are a flexbuffers attribute map that contain the op config
|
// with the key `text_encoder_config` as `TextEncoderConfig`.
|
void* Initialize(TfLiteContext* context, const char* buffer, size_t length) {
|
const flexbuffers::Map& attr_map =
|
flexbuffers::GetRoot(reinterpret_cast<const uint8_t*>(buffer), length)
|
.AsMap();
|
const flexbuffers::Blob serialized_config =
|
attr_map[kTextEncoderConfigAttr].AsBlob();
|
const TextEncoderConfig* config =
|
flatbuffers::GetRoot<TextEncoderConfig>(serialized_config.data());
|
|
std::unique_ptr<TextEncoderOp> encoder_op(new TextEncoderOp());
|
|
// Create normalizer from options.
|
const TrieNode* charsmap_trie_nodes = reinterpret_cast<const TrieNode*>(
|
config->normalization_charsmap()->Data());
|
const int charsmap_trie_nodes_length =
|
config->normalization_charsmap()->Length() / sizeof(TrieNode);
|
encoder_op->normalizer.reset(new SentencePieceNormalizer(
|
DoubleArrayTrie(charsmap_trie_nodes, charsmap_trie_nodes_length),
|
StringPiece(config->normalization_charsmap_values()->data(),
|
config->normalization_charsmap_values()->size()),
|
config->add_dummy_prefix(), config->remove_extra_whitespaces(),
|
config->escape_whitespaces()));
|
|
const int num_pieces = config->pieces_scores()->Length();
|
|
switch (config->matcher_type()) {
|
case SentencePieceMatcherType_MAPPED_TRIE: {
|
const TrieNode* pieces_trie_nodes =
|
reinterpret_cast<const TrieNode*>(config->pieces()->Data());
|
const int pieces_trie_nodes_length =
|
config->pieces()->Length() / sizeof(TrieNode);
|
encoder_op->matcher.reset(
|
new DoubleArrayTrie(pieces_trie_nodes, pieces_trie_nodes_length));
|
break;
|
}
|
case SentencePieceMatcherType_SORTED_STRING_TABLE: {
|
encoder_op->matcher.reset(new SortedStringsTable(
|
num_pieces, config->pieces_offsets()->data(),
|
StringPiece(config->pieces()->data(), config->pieces()->Length())));
|
break;
|
}
|
default: {
|
TC3_LOG(ERROR) << "Unknown sentence piece matcher type.";
|
return nullptr;
|
}
|
}
|
encoder_op->encoder.reset(new Encoder(
|
encoder_op->matcher.get(), num_pieces, config->pieces_scores()->data(),
|
config->start_code(), config->end_code(), config->encoding_offset(),
|
config->unknown_code(), config->unknown_score()));
|
return encoder_op.release();
|
}
|
|
void Free(TfLiteContext* context, void* buffer) {
|
delete reinterpret_cast<TextEncoderOp*>(buffer);
|
}
|
|
TfLiteStatus ResizeOutputTensors(TfLiteContext* context, TfLiteNode* node,
|
int max_output_length) {
|
TF_LITE_ENSURE_OK(
|
context,
|
ResizeOutputTensor(max_output_length,
|
&context->tensors[node->outputs->data[kOutputEncoded]],
|
context));
|
|
TF_LITE_ENSURE_OK(
|
context,
|
ResizeOutputTensor(
|
max_output_length,
|
&context->tensors[node->outputs->data[kOutputPosition]], context));
|
|
const int num_output_attrs = node->outputs->size - kOutputAttr;
|
for (int i = 0; i < num_output_attrs; ++i) {
|
TF_LITE_ENSURE_OK(
|
context,
|
ResizeOutputTensor(
|
max_output_length,
|
&context->tensors[node->outputs->data[kOutputAttr + i]], context));
|
}
|
return kTfLiteOk;
|
}
|
|
TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
|
// Check that the batch dimension is kBatchSize.
|
const TfLiteTensor& input_text =
|
context->tensors[node->inputs->data[kInputTexts]];
|
TF_LITE_ENSURE_EQ(context, input_text.dims->size, kEncoderInputRank);
|
TF_LITE_ENSURE_EQ(context, input_text.dims->data[0], kEncoderBatchSize);
|
|
TfLiteTensor& output_lengths =
|
context->tensors[node->outputs->data[kOutputLengths]];
|
TfLiteTensor& output_encoded =
|
context->tensors[node->outputs->data[kOutputEncoded]];
|
TfLiteTensor& output_positions =
|
context->tensors[node->outputs->data[kOutputPosition]];
|
|
TF_LITE_ENSURE_OK(context,
|
context->ResizeTensor(context, &output_lengths,
|
CreateIntArray({kEncoderBatchSize})));
|
|
// Check that there are enough outputs for attributes.
|
const int num_output_attrs = node->outputs->size - kOutputAttr;
|
TF_LITE_ENSURE_EQ(context, node->inputs->size - kInputAttr, num_output_attrs);
|
|
// Copy attribute types from input to output tensors.
|
for (int i = 0; i < num_output_attrs; ++i) {
|
TfLiteTensor& input = context->tensors[node->inputs->data[kInputAttr + i]];
|
TfLiteTensor& output =
|
context->tensors[node->outputs->data[kOutputAttr + i]];
|
output.type = input.type;
|
}
|
|
const TfLiteTensor& output_length =
|
context->tensors[node->inputs->data[kInputMaxLength]];
|
|
if (tflite::IsConstantTensor(&output_length)) {
|
return ResizeOutputTensors(context, node, output_length.data.i64[0]);
|
} else {
|
tflite::SetTensorToDynamic(&output_encoded);
|
tflite::SetTensorToDynamic(&output_positions);
|
for (int i = 0; i < num_output_attrs; ++i) {
|
TfLiteTensor& output_attr =
|
context->tensors[node->outputs->data[kOutputAttr + i]];
|
tflite::SetTensorToDynamic(&output_attr);
|
}
|
}
|
|
return kTfLiteOk;
|
}
|
|
TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
|
if (node->user_data == nullptr) {
|
return kTfLiteError;
|
}
|
const TextEncoderOp* encoder_op =
|
reinterpret_cast<TextEncoderOp*>(node->user_data);
|
const TfLiteTensor& input_text =
|
context->tensors[node->inputs->data[kInputTexts]];
|
const int num_strings = tflite::GetStringCount(&input_text);
|
// Check that the number of strings matches the length parameter.
|
const int num_strings_param =
|
context->tensors[node->inputs->data[kInputNumInputs]].data.i32[0];
|
TF_LITE_ENSURE_EQ(context, num_strings, num_strings_param);
|
|
TfLiteTensor& output_encoded =
|
context->tensors[node->outputs->data[kOutputEncoded]];
|
if (tflite::IsDynamicTensor(&output_encoded)) {
|
const TfLiteTensor& output_length =
|
context->tensors[node->inputs->data[kInputMaxLength]];
|
TF_LITE_ENSURE_OK(
|
context, ResizeOutputTensors(context, node, output_length.data.i64[0]));
|
}
|
TfLiteTensor& output_positions =
|
context->tensors[node->outputs->data[kOutputPosition]];
|
|
std::vector<int> encoded_total;
|
std::vector<int> encoded_offsets;
|
std::vector<int> encoded_positions;
|
encoded_offsets.reserve(num_strings);
|
const int max_output_length = output_encoded.dims->data[1];
|
const int max_encoded_position = max_output_length;
|
|
for (int i = 0; i < num_strings; ++i) {
|
const auto& strref = tflite::GetString(&input_text, i);
|
std::string normalized;
|
TF_LITE_ENSURE(context,
|
encoder_op->normalizer->Normalize(
|
StringPiece(strref.str, strref.len), &normalized));
|
std::vector<int> encoded;
|
TF_LITE_ENSURE(context, encoder_op->encoder->Encode(normalized, &encoded));
|
encoded_total.insert(encoded_total.end(), encoded.begin(), encoded.end());
|
encoded_offsets.push_back(encoded_total.size());
|
for (int i = 0; i < encoded.size(); i++) {
|
encoded_positions.push_back(std::min(i, max_encoded_position - 1));
|
}
|
}
|
|
const int num_skip = CopyDataToTensorAndPadOrTruncate(
|
max_output_length, encoded_total,
|
/*padding_value=*/encoded_total.back(), &output_encoded);
|
TfLiteTensor& output_lengths =
|
context->tensors[node->outputs->data[kOutputLengths]];
|
output_lengths.data.i32[0] = encoded_total.size() - num_skip;
|
CopyDataToTensorAndPadOrTruncate(max_output_length, encoded_positions,
|
/*padding_value=*/max_encoded_position,
|
&output_positions);
|
|
// Process attributes, all checks of sizes and types are done in Prepare.
|
const int num_output_attrs = node->outputs->size - kOutputAttr;
|
TF_LITE_ENSURE_EQ(context, node->inputs->size - kInputAttr, num_output_attrs);
|
for (int i = 0; i < num_output_attrs; ++i) {
|
TfLiteStatus attr_status = CopyValuesToTensorAndPadOrTruncate(
|
context->tensors[node->inputs->data[kInputAttr + i]], encoded_offsets,
|
num_skip, context,
|
&context->tensors[node->outputs->data[kOutputAttr + i]]);
|
if (attr_status != kTfLiteOk) {
|
return attr_status;
|
}
|
}
|
|
return kTfLiteOk;
|
}
|
|
} // namespace
|
} // namespace libtextclassifier3
|
|
namespace tflite {
|
namespace ops {
|
namespace custom {
|
|
TfLiteRegistration* Register_TEXT_ENCODER() {
|
static TfLiteRegistration registration = {
|
libtextclassifier3::Initialize, libtextclassifier3::Free,
|
libtextclassifier3::Prepare, libtextclassifier3::Eval};
|
return ®istration;
|
}
|
|
} // namespace custom
|
} // namespace ops
|
} // namespace tflite
|
