Update Nn::LabelScorer: allow enabling of specific transition types (#148)

Co-authored-by: Eugen Beck <[email protected]>

Author: SimBe195

src/Core/CollapsedVector.hh

@@ -45,6 +45,7 @@ public:
     inline void                  push_back(const T& value);
     inline const T&              operator[](size_t idx) const;
     inline const T&              at(size_t idx) const;
+    inline void                  set(size_t idx, const T& value);
     inline size_t                size() const noexcept;
     inline void                  clear() noexcept;
     inline void                  reserve(size_t size);
@@ -105,6 +106,21 @@ inline const T& CollapsedVector<T>::at(size_t idx) const {
     return (*this)[idx];
 }
 
+template<typename T>
+inline void CollapsedVector<T>::set(size_t idx, const T& value) {
+    if (idx >= logicalSize_) {
+        throw std::out_of_range("Trying to access illegal index of CollapsedVector");
+    }
+    if (data_.size() != 1ul) {
+        data_[idx] = value;
+        data_.push_back(value);
+    }
+    else if (value != data_.front()) {
+        data_.resize(logicalSize_, data_.front());
+        data_[idx] = value;
+    }
+}
+
 template<typename T>
 inline size_t CollapsedVector<T>::size() const noexcept {
     return logicalSize_;

src/Nn/LabelScorer/BufferedLabelScorer.cc

@@ -17,9 +17,9 @@
 
 namespace Nn {
 
-BufferedLabelScorer::BufferedLabelScorer(Core::Configuration const& config)
+BufferedLabelScorer::BufferedLabelScorer(Core::Configuration const& config, TransitionPresetType defaultPreset)
         : Core::Component(config),
-          Precursor(config),
+          Precursor(config, defaultPreset),
           expectMoreFeatures_(true),
           inputBuffer_(),
           numDeletedInputs_(0ul) {

src/Nn/LabelScorer/BufferedLabelScorer.hh

@@ -31,7 +31,7 @@ class BufferedLabelScorer : public LabelScorer {
 public:
     using Precursor = LabelScorer;
 
-    BufferedLabelScorer(Core::Configuration const& config);
+    BufferedLabelScorer(Core::Configuration const& config, TransitionPresetType defaultPreset);
 
     // Prepares the LabelScorer to receive new inputs by resetting input buffer, timeframe buffer
     // and segment end flag

src/Nn/LabelScorer/CombineLabelScorer.cc

@@ -25,7 +25,8 @@ Core::ParameterFloat CombineLabelScorer::paramScale(
         "scale", "Scores of a sub-label-scorer are scaled by this factor", 1.0);
 
 CombineLabelScorer::CombineLabelScorer(Core::Configuration const& config)
-        : Core::Component(config), Precursor(config), scaledScorers_() {
+        : Core::Component(config),
+          Precursor(config, TransitionPresetType::ALL) {
     size_t numLabelScorers = paramNumLabelScorers(config);
     for (size_t i = 0ul; i < numLabelScorers; ++i) {
         Core::Configuration subConfig = select(std::string("scorer-") + std::to_string(i + 1));
@@ -55,22 +56,6 @@ ScoringContextRef CombineLabelScorer::getInitialScoringContext() {
     return Core::ref(new CombineScoringContext(std::move(scoringContexts)));
 }
 
-ScoringContextRef CombineLabelScorer::extendedScoringContext(Request const& request) {
-    auto combineContext = dynamic_cast<const CombineScoringContext*>(request.context.get());
-
-    std::vector<ScoringContextRef> extScoringContexts;
-    extScoringContexts.reserve(scaledScorers_.size());
-
-    auto scorerIt  = scaledScorers_.begin();
-    auto contextIt = combineContext->scoringContexts.begin();
-
-    for (; scorerIt != scaledScorers_.end(); ++scorerIt, ++contextIt) {
-        Request subRequest{*contextIt, request.nextToken, request.transitionType};
-        extScoringContexts.push_back(scorerIt->scorer->extendedScoringContext(subRequest));
-    }
-    return Core::ref(new CombineScoringContext(std::move(extScoringContexts)));
-}
-
 void CombineLabelScorer::cleanupCaches(Core::CollapsedVector<ScoringContextRef> const& activeContexts) {
     std::vector<const CombineScoringContext*> combineContexts;
     combineContexts.reserve(activeContexts.internalSize());
@@ -101,7 +86,23 @@ void CombineLabelScorer::addInputs(DataView const& input, size_t nTimesteps) {
     }
 }
 
-std::optional<LabelScorer::ScoreWithTime> CombineLabelScorer::computeScoreWithTime(Request const& request) {
+ScoringContextRef CombineLabelScorer::extendedScoringContextInternal(Request const& request) {
+    auto combineContext = dynamic_cast<const CombineScoringContext*>(request.context.get());
+
+    std::vector<ScoringContextRef> extScoringContexts;
+    extScoringContexts.reserve(scaledScorers_.size());
+
+    auto scorerIt  = scaledScorers_.begin();
+    auto contextIt = combineContext->scoringContexts.begin();
+
+    for (; scorerIt != scaledScorers_.end(); ++scorerIt, ++contextIt) {
+        Request subRequest{*contextIt, request.nextToken, request.transitionType};
+        extScoringContexts.push_back(scorerIt->scorer->extendedScoringContext(subRequest));
+    }
+    return Core::ref(new CombineScoringContext(std::move(extScoringContexts)));
+}
+
+std::optional<LabelScorer::ScoreWithTime> CombineLabelScorer::computeScoreWithTimeInternal(Request const& request) {
     // Initialize accumulated result with zero-valued score and timestep
     ScoreWithTime accumResult{0.0, 0};
 
@@ -130,7 +131,11 @@ std::optional<LabelScorer::ScoreWithTime> CombineLabelScorer::computeScoreWithTi
     return accumResult;
 }
 
-std::optional<LabelScorer::ScoresWithTimes> CombineLabelScorer::computeScoresWithTimes(std::vector<Request> const& requests) {
+std::optional<LabelScorer::ScoresWithTimes> CombineLabelScorer::computeScoresWithTimesInternal(std::vector<Request> const& requests) {
+    if (requests.empty()) {
+        return ScoresWithTimes{};
+    }
+
     // Initialize accumulated results with zero-valued scores and timesteps
     ScoresWithTimes accumResult{std::vector<Score>(requests.size(), 0.0), {requests.size(), 0}};
 

src/Nn/LabelScorer/CombineLabelScorer.hh

@@ -46,9 +46,6 @@ public:
     // Combine initial ScoringContexts from all sub-scorers
     ScoringContextRef getInitialScoringContext() override;
 
-    // Combine extended ScoringContexts from all sub-scorers
-    ScoringContextRef extendedScoringContext(Request const& request) override;
-
     // Cleanup all sub-scorers
     void cleanupCaches(Core::CollapsedVector<ScoringContextRef> const& activeContexts) override;
 
@@ -58,19 +55,22 @@ public:
     // Add inputs to all sub-scorers
     virtual void addInputs(DataView const& input, size_t nTimesteps) override;
 
-    // Compute weighted score of request with all sub-scorers
-    std::optional<ScoreWithTime> computeScoreWithTime(Request const& request) override;
-
-    // Compute weighted scores of requests with all sub-scorers
-    std::optional<ScoresWithTimes> computeScoresWithTimes(std::vector<Request> const& requests) override;
-
 protected:
     struct ScaledLabelScorer {
         Core::Ref<LabelScorer> scorer;
         Score                  scale;
     };
 
     std::vector<ScaledLabelScorer> scaledScorers_;
+
+    // Combine extended ScoringContexts from all sub-scorers
+    ScoringContextRef extendedScoringContextInternal(Request const& request) override;
+
+    // Compute weighted score of request with all sub-scorers
+    std::optional<ScoreWithTime> computeScoreWithTimeInternal(Request const& request) override;
+
+    // Compute weighted scores of requests with all sub-scorers
+    std::optional<ScoresWithTimes> computeScoresWithTimesInternal(std::vector<Request> const& requests) override;
 };
 
 }  // namespace Nn

src/Nn/LabelScorer/EncoderDecoderLabelScorer.cc

@@ -19,7 +19,7 @@ namespace Nn {
 
 EncoderDecoderLabelScorer::EncoderDecoderLabelScorer(Core::Configuration const& config, Core::Ref<Encoder> const& encoder, Core::Ref<LabelScorer> const& decoder)
         : Core::Component(config),
-          LabelScorer(config),
+          LabelScorer(config, TransitionPresetType::ALL),
           encoder_(encoder),
           decoder_(decoder) {
 }
@@ -33,10 +33,6 @@ ScoringContextRef EncoderDecoderLabelScorer::getInitialScoringContext() {
     return decoder_->getInitialScoringContext();
 }
 
-ScoringContextRef EncoderDecoderLabelScorer::extendedScoringContext(Request const& request) {
-    return decoder_->extendedScoringContext(request);
-}
-
 void EncoderDecoderLabelScorer::cleanupCaches(Core::CollapsedVector<ScoringContextRef> const& activeContexts) {
     decoder_->cleanupCaches(activeContexts);
 }
@@ -59,11 +55,19 @@ void EncoderDecoderLabelScorer::signalNoMoreFeatures() {
     decoder_->signalNoMoreFeatures();
 }
 
-std::optional<LabelScorer::ScoreWithTime> EncoderDecoderLabelScorer::computeScoreWithTime(LabelScorer::Request const& request) {
+ScoringContextRef EncoderDecoderLabelScorer::extendedScoringContextInternal(Request const& request) {
+    return decoder_->extendedScoringContext(request);
+}
+
+std::optional<LabelScorer::ScoreWithTime> EncoderDecoderLabelScorer::computeScoreWithTimeInternal(LabelScorer::Request const& request) {
     return decoder_->computeScoreWithTime(request);
 }
 
-std::optional<LabelScorer::ScoresWithTimes> EncoderDecoderLabelScorer::computeScoresWithTimes(std::vector<LabelScorer::Request> const& requests) {
+std::optional<LabelScorer::ScoresWithTimes> EncoderDecoderLabelScorer::computeScoresWithTimesInternal(std::vector<LabelScorer::Request> const& requests) {
+    if (requests.empty()) {
+        return ScoresWithTimes{};
+    }
+
     return decoder_->computeScoresWithTimes(requests);
 }
 

src/Nn/LabelScorer/EncoderDecoderLabelScorer.hh

@@ -46,9 +46,6 @@ public:
     // Get start context from decoder component
     ScoringContextRef getInitialScoringContext() override;
 
-    // Get extended context from decoder component
-    ScoringContextRef extendedScoringContext(Request const& request) override;
-
     // Cleanup decoder component. Encoder is "self-cleaning" already in that it only stores outputs until they are
     // retrieved.
     void cleanupCaches(Core::CollapsedVector<ScoringContextRef> const& activeContexts) override;
@@ -60,11 +57,15 @@ public:
     // Same as `addInput` but adds features for multiple timesteps at once
     void addInputs(DataView const& input, size_t nTimesteps) override;
 
+protected:
+    // Get extended context from decoder component
+    ScoringContextRef extendedScoringContextInternal(Request const& request) override;
+
     // Run request through decoder component
-    std::optional<LabelScorer::ScoreWithTime> computeScoreWithTime(LabelScorer::Request const& request) override;
+    std::optional<LabelScorer::ScoreWithTime> computeScoreWithTimeInternal(LabelScorer::Request const& request) override;
 
     // Run requests through decoder component
-    std::optional<LabelScorer::ScoresWithTimes> computeScoresWithTimes(std::vector<LabelScorer::Request> const& requests) override;
+    std::optional<LabelScorer::ScoresWithTimes> computeScoresWithTimesInternal(std::vector<LabelScorer::Request> const& requests) override;
 
 private:
     Core::Ref<Encoder>     encoder_;

src/Nn/LabelScorer/FixedContextOnnxLabelScorer.cc

@@ -72,7 +72,7 @@ static const std::vector<Onnx::IOSpecification> ioSpec = {
 
 FixedContextOnnxLabelScorer::FixedContextOnnxLabelScorer(Core::Configuration const& config)
         : Core::Component(config),
-          Precursor(config),
+          Precursor(config, TransitionPresetType::TRANSDUCER),
           startLabelIndex_(paramStartLabelIndex(config)),
           historyLength_(paramHistoryLength(config)),
           blankUpdatesHistory_(paramBlankUpdatesHistory(config)),
@@ -97,7 +97,32 @@ ScoringContextRef FixedContextOnnxLabelScorer::getInitialScoringContext() {
     return hist;
 }
 
-ScoringContextRef FixedContextOnnxLabelScorer::extendedScoringContext(LabelScorer::Request const& request) {
+size_t FixedContextOnnxLabelScorer::getMinActiveInputIndex(Core::CollapsedVector<ScoringContextRef> const& activeContexts) const {
+    auto minTimeIndex = Core::Type<Speech::TimeframeIndex>::max;
+    for (auto const& context : activeContexts.internalData()) {
+        SeqStepScoringContextRef stepHistory(dynamic_cast<const SeqStepScoringContext*>(context.get()));
+        minTimeIndex = std::min(minTimeIndex, stepHistory->currentStep);
+    }
+
+    return minTimeIndex;
+}
+
+void FixedContextOnnxLabelScorer::cleanupCaches(Core::CollapsedVector<ScoringContextRef> const& activeContexts) {
+    Precursor::cleanupCaches(activeContexts);
+
+    std::unordered_set<ScoringContextRef, ScoringContextHash, ScoringContextEq> activeContextSet(activeContexts.internalData().begin(), activeContexts.internalData().end());
+
+    for (auto it = scoreCache_.begin(); it != scoreCache_.end();) {
+        if (activeContextSet.find(it->first) == activeContextSet.end()) {
+            it = scoreCache_.erase(it);
+        }
+        else {
+            ++it;
+        }
+    }
+}
+
+ScoringContextRef FixedContextOnnxLabelScorer::extendedScoringContextInternal(LabelScorer::Request const& request) {
     SeqStepScoringContextRef context(dynamic_cast<const SeqStepScoringContext*>(request.context.get()));
 
     bool   pushToken     = false;
@@ -144,22 +169,11 @@ ScoringContextRef FixedContextOnnxLabelScorer::extendedScoringContext(LabelScore
     return Core::ref(new SeqStepScoringContext(std::move(newLabelSeq), context->currentStep + timeIncrement));
 }
 
-void FixedContextOnnxLabelScorer::cleanupCaches(Core::CollapsedVector<ScoringContextRef> const& activeContexts) {
-    Precursor::cleanupCaches(activeContexts);
-
-    std::unordered_set<ScoringContextRef, ScoringContextHash, ScoringContextEq> activeContextSet(activeContexts.internalData().begin(), activeContexts.internalData().end());
-
-    for (auto it = scoreCache_.begin(); it != scoreCache_.end();) {
-        if (activeContextSet.find(it->first) == activeContextSet.end()) {
-            it = scoreCache_.erase(it);
-        }
-        else {
-            ++it;
-        }
+std::optional<LabelScorer::ScoresWithTimes> FixedContextOnnxLabelScorer::computeScoresWithTimesInternal(std::vector<LabelScorer::Request> const& requests) {
+    if (requests.empty()) {
+        return ScoresWithTimes{};
     }
-}
 
-std::optional<LabelScorer::ScoresWithTimes> FixedContextOnnxLabelScorer::computeScoresWithTimes(std::vector<LabelScorer::Request> const& requests) {
     ScoresWithTimes result;
     result.scores.reserve(requests.size());
 
@@ -232,24 +246,14 @@ std::optional<LabelScorer::ScoresWithTimes> FixedContextOnnxLabelScorer::compute
     return result;
 }
 
-std::optional<LabelScorer::ScoreWithTime> FixedContextOnnxLabelScorer::computeScoreWithTime(LabelScorer::Request const& request) {
+std::optional<LabelScorer::ScoreWithTime> FixedContextOnnxLabelScorer::computeScoreWithTimeInternal(LabelScorer::Request const& request) {
     auto result = computeScoresWithTimes({request});
     if (not result.has_value()) {
         return {};
     }
     return ScoreWithTime{result->scores.front(), result->timeframes.front()};
 }
 
-size_t FixedContextOnnxLabelScorer::getMinActiveInputIndex(Core::CollapsedVector<ScoringContextRef> const& activeContexts) const {
-    auto minTimeIndex = Core::Type<Speech::TimeframeIndex>::max;
-    for (auto const& context : activeContexts.internalData()) {
-        SeqStepScoringContextRef stepHistory(dynamic_cast<const SeqStepScoringContext*>(context.get()));
-        minTimeIndex = std::min(minTimeIndex, stepHistory->currentStep);
-    }
-
-    return minTimeIndex;
-}
-
 void FixedContextOnnxLabelScorer::forwardBatch(std::vector<SeqStepScoringContextRef> const& contextBatch) {
     if (contextBatch.empty()) {
         return;

src/Nn/LabelScorer/FixedContextOnnxLabelScorer.hh

@@ -47,24 +47,24 @@ public:
     // Initial scoring context contains step 0 and a history vector filled with the start label index
     ScoringContextRef getInitialScoringContext() override;
 
+    // Clean up input buffer as well as cached score vectors that are no longer needed
+    void cleanupCaches(Core::CollapsedVector<ScoringContextRef> const& activeContexts) override;
+
+protected:
+    size_t getMinActiveInputIndex(Core::CollapsedVector<ScoringContextRef> const& activeContexts) const override;
+
     // May increment the step by 1 (except for vertical transitions) and may append the next token to the
     // history label sequence depending on the transition type and whether loops/blanks update the history
     // or not
-    ScoringContextRef extendedScoringContext(LabelScorer::Request const& request) override;
-
-    // Clean up input buffer as well as cached score vectors that are no longer needed
-    void cleanupCaches(Core::CollapsedVector<ScoringContextRef> const& activeContexts) override;
+    ScoringContextRef extendedScoringContextInternal(LabelScorer::Request const& request) override;
 
     // If scores for the given scoring contexts are not yet cached, prepare and run an ONNX session to
     // compute the scores and cache them
     // Then, retreive scores from cache
-    std::optional<LabelScorer::ScoresWithTimes> computeScoresWithTimes(std::vector<LabelScorer::Request> const& requests) override;
+    std::optional<LabelScorer::ScoresWithTimes> computeScoresWithTimesInternal(std::vector<LabelScorer::Request> const& requests) override;
 
     // Uses `getScoresWithTimes` internally with some wrapping for vector packing/expansion
-    std::optional<LabelScorer::ScoreWithTime> computeScoreWithTime(LabelScorer::Request const& request) override;
-
-protected:
-    size_t getMinActiveInputIndex(Core::CollapsedVector<ScoringContextRef> const& activeContexts) const override;
+    std::optional<LabelScorer::ScoreWithTime> computeScoreWithTimeInternal(LabelScorer::Request const& request) override;
 
 private:
     // Forward a batch of histories through the ONNX model and put the resulting scores into the score cache