| /* |
| * Copyright (C) 2017 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 "model-executor.h" |
| |
| #include "util/base/logging.h" |
| |
| namespace libtextclassifier2 { |
| namespace internal { |
| bool FromModelSpec(const tflite::Model* model_spec, |
| std::unique_ptr<tflite::FlatBufferModel>* model, |
| std::unique_ptr<tflite::Interpreter>* interpreter) { |
| *model = tflite::FlatBufferModel::BuildFromModel(model_spec); |
| if (!(*model) || !(*model)->initialized()) { |
| TC_LOG(ERROR) << "Could not build TFLite model from a model spec. "; |
| return false; |
| } |
| |
| tflite::ops::builtin::BuiltinOpResolver builtins; |
| tflite::InterpreterBuilder(**model, builtins)(interpreter); |
| if (!interpreter) { |
| TC_LOG(ERROR) << "Could not build TFLite interpreter."; |
| return false; |
| } |
| return true; |
| } |
| } // namespace internal |
| |
| TFLiteEmbeddingExecutor::TFLiteEmbeddingExecutor( |
| const tflite::Model* model_spec) { |
| internal::FromModelSpec(model_spec, &model_, &interpreter_); |
| if (!interpreter_) { |
| return; |
| } |
| if (interpreter_->tensors_size() != 2) { |
| return; |
| } |
| embeddings_ = interpreter_->tensor(0); |
| if (embeddings_->dims->size != 2) { |
| return; |
| } |
| num_buckets_ = embeddings_->dims->data[0]; |
| scales_ = interpreter_->tensor(1); |
| if (scales_->dims->size != 2 || scales_->dims->data[0] != num_buckets_ || |
| scales_->dims->data[1] != 1) { |
| return; |
| } |
| embedding_size_ = embeddings_->dims->data[1]; |
| initialized_ = true; |
| } |
| |
| bool TFLiteEmbeddingExecutor::AddEmbedding( |
| const TensorView<int>& sparse_features, float* dest, int dest_size) { |
| if (!initialized_ || dest_size != embedding_size_) { |
| return false; |
| } |
| const int num_sparse_features = sparse_features.size(); |
| for (int i = 0; i < num_sparse_features; ++i) { |
| const int bucket_id = sparse_features.data()[i]; |
| if (bucket_id >= num_buckets_) { |
| return false; |
| } |
| const float multiplier = scales_->data.f[bucket_id]; |
| for (int k = 0; k < embedding_size_; ++k) { |
| // Dequantize and add the embedding. |
| dest[k] += |
| 1.0 / num_sparse_features * |
| (static_cast<int>( |
| embeddings_->data.uint8[bucket_id * embedding_size_ + k]) - |
| kQuantBias) * |
| multiplier; |
| } |
| } |
| return true; |
| } |
| |
| TensorView<float> ComputeLogitsHelper(const int input_index_features, |
| const int output_index_logits, |
| const TensorView<float>& features, |
| tflite::Interpreter* interpreter) { |
| interpreter->ResizeInputTensor(input_index_features, features.shape()); |
| if (interpreter->AllocateTensors() != kTfLiteOk) { |
| TC_VLOG(1) << "Allocation failed."; |
| return TensorView<float>::Invalid(); |
| } |
| |
| TfLiteTensor* features_tensor = |
| interpreter->tensor(interpreter->inputs()[input_index_features]); |
| int size = 1; |
| for (int i = 0; i < features_tensor->dims->size; ++i) { |
| size *= features_tensor->dims->data[i]; |
| } |
| features.copy_to(features_tensor->data.f, size); |
| |
| if (interpreter->Invoke() != kTfLiteOk) { |
| TC_VLOG(1) << "Interpreter failed."; |
| return TensorView<float>::Invalid(); |
| } |
| |
| TfLiteTensor* logits_tensor = |
| interpreter->tensor(interpreter->outputs()[output_index_logits]); |
| |
| std::vector<int> output_shape(logits_tensor->dims->size); |
| for (int i = 0; i < logits_tensor->dims->size; ++i) { |
| output_shape[i] = logits_tensor->dims->data[i]; |
| } |
| |
| return TensorView<float>(logits_tensor->data.f, output_shape); |
| } |
| |
| } // namespace libtextclassifier2 |