Add Nn::StatefulOnnxLabelScorer (#140)

Author: SimBe195

src/Nn/LabelScorer/BufferedLabelScorer.cc

@@ -65,4 +65,8 @@ std::optional<DataView> BufferedLabelScorer::getInput(size_t inputIndex) const {
     return inputBuffer_[bufferPosition];
 }
 
+size_t BufferedLabelScorer::bufferSize() const {
+    return inputBuffer_.size();
+}
+
 }  // namespace Nn

src/Nn/LabelScorer/BufferedLabelScorer.hh

@@ -63,6 +63,9 @@ protected:
     //  - getInput(3) will return None since no fourth input was added yet
     std::optional<DataView> getInput(size_t inputIndex) const;
 
+    // Get number of currently buffered elements
+    size_t bufferSize() const;
+
 private:
     std::deque<DataView> inputBuffer_;       // Buffer that contains all the feature data for the current segment
     size_t               numDeletedInputs_;  // Count deleted inputs in order to address the correct index in inputBuffer_

src/Nn/LabelScorer/Makefile

@@ -21,7 +21,8 @@ LIBSPRINTLABELSCORER_O =  \
     $(OBJDIR)/FixedContextOnnxLabelScorer.o \
     $(OBJDIR)/NoContextOnnxLabelScorer.o \
     $(OBJDIR)/NoOpLabelScorer.o \
-    $(OBJDIR)/ScoringContext.o
+    $(OBJDIR)/ScoringContext.o \
+    $(OBJDIR)/StatefulOnnxLabelScorer.o
 
 # -----------------------------------------------------------------------------
 

src/Nn/LabelScorer/ScoringContext.cc

@@ -108,4 +108,32 @@ bool SeqStepScoringContext::isEqual(ScoringContextRef const& other) const {
     return true;
 }
 
+/*
+ * =================================
+ * = OnnxHiddenStateScoringContext =
+ * =================================
+ */
+size_t OnnxHiddenStateScoringContext::hash() const {
+    return Core::MurmurHash3_x64_64(reinterpret_cast<void const*>(labelSeq.data()), labelSeq.size() * sizeof(LabelIndex), 0x78b174eb);
+}
+
+bool OnnxHiddenStateScoringContext::isEqual(ScoringContextRef const& other) const {
+    auto* otherPtr = dynamic_cast<const OnnxHiddenStateScoringContext*>(other.get());
+    if (otherPtr == nullptr) {
+        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

@@ -18,6 +18,7 @@
 
 #include <Core/ReferenceCounting.hh>
 #include <Mm/Types.hh>
+#include <Onnx/Value.hh>
 #include <Speech/Types.hh>
 
 namespace Nn {
@@ -105,6 +106,47 @@ struct SeqStepScoringContext : public ScoringContext {
 
 typedef Core::Ref<const SeqStepScoringContext> SeqStepScoringContextRef;
 
+/*
+ * Hidden state represented by a dictionary of named ONNX values
+ */
+struct OnnxHiddenState : public Core::ReferenceCounted {
+    std::unordered_map<std::string, Onnx::Value> stateValueMap;
+
+    OnnxHiddenState()
+            : stateValueMap() {}
+
+    OnnxHiddenState(std::vector<std::string>&& names, std::vector<Onnx::Value>&& values) {
+        verify(names.size() == values.size());
+        stateValueMap.reserve(names.size());
+        for (size_t i = 0ul; i < names.size(); ++i) {
+            stateValueMap.emplace(std::move(names[i]), std::move(values[i]));
+        }
+    }
+};
+
+typedef Core::Ref<OnnxHiddenState> OnnxHiddenStateRef;
+
+/*
+ * Scoring context consisting of a hidden state.
+ * Assumes that two hidden states are equal if and only if they were created
+ * from the same label history.
+ */
+struct OnnxHiddenStateScoringContext : public ScoringContext {
+    std::vector<LabelIndex> labelSeq;  // Used for hashing
+    OnnxHiddenStateRef      hiddenState;
+
+    OnnxHiddenStateScoringContext()
+            : labelSeq(), hiddenState() {}
+
+    OnnxHiddenStateScoringContext(std::vector<LabelIndex> const& labelSeq, OnnxHiddenStateRef state)
+            : labelSeq(labelSeq), hiddenState(state) {}
+
+    bool   isEqual(ScoringContextRef const& other) const;
+    size_t hash() const;
+};
+
+typedef Core::Ref<const OnnxHiddenStateScoringContext> OnnxHiddenStateScoringContextRef;
+
 }  // namespace Nn
 
 #endif  // SCORING_CONTEXT_HH