// Copyright 2020 Jiayu DU
|
|
// See ../../COPYING for clarification regarding multiple authors
|
//
|
// 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
|
//
|
// THIS CODE IS PROVIDED *AS IS* BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
// KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION ANY IMPLIED
|
// WARRANTIES OR CONDITIONS OF TITLE, FITNESS FOR A PARTICULAR PURPOSE,
|
// MERCHANTABLITY OR NON-INFRINGEMENT.
|
// See the Apache 2 License for the specific language governing permissions and
|
// limitations under the License.
|
#ifdef HAVE_KENLM
|
#ifndef KALDI_LM_KENLM_H
|
#define KALDI_LM_KENLM_H
|
|
#include <base/kaldi-common.h>
|
#include <fst/fstlib.h>
|
#include <fst/fst-decl.h>
|
#include <fstext/deterministic-fst.h>
|
|
#include "lm/model.hh"
|
#include "util/murmur_hash.hh"
|
|
namespace kaldi {
|
|
// KenLm class wraps kenlm model(supporting both "trie" or "probing" models):
|
// 1. provides interface for loading binary LM, and holds it with ownership
|
// 2. provides interface for ngram score query at runtime
|
// 3. handles the index mapping between kaldi's symbols & kenlm's words
|
// KenLm object is heavy, stateless and thread-safe,
|
// can be shared by Fst wrapper class(i.e. KenLmDeterministicOnDemandFst)
|
class KenLm {
|
public:
|
typedef lm::WordIndex WordIndex;
|
typedef lm::ngram::State State;
|
|
public:
|
KenLm() :
|
model_(nullptr), vocab_(nullptr),
|
bos_sym_("<s>"), eos_sym_("</s>"), unk_sym_("<unk>"),
|
bos_symid_(0), eos_symid_(0), unk_symid_(0)
|
{ }
|
|
~KenLm() {
|
if (model_ != nullptr) {
|
delete model_;
|
}
|
model_ = nullptr;
|
vocab_ = nullptr;
|
symid_to_wid_.clear();
|
}
|
|
// If you have big LM on SSD hard-drive,
|
// you can set load_method to util::LoadMethod::LAZY,
|
// which enables "on-demand" model reading(via POSIX mmap) at runtime.
|
// Refer to tools/kenlm/util/mmap.hh for more load methods.
|
int Load(std::string kenlm_filename,
|
std::string kaldi_symbol_table_filename,
|
util::LoadMethod load_method = util::LoadMethod::POPULATE_OR_READ);
|
|
inline WordIndex GetWordIndex(std::string word) const {
|
return vocab_->Index(word.c_str());
|
}
|
|
inline WordIndex GetWordIndex(int32 symbol_id) const {
|
return symid_to_wid_[symbol_id];
|
}
|
|
void SetStateToBeginOfSentence(State *s) const { model_->BeginSentenceWrite(s); }
|
void SetStateToNull(State *s) const { model_->NullContextWrite(s); }
|
|
int32 BosSymbolIndex() const { return bos_symid_; }
|
int32 EosSymbolIndex() const { return eos_symid_; }
|
int32 UnkSymbolIndex() const { return unk_symid_; }
|
|
inline BaseFloat Score(const State *in_state,
|
WordIndex word,
|
State *out_state) const {
|
return model_->BaseScore(in_state, word, out_state);
|
}
|
|
// This provides a fast state hashing,
|
// KenLmDeterministicOnDemandFst needs this for Fst states managing.
|
struct StateHasher {
|
inline size_t operator()(const State &s) const noexcept {
|
return util::MurmurHashNative(s.words, sizeof(WordIndex) * s.Length());
|
}
|
};
|
|
private:
|
void ComputeSymbolToWordIndexMapping(std::string symbol_table);
|
|
private:
|
lm::base::Model *model_; // with ownership
|
|
// without ownership, points to internal vocabulary of model_
|
const lm::base::Vocabulary* vocab_;
|
|
// There are two integerized indexing systems here:
|
// 1. Kaldi's fst output *symbol index*(defined in words.txt),
|
// 2. KenLm's *word index*(defined by word string hashing).
|
// In order to rescore kaldi hypotheses with kenlm ngrams,
|
// we need to know the index mapping from symbol to word.
|
// KenLm class precomputes (during model loading) and stores this mapping,
|
// and apply the mapping at runtime.
|
// This is slower, but at least we don't need
|
// to modify/convert runtime resources.(e.g. HCLG/lattices or kenlm models)
|
//
|
// In the mapping, <eps> and #0 symbols are special:
|
// They do not correspond to any word in KenLm,
|
// so the mapping of these two symbols are logically undefined,
|
// we just map them to KenLm's <unk> to avoid random invalid mapping.
|
|
// symid_to_wid_[kaldi_symbol_index] -> kenlm word index
|
std::vector<WordIndex> symid_to_wid_;
|
|
// special lm symbols
|
std::string bos_sym_;
|
std::string eos_sym_;
|
std::string unk_sym_;
|
|
int32 bos_symid_;
|
int32 eos_symid_;
|
int32 unk_symid_;
|
}; // class KenLm
|
|
|
// DeterministicOnDemandFst wraps a KenLm object as a deteministic Fst.
|
// Internally, it manages dynamically expanded Fst states(so not thread-safe),
|
// different threads should create their own instances of this class.
|
// They are lightweight and can share the same KenLm object.
|
template<class Arc>
|
class KenLmDeterministicOnDemandFst : public fst::DeterministicOnDemandFst<Arc> {
|
public:
|
typedef typename Arc::Weight Weight;
|
typedef typename Arc::StateId StateId;
|
typedef typename Arc::Label Label;
|
typedef typename KenLm::State State;
|
typedef typename KenLm::WordIndex WordIndex;
|
|
explicit KenLmDeterministicOnDemandFst(const KenLm *lm)
|
: lm_(lm), num_states_(0), bos_state_id_(0)
|
{
|
// create bos to be FST start state
|
MapElem e;
|
lm->SetStateToBeginOfSentence(&e.first);
|
e.second = bos_state_id_;
|
std::pair<IterType, bool> r = state_map_.insert(e);
|
KALDI_ASSERT(r.second == true); // bos successfully inserted into state map
|
state_vec_.push_back(&r.first->first);
|
num_states_++;
|
|
eos_symbol_id_ = lm_->EosSymbolIndex();
|
}
|
virtual ~KenLmDeterministicOnDemandFst() { }
|
|
virtual StateId Start() {
|
return bos_state_id_;
|
}
|
|
virtual bool GetArc(StateId s, Label label, Arc *oarc) {
|
KALDI_ASSERT(s < static_cast<StateId>(state_vec_.size()));
|
const State* istate = state_vec_[s];
|
MapElem e;
|
WordIndex word = lm_->GetWordIndex(label);
|
BaseFloat log_10_prob = lm_->Score(istate, word, &e.first);
|
e.second = num_states_;
|
std::pair<IterType, bool> r = state_map_.insert(e);
|
if (r.second == true) { // new state
|
state_vec_.push_back(&(r.first->first));
|
num_states_++;
|
}
|
|
oarc->ilabel = label;
|
oarc->olabel = oarc->ilabel;
|
oarc->nextstate = r.first->second;
|
oarc->weight = Weight(-log_10_prob * M_LN10); // KenLm log10() -> Kaldi ln()
|
|
return true;
|
}
|
|
virtual Weight Final(StateId s) {
|
Arc oarc;
|
GetArc(s, eos_symbol_id_, &oarc);
|
return oarc.weight;
|
}
|
|
private:
|
typedef std::pair<State, StateId> MapElem;
|
typedef unordered_map<State, StateId, KenLm::StateHasher> MapType;
|
typedef typename MapType::iterator IterType;
|
|
const KenLm *lm_; // no ownership
|
MapType state_map_;
|
std::vector<const State*> state_vec_;
|
StateId num_states_; // state vector index range, [0, num_states_)
|
StateId bos_state_id_; // fst start state id
|
Label eos_symbol_id_;
|
}; // class KenLmDeterministicOnDemandFst
|
} // namespace kaldi
|
#endif
|
#endif
|
