/*
* 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.
*/
#ifndef LIBTEXTCLASSIFIER_COMMON_EMBEDDING_NETWORK_PARAMS_FROM_PROTO_H_
#define LIBTEXTCLASSIFIER_COMMON_EMBEDDING_NETWORK_PARAMS_FROM_PROTO_H_
#include <algorithm>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "common/embedding-network-package.pb.h"
#include "common/embedding-network-params.h"
#include "common/embedding-network.pb.h"
#include "common/float16.h"
#include "common/little-endian-data.h"
#include "common/task-context.h"
#include "common/task-spec.pb.h"
#include "util/base/integral_types.h"
#include "util/base/logging.h"
namespace libtextclassifier {
namespace nlp_core {
// A wrapper class that owns and exposes an EmbeddingNetworkProto message via
// the EmbeddingNetworkParams interface.
//
// The EmbeddingNetworkParams interface encapsulates the weight matrices of the
// embeddings, hidden and softmax layers as transposed versions of their
// counterparts in the original EmbeddingNetworkProto. The matrices in the proto
// passed to this class' constructor must likewise already have been transposed.
// See embedding-network-params.h for details.
class EmbeddingNetworkParamsFromProto : public EmbeddingNetworkParams {
public:
// Constructor that takes ownership of the provided proto. See class-comment
// for the requirements that certain weight matrices must satisfy.
explicit EmbeddingNetworkParamsFromProto(
std::unique_ptr<EmbeddingNetworkProto> proto)
: proto_(std::move(proto)) {
valid_ = true;
// Initialize these vectors to have the required number of elements
// regardless of quantization status. This is to support the unlikely case
// where only some embeddings are quantized, along with the fact that
// EmbeddingNetworkParams interface accesses them by index.
embeddings_quant_scales_.resize(proto_->embeddings_size());
embeddings_quant_weights_.resize(proto_->embeddings_size());
for (int i = 0; i < proto_->embeddings_size(); ++i) {
MatrixParams *embedding = proto_->mutable_embeddings()->Mutable(i);
if (!embedding->is_quantized()) {
continue;
}
bool success = FillVectorFromDataBytesInLittleEndian(
embedding->bytes_for_quantized_values(),
embedding->rows() * embedding->cols(),
&(embeddings_quant_weights_[i]));
if (!success) {
TC_LOG(ERROR) << "Problem decoding quant_weights for embeddings #" << i;
valid_ = false;
}
// The repeated field bytes_for_quantized_values uses a lot of memory.
// Since it's no longer necessary (and we own the proto), we clear it.
embedding->clear_bytes_for_quantized_values();
success = FillVectorFromDataBytesInLittleEndian(
embedding->bytes_for_col_scales(),
embedding->rows(),
&(embeddings_quant_scales_[i]));
if (!success) {
TC_LOG(ERROR) << "Problem decoding col_scales for embeddings #" << i;
valid_ = false;
}
// See comments for clear_bytes_for_quantized_values().
embedding->clear_bytes_for_col_scales();
}
}
const TaskSpec *GetTaskSpec() override {
if (!proto_) {
return nullptr;
}
auto extension_id = task_spec_in_embedding_network_proto;
if (proto_->HasExtension(extension_id)) {
return &(proto_->GetExtension(extension_id));
} else {
TC_LOG(ERROR) << "Unable to get TaskSpec from EmbeddingNetworkProto";
return nullptr;
}
}
// Returns true if these params are valid. False otherwise (e.g., if the
// original proto data was corrupted).
bool is_valid() { return valid_; }
protected:
int embeddings_size() const override { return proto_->embeddings_size(); }
int embeddings_num_rows(int i) const override {
TC_DCHECK(InRange(i, embeddings_size()));
return proto_->embeddings(i).rows();
}
int embeddings_num_cols(int i) const override {
TC_DCHECK(InRange(i, embeddings_size()));
return proto_->embeddings(i).cols();
}
const void *embeddings_weights(int i) const override {
TC_DCHECK(InRange(i, embeddings_size()));
if (proto_->embeddings(i).is_quantized()) {
return static_cast<const void *>(embeddings_quant_weights_.at(i).data());
} else {
return static_cast<const void *>(proto_->embeddings(i).value().data());
}
}
QuantizationType embeddings_quant_type(int i) const override {
TC_DCHECK(InRange(i, embeddings_size()));
return proto_->embeddings(i).is_quantized() ? QuantizationType::UINT8
: QuantizationType::NONE;
}
const float16 *embeddings_quant_scales(int i) const override {
TC_DCHECK(InRange(i, embeddings_size()));
return proto_->embeddings(i).is_quantized()
? embeddings_quant_scales_.at(i).data()
: nullptr;
}
int hidden_size() const override { return proto_->hidden_size(); }
int hidden_num_rows(int i) const override {
TC_DCHECK(InRange(i, hidden_size()));
return proto_->hidden(i).rows();
}
int hidden_num_cols(int i) const override {
TC_DCHECK(InRange(i, hidden_size()));
return proto_->hidden(i).cols();
}
const void *hidden_weights(int i) const override {
TC_DCHECK(InRange(i, hidden_size()));
return proto_->hidden(i).value().data();
}
int hidden_bias_size() const override { return proto_->hidden_bias_size(); }
int hidden_bias_num_rows(int i) const override {
TC_DCHECK(InRange(i, hidden_bias_size()));
return proto_->hidden_bias(i).rows();
}
int hidden_bias_num_cols(int i) const override {
TC_DCHECK(InRange(i, hidden_bias_size()));
return proto_->hidden_bias(i).cols();
}
const void *hidden_bias_weights(int i) const override {
TC_DCHECK(InRange(i, hidden_bias_size()));
return proto_->hidden_bias(i).value().data();
}
int softmax_size() const override { return proto_->has_softmax() ? 1 : 0; }
int softmax_num_rows(int i) const override {
TC_DCHECK(InRange(i, softmax_size()));
return proto_->has_softmax() ? proto_->softmax().rows() : 0;
}
int softmax_num_cols(int i) const override {
TC_DCHECK(InRange(i, softmax_size()));
return proto_->has_softmax() ? proto_->softmax().cols() : 0;
}
const void *softmax_weights(int i) const override {
TC_DCHECK(InRange(i, softmax_size()));
return proto_->has_softmax() ? proto_->softmax().value().data() : nullptr;
}
int softmax_bias_size() const override {
return proto_->has_softmax_bias() ? 1 : 0;
}
int softmax_bias_num_rows(int i) const override {
TC_DCHECK(InRange(i, softmax_bias_size()));
return proto_->has_softmax_bias() ? proto_->softmax_bias().rows() : 0;
}
int softmax_bias_num_cols(int i) const override {
TC_DCHECK(InRange(i, softmax_bias_size()));
return proto_->has_softmax_bias() ? proto_->softmax_bias().cols() : 0;
}
const void *softmax_bias_weights(int i) const override {
TC_DCHECK(InRange(i, softmax_bias_size()));
return proto_->has_softmax_bias() ? proto_->softmax_bias().value().data()
: nullptr;
}
int embedding_num_features_size() const override {
return proto_->embedding_num_features_size();
}
int embedding_num_features(int i) const override {
TC_DCHECK(InRange(i, embedding_num_features_size()));
return proto_->embedding_num_features(i);
}
private:
std::unique_ptr<EmbeddingNetworkProto> proto_;
// True if these params are valid. May be false if the original proto was
// corrupted. We prefer to set this to false to CHECK-failing.
bool valid_;
// When the embeddings are quantized, these members are used to store their
// numeric values using the types expected by the rest of the class. Due to
// technical reasons, the proto stores this info using larger types (i.e.,
// more bits).
std::vector<std::vector<float16>> embeddings_quant_scales_;
std::vector<std::vector<uint8>> embeddings_quant_weights_;
};
} // namespace nlp_core
} // namespace libtextclassifier
#endif // LIBTEXTCLASSIFIER_COMMON_EMBEDDING_NETWORK_PARAMS_FROM_PROTO_H_