summaryrefslogtreecommitdiff
path: root/klm/lm/search_trie.cc
diff options
context:
space:
mode:
Diffstat (limited to 'klm/lm/search_trie.cc')
-rw-r--r--klm/lm/search_trie.cc369
1 files changed, 272 insertions, 97 deletions
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 <algorithm>
+#include <cmath>
#include <cstring>
#include <cstdio>
#include <deque>
-#include <iostream>
#include <limits>
//#include <parallel/algorithm>
#include <vector>
@@ -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_t>(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<const uint8_t*>(mem_begin); group_begin != static_cast<const uint8_t*>(mem_end);) {
- const uint8_t *group_end = group_begin;
- for (group_end += entry_size; (group_end != static_cast<const uint8_t*>(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<const uint8_t*>(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 <class Weights> void ReadWeights(Weights &to) {
- ReadOrThrow(file_.get(), &to, sizeof(Weights));
+ template <class Weights> 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<WordIndex> 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<uint64_t> &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<ProbBackoff*>(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 &params, 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<uint64_t> &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<BitPacked&>(longest_) : static_cast<BitPacked&>(*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 Doing> 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<WordIndex>::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<WordIndex>::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<WordIndex>::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<WordIndex> words_;
+
+ const unsigned char order_minus_2_;
+};
+
+void SanityCheckCounts(const std::vector<uint64_t> &initial, const std::vector<uint64_t> &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<uint64_t> &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<unsigned int>(i) << "_merged";
@@ -434,36 +594,49 @@ void BuildTrie(const std::string &file_prefix, const std::vector<uint64_t> &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<unsigned char>(counts.size());
- params.middle = &*out.middle.begin();
- params.longest = &out.longest;
+ std::vector<uint64_t> fixed_counts(counts.size());
+ {
+ RecursiveInsert<JustCount> 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<WriteEntries> 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<uint64_t> &counts, const Config &config, SortedVocabulary &vocab) {
+void TrieSearch::InitializeFromARPA(const char *file, util::FilePiece &f, std::vector<uint64_t> &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<size_t>(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;
}
}
generated by cgit v1.2.3 (git 2.39.1) at 2024-09-21 17:21:47 +0000