xls-r-300m-sv-robust / kenlm /lm /interpolate /merge_probabilities.hh
marinone94's picture
Training in progress, epoch 0
1ce325b
raw
history blame
3.39 kB
#ifndef LM_INTERPOLATE_MERGE_PROBABILITIES_H
#define LM_INTERPOLATE_MERGE_PROBABILITIES_H
#include "../common/ngram.hh"
#include "bounded_sequence_encoding.hh"
#include "../../util/fixed_array.hh"
#include "../../util/stream/multi_stream.hh"
#include <stdint.h>
namespace lm {
namespace interpolate {
struct InterpolateInfo;
/**
* Make the encoding of backoff values for a given order. This stores values
* in [PartialProbGamma::FromBegin(), PartialProbGamma::FromEnd())
*/
BoundedSequenceEncoding MakeEncoder(const InterpolateInfo &info, uint8_t order);
/**
* The first pass for the offline log-linear interpolation algorithm. This
* reads K **suffix-ordered** streams for each model, for each order, of
* ngram records (ngram-id, prob, backoff). It further assumes that the
* ngram-ids have been unified over all of the stream inputs.
*
* Its output is records of (ngram-id, prob-prod, backoff-level,
* backoff-level, ...) where the backoff-levels (of which there are K) are
* the context length (0 for unigrams) that the corresponding model had to
* back off to in order to obtain a probability for that ngram-id. Each of
* these streams is terminated with a record whose ngram-id is all
* maximum-integers for simplicity in implementation here.
*
* @param model_by_order An array of length N (max_i N_i) containing at
* the ChainPositions for the streams for order (i + 1).
* The Rus attached to output chains for each order (of length K)
*/
class MergeProbabilities {
public:
MergeProbabilities(const InterpolateInfo &info, util::FixedArray<util::stream::ChainPositions> &models_by_order)
: info_(info), models_by_order_(models_by_order) {}
void Run(const util::stream::ChainPositions &outputs);
private:
const InterpolateInfo &info_;
util::FixedArray<util::stream::ChainPositions> &models_by_order_;
};
/**
* This class represents the output payload for this pass, which consists
* of an ngram-id, a probability, and then a vector of orders from which
* each of the component models backed off to for this ngram, encoded
* using the BoundedSequenceEncoding class.
*/
class PartialProbGamma : public lm::NGramHeader {
public:
PartialProbGamma(std::size_t order, std::size_t backoff_bytes)
: lm::NGramHeader(NULL, order), backoff_bytes_(backoff_bytes) {
// nothing
}
std::size_t TotalSize() const {
return sizeof(WordIndex) * Order() + sizeof(After) + backoff_bytes_;
}
// TODO: cache bounded sequence encoding in the pipeline?
static std::size_t TotalSize(const InterpolateInfo &info, uint8_t order) {
return sizeof(WordIndex) * order + sizeof(After) + MakeEncoder(info, order).EncodedLength();
}
float &Prob() { return Pay().prob; }
float Prob() const { return Pay().prob; }
float &LowerProb() { return Pay().lower_prob; }
float LowerProb() const { return Pay().lower_prob; }
const uint8_t *FromBegin() const { return Pay().from; }
uint8_t *FromBegin() { return Pay().from; }
private:
struct After {
// Note that backoff_and_normalize assumes this comes first.
float prob;
float lower_prob;
uint8_t from[];
};
const After &Pay() const { return *reinterpret_cast<const After *>(end()); }
After &Pay() { return *reinterpret_cast<After*>(end()); }
std::size_t backoff_bytes_;
};
}} // namespaces
#endif // LM_INTERPOLATE_MERGE_PROBABILITIES_H