Merge branch 'master'

qindazhu · qindazhu · commit 37aafd844d80 · 2020-06-05T16:04:42.000+08:00
diff --git a/k2/csrc/aux_labels.cc b/k2/csrc/aux_labels.cc
@@ -6,14 +6,94 @@
 
 #include "k2/csrc/aux_labels.h"
 
+#include <algorithm>
 #include <numeric>
+#include <utility>
 #include <vector>
 
 #include "glog/logging.h"
 #include "k2/csrc/fsa.h"
+#include "k2/csrc/fsa_util.h"
+#include "k2/csrc/properties.h"
+
+namespace {
+
+/*
+  This function counts how many extra states we need to create for each state in
+  the input FSA when we invert an FSA. Generally, if an entering arc of state
+  `i` in the input FSA has n olabels, then we need to create n-1 extra states
+  for state `i`.
+
+  @param [in] fsa_in      Input FSA
+  @param [in] labels_in   Aux-label sequences for the input FSA
+  @param [out] num_extra_states For state `i` in `fsa_in`, we need to create
+                                extra `num_extra_states[i]` states in the output
+                                inverted FSA.
+*/
+static void CountExtraStates(const k2::Fsa &fsa_in,
+                             const k2::AuxLabels &labels_in,
+                             std::vector<int32_t> *num_extra_states) {
+  CHECK_EQ(num_extra_states->size(), fsa_in.NumStates());
+  auto &states = *num_extra_states;
+  for (int32_t i = 0; i != fsa_in.arcs.size(); ++i) {
+    const auto &arc = fsa_in.arcs[i];
+    int32_t pos_start = labels_in.start_pos[i];
+    int32_t pos_end = labels_in.start_pos[i + 1];
+    states[arc.dest_state] += std::max(0, pos_end - pos_start - 1);
+  }
+}
+
+/*
+  Map the state in the input FSA to state in the output inverted FSA.
+
+  @param [in] num_extra_states Output of function `CountExtraStates`
+                               which gives how many extra states we need
+                               to create for each state in the input FSA.
+  @param [out] state_map       Map state `i` in the input FSA to state
+                               `state_map[i]` in the output FSA.
+                               At exit, it will be
+                               state_map[0] = 0,
+                               state_map[i] = state_map[i-1]
+                                            + num_extra_states[i]
+                                            + 1, for any i >=1
+  @param [out] state_ids       At exit, it will be
+                               state_ids[0] = 0,
+                               state_ids[i] = state_map[i-1], for any i >= 1.
+*/
+static void MapStates(const std::vector<int32_t> &num_extra_states,
+                      std::vector<int32_t> *state_map,
+                      std::vector<int32_t> *state_ids) {
+  CHECK_EQ(state_map->size(), num_extra_states.size());
+  CHECK_EQ(state_ids->size(), num_extra_states.size());
+  auto &s_map = *state_map;
+  auto &s_ids = *state_ids;
+  // we suppose there's no arcs entering the start state (i.e. state id of the
+  // start state in output FSA will be 0), otherwise we may need to create a new
+  // state as the real start state.
+  CHECK_EQ(num_extra_states[0], 0);
+  auto num_states_in = num_extra_states.size();
+  // process from the second state
+  s_map[0] = 0;
+  s_ids[0] = 0;
+  int32_t num_states_out = 0;
+  for (auto i = 1; i != num_states_in; ++i) {
+    s_ids[i] = num_states_out;
+    // `+1` as we did not count state `i` itself in `num_extra_states`
+    num_states_out += num_extra_states[i] + 1;
+    s_map[i] = num_states_out;
+  }
+}
+}  // namespace
 
 namespace k2 {
 
+void Swap(AuxLabels *labels1, AuxLabels *labels2) {
+  CHECK_NOTNULL(labels1);
+  CHECK_NOTNULL(labels2);
+  std::swap(labels1->start_pos, labels2->start_pos);
+  std::swap(labels1->labels, labels2->labels);
+}
+
 void MapAuxLabels1(const AuxLabels &labels_in,
                    const std::vector<int32_t> &arc_map, AuxLabels *labels_out) {
   CHECK_NOTNULL(labels_out);
@@ -61,4 +141,93 @@ void MapAuxLabels2(const AuxLabels &labels_in,
   start_pos.push_back(num_labels);
 }
 
+void InvertFst(const Fsa &fsa_in, const AuxLabels &labels_in, Fsa *fsa_out,
+               AuxLabels *aux_labels_out) {
+  CHECK_NOTNULL(fsa_out);
+  CHECK_NOTNULL(aux_labels_out);
+  fsa_out->arc_indexes.clear();
+  fsa_out->arcs.clear();
+  aux_labels_out->start_pos.clear();
+  aux_labels_out->labels.clear();
+
+  if (IsEmpty(fsa_in)) {
+    aux_labels_out->start_pos.push_back(0);
+    return;
+  }
+
+  auto num_states_in = fsa_in.NumStates();
+  // get the number of extra states we need to create for each state
+  // in fsa_in when inverting
+  std::vector<int32_t> num_extra_states(num_states_in, 0);
+  CountExtraStates(fsa_in, labels_in, &num_extra_states);
+
+  // map state in fsa_in to state in fsa_out
+  std::vector<int32_t> state_map(num_states_in, 0);
+  std::vector<int32_t> state_ids(num_states_in, 0);
+  MapStates(num_extra_states, &state_map, &state_ids);
+
+  // a maximal approximation
+  int32_t num_arcs_out = labels_in.labels.size() + fsa_in.arcs.size();
+  std::vector<Arc> arcs;
+  arcs.reserve(num_arcs_out);
+  // `+1` for the end position of the last arc's olabel sequence
+  std::vector<int32_t> start_pos;
+  start_pos.reserve(num_arcs_out + 1);
+  std::vector<int32_t> labels;
+  labels.reserve(fsa_in.arcs.size());
+  int32_t final_state_in = fsa_in.FinalState();
+
+  int32_t num_non_eps_ilabel_processed = 0;
+  start_pos.push_back(0);
+  for (auto i = 0; i != fsa_in.arcs.size(); ++i) {
+    const auto &arc = fsa_in.arcs[i];
+    int32_t pos_start = labels_in.start_pos[i];
+    int32_t pos_end = labels_in.start_pos[i + 1];
+    int32_t src_state = arc.src_state;
+    int32_t dest_state = arc.dest_state;
+    if (dest_state == final_state_in) {
+      // every arc entering the final state must have exactly
+      // one olabel == kFinalSymbol
+      CHECK_EQ(pos_start + 1, pos_end);
+      CHECK_EQ(labels_in.labels[pos_start], kFinalSymbol);
+    }
+    if (pos_end - pos_start <= 1) {
+      int32_t curr_label =
+          (pos_end - pos_start == 0) ? kEpsilon : labels_in.labels[pos_start];
+      arcs.emplace_back(state_map[src_state], state_map[dest_state],
+                        curr_label);
+    } else {
+      // expand arcs with olabels
+      arcs.emplace_back(state_map[src_state], state_ids[dest_state] + 1,
+                        labels_in.labels[pos_start]);
+      start_pos.push_back(num_non_eps_ilabel_processed);
+      for (int32_t pos = pos_start + 1; pos < pos_end - 1; ++pos) {
+        ++state_ids[dest_state];
+        arcs.emplace_back(state_ids[dest_state], state_ids[dest_state] + 1,
+                          labels_in.labels[pos]);
+        start_pos.push_back(num_non_eps_ilabel_processed);
+      }
+      ++state_ids[dest_state];
+      arcs.emplace_back(state_ids[dest_state], state_map[arc.dest_state],
+                        labels_in.labels[pos_end - 1]);
+    }
+    // push non-epsilon ilabel in fsa_in as olabel of fsa_out
+    if (arc.label != kEpsilon) {
+      labels.push_back(arc.label);
+      ++num_non_eps_ilabel_processed;
+    }
+    start_pos.push_back(num_non_eps_ilabel_processed);
+  }
+
+  labels.resize(labels.size());
+  arcs.resize(arcs.size());
+  start_pos.resize(start_pos.size());
+
+  std::vector<int32_t> arc_map;
+  ReorderArcs(arcs, fsa_out, &arc_map);
+  AuxLabels labels_tmp;
+  labels_tmp.start_pos = std::move(start_pos);
+  labels_tmp.labels = std::move(labels);
+  MapAuxLabels1(labels_tmp, arc_map, aux_labels_out);
+}
 }  // namespace k2
diff --git a/k2/csrc/aux_labels.h b/k2/csrc/aux_labels.h
@@ -46,6 +46,9 @@ struct AuxLabels {
   std::vector<int32_t> labels;
 };
 
+// Swap AuxLabels; it's cheap to to this as we are actually doing shallow swap.
+void Swap(AuxLabels *labels1, AuxLabels *labels2);
+
 /*
   Maps auxiliary labels after an FSA operation where each arc in the output
   FSA corresponds to exactly one arc in the input FSA.
diff --git a/k2/csrc/aux_labels_test.cc b/k2/csrc/aux_labels_test.cc
@@ -12,6 +12,7 @@
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include "k2/csrc/fsa.h"
+#include "k2/csrc/properties.h"
 
 namespace k2 {
 
@@ -97,4 +98,106 @@ TEST_F(AuxLablesTest, MapAuxLabels2) {
   }
 }
 
+TEST(AuxLabels, InvertFst) {
+  {
+    // empty input FSA
+    Fsa fsa_in;
+    AuxLabels labels_in;
+    std::vector<int32_t> start_pos = {0, 1, 3, 6, 7};
+    std::vector<int32_t> labels = {1, 2, 3, 4, 5, 6, 7};
+    labels_in.start_pos = std::move(start_pos);
+    labels_in.labels = std::move(labels);
+
+    std::vector<Arc> arcs = {{0, 1, 1}, {1, 2, -1}};
+    Fsa fsa_out(std::move(arcs), 2);
+    AuxLabels labels_out;
+    // some dirty data
+    labels_out.start_pos = {1, 2, 3};
+    labels_out.labels = {4, 5};
+    InvertFst(fsa_in, labels_in, &fsa_out, &labels_out);
+
+    EXPECT_TRUE(IsEmpty(fsa_out));
+    EXPECT_TRUE(labels_out.labels.empty());
+    ASSERT_EQ(labels_out.start_pos.size(), 1);
+    EXPECT_EQ(labels_out.start_pos[0], 0);
+  }
+
+  {
+    // top-sorted input FSA
+    std::vector<Arc> arcs = {{0, 1, 1}, {0, 1, 0},  {0, 3, 2},
+                             {1, 2, 3}, {1, 3, 4},  {1, 5, -1},
+                             {2, 3, 0}, {2, 5, -1}, {4, 5, -1}};
+    Fsa fsa_in(std::move(arcs), 5);
+    EXPECT_TRUE(IsTopSorted(fsa_in));
+    AuxLabels labels_in;
+    std::vector<int32_t> start_pos = {0, 2, 3, 3, 6, 6, 7, 7, 8, 9};
+    EXPECT_EQ(start_pos.size(), fsa_in.arcs.size() + 1);
+    std::vector<int32_t> labels = {1, 2, 3, 5, 6, 7, -1, -1, -1};
+    labels_in.start_pos = std::move(start_pos);
+    labels_in.labels = std::move(labels);
+
+    Fsa fsa_out;
+    AuxLabels labels_out;
+    InvertFst(fsa_in, labels_in, &fsa_out, &labels_out);
+
+    EXPECT_TRUE(IsTopSorted(fsa_out));
+    std::vector<Arc> arcs_out = {
+        {0, 1, 1},  {0, 2, 3}, {0, 6, 0}, {1, 2, 2}, {2, 3, 5},  {2, 6, 0},
+        {2, 8, -1}, {3, 4, 6}, {4, 5, 7}, {5, 6, 0}, {5, 8, -1}, {7, 8, -1},
+    };
+    ASSERT_EQ(fsa_out.arcs.size(), arcs_out.size());
+    for (auto i = 0; i != arcs_out.size(); ++i) {
+      EXPECT_EQ(fsa_out.arcs[i], arcs_out[i]);
+    }
+    ASSERT_EQ(fsa_out.arc_indexes.size(), 10);
+    EXPECT_THAT(fsa_out.arc_indexes,
+                ::testing::ElementsAre(0, 3, 4, 7, 8, 9, 11, 11, 12, 12));
+    ASSERT_EQ(labels_out.labels.size(), 7);
+    EXPECT_THAT(labels_out.labels,
+                ::testing::ElementsAre(2, 1, 4, -1, 3, -1, -1));
+    ASSERT_EQ(labels_out.start_pos.size(), 13);
+    EXPECT_THAT(labels_out.start_pos,
+                ::testing::ElementsAre(0, 0, 0, 1, 2, 2, 3, 4, 4, 5, 5, 6, 7));
+  }
+
+  {
+    // non-top-sorted input FSA
+    std::vector<Arc> arcs = {{0, 1, 1},  {0, 1, 0}, {0, 3, 2},
+                             {1, 2, 3},  {1, 3, 4}, {2, 1, 5},
+                             {2, 5, -1}, {3, 1, 6}, {4, 5, -1}};
+    Fsa fsa_in(std::move(arcs), 5);
+    EXPECT_FALSE(IsTopSorted(fsa_in));
+    AuxLabels labels_in;
+    std::vector<int32_t> start_pos = {0, 2, 3, 3, 6, 6, 7, 8, 10, 11};
+    EXPECT_EQ(start_pos.size(), fsa_in.arcs.size() + 1);
+    std::vector<int32_t> labels = {1, 2, 3, 5, 6, 7, 8, -1, 9, 10, -1};
+    labels_in.start_pos = std::move(start_pos);
+    labels_in.labels = std::move(labels);
+
+    Fsa fsa_out;
+    AuxLabels labels_out;
+    InvertFst(fsa_in, labels_in, &fsa_out, &labels_out);
+
+    EXPECT_FALSE(IsTopSorted(fsa_out));
+    std::vector<Arc> arcs_out = {{0, 1, 1},  {0, 3, 3}, {0, 7, 0},  {1, 3, 2},
+                                 {2, 3, 10}, {3, 4, 5}, {3, 7, 0},  {4, 5, 6},
+                                 {5, 6, 7},  {6, 3, 8}, {6, 9, -1}, {7, 2, 9},
+                                 {8, 9, -1}};
+    ASSERT_EQ(fsa_out.arcs.size(), arcs_out.size());
+    for (auto i = 0; i != arcs_out.size(); ++i) {
+      EXPECT_EQ(fsa_out.arcs[i], arcs_out[i]);
+    }
+    ASSERT_EQ(fsa_out.arc_indexes.size(), 11);
+    EXPECT_THAT(fsa_out.arc_indexes,
+                ::testing::ElementsAre(0, 3, 4, 5, 7, 8, 9, 11, 12, 13, 13));
+    ASSERT_EQ(labels_out.labels.size(), 8);
+    EXPECT_THAT(labels_out.labels,
+                ::testing::ElementsAre(2, 1, 6, 4, 3, 5, -1, -1));
+    ASSERT_EQ(labels_out.start_pos.size(), 14);
+    EXPECT_THAT(
+        labels_out.start_pos,
+        ::testing::ElementsAre(0, 0, 0, 1, 2, 3, 3, 4, 4, 5, 6, 7, 7, 8));
+  }
+}
+
 }  // namespace k2
diff --git a/k2/csrc/fsa_util.cc b/k2/csrc/fsa_util.cc
@@ -177,6 +177,43 @@ void GetArcWeights(const float *arc_weights_in,
   }
 }
 
+void ReorderArcs(const std::vector<Arc> &arcs, Fsa *fsa,
+                 std::vector<int32_t> *arc_map /*= nullptr*/) {
+  CHECK_NOTNULL(fsa);
+  fsa->arc_indexes.clear();
+  fsa->arcs.clear();
+  if (arc_map != nullptr) arc_map->clear();
+
+  if (arcs.empty()) return;
+
+  using ArcWithIndex = std::pair<Arc, int32_t>;
+  int arc_id = 0;
+  std::vector<std::vector<ArcWithIndex>> vec;
+  for (const auto &arc : arcs) {
+    auto src_state = arc.src_state;
+    auto dest_state = arc.dest_state;
+    auto new_size = std::max(src_state, dest_state);
+    if (new_size >= vec.size()) vec.resize(new_size + 1);
+    vec[src_state].push_back({arc, arc_id++});
+  }
+
+  std::size_t num_states = vec.size();
+  fsa->arc_indexes.resize(num_states + 1);
+  fsa->arcs.reserve(arcs.size());
+  std::vector<int32_t> arc_map_out;
+  arc_map_out.reserve(arcs.size());
+
+  for (auto i = 0; i != num_states; ++i) {
+    fsa->arc_indexes[i] = static_cast<int32_t>(fsa->arcs.size());
+    for (auto arc_with_index : vec[i]) {
+      fsa->arcs.emplace_back(arc_with_index.first);
+      arc_map_out.push_back(arc_with_index.second);
+    }
+  }
+  fsa->arc_indexes.back() = static_cast<int32_t>(fsa->arcs.size());
+  if (arc_map != nullptr) arc_map->swap(arc_map_out);
+}
+
 void Swap(Fsa *a, Fsa *b) {
   CHECK_NOTNULL(a);
   CHECK_NOTNULL(b);
diff --git a/k2/csrc/fsa_util.h b/k2/csrc/fsa_util.h
@@ -59,10 +59,27 @@ void GetArcWeights(const float *arc_weights_in,
                    const std::vector<std::vector<int32_t>> &arc_map,
                    float *arc_weights_out);
 
-// Version of GetArcWeights where arc_map maps each arc in the output FSA to one
-// arc (instead of a sequence of arcs) in the input FSA; see its documentation.
+// Version of GetArcWeights where arc_map maps each arc in the output FSA to
+// one arc (instead of a sequence of arcs) in the input FSA; see its
+// documentation.
 void GetArcWeights(const float *arc_weights_in,
                    const std::vector<int32_t> &arc_map, float *arc_weights_out);
+
+/* Reorder a list of arcs to get a valid FSA. This function will be used in a
+   situation that the input list of arcs is not sorted by src_state, we'll
+   reorder the arcs and generate the corresponding valid FSA. Note that we don't
+   remap any state index here, it is supposed that the start state is 0 and the
+   final state is the largest state number in the input arcs.
+
+   @param [in] arcs  A list of arcs.
+   @param [out] fsa  Output fsa.
+   @param [out] arc_map   If non-NULL, this function will
+                            output a map from the arc-index in `fsa` to
+                            the corresponding arc-index in input `arcs`.
+*/
+void ReorderArcs(const std::vector<Arc> &arcs, Fsa *fsa,
+                 std::vector<int32_t> *arc_map = nullptr);
+
 /*
   Convert indexes (typically arc-mapping indexes, e.g. as output by Compose())
   from int32 to int64; this will be needed for conversion to LongTensor.
diff --git a/k2/csrc/fsa_util_test.cc b/k2/csrc/fsa_util_test.cc
