experiments/STT
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Write a CTC beam search decoder TF op that scores beams with our LM

Browse Source
This commit is contained in:
Reuben Morais 2017-09-01 11:45:35 +02:00
parent af71da0d4d
commit fc91e3d7b8
5 changed files with 774 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

32
native_client/BUILD
View File

@ -34,3 +34,35 @@ cc_library(
copts = [] + if_linux_x86_64(["-mno-fma", "-mno-avx", "-mno-avx2"]),
nocopts = "(-fstack-protector|-fno-omit-frame-pointer)",
)
cc_library(
name = "ctc_decoder_with_kenlm",
srcs = ["beam_search.cc",
"alphabet.h",
"trie_node.h"] +
glob(["kenlm/lm/*.cc", "kenlm/util/*.cc", "kenlm/util/double-conversion/*.cc",
"kenlm/lm/*.hh", "kenlm/util/*.hh", "kenlm/util/double-conversion/*.h"],
exclude = ["kenlm/*/*test.cc", "kenlm/*/*main.cc"]),
includes = ["kenlm"],
defines = ["KENLM_MAX_ORDER=6"],
deps = ["//tensorflow/core:core",
"//tensorflow/core/util/ctc",
"//third_party/eigen3",
],
)
cc_binary(
name = "generate_trie",
srcs = [
"generate_trie.cpp",
"trie_node.h",
"alphabet.h",
] + glob(["kenlm/lm/*.cc", "kenlm/util/*.cc", "kenlm/util/double-conversion/*.cc",
"kenlm/lm/*.hh", "kenlm/util/*.hh", "kenlm/util/double-conversion/*.h"],
exclude = ["kenlm/*/*test.cc", "kenlm/*/*main.cc"]),
includes = ["kenlm"],
copts = ['-std=c++11'],
linkopts = ['-lm'],
defines = ["KENLM_MAX_ORDER=6"],
)

5
native_client/alphabet.h
View File

@ -54,6 +54,11 @@ public:
return size_;
}
bool IsSpace(unsigned int label) const {
const std::string& str = StringFromLabel(label);
return str.size() == 1 && str[0] == ' ';
}
private:
size_t size_;
std::unordered_map<unsigned int, std::string> label_to_str_;

546
native_client/beam_search.cc Normal file
View File

@ -0,0 +1,546 @@
/* Copyright 2016 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.
==============================================================================*/
// This test illustrates how to make use of the CTCBeamSearchDecoder using a
// custom BeamScorer and BeamState based on a dictionary with a few artificial
// words.
#include "tensorflow/core/util/ctc/ctc_beam_search.h"
#include <algorithm>
#include <cmath>
#include <vector>
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/shape_inference.h"
#include "tensorflow/core/kernels/bounds_check.h"
#include "kenlm/lm/model.hh"
#include "alphabet.h"
#include "trie_node.h"
namespace tf = tensorflow;
using tf::shape_inference::DimensionHandle;
using tf::shape_inference::InferenceContext;
using tf::shape_inference::ShapeHandle;
REGISTER_OP("CTCBeamSearchDecoderWithLM")
.Input("inputs: float")
.Input("sequence_length: int32")
.Attr("model_path: string")
.Attr("trie_path: string")
.Attr("alphabet_path: string")
.Attr("beam_width: int >= 1 = 100")
.Attr("top_paths: int >= 1 = 1")
.Attr("merge_repeated: bool = true")
.Output("decoded_indices: top_paths * int64")
.Output("decoded_values: top_paths * int64")
.Output("decoded_shape: top_paths * int64")
.Output("log_probability: float")
.SetShapeFn([](InferenceContext *c) {
ShapeHandle inputs;
ShapeHandle sequence_length;
TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 3, &inputs));
TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 1, &sequence_length));
// Get batch size from inputs and sequence_length.
DimensionHandle batch_size;
TF_RETURN_IF_ERROR(
c->Merge(c->Dim(inputs, 1), c->Dim(sequence_length, 0), &batch_size));
tf::int32 top_paths;
TF_RETURN_IF_ERROR(c->GetAttr("top_paths", &top_paths));
// Outputs.
int out_idx = 0;
for (int i = 0; i < top_paths; ++i) { // decoded_indices
c->set_output(out_idx++, c->Matrix(InferenceContext::kUnknownDim, 2));
}
for (int i = 0; i < top_paths; ++i) { // decoded_values
c->set_output(out_idx++, c->Vector(InferenceContext::kUnknownDim));
}
ShapeHandle shape_v = c->Vector(2);
for (int i = 0; i < top_paths; ++i) { // decoded_shape
c->set_output(out_idx++, shape_v);
}
c->set_output(out_idx++, c->Matrix(batch_size, top_paths));
return tf::Status::OK();
})
.Doc(R"doc(
Performs beam search decoding on the logits given in input.
A note about the attribute merge_repeated: For the beam search decoder,
this means that if consecutive entries in a beam are the same, only
the first of these is emitted. That is, when the top path is "A B B B B",
"A B" is returned if merge_repeated = True but "A B B B B" is
returned if merge_repeated = False.
inputs: 3-D, shape: `(max_time x batch_size x num_classes)`, the logits.
sequence_length: A vector containing sequence lengths, size `(batch)`.
beam_width: A scalar >= 0 (beam search beam width).
top_paths: A scalar >= 0, <= beam_width (controls output size).
merge_repeated: If true, merge repeated classes in output.
decoded_indices: A list (length: top_paths) of indices matrices. Matrix j,
size `(total_decoded_outputs[j] x 2)`, has indices of a
`SparseTensor<int64, 2>`. The rows store: [batch, time].
decoded_values: A list (length: top_paths) of values vectors. Vector j,
size `(length total_decoded_outputs[j])`, has the values of a
`SparseTensor<int64, 2>`. The vector stores the decoded classes for beam j.
decoded_shape: A list (length: top_paths) of shape vector. Vector j,
size `(2)`, stores the shape of the decoded `SparseTensor[j]`.
Its values are: `[batch_size, max_decoded_length[j]]`.
log_probability: A matrix, shaped: `(batch_size x top_paths)`. The
sequence log-probabilities.
)doc");
struct KenLMBeamState {
float language_model_score;
float score;
float delta_score;
std::string incomplete_word;
TrieNode *incomplete_word_trie_node;
lm::ngram::ProbingModel::State model_state;
};
class KenLMBeamScorer : public tf::ctc::BaseBeamScorer<KenLMBeamState> {
public:
typedef lm::ngram::ProbingModel Model;
KenLMBeamScorer(const std::string &kenlm_path, const std::string &trie_path, const std::string &alphabet_path)
: lm_weight(1.0f)
, word_count_weight(-0.1f)
, valid_word_count_weight(1.0f)
{
lm::ngram::Config config;
config.load_method = util::POPULATE_OR_READ;
model = new Model(kenlm_path.c_str(), config);
alphabet = new Alphabet(alphabet_path.c_str());
std::ifstream in;
in.open(trie_path, std::ios::in);
TrieNode::ReadFromStream(in, trieRoot, alphabet->GetSize());
in.close();
}
virtual ~KenLMBeamScorer() {
delete model;
delete trieRoot;
}
// State initialization.
void InitializeState(KenLMBeamState* root) const {
root->language_model_score = 0.0f;
root->score = 0.0f;
root->delta_score = 0.0f;
root->incomplete_word.clear();
root->incomplete_word_trie_node = trieRoot;
root->model_state = model->BeginSentenceState();
}
// ExpandState is called when expanding a beam to one of its children.
// Called at most once per child beam. In the simplest case, no state
// expansion is done.
void ExpandState(const KenLMBeamState& from_state, int from_label,
KenLMBeamState* to_state, int to_label) const {
CopyState(from_state, to_state);
if (!alphabet->IsSpace(to_label)) {
to_state->incomplete_word += alphabet->StringFromLabel(to_label);
TrieNode *trie_node = from_state.incomplete_word_trie_node;
// TODO replace with OOV unigram prob?
// If we have no valid prefix we assume a very low log probability
float min_unigram_score = -10.0f;
// If prefix does exist
if (trie_node != nullptr) {
trie_node = trie_node->GetChildAt(to_label);
to_state->incomplete_word_trie_node = trie_node;
if (trie_node != nullptr) {
min_unigram_score = trie_node->GetMinUnigramScore();
}
}
// TODO try two options
// 1) unigram score added up to language model scare
// 2) langugage model score of (preceding_words + unigram_word)
to_state->score = min_unigram_score + to_state->language_model_score;
to_state->delta_score = to_state->score - from_state.score;
} else {
float lm_score_delta = ScoreIncompleteWord(from_state.model_state,
to_state->incomplete_word,
to_state->model_state);
// Give fixed word bonus
if (!IsOOV(to_state->incomplete_word)) {
to_state->language_model_score += valid_word_count_weight;
}
to_state->language_model_score += word_count_weight;
UpdateWithLMScore(to_state, lm_score_delta);
ResetIncompleteWord(to_state);
}
}
// ExpandStateEnd is called after decoding has finished. Its purpose is to
// allow a final scoring of the beam in its current state, before resorting
// and retrieving the TopN requested candidates. Called at most once per beam.
void ExpandStateEnd(KenLMBeamState* state) const {
float lm_score_delta = 0.0f;
Model::State out;
if (state->incomplete_word.size() > 0) {
lm_score_delta += ScoreIncompleteWord(state->model_state,
state->incomplete_word,
out);
ResetIncompleteWord(state);
state->model_state = out;
}
lm_score_delta += model->FullScore(state->model_state,
model->GetVocabulary().EndSentence(),
out).prob;
UpdateWithLMScore(state, lm_score_delta);
}
// GetStateExpansionScore should be an inexpensive method to retrieve the
// (cached) expansion score computed within ExpandState. The score is
// multiplied (log-addition) with the input score at the current step from
// the network.
//
// The score returned should be a log-probability. In the simplest case, as
// there's no state expansion logic, the expansion score is zero.
float GetStateExpansionScore(const KenLMBeamState& state,
float previous_score) const {
return lm_weight * state.delta_score + previous_score;
}
// GetStateEndExpansionScore should be an inexpensive method to retrieve the
// (cached) expansion score computed within ExpandStateEnd. The score is
// multiplied (log-addition) with the final probability of the beam.
//
// The score returned should be a log-probability.
float GetStateEndExpansionScore(const KenLMBeamState& state) const {
return lm_weight * state.delta_score;
}
void SetLMWeight(float lm_weight) {
this->lm_weight = lm_weight;
}
void SetWordCountWeight(float word_count_weight) {
this->word_count_weight = word_count_weight;
}
void SetValidWordCountWeight(float valid_word_count_weight) {
this->valid_word_count_weight = valid_word_count_weight;
}
private:
Model *model;
Alphabet *alphabet;
TrieNode *trieRoot;
float lm_weight;
float word_count_weight;
float valid_word_count_weight;
void UpdateWithLMScore(KenLMBeamState *state, float lm_score_delta) const {
float previous_score = state->score;
state->language_model_score += lm_score_delta;
state->score = state->language_model_score;
state->delta_score = state->language_model_score - previous_score;
}
void ResetIncompleteWord(KenLMBeamState *state) const {
state->incomplete_word.clear();
state->incomplete_word_trie_node = trieRoot;
}
bool IsOOV(const std::string& word) const {
auto &vocabulary = model->GetVocabulary();
return vocabulary.Index(word) == vocabulary.NotFound();
}
float ScoreIncompleteWord(const Model::State& model_state,
const std::string& word,
Model::State& out) const {
lm::FullScoreReturn full_score_return;
lm::WordIndex vocab = model->GetVocabulary().Index(word);
full_score_return = model->FullScore(model_state, vocab, out);
return full_score_return.prob;
}
void CopyState(const KenLMBeamState& from, KenLMBeamState* to) const {
to->language_model_score = from.language_model_score;
to->score = from.score;
to->delta_score = from.delta_score;
to->incomplete_word = from.incomplete_word;
to->incomplete_word_trie_node = from.incomplete_word_trie_node;
to->model_state = from.model_state;
}
};
class CTCDecodeHelper {
public:
CTCDecodeHelper() : top_paths_(1) {}
inline int GetTopPaths() const { return top_paths_; }
void SetTopPaths(int tp) { top_paths_ = tp; }
tf::Status ValidateInputsGenerateOutputs(
tf::OpKernelContext *ctx, const tf::Tensor **inputs, const tf::Tensor **seq_len,
std::string *model_path, std::string *trie_path, std::string *alphabet_path,
tf::Tensor **log_prob, tf::OpOutputList *decoded_indices,
tf::OpOutputList *decoded_values, tf::OpOutputList *decoded_shape) const {
tf::Status status = ctx->input("inputs", inputs);
if (!status.ok()) return status;
status = ctx->input("sequence_length", seq_len);
if (!status.ok()) return status;
const tf::TensorShape &inputs_shape = (*inputs)->shape();
if (inputs_shape.dims() != 3) {
return tf::errors::InvalidArgument("inputs is not a 3-Tensor");
}
const tf::int64 max_time = inputs_shape.dim_size(0);
const tf::int64 batch_size = inputs_shape.dim_size(1);
if (max_time == 0) {
return tf::errors::InvalidArgument("max_time is 0");
}
if (!tf::TensorShapeUtils::IsVector((*seq_len)->shape())) {
return tf::errors::InvalidArgument("sequence_length is not a vector");
}
if (!(batch_size == (*seq_len)->dim_size(0))) {
return tf::errors::FailedPrecondition(
"len(sequence_length) != batch_size. ", "len(sequence_length): ",
(*seq_len)->dim_size(0), " batch_size: ", batch_size);
}
auto seq_len_t = (*seq_len)->vec<tf::int32>();
for (int b = 0; b < batch_size; ++b) {
if (!(seq_len_t(b) <= max_time)) {
return tf::errors::FailedPrecondition("sequence_length(", b, ") <= ",
max_time);
}
}
tf::Status s = ctx->allocate_output(
"log_probability", tf::TensorShape({batch_size, top_paths_}), log_prob);
if (!s.ok()) return s;
s = ctx->output_list("decoded_indices", decoded_indices);
if (!s.ok()) return s;
s = ctx->output_list("decoded_values", decoded_values);
if (!s.ok()) return s;
s = ctx->output_list("decoded_shape", decoded_shape);
if (!s.ok()) return s;
return tf::Status::OK();
}
// sequences[b][p][ix] stores decoded value "ix" of path "p" for batch "b".
tf::Status StoreAllDecodedSequences(
const std::vector<std::vector<std::vector<int>>> &sequences,
tf::OpOutputList *decoded_indices, tf::OpOutputList *decoded_values,
tf::OpOutputList *decoded_shape) const {
// Calculate the total number of entries for each path
const tf::int64 batch_size = sequences.size();
std::vector<tf::int64> num_entries(top_paths_, 0);
// Calculate num_entries per path
for (const auto &batch_s : sequences) {
CHECK_EQ(batch_s.size(), top_paths_);
for (int p = 0; p < top_paths_; ++p) {
num_entries[p] += batch_s[p].size();
}
}
for (int p = 0; p < top_paths_; ++p) {
tf::Tensor *p_indices = nullptr;
tf::Tensor *p_values = nullptr;
tf::Tensor *p_shape = nullptr;
const tf::int64 p_num = num_entries[p];
tf::Status s =
decoded_indices->allocate(p, tf::TensorShape({p_num, 2}), &p_indices);
if (!s.ok()) return s;
s = decoded_values->allocate(p, tf::TensorShape({p_num}), &p_values);
if (!s.ok()) return s;
s = decoded_shape->allocate(p, tf::TensorShape({2}), &p_shape);
if (!s.ok()) return s;
auto indices_t = p_indices->matrix<tf::int64>();
auto values_t = p_values->vec<tf::int64>();
auto shape_t = p_shape->vec<tf::int64>();
tf::int64 max_decoded = 0;
tf::int64 offset = 0;
for (tf::int64 b = 0; b < batch_size; ++b) {
auto &p_batch = sequences[b][p];
tf::int64 num_decoded = p_batch.size();
max_decoded = std::max(max_decoded, num_decoded);
std::copy_n(p_batch.begin(), num_decoded, &values_t(offset));
for (tf::int64 t = 0; t < num_decoded; ++t, ++offset) {
indices_t(offset, 0) = b;
indices_t(offset, 1) = t;
}
}
shape_t(0) = batch_size;
shape_t(1) = max_decoded;
}
return tf::Status::OK();
}
private:
int top_paths_;
TF_DISALLOW_COPY_AND_ASSIGN(CTCDecodeHelper);
};
// CTC beam search
class CTCBeamSearchDecoderOp : public tf::OpKernel {
public:
explicit CTCBeamSearchDecoderOp(tf::OpKernelConstruction *ctx)
: tf::OpKernel(ctx)
, beam_scorer_(GetModelPath(ctx),
GetTriePath(ctx),
GetAlphabetPath(ctx))
{
OP_REQUIRES_OK(ctx, ctx->GetAttr("merge_repeated", &merge_repeated_));
OP_REQUIRES_OK(ctx, ctx->GetAttr("beam_width", &beam_width_));
int top_paths;
OP_REQUIRES_OK(ctx, ctx->GetAttr("top_paths", &top_paths));
decode_helper_.SetTopPaths(top_paths);
// const tf::Tensor* model_tensor;
// tf::Status status = ctx->input("model_path", &model_tensor);
// if (!status.ok()) return status;
// auto model_vec = model_tensor->flat<std::string>();
// *model_path = model_vec(0);
// const tf::Tensor* trie_tensor;
// status = ctx->input("trie_path", &trie_tensor);
// if (!status.ok()) return status;
// auto trie_vec = trie_tensor->flat<std::string>();
// *trie_path = model_vec(0);
// const tf::Tensor* alphabet_tensor;
// status = ctx->input("alphabet_path", &alphabet_tensor);
// if (!status.ok()) return status;
// auto alphabet_vec = alphabet_tensor->flat<std::string>();
// *alphabet_path = alphabet_vec(0);
}
std::string GetModelPath(tf::OpKernelConstruction *ctx) {
std::string model_path;
ctx->GetAttr("model_path", &model_path);
return model_path;
}
std::string GetTriePath(tf::OpKernelConstruction *ctx) {
std::string trie_path;
ctx->GetAttr("trie_path", &trie_path);
return trie_path;
}
std::string GetAlphabetPath(tf::OpKernelConstruction *ctx) {
std::string alphabet_path;
ctx->GetAttr("alphabet_path", &alphabet_path);
return alphabet_path;
}
void Compute(tf::OpKernelContext *ctx) override {
const tf::Tensor *inputs;
const tf::Tensor *seq_len;
std::string model_path;
std::string trie_path;
std::string alphabet_path;
tf::Tensor *log_prob = nullptr;
tf::OpOutputList decoded_indices;
tf::OpOutputList decoded_values;
tf::OpOutputList decoded_shape;
OP_REQUIRES_OK(ctx, decode_helper_.ValidateInputsGenerateOutputs(
ctx, &inputs, &seq_len, &model_path, &trie_path,
&alphabet_path, &log_prob, &decoded_indices,
&decoded_values, &decoded_shape));
auto inputs_t = inputs->tensor<float, 3>();
auto seq_len_t = seq_len->vec<tf::int32>();
auto log_prob_t = log_prob->matrix<float>();
const tf::TensorShape &inputs_shape = inputs->shape();
const tf::int64 max_time = inputs_shape.dim_size(0);
const tf::int64 batch_size = inputs_shape.dim_size(1);
const tf::int64 num_classes_raw = inputs_shape.dim_size(2);
OP_REQUIRES(
ctx, tf::FastBoundsCheck(num_classes_raw, std::numeric_limits<int>::max()),
tf::errors::InvalidArgument("num_classes cannot exceed max int"));
const int num_classes = static_cast<const int>(num_classes_raw);
log_prob_t.setZero();
std::vector<tf::TTypes<float>::UnalignedConstMatrix> input_list_t;
for (std::size_t t = 0; t < max_time; ++t) {
input_list_t.emplace_back(inputs_t.data() + t * batch_size * num_classes,
batch_size, num_classes);
}
tf::ctc::CTCBeamSearchDecoder<KenLMBeamState> beam_search(num_classes, beam_width_,
&beam_scorer_, 1 /* batch_size */,
merge_repeated_);
tf::Tensor input_chip(tf::DT_FLOAT, tf::TensorShape({num_classes}));
auto input_chip_t = input_chip.flat<float>();
std::vector<std::vector<std::vector<int>>> best_paths(batch_size);
std::vector<float> log_probs;
// Assumption: the blank index is num_classes - 1
for (int b = 0; b < batch_size; ++b) {
auto &best_paths_b = best_paths[b];
best_paths_b.resize(decode_helper_.GetTopPaths());
for (int t = 0; t < seq_len_t(b); ++t) {
input_chip_t = input_list_t[t].chip(b, 0);
auto input_bi =
Eigen::Map<const Eigen::ArrayXf>(input_chip_t.data(), num_classes);
beam_search.Step(input_bi);
}
OP_REQUIRES_OK(
ctx, beam_search.TopPaths(decode_helper_.GetTopPaths(), &best_paths_b,
&log_probs, merge_repeated_));
beam_search.Reset();
for (int bp = 0; bp < decode_helper_.GetTopPaths(); ++bp) {
log_prob_t(b, bp) = log_probs[bp];
}
}
OP_REQUIRES_OK(ctx, decode_helper_.StoreAllDecodedSequences(
best_paths, &decoded_indices, &decoded_values,
&decoded_shape));
}
private:
CTCDecodeHelper decode_helper_;
KenLMBeamScorer beam_scorer_;
bool merge_repeated_;
int beam_width_;
TF_DISALLOW_COPY_AND_ASSIGN(CTCBeamSearchDecoderOp);
};
REGISTER_KERNEL_BUILDER(Name("CTCBeamSearchDecoderWithLM").Device(tf::DEVICE_CPU),
CTCBeamSearchDecoderOp);

66
native_client/generate_trie.cpp Normal file
View File

@ -0,0 +1,66 @@
#include <algorithm>
#include <iostream>
#include <string>
using namespace std;
#include "lm/model.hh"
#include "trie_node.h"
#include "alphabet.h"
typedef lm::ngram::ProbingModel Model;
lm::WordIndex GetWordIndex(const Model& model, const std::string& word) {
lm::WordIndex vocab;
vocab = model.GetVocabulary().Index(word);
return vocab;
}
float ScoreWord(const Model& model, lm::WordIndex vocab) {
Model::State in_state = model.NullContextState();
Model::State out;
lm::FullScoreReturn full_score_return;
full_score_return = model.FullScore(in_state, vocab, out);
return full_score_return.prob;
}
int generate_trie(const char* alphabet_path, const char* kenlm_path, const char* vocab_path, const char* trie_path) {
Alphabet a(alphabet_path);
lm::ngram::Config config;
config.load_method = util::POPULATE_OR_READ;
Model model(kenlm_path, config);
TrieNode root(a.GetSize());
std::ifstream ifs;
ifs.open(vocab_path, std::ifstream::in);
if (!ifs.is_open()) {
std::cout << "unable to open vocabulary" << std::endl;
return -1;
}
std::ofstream ofs;
ofs.open(trie_path);
std::string word;
while (ifs >> word) {
for_each(word.begin(), word.end(), [](char& a) { a = tolower(a); });
lm::WordIndex vocab = GetWordIndex(model, word);
float unigram_score = ScoreWord(model, vocab);
root.Insert(word.c_str(), [&a](char c) {
return a.LabelFromString(string(1, c));
}, vocab, unigram_score);
}
root.WriteToStream(ofs);
ifs.close();
ofs.close();
return 0;
}
int main(void) {
return generate_trie("/Users/remorais/Development/DeepSpeech/data/alphabet.txt",
"/Users/remorais/Development/DeepSpeech/data/lm/lm.binary",
"/Users/remorais/Development/DeepSpeech/data/lm/vocab.txt",
"/Users/remorais/Development/DeepSpeech/data/lm/trie");
}

125
native_client/trie_node.h Normal file
View File

@ -0,0 +1,125 @@
/* 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.
==============================================================================*/
#ifndef TRIE_NODE_H
#define TRIE_NODE_H
#include "lm/model.hh"
#include <functional>
#include <istream>
#include <iostream>
#include <limits>
class TrieNode {
public:
TrieNode(int vocab_size)
: vocab_size(vocab_size)
, prefixCount(0)
, min_score_word(0)
, min_unigram_score(std::numeric_limits<float>::max())
{
children = new TrieNode*[vocab_size]();
}
~TrieNode() {
for (int i = 0; i < vocab_size; i++) {
delete children[i];
}
delete children;
}
void WriteToStream(std::ostream& os) {
WriteNode(os);
for (int i = 0; i < vocab_size; i++) {
if (children[i] == nullptr) {
os << -1 << std::endl;
} else {
// Recursive call
children[i]->WriteToStream(os);
}
}
}
static void ReadFromStream(std::istream& is, TrieNode* &obj, int vocab_size) {
int prefixCount;
is >> prefixCount;
if (prefixCount == -1) {
// This is an undefined child
obj = nullptr;
return;
}
obj = new TrieNode(vocab_size);
obj->ReadNode(is, prefixCount);
for (int i = 0; i < vocab_size; i++) {
// Recursive call
ReadFromStream(is, obj->children[i], vocab_size);
}
}
void Insert(const char* word, std::function<int (char)> translator,
lm::WordIndex lm_word, float unigram_score) {
char wordCharacter = *word;
prefixCount++;
if (unigram_score < min_unigram_score) {
min_unigram_score = unigram_score;
min_score_word = lm_word;
}
if (wordCharacter != '\0') {
int vocabIndex = translator(wordCharacter);
TrieNode *child = children[vocabIndex];
if (child == nullptr)
child = children[vocabIndex] = new TrieNode(vocab_size);
child->Insert(word + 1, translator, lm_word, unigram_score);
}
}
int GetFrequency() {
return prefixCount;
}
lm::WordIndex GetMinScoreWordIndex() {
return min_score_word;
}
float GetMinUnigramScore() {
return min_unigram_score;
}
TrieNode *GetChildAt(int vocabIndex) {
return children[vocabIndex];
}
private:
int vocab_size;
int prefixCount;
lm::WordIndex min_score_word;
float min_unigram_score;
TrieNode **children;
void WriteNode(std::ostream& os) const {
os << prefixCount << std::endl;
os << min_score_word << std::endl;
os << min_unigram_score << std::endl;
}
void ReadNode(std::istream& is, int first_input) {
prefixCount = first_input;
is >> min_score_word;
is >> min_unigram_score;
}
};
#endif //TRIE_NODE_H