Extending 'lattice-compose.cc' to compose with ark of fsts,

- This is a follow-up of kaldi-asr#4571

- Refactoring 'lattice-compose.cc' to support composition with ark
  of fsts, so that it is done as Dan suggested before:

  I am thinking this can be done with a string arg called e.g.
  "--compose-with-fst", defaulting to "auto" which is the old behavior,
  meaning: rspecifier=lats, rxfilename=FST; and true/True or false/False
  is FST or lattice respectively.

- I added there possibility of rho-composition, which is useful for
  biasing lattices with word-sequences. Thanks to rho-composition,
  the biasing graph does not need to contain all words from lexicon.

- Would you be interested in an example how to use this?
  (i.e. create graphs from text file with python script
   using openfst as library, but that would need to change
   build of openfst to enable python extensions)

- Also which 'egs' recipe would be convenient to use it with?

Author: KarelVesely84

src/fstext/fstext-utils.h

@@ -381,7 +381,7 @@ template<class Arc>
 void PropagateFinal(typename Arc::Label phi_label,
                     MutableFst<Arc> *fst);
 
-// PhiCompose is a version of composition where
+// RhoCompose is a version of composition where
 // the right hand FST (fst2) has speciall "rho transitions"
 // which are taken whenever no normal transition matches; these
 // transitions will be rewritten with whatever symbol was on

src/latbin/lattice-compose.cc

@@ -1,6 +1,7 @@
 // latbin/lattice-compose.cc
 
 // Copyright 2009-2011  Microsoft Corporation;  Saarland University
+//           2022  Brno University of Technology
 
 // See ../../COPYING for clarification regarding multiple authors
 //
@@ -17,7 +18,6 @@
 // See the Apache 2 License for the specific language governing permissions and
 // limitations under the License.
 
-
 #include "base/kaldi-common.h"
 #include "util/common-utils.h"
 #include "fstext/fstext-lib.h"
@@ -39,22 +39,34 @@ int main(int argc, char *argv[]) {
         "or lattices with FSTs (rspecifiers are assumed to be lattices, and\n"
         "rxfilenames are assumed to be FSTs, which have their weights interpreted\n"
         "as \"graph weights\" when converted into the Lattice format.\n"
+        "Or, rspecifier can be ark of biasing FSTs, see --compose-with-fst=true.\n"
         "\n"
         "Usage: lattice-compose [options] lattice-rspecifier1 "
         "(lattice-rspecifier2|fst-rxfilename2) lattice-wspecifier\n"
         " e.g.: lattice-compose ark:1.lats ark:2.lats ark:composed.lats\n"
-        " or: lattice-compose ark:1.lats G.fst ark:composed.lats\n";
+        " or: lattice-compose ark:1.lats G.fst ark:composed.lats\n"
+        " or: lattice-compose --compose-with-fst=true ark:1.lats\n"
+        "       ark:biasing.fsts ark:composed.lats\n";
 
     ParseOptions po(usage);
 
     bool write_compact = true;
     int32 num_states_cache = 50000;
     int32 phi_label = fst::kNoLabel; // == -1
+    int32 rho_label = fst::kNoLabel; // == -1
+    std::string compose_with_fst = "auto";
+
     po.Register("write-compact", &write_compact, "If true, write in normal (compact) form.");
     po.Register("phi-label", &phi_label, "If >0, the label on backoff arcs of the LM");
+    po.Register("rho-label", &rho_label,
+                "If >0, the label to forward fst1 paths not present biasing graph fst2. "
+                "(rho is input and output symbol on special arc in biasing graph)");
     po.Register("num-states-cache", &num_states_cache,
                 "Number of states we cache when mapping LM FST to lattice type. "
                 "More -> more memory but faster.");
+    po.Register("compose-with-fst", &compose_with_fst,
+                "(true|false|auto) For auto arg2 is: rspecifier=lats, rxfilename=fst "
+                "(old behavior), for true/false rspecifier is fst/lattice.");
     po.Read(argc, argv);
 
     if (po.NumArgs() != 3) {
@@ -63,14 +75,30 @@ int main(int argc, char *argv[]) {
     }
 
     KALDI_ASSERT(phi_label > 0 || phi_label == fst::kNoLabel); // e.g. 0 not allowed.
+    KALDI_ASSERT(rho_label > 0 || rho_label == fst::kNoLabel); // e.g. 0 not allowed.
+    if (phi_label > 0 && rho_label > 0) {
+      KALDI_ERR << "You cannot set both 'phi_label' and 'rho_label' at the same time.";
+    }
+
+    { // convert 'compose_with_fst' to lowercase to support: true, True, TRUE
+      std::string tmp_lc(compose_with_fst);
+      std::transform(compose_with_fst.begin(), compose_with_fst.end(),
+                     tmp_lc.begin(), ::tolower);  // lc
+      compose_with_fst.swap(tmp_lc);
+    }
+    if (compose_with_fst != "auto" && compose_with_fst != "true" &&
+                                                 compose_with_fst != "false") {
+      KALDI_ERR << "Unkown 'compose_with_fst' value : " << compose_with_fst
+                << " , values are (auto|true|false)";
+    }
 
     std::string lats_rspecifier1 = po.GetArg(1),
         arg2 = po.GetArg(2),
         lats_wspecifier = po.GetArg(3);
     int32 n_done = 0, n_fail = 0;
 
     SequentialLatticeReader lattice_reader1(lats_rspecifier1);
-    
+
     CompactLatticeWriter compact_lattice_writer;
     LatticeWriter lattice_writer;
 
@@ -79,33 +107,48 @@ int main(int argc, char *argv[]) {
     else
       lattice_writer.Open(lats_wspecifier);
 
-    if (ClassifyRspecifier(arg2, NULL, NULL) == kNoRspecifier) {
+    bool arg2_is_rxfilename = (ClassifyRspecifier(arg2, NULL, NULL) == kNoRspecifier);
+
+    /**
+     * arg2 is rxfilename that contains a single fst
+     * - compose arg1 lattices with single fst in arg2
+     */
+    if (arg2_is_rxfilename && (compose_with_fst == "auto" || compose_with_fst == "true")) {
       std::string fst_rxfilename = arg2;
-      VectorFst<StdArc> *fst2 = fst::ReadFstKaldi(fst_rxfilename);
-      // mapped_fst2 is fst2 interpreted using the LatticeWeight semiring,
-      // with all the cost on the first member of the pair (since we're
-      // assuming it's a graph weight).
+      VectorFst<StdArc>* fst2 = fst::ReadFstKaldi(fst_rxfilename);
+
+      // Make sure fst2 is sorted on ilabel
       if (fst2->Properties(fst::kILabelSorted, true) == 0) {
-        // Make sure fst2 is sorted on ilabel.
         fst::ILabelCompare<StdArc> ilabel_comp;
         ArcSort(fst2, ilabel_comp);
       }
+
       if (phi_label > 0)
         PropagateFinal(phi_label, fst2);
 
+      // mapped_fst2 is fst2 interpreted using the LatticeWeight semiring,
+      // with all the cost on the first member of the pair (since we're
+      // assuming it's a graph weight).
       fst::CacheOptions cache_opts(true, num_states_cache);
       fst::MapFstOptions mapfst_opts(cache_opts);
       fst::StdToLatticeMapper<BaseFloat> mapper;
       fst::MapFst<StdArc, LatticeArc, fst::StdToLatticeMapper<BaseFloat> >
           mapped_fst2(*fst2, mapper, mapfst_opts);
+
       for (; !lattice_reader1.Done(); lattice_reader1.Next()) {
         std::string key = lattice_reader1.Key();
         KALDI_VLOG(1) << "Processing lattice for key " << key;
         Lattice lat1 = lattice_reader1.Value();
         ArcSort(&lat1, fst::OLabelCompare<LatticeArc>());
+
         Lattice composed_lat;
-        if (phi_label > 0) PhiCompose(lat1, mapped_fst2, phi_label, &composed_lat);
-        else Compose(lat1, mapped_fst2, &composed_lat);
+        if (phi_label > 0) {
+          PhiCompose(lat1, mapped_fst2, phi_label, &composed_lat);
+        } else if (rho_label > 0) {
+          RhoCompose(lat1, mapped_fst2, rho_label, &composed_lat);
+        } else {
+          Compose(lat1, mapped_fst2, &composed_lat);
+        }
         if (composed_lat.Start() == fst::kNoStateId) {
           KALDI_WARN << "Empty lattice for utterance " << key << " (incompatible LM?)";
           n_fail++;
@@ -121,7 +164,23 @@ int main(int argc, char *argv[]) {
         }
       }
       delete fst2;
-    } else {
+    }
+
+    /**
+     * arg2 is rxfilename that contains a single lattice
+     */
+    else if (arg2_is_rxfilename && compose_with_fst == "false") {
+      // Would it make sense to do this? Not implementing...
+      KALDI_ERR << "Unimplemented...";
+    }
+
+    /**
+     * arg2 is rspecifier that contains a table of lattices
+     * - composing arg1 lattices with arg2 lattices
+     */
+    else if (not arg2_is_rxfilename &&
+                  (compose_with_fst == "auto" || compose_with_fst == "false")) {
+      //
       std::string lats_rspecifier2 = arg2;
       // This is the case similar to lattice-interp.cc, where we
       // read in another set of lattices and compose them.  But in this
@@ -139,6 +198,7 @@ int main(int argc, char *argv[]) {
           n_fail++;
           continue;
         }
+
         Lattice lat2 = lattice_reader2.Value(key);
         // Make sure that either lat2 is ilabel sorted
         // or lat1 is olabel sorted, to ensure that
@@ -150,29 +210,108 @@ int main(int argc, char *argv[]) {
           fst::ArcSort(&lat2, ilabel_comp);
         }
 
-        Lattice lat_out;
+        Lattice composed_lat;
+        // Btw, can the lat2 lattice contin phi/rho symbols ?
         if (phi_label > 0) {
           PropagateFinal(phi_label, &lat2);
-          PhiCompose(lat1, lat2, phi_label, &lat_out);
+          PhiCompose(lat1, lat2, phi_label, &composed_lat);
+        } else if (rho_label > 0) {
+          RhoCompose(lat1, lat2, rho_label, &composed_lat);
+        } else {
+          Compose(lat1, lat2, &composed_lat);
+        }
+        if (composed_lat.Start() == fst::kNoStateId) {
+          KALDI_WARN << "Empty lattice for utterance " << key << " (incompatible LM?)";
+          n_fail++;
+        } else {
+          if (write_compact) {
+            CompactLattice clat;
+            ConvertLattice(composed_lat, &clat);
+            compact_lattice_writer.Write(key, clat);
+          } else {
+            lattice_writer.Write(key, composed_lat);
+          }
+          n_done++;
+        }
+      }
+    }
+
+    /**
+     * arg2 is rspecifier that contains a table of fsts
+     * - composing arg1 lattices with arg2 fsts
+     */
+    else if (not arg2_is_rxfilename && compose_with_fst == "true") {
+      std::string fst_rspecifier2 = arg2;
+      RandomAccessTableReader<fst::VectorFstHolder> fst_reader2(fst_rspecifier2);
+
+      for (; !lattice_reader1.Done(); lattice_reader1.Next()) {
+        std::string key = lattice_reader1.Key();
+        KALDI_VLOG(1) << "Processing lattice for key " << key;
+        Lattice lat1 = lattice_reader1.Value();
+        lattice_reader1.FreeCurrent();
+
+        if (!fst_reader2.HasKey(key)) {
+          KALDI_WARN << "Not producing output for utterance " << key
+                     << " because not present in second table.";
+          n_fail++;
+          continue;
+        }
+
+        VectorFst<StdArc> fst2 = fst_reader2.Value(key);
+        // Make sure fst2 is sorted on ilabel
+        if (fst2.Properties(fst::kILabelSorted, true) == 0) {
+          fst::ILabelCompare<StdArc> ilabel_comp;
+          fst::ArcSort(&fst2, ilabel_comp);
+        }
+
+        // for composing with LM-fsts, it makes all fst2 states final
+        if (phi_label > 0)
+          PropagateFinal(phi_label, &fst2);
+
+        // mapped_fst2 is fst2 interpreted using the LatticeWeight semiring,
+        // with all the cost on the first member of the pair (since we're
+        // assuming it's a graph weight).
+        fst::CacheOptions cache_opts(true, num_states_cache);
+        fst::MapFstOptions mapfst_opts(cache_opts);
+        fst::StdToLatticeMapper<BaseFloat> mapper;
+        fst::MapFst<StdArc, LatticeArc, fst::StdToLatticeMapper<BaseFloat> >
+            mapped_fst2(fst2, mapper, mapfst_opts);
+
+        // sort lat1 on olabel.
+        ArcSort(&lat1, fst::OLabelCompare<LatticeArc>());
+
+        Lattice composed_lat;
+        if (phi_label > 0) {
+          PhiCompose(lat1, mapped_fst2, phi_label, &composed_lat);
+        } else if (rho_label > 0) {
+          RhoCompose(lat1, mapped_fst2, rho_label, &composed_lat);
         } else {
-          Compose(lat1, lat2, &lat_out);
+          Compose(lat1, mapped_fst2, &composed_lat);
         }
-        if (lat_out.Start() == fst::kNoStateId) {
+
+        if (composed_lat.Start() == fst::kNoStateId) {
           KALDI_WARN << "Empty lattice for utterance " << key << " (incompatible LM?)";
           n_fail++;
         } else {
           if (write_compact) {
-            CompactLattice clat_out;
-            ConvertLattice(lat_out, &clat_out);
-            compact_lattice_writer.Write(key, clat_out);
+            CompactLattice clat;
+            ConvertLattice(composed_lat, &clat);
+            compact_lattice_writer.Write(key, clat);
           } else {
-            lattice_writer.Write(key, lat_out);
+            lattice_writer.Write(key, composed_lat);
           }
           n_done++;
         }
       }
     }
 
+    /**
+     * none of the 'if-else-if' applied...
+     */
+    else {
+      KALDI_ERR << "You should never reach here...";
+    }
+
     KALDI_LOG << "Done " << n_done << " lattices; failed for "
               << n_fail;
     return (n_done != 0 ? 0 : 1);