Modify AdvancedTreeSearch for CtcTreeBuilder support

src/Search/AdvancedTreeSearch/SearchSpace.cc

@@ -461,7 +461,7 @@ void StaticSearchAutomaton::buildDepths(bool onlyFromRoot) {
     for (u32 a = 1; a < stateDepths.size(); ++a) {
         if (stateDepths[a] != Core::Type<int>::min && stateDepths[a] != Core::Type<int>::max) {
             for (HMMStateNetwork::SuccessorIterator it = network.structure.successors(a); it; ++it) {
-                if (!it.isLabel()) {
+                if (!it.isLabel() and *it != a) {
                     verify(stateDepths[*it] > stateDepths[a]);
                 }
             }
@@ -498,7 +498,7 @@ int StaticSearchAutomaton::fillStateDepths(StateId state, int depth) {
     localDepth = 0;
 
     for (HMMStateNetwork::SuccessorIterator it = network.structure.successors(state); it; ++it) {
-        if (not it.isLabel()) {
+        if (not it.isLabel() and *it != state) {
             int d = fillStateDepths(*it, depth + 1);
 
             if (d > localDepth) {
@@ -723,7 +723,7 @@ void StaticSearchAutomaton::buildBatches() {
         network.dumpDotGraph(paramDumpDotGraph(config), stateDepths);
 
     // Print some useful statistics about pushed and unpushed labels
-    verify(!network.unpushedCoarticulatedRootStates.empty());
+    verify(network.coarticulatedRootStates.empty() || !network.unpushedCoarticulatedRootStates.empty());
 
     u32 unpushedLabels = 0;
     u32 pushedLabels   = 0;
@@ -734,8 +734,10 @@ void StaticSearchAutomaton::buildBatches() {
                 bool    isUnpushed = network.unpushedCoarticulatedRootStates.count(transit) || transit == network.ciRootState || transit == network.rootState;
                 if (isUnpushed) {
                     ++unpushedLabels;
-                    std::map<StateId, std::pair<Bliss::Phoneme::Id, Bliss::Phoneme::Id>>::iterator it = network.rootTransitDescriptions.find(transit);
-                    verify(it != network.rootTransitDescriptions.end());
+                    if (!network.rootTransitDescriptions.empty()) {
+                        std::map<StateId, std::pair<Bliss::Phoneme::Id, Bliss::Phoneme::Id>>::iterator it = network.rootTransitDescriptions.find(transit);
+                        verify(it != network.rootTransitDescriptions.end());
+                    }
                 }
                 else {
                     ++pushedLabels;

src/Search/TreeBuilder.cc

@@ -1210,9 +1210,21 @@ const Core::ParameterBool CtcTreeBuilder::paramLabelLoop(
         "allow label loops in the search tree",
         true);
 
+const Core::ParameterBool CtcTreeBuilder::paramBlankLoop(
+        "allow-blank-loop",
+        "allow loops on the blank nodes in the search tree",
+        true);
+
+const Core::ParameterBool CtcTreeBuilder::paramForceBlank(
+        "force-blank-between-repeated-labels",
+        "require a blank label between two identical labels (only works if label-loops are disabled)",
+        true);
+
 CtcTreeBuilder::CtcTreeBuilder(Core::Configuration config, const Bliss::Lexicon& lexicon, const Am::AcousticModel& acousticModel, Search::PersistentStateTree& network, bool initialize)
         : AbstractTreeBuilder(config, lexicon, acousticModel, network),
-          labelLoop_(paramLabelLoop(config)) {
+          labelLoop_(paramLabelLoop(config)),
+          blankLoop_(paramBlankLoop(config)),
+          forceBlank_(paramForceBlank(config)) {
     auto iters = lexicon.phonemeInventory()->phonemes();
     for (auto it = iters.first; it != iters.second; ++it) {
         require(not(*it)->isContextDependent());  // Context dependent labels are not supported
@@ -1309,7 +1321,8 @@ StateId CtcTreeBuilder::extendState(StateId predecessor, StateTree::StateDesc de
 }
 
 void CtcTreeBuilder::addTransition(StateId predecessor, StateId successor) {
-    auto const& successorStateDesc = network_.structure.state(successor).stateDesc;
+    auto const& predecessorStateDesc = network_.structure.state(predecessor).stateDesc;
+    auto const& successorStateDesc   = network_.structure.state(successor).stateDesc;
 
     for (HMMStateNetwork::SuccessorIterator target = network_.structure.successors(predecessor); target; ++target) {
         if (!target.isLabel() && network_.structure.state(*target).stateDesc == successorStateDesc) {
@@ -1361,8 +1374,10 @@ StateId CtcTreeBuilder::extendPronunciation(StateId startState, Bliss::Pronuncia
                     addTransition(currentState, currentState);
                 }
 
-                // Add transition from previous non-blank state to this state, allowing to skip the blank state in-between these two
-                if (prevNonBlankState != invalidTreeNodeIndex) {
+                bool label_repetition = prevNonBlankState != currentState and prevNonBlankState != invalidTreeNodeIndex and network_.structure.state(prevNonBlankState).stateDesc == network_.structure.state(currentState).stateDesc;
+                if (prevNonBlankState != invalidTreeNodeIndex and not(label_repetition and forceBlank_)) {
+                    // Add transition from previous non-blank state to this state, allowing to skip the blank state in-between these two
+                    // If we want to enforce blank between repeated labels, don't add a transition between two distinct states of equal description
                     addTransition(prevNonBlankState, currentState);
                 }
                 prevNonBlankState = currentState;
@@ -1371,8 +1386,11 @@ StateId CtcTreeBuilder::extendPronunciation(StateId startState, Bliss::Pronuncia
                 if (not allophoneIsBlank and not isLastStateInLemma) {
                     // Add blank state after the newly created state
                     currentState = extendState(currentState, blankDesc_);
-                    // Add loop for this blank state
-                    addTransition(currentState, currentState);
+
+                    if (blankLoop_) {
+                        // Add loop for this blank state
+                        addTransition(currentState, currentState);
+                    }
                 }
             }
         }
@@ -1411,8 +1429,11 @@ void CtcTreeBuilder::addWordBoundaryStates() {
     for (StateId wbs : wordBoundaryLemmaStartStates) {
         network_.structure.addTargetToNode(blankBefore, wbs);
     }
-    // Add loop for this blank state
-    addTransition(blankBefore, blankBefore);
+
+    if (blankLoop_) {
+        // Add loop for this blank state
+        addTransition(blankBefore, blankBefore);
+    }
 }
 
 // -------------------- RnaTreeBuilder --------------------
@@ -1422,7 +1443,13 @@ const Core::ParameterBool RnaTreeBuilder::paramLabelLoop(
         "allow label loops in the search tree",
         false);
 
+const Core::ParameterBool RnaTreeBuilder::paramForceBlank(
+        "force-blank-between-repeated-labels",
+        "require a blank label between two identical labels (only works if label-loops are disabled)",
+        false);
+
 RnaTreeBuilder::RnaTreeBuilder(Core::Configuration config, const Bliss::Lexicon& lexicon, const Am::AcousticModel& acousticModel, Search::PersistentStateTree& network, bool initialize)
         : CtcTreeBuilder(config, lexicon, acousticModel, network, initialize) {
-    this->labelLoop_ = paramLabelLoop(config);
+    this->labelLoop_  = paramLabelLoop(config);
+    this->forceBlank_ = paramForceBlank(config);
 }

src/Search/TreeBuilder.hh

@@ -250,6 +250,8 @@ protected:
 class CtcTreeBuilder : public AbstractTreeBuilder {
 public:
     static const Core::ParameterBool paramLabelLoop;
+    static const Core::ParameterBool paramBlankLoop;
+    static const Core::ParameterBool paramForceBlank;
 
     CtcTreeBuilder(Core::Configuration config, const Bliss::Lexicon& lexicon, const Am::AcousticModel& acousticModel, Search::PersistentStateTree& network, bool initialize = true);
     virtual ~CtcTreeBuilder() = default;
@@ -261,6 +263,8 @@ public:
 
 protected:
     bool labelLoop_;
+    bool blankLoop_;
+    bool forceBlank_;
 
     StateId                      wordBoundaryRoot_;
     Search::StateTree::StateDesc blankDesc_;
@@ -292,6 +296,7 @@ protected:
 class RnaTreeBuilder : public CtcTreeBuilder {
 public:
     static const Core::ParameterBool paramLabelLoop;
+    static const Core::ParameterBool paramForceBlank;
 
     RnaTreeBuilder(Core::Configuration config, const Bliss::Lexicon& lexicon, const Am::AcousticModel& acousticModel, Search::PersistentStateTree& network, bool initialize = true);
     virtual ~RnaTreeBuilder() = default;