1 /* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15
16 /**
17 * Sentencepiece tflite tokenizer implementation.
18 */
19 #include "tensorflow_lite_support/custom_ops/kernel/sentencepiece/optimized_encoder.h"
20 #include "tensorflow_lite_support/custom_ops/kernel/sentencepiece/sentencepiece_tokenizer.h"
21 #include "flatbuffers/flexbuffers.h" // from @flatbuffers
22 #include "tensorflow/lite/c/common.h"
23 #include "tensorflow/lite/context.h"
24 #include "tensorflow/lite/kernels/internal/tensor.h"
25 #include "tensorflow/lite/kernels/kernel_util.h"
26 #include "tensorflow/lite/model.h"
27 #include "tensorflow/lite/string_util.h"
28
29 namespace tflite {
30 namespace ops {
31 namespace custom {
32 namespace sentencepiece {
33 namespace tokenizer {
34
35 constexpr int kOutputValuesInd = 0;
36 constexpr int kOutputSplitsInd = 1;
37
38 namespace {
CreateSizeArray(const std::initializer_list<int> & sizes)39 TfLiteIntArray* CreateSizeArray(const std::initializer_list<int>& sizes) {
40 TfLiteIntArray* array_size = TfLiteIntArrayCreate(sizes.size());
41 int index = 0;
42 for (const int size : sizes) {
43 array_size->data[index++] = size;
44 }
45 return array_size;
46 }
47 } // namespace
48
49 // Initializes text encoder object from serialized parameters.
Initialize(TfLiteContext *,const char *,size_t)50 void* Initialize(TfLiteContext* /*context*/, const char* /*buffer*/,
51 size_t /*length*/) {
52 return nullptr;
53 }
Free(TfLiteContext *,void *)54 void Free(TfLiteContext* /*context*/, void* /*buffer*/) {}
55
Prepare(TfLiteContext * context,TfLiteNode * node)56 TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
57 // TODO(mgubin): Add checks for input and output tensors.
58 TfLiteTensor& output_values =
59 context->tensors[node->outputs->data[kOutputValuesInd]];
60 SetTensorToDynamic(&output_values);
61
62 TfLiteTensor& output_splits =
63 context->tensors[node->outputs->data[kOutputSplitsInd]];
64 SetTensorToDynamic(&output_splits);
65 return kTfLiteOk;
66 }
67
Eval(TfLiteContext * context,TfLiteNode * node)68 TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
69 const TfLiteTensor& model_tensor =
70 context->tensors[node->inputs->data[tensorflow::ops::kSPModelIndex]];
71 const auto model_buffer_data = model_tensor.data.data;
72 const TfLiteTensor& input_text =
73 context->tensors[node->inputs->data[tensorflow::ops::kInputIndex]];
74
75 const TfLiteTensor add_bos_tensor =
76 context->tensors[node->inputs->data[tensorflow::ops::kAddBOSInput]];
77 const bool add_bos = add_bos_tensor.data.b[0];
78 const TfLiteTensor add_eos_tensor =
79 context->tensors[node->inputs->data[tensorflow::ops::kAddEOSInput]];
80 const bool add_eos = add_eos_tensor.data.b[0];
81 const TfLiteTensor reverse_tensor =
82 context->tensors[node->inputs->data[tensorflow::ops::kReverseInput]];
83 const bool reverse = reverse_tensor.data.b[0];
84
85 std::vector<int32> encoded;
86 std::vector<int32> splits;
87 const int num_strings = tflite::GetStringCount(&input_text);
88 for (int i = 0; i < num_strings; ++i) {
89 const auto strref = tflite::GetString(&input_text, i);
90 const auto res = EncodeString(std::string(strref.str, strref.len),
91 model_buffer_data, add_bos, add_eos, reverse);
92 TF_LITE_ENSURE_MSG(context, res.type == EncoderResultType::SUCCESS,
93 "Sentencepiece conversion failed");
94 std::copy(res.codes.begin(), res.codes.end(), std::back_inserter(encoded));
95 splits.emplace_back(encoded.size());
96 }
97
98 TfLiteTensor& output_values =
99 context->tensors[node->outputs->data[kOutputValuesInd]];
100 TF_LITE_ENSURE_OK(context,
101 context->ResizeTensor(
102 context, &output_values,
103 CreateSizeArray({static_cast<int>(encoded.size())})));
104 int32_t* output_values_flat = output_values.data.i32;
105 std::copy(encoded.begin(), encoded.end(), output_values_flat);
106 TfLiteTensor& output_splits =
107 context->tensors[node->outputs->data[kOutputSplitsInd]];
108 TF_LITE_ENSURE_OK(
109 context, context->ResizeTensor(
110 context, &output_splits,
111 CreateSizeArray({static_cast<int>(splits.size() + 1)})));
112 int32_t* output_splits_flat = output_splits.data.i32;
113 *output_splits_flat = 0;
114 std::copy(splits.begin(), splits.end(), output_splits_flat + 1);
115 return kTfLiteOk;
116 }
117 } // namespace tokenizer
118 } // namespace sentencepiece
119
Register_SENTENCEPIECE_TOKENIZER()120 TfLiteRegistration* Register_SENTENCEPIECE_TOKENIZER() {
121 static TfLiteRegistration r = {
122 sentencepiece::tokenizer::Initialize, sentencepiece::tokenizer::Free,
123 sentencepiece::tokenizer::Prepare, sentencepiece::tokenizer::Eval};
124 return &r;
125 }
126
127 } // namespace custom
128 } // namespace ops
129 } // namespace tflite
130