1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
|
/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
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 "tensorflow/contrib/lite/models/smartreply/predictor.h"
#include "absl/strings/str_split.h"
#include "re2/re2.h"
#include "tensorflow/contrib/lite/interpreter.h"
#include "tensorflow/contrib/lite/kernels/register.h"
#include "tensorflow/contrib/lite/model.h"
#include "tensorflow/contrib/lite/string_util.h"
#include "tensorflow/contrib/lite/tools/mutable_op_resolver.h"
void RegisterSelectedOps(::tflite::MutableOpResolver* resolver);
namespace tflite {
namespace custom {
namespace smartreply {
// Split sentence into segments (using punctuation).
std::vector<string> SplitSentence(const string& input) {
string result(input);
RE2::GlobalReplace(&result, "([?.!,])+", " \\1");
RE2::GlobalReplace(&result, "([?.!,])+\\s+", "\\1\t");
RE2::GlobalReplace(&result, "[ ]+", " ");
RE2::GlobalReplace(&result, "\t+$", "");
return strings::Split(result, '\t');
}
// Predict with TfLite model.
void ExecuteTfLite(const string& sentence, ::tflite::Interpreter* interpreter,
std::map<string, float>* response_map) {
{
TfLiteTensor* input = interpreter->tensor(interpreter->inputs()[0]);
tflite::DynamicBuffer buf;
buf.AddString(sentence.data(), sentence.length());
buf.WriteToTensor(input);
interpreter->AllocateTensors();
interpreter->Invoke();
TfLiteTensor* messages = interpreter->tensor(interpreter->outputs()[0]);
TfLiteTensor* confidence = interpreter->tensor(interpreter->outputs()[1]);
for (int i = 0; i < confidence->dims->data[0]; i++) {
float weight = confidence->data.f[i];
auto response_text = tflite::GetString(messages, i);
if (response_text.len > 0) {
(*response_map)[string(response_text.str, response_text.len)] += weight;
}
}
}
}
void GetSegmentPredictions(
const std::vector<string>& input, const ::tflite::FlatBufferModel& model,
const SmartReplyConfig& config,
std::vector<PredictorResponse>* predictor_responses) {
// Initialize interpreter
std::unique_ptr<::tflite::Interpreter> interpreter;
::tflite::MutableOpResolver resolver;
RegisterSelectedOps(&resolver);
::tflite::InterpreterBuilder(model, resolver)(&interpreter);
if (!model.initialized()) {
fprintf(stderr, "Failed to mmap model \n");
return;
}
// Execute Tflite Model
std::map<string, float> response_map;
std::vector<string> sentences;
for (const string& str : input) {
std::vector<string> splitted_str = SplitSentence(str);
sentences.insert(sentences.end(), splitted_str.begin(), splitted_str.end());
}
for (const auto& sentence : sentences) {
ExecuteTfLite(sentence, interpreter.get(), &response_map);
}
// Generate the result.
for (const auto& iter : response_map) {
PredictorResponse prediction(iter.first, iter.second);
predictor_responses->emplace_back(prediction);
}
std::sort(predictor_responses->begin(), predictor_responses->end(),
[](const PredictorResponse& a, const PredictorResponse& b) {
return a.GetScore() > b.GetScore();
});
// Add backoff response.
for (const string& backoff : config.backoff_responses) {
if (predictor_responses->size() >= config.num_response) {
break;
}
predictor_responses->push_back({backoff, config.backoff_confidence});
}
}
} // namespace smartreply
} // namespace custom
} // namespace tflite
|
