diff options
author | Kenneth Heafield <github@kheafield.com> | 2012-06-28 18:58:59 +0400 |
---|---|---|
committer | Kenneth Heafield <github@kheafield.com> | 2012-06-28 18:58:59 +0400 |
commit | c7e02eb67743b31f4cfab752c60fe96f8e2ae630 (patch) | |
tree | e25d3c43e239e02ae3adff5262b8f674decb826a /lm/search_trie.hh | |
parent | f6b3354f396f8d90e273a3a8255b68cfd5f33f70 (diff) |
KenLM e3b5c55910 including rest costs for probing
Diffstat (limited to 'lm/search_trie.hh')
-rw-r--r-- | lm/search_trie.hh | 71 |
1 files changed, 34 insertions, 37 deletions
diff --git a/lm/search_trie.hh b/lm/search_trie.hh index 5155ca020..10b22ab18 100644 --- a/lm/search_trie.hh +++ b/lm/search_trie.hh @@ -28,13 +28,11 @@ template <class Quant, class Bhiksha> class TrieSearch { public: typedef NodeRange Node; - typedef ::lm::ngram::trie::Unigram Unigram; - Unigram unigram; - - typedef trie::BitPackedMiddle<typename Quant::Middle, Bhiksha> Middle; + typedef ::lm::ngram::trie::UnigramPointer UnigramPointer; + typedef typename Quant::MiddlePointer MiddlePointer; + typedef typename Quant::LongestPointer LongestPointer; - typedef trie::BitPackedLongest<typename Quant::Longest> Longest; - Longest longest; + static const bool kDifferentRest = false; static const ModelType kModelType = static_cast<ModelType>(TRIE_SORTED + Quant::kModelTypeAdd + Bhiksha::kModelTypeAdd); @@ -62,55 +60,46 @@ template <class Quant, class Bhiksha> class TrieSearch { void LoadedBinary(); - typedef const Middle *MiddleIter; + void InitializeFromARPA(const char *file, util::FilePiece &f, std::vector<uint64_t> &counts, const Config &config, SortedVocabulary &vocab, Backing &backing); - const Middle *MiddleBegin() const { return middle_begin_; } - const Middle *MiddleEnd() const { return middle_end_; } + unsigned char Order() const { + return middle_end_ - middle_begin_ + 2; + } - void InitializeFromARPA(const char *file, util::FilePiece &f, std::vector<uint64_t> &counts, const Config &config, SortedVocabulary &vocab, Backing &backing); + ProbBackoff &UnknownUnigram() { return unigram_.Unknown(); } - void LookupUnigram(WordIndex word, float &backoff, Node &node, FullScoreReturn &ret) const { - unigram.Find(word, ret.prob, backoff, node); - ret.independent_left = (node.begin == node.end); - ret.extend_left = static_cast<uint64_t>(word); + UnigramPointer LookupUnigram(WordIndex word, Node &next, bool &independent_left, uint64_t &extend_left) const { + extend_left = static_cast<uint64_t>(word); + UnigramPointer ret(unigram_.Find(word, next)); + independent_left = (next.begin == next.end); + return ret; } - bool LookupMiddle(const Middle &mid, WordIndex word, float &backoff, Node &node, FullScoreReturn &ret) const { - if (!mid.Find(word, ret.prob, backoff, node, ret.extend_left)) return false; - ret.independent_left = (node.begin == node.end); - return true; + MiddlePointer Unpack(uint64_t extend_pointer, unsigned char extend_length, Node &node) const { + return MiddlePointer(quant_, extend_length - 2, middle_begin_[extend_length - 2].ReadEntry(extend_pointer, node)); } - bool LookupMiddleNoProb(const Middle &mid, WordIndex word, float &backoff, Node &node) const { - return mid.FindNoProb(word, backoff, node); + MiddlePointer LookupMiddle(unsigned char order_minus_2, WordIndex word, Node &node, bool &independent_left, uint64_t &extend_left) const { + util::BitAddress address(middle_begin_[order_minus_2].Find(word, node, extend_left)); + independent_left = (address.base == NULL) || (node.begin == node.end); + return MiddlePointer(quant_, order_minus_2, address); } - bool LookupLongest(WordIndex word, float &prob, const Node &node) const { - return longest.Find(word, prob, node); + LongestPointer LookupLongest(WordIndex word, const Node &node) const { + return LongestPointer(quant_, longest_.Find(word, node)); } bool FastMakeNode(const WordIndex *begin, const WordIndex *end, Node &node) const { - // TODO: don't decode backoff. assert(begin != end); - FullScoreReturn ignored; - float ignored_backoff; - LookupUnigram(*begin, ignored_backoff, node, ignored); + bool independent_left; + uint64_t ignored; + LookupUnigram(*begin, node, independent_left, ignored); for (const WordIndex *i = begin + 1; i < end; ++i) { - if (!LookupMiddleNoProb(middle_begin_[i - begin - 1], *i, ignored_backoff, node)) return false; + if (independent_left || !LookupMiddle(i - begin - 1, *i, node, independent_left, ignored).Found()) return false; } return true; } - Node Unpack(uint64_t extend_pointer, unsigned char extend_length, float &prob) const { - if (extend_length == 1) { - float ignored; - Node ret; - unigram.Find(static_cast<WordIndex>(extend_pointer), prob, ignored, ret); - return ret; - } - return middle_begin_[extend_length - 2].ReadEntry(extend_pointer, prob); - } - private: friend void BuildTrie<Quant, Bhiksha>(SortedFiles &files, std::vector<uint64_t> &counts, const Config &config, TrieSearch<Quant, Bhiksha> &out, Quant &quant, const SortedVocabulary &vocab, Backing &backing); @@ -122,8 +111,16 @@ template <class Quant, class Bhiksha> class TrieSearch { free(middle_begin_); } + typedef trie::BitPackedMiddle<Bhiksha> Middle; + + typedef trie::BitPackedLongest Longest; + Longest longest_; + Middle *middle_begin_, *middle_end_; Quant quant_; + + typedef ::lm::ngram::trie::Unigram Unigram; + Unigram unigram_; }; } // namespace trie |