Final lattice can be rescored with a bigger LM.

Author: ondrejklejch

Makefile

@@ -47,7 +47,8 @@ ADDLIBS = $(FSTROOT)/src/lib/.libs/libfst.a \
           $(KALDI_DIR)/src/tree/kaldi-tree.a \
           $(KALDI_DIR)/src/matrix/kaldi-matrix.a \
           $(KALDI_DIR)/src/util/kaldi-util.a \
-          $(KALDI_DIR)/src/base/kaldi-base.a
+          $(KALDI_DIR)/src/base/kaldi-base.a \
+          $(KALDI_DIR)/src/fstext/kaldi-fstext.a
 
 LDFLAGS = $(ADDLIBS) -llapack_atlas -lcblas -latlas -lf77blas -lm -lpthread -ldl
 

README.md

@@ -70,6 +70,10 @@ Example of `alex_asr.conf` that should reside in `model_dir`:
                        # with configuration of the pitch extractor.
 --bits_per_sample=16   # 8/16; How many bits per sample frame?
 
+--rescore=True         # Rescore lattice with a bigger LM?
+--lm_small=G_small.fst # Original G.fst
+--lm_big=G_big.fst     # G.fst corresponding to the bigger LM
+
 # These parameters specify filenames of configuration of the particular parts of the decoder. Detailed below.
 --cfg_decoder=decoder.cfg
 --cfg_decodable=decodable.cfg

setup.py

@@ -58,7 +58,7 @@ def run(self):
     packages=find_packages(exclude=["alex_asr/decoder.cpp"]),
     include_package_data=True,
     cmdclass={'build_ext': build_ext_with_make},
-    version='1.0.4',
+    version='1.0.5',
     install_requires=install_requires,
     setup_requires=['cython>=0.19.1', 'nose>=1.0'],
     ext_modules=[

src/decoder.cc

@@ -2,6 +2,9 @@
 #include "src/utils.h"
 
 #include "online2/onlinebin-util.h"
+#include "fst/fstlib.h"
+#include "fstext/table-matcher.h"
+#include "fstext/kaldi-fst-io.h"
 #include "lat/kaldi-lattice.h"
 #include "lat/sausages.h"
 
@@ -11,6 +14,8 @@ namespace alex_asr {
     Decoder::Decoder(const string model_path) :
             feature_pipeline_(NULL),
             hclg_(NULL),
+            lm_small_(NULL),
+            lm_big_(NULL),
             decoder_(NULL),
             trans_model_(NULL),
             am_nnet2_(NULL),
@@ -19,7 +24,6 @@ namespace alex_asr {
             words_(NULL),
             config_(NULL),
             decodable_(NULL)
-
     {
         // Change dir to model_path. Change back when leaving the scope.
         local_cwd cwd_to_model_path(model_path);
@@ -35,6 +39,8 @@ namespace alex_asr {
     Decoder::~Decoder() {
         delete feature_pipeline_;
         delete hclg_;
+        delete lm_small_;
+        delete lm_big_;
         delete decoder_;
         delete trans_model_;
         delete am_nnet2_;
@@ -109,6 +115,37 @@ namespace alex_asr {
             WordBoundaryInfoNewOpts word_boundary_info_opts;
             word_boundary_info_ = new WordBoundaryInfo(word_boundary_info_opts, config_->word_boundary_rxfilename);
         }
+
+        if(config_->rescore == true) {
+            LoadLM(config_->lm_small_rxfilename, &lm_small_);
+            LoadLM(config_->lm_big_rxfilename, &lm_big_);
+        }
+    }
+
+    void Decoder::LoadLM(
+        const string path,
+        fst::MapFst<fst::StdArc, LatticeArc, fst::StdToLatticeMapper<BaseFloat> > **lm_fst
+    ) {
+        int num_states_cache = 50000;
+
+        if (FileExists(path)) {
+            fst::VectorFst<fst::StdArc> *std_lm_fst = fst::ReadFstKaldi(path);
+            fst::Project(std_lm_fst, fst::PROJECT_OUTPUT);
+            if (std_lm_fst->Properties(fst::kILabelSorted, true) == 0) {
+                fst::ILabelCompare<fst::StdArc> ilabel_comp;
+                fst::ArcSort(std_lm_fst, ilabel_comp);
+            }
+
+            fst::CacheOptions cache_opts(true, num_states_cache);
+            fst::MapFstOptions mapfst_opts(cache_opts);
+            fst::StdToLatticeMapper<BaseFloat> mapper;
+            *lm_fst = new fst::MapFst<fst::StdArc, LatticeArc, fst::StdToLatticeMapper<BaseFloat> >(*std_lm_fst, mapper, mapfst_opts);
+            delete std_lm_fst;
+
+            KALDI_VLOG(2) << "LM loaded: " << path;
+        } else {
+            KALDI_ERR << "LM" << path << "doesn't exist.";
+        }
     }
 
     void Decoder::Reset() {
@@ -213,9 +250,7 @@ namespace alex_asr {
         return ok;
     }
 
-    bool Decoder::GetLattice(fst::VectorFst<fst::LogArc> *fst_out,
-                                     double *tot_lik, bool end_of_utterance) {
-        CompactLattice lat;
+    bool Decoder::GetPrunedLattice(CompactLattice *lat) {
         Lattice raw_lat;
 
         if (decoder_->NumFramesDecoded() == 0)
@@ -224,27 +259,76 @@ namespace alex_asr {
         if (!config_->decoder_opts.determinize_lattice)
             KALDI_ERR << "--determinize-lattice=false option is not supported at the moment";
 
-        bool ok = decoder_->GetRawLattice(&raw_lat, end_of_utterance);
+        bool ok = decoder_->GetRawLattice(&raw_lat);
+
 
         BaseFloat lat_beam = config_->decoder_opts.lattice_beam;
-        DeterminizeLatticePhonePrunedWrapper(
-                *trans_model_, &raw_lat, lat_beam, &lat, config_->decoder_opts.det_opts);
+        if(!config_->rescore) {
+            DeterminizeLatticePhonePrunedWrapper(*trans_model_, &raw_lat, lat_beam, lat, config_->decoder_opts.det_opts);
+        } else {
+            CompactLattice pruned_lat;
+
+            DeterminizeLatticePhonePrunedWrapper(*trans_model_, &raw_lat, lat_beam, &pruned_lat, config_->decoder_opts.det_opts);
+            ok = ok && RescoreLattice(pruned_lat, lat);
+        }
+
+        return ok;
+    }
+
+    bool Decoder::RescoreLattice(CompactLattice lat, CompactLattice *rescored_lattice) {
+        CompactLattice intermidiate_lattice;
+        bool ok = true;
+
+        ok = ok && RescoreLatticeWithLM(lat, -1.0, lm_small_, &intermidiate_lattice);
+        ok = ok && RescoreLatticeWithLM(intermidiate_lattice, 1.0, lm_big_, rescored_lattice);
+
+        return ok;
+    }
+
+    bool Decoder::RescoreLatticeWithLM(
+        CompactLattice lat,
+        float lm_scale,
+        fst::MapFst<fst::StdArc, LatticeArc, fst::StdToLatticeMapper<BaseFloat> > *lm_fst,
+        CompactLattice *rescored_lattice) {
+
+        // Taken from https://github.com/kaldi-asr/kaldi/blob/9b9b561e2b3d2bdf64c84f8626175953f0885264/src/latbin/lattice-lmrescore.cc
+
+        Lattice lattice;
+        ConvertLattice(lat, &lattice);
+
+        fst::ScaleLattice(fst::GraphLatticeScale(1.0 / lm_scale), &lattice);
+        ArcSort(&lattice, fst::OLabelCompare<LatticeArc>());
+
+        Lattice composed_lat;
+        fst::TableComposeOptions compose_opts(fst::TableMatcherOptions(), true, fst::SEQUENCE_FILTER, fst::MATCH_INPUT);
+        fst::TableComposeCache<fst::Fst<LatticeArc> > lm_compose_cache(compose_opts);
+        TableCompose(lattice, *lm_fst, &composed_lat, &lm_compose_cache);
+        Invert(&composed_lat);
+
+        DeterminizeLattice(composed_lat, rescored_lattice);
+        fst::ScaleLattice(fst::GraphLatticeScale(lm_scale), rescored_lattice);
 
+        return rescored_lattice->Start() != fst::kNoStateId;
+    }
+
+    bool Decoder::GetLattice(fst::VectorFst<fst::LogArc> *fst_out,
+                                     double *tot_lik, bool end_of_utterance) {
+        CompactLattice lat;
+        bool ok = true;
+
+        ok = this->GetPrunedLattice(&lat);
         *tot_lik = CompactLatticeToWordsPost(lat, fst_out);
 
         return ok;
     }
 
     bool Decoder::GetTimeAlignment(std::vector<int> *words, std::vector<int> *times, std::vector<int> *lengths) {
-        Lattice lat;
         CompactLattice compact_lat;
         CompactLattice best_path;
         CompactLattice aligned_best_path;
         bool ok = true;
 
-        ok = ok && decoder_->GetRawLattice(&lat);
-        BaseFloat lat_beam = config_->decoder_opts.lattice_beam;
-        DeterminizeLatticePhonePrunedWrapper(*trans_model_, &lat, lat_beam, &compact_lat, config_->decoder_opts.det_opts);
+        ok = this->GetPrunedLattice(&compact_lat);
         CompactLatticeShortestPath(compact_lat, &best_path);
 
         if(config_->word_boundary_rxfilename == "") {
@@ -264,9 +348,7 @@ namespace alex_asr {
         CompactLattice aligned_best_path;
         bool ok = true;
 
-        ok = ok && decoder_->GetRawLattice(&lat);
-        BaseFloat lat_beam = config_->decoder_opts.lattice_beam;
-        DeterminizeLatticePhonePrunedWrapper(*trans_model_, &lat, lat_beam, &compact_lat, config_->decoder_opts.det_opts);
+        ok = this->GetPrunedLattice(&compact_lat);
         CompactLatticeShortestPath(compact_lat, &best_path);
 
         if(config_->word_boundary_rxfilename != "") {

src/decoder.h

@@ -32,7 +32,7 @@ namespace alex_asr {
         void FrameIn(unsigned char *buffer, int32 buffer_length);
         void FrameIn(VectorBase<BaseFloat> *waveform_in);
         bool GetBestPath(std::vector<int> *v_out, BaseFloat *prob);
-        bool GetLattice(fst::VectorFst<fst::LogArc> * out_fst, double *tot_lik, bool end_of_utt=true);
+        bool GetLattice(fst::VectorFst<fst::LogArc> *out_fst, double *tot_lik, bool end_of_utt=true);
         bool GetTimeAlignment(std::vector<int> *words, std::vector<int> *times, std::vector<int> *lengths);
         bool GetTimeAlignmentWithWordConfidence(std::vector<int> *words, std::vector<int> *times, std::vector<int> *lengths, std::vector<float> *confs);
         string GetWord(int word_id);
@@ -51,6 +51,8 @@ namespace alex_asr {
         FeaturePipeline *feature_pipeline_;
 
         fst::StdFst *hclg_;
+        fst::MapFst<fst::StdArc, LatticeArc, fst::StdToLatticeMapper<BaseFloat> > *lm_small_;
+        fst::MapFst<fst::StdArc, LatticeArc, fst::StdToLatticeMapper<BaseFloat> > *lm_big_;
         LatticeFasterOnlineDecoder *decoder_;
         TransitionModel *trans_model_;
         nnet2::AmNnet *am_nnet2_;
@@ -63,8 +65,19 @@ namespace alex_asr {
 
         void InitTransformMatrices();
         void LoadDecoder();
+        void LoadLM(
+            const string path,
+            fst::MapFst<fst::StdArc, LatticeArc, fst::StdToLatticeMapper<BaseFloat> > **lm_fst
+        );
         void ParseConfig();
         void Deallocate();
+        bool GetPrunedLattice(CompactLattice *lat);
+        bool RescoreLattice(CompactLattice lat, CompactLattice *rescored_lattice);
+        bool RescoreLatticeWithLM(
+            CompactLattice lat,
+            float lm_scale,
+            fst::MapFst<fst::StdArc, LatticeArc, fst::StdToLatticeMapper<BaseFloat> > *lm_fst,
+            CompactLattice *rescored_lattice);
         bool FileExists(const std::string& name);
     };
 

src/decoder_config.cc

@@ -11,6 +11,7 @@ namespace alex_asr {
             use_ivectors(false),
             use_cmvn(false),
             use_pitch(false),
+            rescore(false),
             cfg_decoder(""),
             cfg_decodable(""),
             cfg_mfcc(""),
@@ -48,6 +49,9 @@ namespace alex_asr {
         po->Register("use_cmvn", &use_cmvn, "Are we using cmvn transform?");
         po->Register("use_pitch", &use_pitch, "Are we using pitch feature?");
         po->Register("bits_per_sample", &bits_per_sample, "Bits per sample for input.");
+        po->Register("rescore", &rescore, "Rescore lattice with bigger LM?");
+        po->Register("lm_small", &lm_small_rxfilename, "G.fst of the LM used in HCLG.");
+        po->Register("lm_big", &lm_big_rxfilename, "G.fst of the LM to rescore with.");
 
         po->Register("cfg_decoder", &cfg_decoder, "");
         po->Register("cfg_decodable", &cfg_decodable, "");

src/decoder_config.h

@@ -58,6 +58,7 @@ namespace alex_asr {
         bool use_ivectors;
         bool use_cmvn;
         bool use_pitch;
+        bool rescore;
 
         std::string cfg_decoder;
         std::string cfg_decodable;
@@ -72,6 +73,8 @@ namespace alex_asr {
 
         std::string model_rxfilename;
         std::string fst_rxfilename;
+        std::string lm_small_rxfilename;
+        std::string lm_big_rxfilename;
         std::string words_rxfilename;
         std::string word_boundary_rxfilename;
         std::string lda_mat_rspecifier;

test/asr_model_digits/alex_asr.conf

@@ -8,3 +8,4 @@
 --cfg_decoder=decoder.conf
 --cfg_decodable=decodable.conf
 --cfg_endpoint=endpoint.conf
+--rescore=false