diff options
Diffstat (limited to 'klm/lm/search_trie.hh')
| -rw-r--r-- | klm/lm/search_trie.hh | 10 | 
1 files changed, 7 insertions, 3 deletions
| diff --git a/klm/lm/search_trie.hh b/klm/lm/search_trie.hh index 33ae8cff..5155ca02 100644 --- a/klm/lm/search_trie.hh +++ b/klm/lm/search_trie.hh @@ -7,6 +7,7 @@  #include "lm/trie.hh"  #include "lm/weights.hh" +#include "util/file.hh"  #include "util/file_piece.hh"  #include <vector> @@ -20,7 +21,8 @@ class SortedVocabulary;  namespace trie {  template <class Quant, class Bhiksha> class TrieSearch; -template <class Quant, class Bhiksha> void BuildTrie(const std::string &file_prefix, std::vector<uint64_t> &counts, const Config &config, TrieSearch<Quant, Bhiksha> &out, Quant &quant, const SortedVocabulary &vocab, Backing &backing); +class SortedFiles; +template <class Quant, class Bhiksha> void BuildTrie(SortedFiles &files, std::vector<uint64_t> &counts, const Config &config, TrieSearch<Quant, Bhiksha> &out, Quant &quant, const SortedVocabulary &vocab, Backing &backing);  template <class Quant, class Bhiksha> class TrieSearch {    public: @@ -40,7 +42,7 @@ template <class Quant, class Bhiksha> class TrieSearch {      static void UpdateConfigFromBinary(int fd, const std::vector<uint64_t> &counts, Config &config) {        Quant::UpdateConfigFromBinary(fd, counts, config); -      AdvanceOrThrow(fd, Quant::Size(counts.size(), config) + Unigram::Size(counts[0])); +      util::AdvanceOrThrow(fd, Quant::Size(counts.size(), config) + Unigram::Size(counts[0]));        Bhiksha::UpdateConfigFromBinary(fd, config);      } @@ -60,6 +62,8 @@ template <class Quant, class Bhiksha> class TrieSearch {      void LoadedBinary(); +    typedef const Middle *MiddleIter; +      const Middle *MiddleBegin() const { return middle_begin_; }      const Middle *MiddleEnd() const { return middle_end_; } @@ -108,7 +112,7 @@ template <class Quant, class Bhiksha> class TrieSearch {      }    private: -    friend void BuildTrie<Quant, Bhiksha>(const std::string &file_prefix, std::vector<uint64_t> &counts, const Config &config, TrieSearch<Quant, Bhiksha> &out, Quant &quant, const SortedVocabulary &vocab, Backing &backing); +    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);      // Middles are managed manually so we can delay construction and they don't have to be copyable.        void FreeMiddles() { | 
