Add stateful-transducer-onnx label scorer

src/Nn/LabelScorer/Makefile

@@ -27,6 +27,7 @@ LIBSPRINTLABELSCORER_O =  \
     $(OBJDIR)/ScoringContext.o \
     $(OBJDIR)/ScaledLabelScorer.o \
     $(OBJDIR)/StatefulOnnxLabelScorer.o \
+    $(OBJDIR)/StatefulTransducerOnnxLabelScorer.o \
     $(OBJDIR)/TransitionLabelScorer.o \
     $(OBJDIR)/TransitionTypes.o
 

src/Nn/LabelScorer/ScoringContext.cc

@@ -136,4 +136,36 @@ bool OnnxHiddenStateScoringContext::isEqual(ScoringContextRef const& other) cons
     return true;
 }
 
+/*
+ * =====================================
+ * = StepOnnxHiddenStateScoringContext =
+ * =====================================
+ */
+size_t StepOnnxHiddenStateScoringContext::hash() const {
+    return Core::combineHashes(currentStep, Core::MurmurHash3_x64_64(reinterpret_cast<void const*>(labelSeq.data()), labelSeq.size() * sizeof(LabelIndex), 0x78b174eb));
+}
+
+bool StepOnnxHiddenStateScoringContext::isEqual(ScoringContextRef const& other) const {
+    auto* otherPtr = dynamic_cast<const StepOnnxHiddenStateScoringContext*>(other.get());
+    if (otherPtr == nullptr) {
+        return false;
+    }
+
+    if (currentStep != otherPtr->currentStep) {
+        return false;
+    }
+
+    if (labelSeq.size() != otherPtr->labelSeq.size()) {
+        return false;
+    }
+
+    for (auto it_l = labelSeq.begin(), it_r = otherPtr->labelSeq.begin(); it_l != labelSeq.end(); ++it_l, ++it_r) {
+        if (*it_l != *it_r) {
+            return false;
+        }
+    }
+
+    return true;
+}
+
 }  // namespace Nn

src/Nn/LabelScorer/ScoringContext.hh

@@ -145,6 +145,29 @@ struct OnnxHiddenStateScoringContext : public ScoringContext {
 
 typedef Core::Ref<OnnxHiddenStateScoringContext const> OnnxHiddenStateScoringContextRef;
 
+/*
+ * Scoring context consisting of a hidden state and a step.
+ * Assumes that two hidden states are equal if and only if they were created
+ * from the same label history.
+ */
+struct StepOnnxHiddenStateScoringContext : public ScoringContext {
+    Speech::TimeframeIndex     currentStep;
+    std::vector<LabelIndex>    labelSeq;  // Used for hashing
+    mutable OnnxHiddenStateRef hiddenState;
+    mutable bool               requiresFinalize;
+
+    StepOnnxHiddenStateScoringContext()
+            : currentStep(0u), labelSeq(), hiddenState(), requiresFinalize(false) {}
+
+    StepOnnxHiddenStateScoringContext(Speech::TimeframeIndex step, std::vector<LabelIndex> const& labelSeq, OnnxHiddenStateRef state)
+            : currentStep(step), labelSeq(labelSeq), hiddenState(state), requiresFinalize(false) {}
+
+    bool   isEqual(ScoringContextRef const& other) const;
+    size_t hash() const;
+};
+
+typedef Core::Ref<const StepOnnxHiddenStateScoringContext> StepOnnxHiddenStateScoringContextRef;
+
 }  // namespace Nn
 
 #endif  // SCORING_CONTEXT_HH

src/Nn/LabelScorer/StatefulOnnxLabelScorer.hh

@@ -56,6 +56,10 @@ namespace Nn {
  *
  * A common use case for this Label Scorer would be an AED model with cross-attention over the encoder output.
  * Since the encoder state inputs are optional, it can also be used for stateful language models without acoustic input.
+ *
+ * Note: This LabelScorer is similar to the `StatefulTransducerOnnxLabelScorer`. The difference is that in this it is assumed that the
+ * input features are processed into the hidden states and they are not directly fed into the scorer. For this, the state initializer
+ * and updater here also take input features in addition to tokens.
  */
 class StatefulOnnxLabelScorer : public BufferedLabelScorer {
     using Precursor = BufferedLabelScorer;
@@ -99,6 +103,9 @@ private:
     // Compute updated states for all non-finalized scoring contexts and put them into the state cache
     void cacheStates(std::vector<OnnxHiddenStateScoringContextRef> const& scoringContextBatch);
 
+    // Replace hidden-state in scoringContext with an updated version that includes the last label
+    void finalizeScoringContext(OnnxHiddenStateScoringContextRef const& scoringContext);
+
     // Since the hidden-state matrix depends on the encoder time axis, we cannot create properly create hidden-states until all encoder states have been passed.
     // So getInitialScoringContext sets the initial hidden-state to a sentinel value (empty Ref) and when other functions such as `extendedScoringContext` and `getScoresWithTime`
     // encounter this sentinel value they call `computeInitialHiddenState` instead to get a usable hidden-state.