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.cc

@@ -55,7 +55,7 @@ const Core::ParameterInt StatefulOnnxLabelScorer::paramMaxBatchSize(
 const Core::ParameterInt StatefulOnnxLabelScorer::paramMaxCachedScores(
         "max-cached-score-vectors",
         "Maximum size of cache that maps scoring contexts to scores. This prevents memory overflow in case of very long audio segments.",
-        1000);
+        10000);
 
 // Scorer only takes hidden states as input which are not part of the IO spec
 const std::vector<Onnx::IOSpecification> scorerModelIoSpec = {

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;
@@ -117,7 +121,7 @@ private:
 
     OnnxHiddenStateRef initialHiddenState_;
 
-    // Map input/output names of onnx models to hidden state names taken from state initializer model
+    // Map input/output names of onnx models to hidden state names
     std::unordered_map<std::string, std::string> initializerOutputToStateNameMap_;
     std::unordered_map<std::string, std::string> updaterInputToStateNameMap_;
     std::unordered_map<std::string, std::string> updaterOutputToStateNameMap_;