From bee6a3c3f6c54cf7449229488c6124dddc7e2f31 Mon Sep 17 00:00:00 2001 From: Chris Dyer Date: Tue, 18 Jan 2011 15:55:40 -0500 Subject: new version of klm --- klm/lm/search_trie.cc | 369 +++++++++++++++++++++++++++++++++++++------------- 1 file changed, 272 insertions(+), 97 deletions(-) (limited to 'klm/lm/search_trie.cc') diff --git a/klm/lm/search_trie.cc b/klm/lm/search_trie.cc index 12294682..3aeeeca3 100644 --- a/klm/lm/search_trie.cc +++ b/klm/lm/search_trie.cc @@ -1,6 +1,7 @@ /* This is where the trie is built. It's on-disk. */ #include "lm/search_trie.hh" +#include "lm/blank.hh" #include "lm/lm_exception.hh" #include "lm/read_arpa.hh" #include "lm/trie.hh" @@ -13,10 +14,10 @@ #include "util/scoped.hh" #include +#include #include #include #include -#include #include //#include #include @@ -152,11 +153,11 @@ void ReadOrThrow(FILE *from, void *data, size_t size) { if (1 != std::fread(data, size, 1, from)) UTIL_THROW(util::ErrnoException, "Short read; requested size " << size); } +const std::size_t kCopyBufSize = 512; void CopyOrThrow(FILE *from, FILE *to, size_t size) { - const size_t kBufSize = 512; - char buf[kBufSize]; - for (size_t i = 0; i < size; i += kBufSize) { - std::size_t amount = std::min(size - i, kBufSize); + char buf[std::min(size, kCopyBufSize)]; + for (size_t i = 0; i < size; i += kCopyBufSize) { + std::size_t amount = std::min(size - i, kCopyBufSize); ReadOrThrow(from, buf, amount); WriteOrThrow(to, buf, amount); } @@ -172,8 +173,10 @@ std::string DiskFlush(const void *mem_begin, const void *mem_end, const std::str if (!out.get()) UTIL_THROW(util::ErrnoException, "Couldn't open " << assembled.str().c_str() << " for writing"); // Compress entries that being with the same (order-1) words. for (const uint8_t *group_begin = static_cast(mem_begin); group_begin != static_cast(mem_end);) { - const uint8_t *group_end = group_begin; - for (group_end += entry_size; (group_end != static_cast(mem_end)) && !memcmp(group_begin, group_end, prefix_size); group_end += entry_size) {} + const uint8_t *group_end; + for (group_end = group_begin + entry_size; + (group_end != static_cast(mem_end)) && !memcmp(group_begin, group_end, prefix_size); + group_end += entry_size) {} WriteOrThrow(out.get(), group_begin, prefix_size); WordIndex group_size = (group_end - group_begin) / entry_size; WriteOrThrow(out.get(), &group_size, sizeof(group_size)); @@ -188,7 +191,7 @@ std::string DiskFlush(const void *mem_begin, const void *mem_end, const std::str class SortedFileReader { public: - SortedFileReader() {} + SortedFileReader() : ended_(false) {} void Init(const std::string &name, unsigned char order) { file_.reset(fopen(name.c_str(), "r")); @@ -206,25 +209,39 @@ class SortedFileReader { std::size_t HeaderBytes() const { return header_.size() * sizeof(WordIndex); } void NextHeader() { - if (1 != fread(&*header_.begin(), HeaderBytes(), 1, file_.get()) && !Ended()) { - UTIL_THROW(util::ErrnoException, "Short read of counts"); + if (1 != fread(&*header_.begin(), HeaderBytes(), 1, file_.get())) { + if (feof(file_.get())) { + ended_ = true; + } else { + UTIL_THROW(util::ErrnoException, "Short read of counts"); + } } } - void ReadCount(WordIndex &to) { - ReadOrThrow(file_.get(), &to, sizeof(WordIndex)); + WordIndex ReadCount() { + WordIndex ret; + ReadOrThrow(file_.get(), &ret, sizeof(WordIndex)); + return ret; } - void ReadWord(WordIndex &to) { - ReadOrThrow(file_.get(), &to, sizeof(WordIndex)); + WordIndex ReadWord() { + WordIndex ret; + ReadOrThrow(file_.get(), &ret, sizeof(WordIndex)); + return ret; } - template void ReadWeights(Weights &to) { - ReadOrThrow(file_.get(), &to, sizeof(Weights)); + template void ReadWeights(Weights &weights) { + ReadOrThrow(file_.get(), &weights, sizeof(Weights)); } bool Ended() { - return feof(file_.get()); + return ended_; + } + + void Rewind() { + rewind(file_.get()); + ended_ = false; + NextHeader(); } FILE *File() { return file_.get(); } @@ -233,12 +250,13 @@ class SortedFileReader { util::scoped_FILE file_; std::vector header_; + + bool ended_; }; void CopyFullRecord(SortedFileReader &from, FILE *to, std::size_t weights_size) { WriteOrThrow(to, from.Header(), from.HeaderBytes()); - WordIndex count; - from.ReadCount(count); + WordIndex count = from.ReadCount(); WriteOrThrow(to, &count, sizeof(WordIndex)); CopyOrThrow(from.File(), to, (weights_size + sizeof(WordIndex)) * count); @@ -263,25 +281,23 @@ void MergeSortedFiles(const char *first_name, const char *second_name, const cha } // Merge at the entry level. WriteOrThrow(out_file.get(), first.Header(), first.HeaderBytes()); - WordIndex first_count, second_count; - first.ReadCount(first_count); second.ReadCount(second_count); + WordIndex first_count = first.ReadCount(), second_count = second.ReadCount(); WordIndex total_count = first_count + second_count; WriteOrThrow(out_file.get(), &total_count, sizeof(WordIndex)); - WordIndex first_word, second_word; - first.ReadWord(first_word); second.ReadWord(second_word); + WordIndex first_word = first.ReadWord(), second_word = second.ReadWord(); WordIndex first_index = 0, second_index = 0; while (true) { if (first_word < second_word) { WriteOrThrow(out_file.get(), &first_word, sizeof(WordIndex)); CopyOrThrow(first.File(), out_file.get(), weights_size); if (++first_index == first_count) break; - first.ReadWord(first_word); + first_word = first.ReadWord(); } else { WriteOrThrow(out_file.get(), &second_word, sizeof(WordIndex)); CopyOrThrow(second.File(), out_file.get(), weights_size); if (++second_index == second_count) break; - second.ReadWord(second_word); + second_word = second.ReadWord(); } } if (first_index == first_count) { @@ -358,75 +374,219 @@ void ARPAToSortedFiles(util::FilePiece &f, const std::vector &counts, { std::string unigram_name = file_prefix + "unigrams"; util::scoped_fd unigram_file; - util::scoped_mmap unigram_mmap; - unigram_mmap.reset(util::MapZeroedWrite(unigram_name.c_str(), counts[0] * sizeof(ProbBackoff), unigram_file), counts[0] * sizeof(ProbBackoff)); + util::scoped_mmap unigram_mmap(util::MapZeroedWrite(unigram_name.c_str(), counts[0] * sizeof(ProbBackoff), unigram_file), counts[0] * sizeof(ProbBackoff)); Read1Grams(f, counts[0], vocab, reinterpret_cast(unigram_mmap.get())); } util::scoped_memory mem; - mem.reset(malloc(buffer), buffer, util::scoped_memory::ARRAY_ALLOCATED); + mem.reset(malloc(buffer), buffer, util::scoped_memory::MALLOC_ALLOCATED); if (!mem.get()) UTIL_THROW(util::ErrnoException, "malloc failed for sort buffer size " << buffer); ConvertToSorted(f, vocab, counts, mem, file_prefix, counts.size()); ReadEnd(f); } -struct RecursiveInsertParams { - WordIndex *words; - SortedFileReader *files; - unsigned char max_order; - // This is an array of size order - 2. - BitPackedMiddle *middle; - // This has exactly one entry. - BitPackedLongest *longest; -}; - -uint64_t RecursiveInsert(RecursiveInsertParams ¶ms, unsigned char order) { - SortedFileReader &file = params.files[order - 2]; - const uint64_t ret = (order == params.max_order) ? params.longest->InsertIndex() : params.middle[order - 2].InsertIndex(); - if (std::memcmp(params.words, file.Header(), sizeof(WordIndex) * (order - 1))) - return ret; - WordIndex count; - file.ReadCount(count); - WordIndex key; - if (order == params.max_order) { - Prob value; - for (WordIndex i = 0; i < count; ++i) { - file.ReadWord(key); - file.ReadWeights(value); - params.longest->Insert(key, value.prob); - } - file.NextHeader(); - return ret; - } - ProbBackoff value; - for (WordIndex i = 0; i < count; ++i) { - file.ReadWord(params.words[order - 1]); - file.ReadWeights(value); - params.middle[order - 2].Insert( - params.words[order - 1], - value.prob, - value.backoff, - RecursiveInsert(params, order + 1)); +bool HeadMatch(const WordIndex *words, const WordIndex *const words_end, const WordIndex *header) { + for (; words != words_end; ++words, ++header) { + if (*words != *header) { + assert(*words <= *header); + return false; + } } - file.NextHeader(); - return ret; + return true; } -void BuildTrie(const std::string &file_prefix, const std::vector &counts, std::ostream *messages, TrieSearch &out) { - UnigramValue *unigrams = out.unigram.Raw(); - // Load unigrams. Leave the next pointers uninitialized. - { - std::string name(file_prefix + "unigrams"); - util::scoped_FILE file(fopen(name.c_str(), "r")); - if (!file.get()) UTIL_THROW(util::ErrnoException, "Opening " << name << " failed"); - for (WordIndex i = 0; i < counts[0]; ++i) { - ReadOrThrow(file.get(), &unigrams[i].weights, sizeof(ProbBackoff)); +class JustCount { + public: + JustCount(UnigramValue * /*unigrams*/, BitPackedMiddle * /*middle*/, BitPackedLongest &/*longest*/, uint64_t *counts, unsigned char order) + : counts_(counts), longest_counts_(counts + order - 1) {} + + void Unigrams(WordIndex begin, WordIndex end) { + counts_[0] += end - begin; } - unlink(name.c_str()); + + void MiddleBlank(const unsigned char mid_idx, WordIndex /* idx */) { + ++counts_[mid_idx + 1]; + } + + void Middle(const unsigned char mid_idx, WordIndex /*key*/, const ProbBackoff &/*weights*/) { + ++counts_[mid_idx + 1]; + } + + void Longest(WordIndex /*key*/, Prob /*prob*/) { + ++*longest_counts_; + } + + // Unigrams wrote one past. + void Cleanup() { + --counts_[0]; + } + + private: + uint64_t *const counts_, *const longest_counts_; +}; + +class WriteEntries { + public: + WriteEntries(UnigramValue *unigrams, BitPackedMiddle *middle, BitPackedLongest &longest, const uint64_t * /*counts*/, unsigned char order) : + unigrams_(unigrams), + middle_(middle), + longest_(longest), + bigram_pack_((order == 2) ? static_cast(longest_) : static_cast(*middle_)) {} + + void Unigrams(WordIndex begin, WordIndex end) { + uint64_t next = bigram_pack_.InsertIndex(); + for (UnigramValue *i = unigrams_ + begin; i < unigrams_ + end; ++i) { + i->next = next; + } + } + + void MiddleBlank(const unsigned char mid_idx, WordIndex key) { + middle_[mid_idx].Insert(key, kBlankProb, kBlankBackoff); + } + + void Middle(const unsigned char mid_idx, WordIndex key, const ProbBackoff &weights) { + middle_[mid_idx].Insert(key, weights.prob, weights.backoff); + } + + void Longest(WordIndex key, Prob prob) { + longest_.Insert(key, prob.prob); + } + + void Cleanup() {} + + private: + UnigramValue *const unigrams_; + BitPackedMiddle *const middle_; + BitPackedLongest &longest_; + BitPacked &bigram_pack_; +}; + +template class RecursiveInsert { + public: + RecursiveInsert(SortedFileReader *inputs, UnigramValue *unigrams, BitPackedMiddle *middle, BitPackedLongest &longest, uint64_t *counts, unsigned char order) : + doing_(unigrams, middle, longest, counts, order), inputs_(inputs), inputs_end_(inputs + order - 1), words_(new WordIndex[order]), order_minus_2_(order - 2) { + } + + // Outer unigram loop. + void Apply(std::ostream *progress_out, const char *message, WordIndex unigram_count) { + util::ErsatzProgress progress(progress_out, message, unigram_count + 1); + for (words_[0] = 0; ; ++words_[0]) { + WordIndex min_continue = unigram_count; + for (SortedFileReader *other = inputs_; other != inputs_end_; ++other) { + if (other->Ended()) continue; + min_continue = std::min(min_continue, other->Header()[0]); + } + // This will write at unigram_count. This is by design so that the next pointers will make sense. + doing_.Unigrams(words_[0], min_continue + 1); + if (min_continue == unigram_count) break; + progress += min_continue - words_[0]; + words_[0] = min_continue; + Middle(0); + } + doing_.Cleanup(); + } + + private: + void Middle(const unsigned char mid_idx) { + // (mid_idx + 2)-gram. + if (mid_idx == order_minus_2_) { + Longest(); + return; + } + // Orders [2, order) + + SortedFileReader &reader = inputs_[mid_idx]; + + if (reader.Ended() || !HeadMatch(words_.get(), words_.get() + mid_idx + 1, reader.Header())) { + // This order doesn't have a header match, but longer ones might. + MiddleAllBlank(mid_idx); + return; + } + + // There is a header match. + WordIndex count = reader.ReadCount(); + WordIndex current = reader.ReadWord(); + while (count) { + WordIndex min_continue = std::numeric_limits::max(); + for (SortedFileReader *other = inputs_ + mid_idx + 1; other < inputs_end_; ++other) { + if (!other->Ended() && HeadMatch(words_.get(), words_.get() + mid_idx + 1, other->Header())) + min_continue = std::min(min_continue, other->Header()[mid_idx + 1]); + } + while (true) { + if (current > min_continue) { + doing_.MiddleBlank(mid_idx, min_continue); + words_[mid_idx + 1] = min_continue; + Middle(mid_idx + 1); + break; + } + ProbBackoff weights; + reader.ReadWeights(weights); + doing_.Middle(mid_idx, current, weights); + --count; + if (current == min_continue) { + words_[mid_idx + 1] = min_continue; + Middle(mid_idx + 1); + if (count) current = reader.ReadWord(); + break; + } + if (!count) break; + current = reader.ReadWord(); + } + } + // Count is now zero. Finish off remaining blanks. + MiddleAllBlank(mid_idx); + reader.NextHeader(); + } + + void MiddleAllBlank(const unsigned char mid_idx) { + while (true) { + WordIndex min_continue = std::numeric_limits::max(); + for (SortedFileReader *other = inputs_ + mid_idx + 1; other < inputs_end_; ++other) { + if (!other->Ended() && HeadMatch(words_.get(), words_.get() + mid_idx + 1, other->Header())) + min_continue = std::min(min_continue, other->Header()[mid_idx + 1]); + } + if (min_continue == std::numeric_limits::max()) return; + doing_.MiddleBlank(mid_idx, min_continue); + words_[mid_idx + 1] = min_continue; + Middle(mid_idx + 1); + } + } + + void Longest() { + SortedFileReader &reader = *(inputs_end_ - 1); + if (reader.Ended() || !HeadMatch(words_.get(), words_.get() + order_minus_2_ + 1, reader.Header())) return; + WordIndex count = reader.ReadCount(); + for (WordIndex i = 0; i < count; ++i) { + WordIndex word = reader.ReadWord(); + Prob prob; + reader.ReadWeights(prob); + doing_.Longest(word, prob); + } + reader.NextHeader(); + return; + } + + Doing doing_; + + SortedFileReader *inputs_; + SortedFileReader *inputs_end_; + + util::scoped_array words_; + + const unsigned char order_minus_2_; +}; + +void SanityCheckCounts(const std::vector &initial, const std::vector &fixed) { + if (fixed[0] != initial[0]) UTIL_THROW(util::Exception, "Unigram count should be constant but initial is " << initial[0] << " and recounted is " << fixed[0]); + if (fixed.back() != initial.back()) UTIL_THROW(util::Exception, "Longest count should be constant"); + for (unsigned char i = 0; i < initial.size(); ++i) { + if (fixed[i] < initial[i]) UTIL_THROW(util::Exception, "Counts came out lower than expected. This shouldn't happen"); } +} - // inputs[0] is bigrams. +void BuildTrie(const std::string &file_prefix, const std::vector &counts, const Config &config, TrieSearch &out, Backing &backing) { SortedFileReader inputs[counts.size() - 1]; + for (unsigned char i = 2; i <= counts.size(); ++i) { std::stringstream assembled; assembled << file_prefix << static_cast(i) << "_merged"; @@ -434,36 +594,49 @@ void BuildTrie(const std::string &file_prefix, const std::vector &coun unlink(assembled.str().c_str()); } - // words[0] is unigrams. - WordIndex words[counts.size()]; - RecursiveInsertParams params; - params.words = words; - params.files = inputs; - params.max_order = static_cast(counts.size()); - params.middle = &*out.middle.begin(); - params.longest = &out.longest; + std::vector fixed_counts(counts.size()); + { + RecursiveInsert counter(inputs, NULL, &*out.middle.begin(), out.longest, &*fixed_counts.begin(), counts.size()); + counter.Apply(config.messages, "Counting n-grams that should not have been pruned", counts[0]); + } + SanityCheckCounts(counts, fixed_counts); + + out.SetupMemory(GrowForSearch(config, TrieSearch::kModelType, fixed_counts, TrieSearch::Size(fixed_counts, config), backing), fixed_counts, config); + + for (unsigned char i = 2; i <= counts.size(); ++i) { + inputs[i-2].Rewind(); + } + UnigramValue *unigrams = out.unigram.Raw(); + // Fill entries except unigram probabilities. { - util::ErsatzProgress progress(messages, "Building trie", counts[0]); - for (words[0] = 0; words[0] < counts[0]; ++words[0], ++progress) { - unigrams[words[0]].next = RecursiveInsert(params, 2); + RecursiveInsert inserter(inputs, unigrams, &*out.middle.begin(), out.longest, &*fixed_counts.begin(), counts.size()); + inserter.Apply(config.messages, "Building trie", fixed_counts[0]); + } + + // Fill unigram probabilities. + { + std::string name(file_prefix + "unigrams"); + util::scoped_FILE file(fopen(name.c_str(), "r")); + if (!file.get()) UTIL_THROW(util::ErrnoException, "Opening " << name << " failed"); + for (WordIndex i = 0; i < counts[0]; ++i) { + ReadOrThrow(file.get(), &unigrams[i].weights, sizeof(ProbBackoff)); } + unlink(name.c_str()); } /* Set ending offsets so the last entry will be sized properly */ + // Last entry for unigrams was already set. if (!out.middle.empty()) { - unigrams[counts[0]].next = out.middle.front().InsertIndex(); for (size_t i = 0; i < out.middle.size() - 1; ++i) { out.middle[i].FinishedLoading(out.middle[i+1].InsertIndex()); } out.middle.back().FinishedLoading(out.longest.InsertIndex()); - } else { - unigrams[counts[0]].next = out.longest.InsertIndex(); } } } // namespace -void TrieSearch::InitializeFromARPA(const char *file, util::FilePiece &f, const std::vector &counts, const Config &config, SortedVocabulary &vocab) { +void TrieSearch::InitializeFromARPA(const char *file, util::FilePiece &f, std::vector &counts, const Config &config, SortedVocabulary &vocab, Backing &backing) { std::string temporary_directory; if (config.temporary_directory_prefix) { temporary_directory = config.temporary_directory_prefix; @@ -473,7 +646,8 @@ void TrieSearch::InitializeFromARPA(const char *file, util::FilePiece &f, const temporary_directory = file; } // Null on end is kludge to ensure null termination. - temporary_directory += "-tmp-XXXXXX\0"; + temporary_directory += "-tmp-XXXXXX"; + temporary_directory += '\0'; if (!mkdtemp(&temporary_directory[0])) { UTIL_THROW(util::ErrnoException, "Failed to make a temporary directory based on the name " << temporary_directory.c_str()); } @@ -483,9 +657,10 @@ void TrieSearch::InitializeFromARPA(const char *file, util::FilePiece &f, const temporary_directory += '/'; // At least 1MB sorting memory. ARPAToSortedFiles(f, counts, std::max(config.building_memory, 1048576), temporary_directory.c_str(), vocab); - BuildTrie(temporary_directory.c_str(), counts, config.messages, *this); - if (rmdir(temporary_directory.c_str())) { - std::cerr << "Failed to delete " << temporary_directory << std::endl; + + BuildTrie(temporary_directory.c_str(), counts, config, *this, backing); + if (rmdir(temporary_directory.c_str()) && config.messages) { + *config.messages << "Failed to delete " << temporary_directory << std::endl; } } -- cgit v1.2.3