Update Nn::LabelScorer: allow enabling of specific transition types

src/Nn/LabelScorer/CombineLabelScorer.hh

@@ -71,6 +71,11 @@ protected:
 
     // Compute weighted scores of requests with all sub-scorers
     std::optional<ScoresWithTimes> computeScoresWithTimesInternal(std::vector<Request> const& requests) override;
+
+    // Request filtering should happen in the sub-scorers, this one should let everything through
+    virtual TransitionPresetType defaultPreset() const override {
+        return TransitionPresetType::ALL;
+    }
 };
 
 }  // namespace Nn

src/Nn/LabelScorer/EncoderDecoderLabelScorer.hh

@@ -67,6 +67,11 @@ protected:
     // Run requests through decoder component
     std::optional<LabelScorer::ScoresWithTimes> computeScoresWithTimesInternal(std::vector<LabelScorer::Request> const& requests) override;
 
+    // Request filtering should happen in the decoder, this one should let everything through
+    virtual TransitionPresetType defaultPreset() const override {
+        return TransitionPresetType::ALL;
+    }
+
 private:
     Core::Ref<Encoder>     encoder_;
     Core::Ref<LabelScorer> decoder_;

src/Nn/LabelScorer/FixedContextOnnxLabelScorer.hh

@@ -66,6 +66,10 @@ protected:
     // Uses `getScoresWithTimes` internally with some wrapping for vector packing/expansion
     std::optional<LabelScorer::ScoreWithTime> computeScoreWithTimeInternal(LabelScorer::Request const& request) override;
 
+    virtual TransitionPresetType defaultPreset() const override {
+        return TransitionPresetType::TRANSDUCER;
+    }
+
 private:
     // Forward a batch of histories through the ONNX model and put the resulting scores into the score cache
     // Assumes that all histories in the batch are based on the same timestep

src/Nn/LabelScorer/LabelScorer.cc

@@ -23,26 +23,29 @@ namespace Nn {
  * =============================
  */
 
-const Core::ParameterStringVector LabelScorer::paramIgnoredTransitionTypes(
-        "ignored-transition-types",
-        "Transition types that should be ignored by the label scorer (i.e. get assigned score 0 and do not affect the ScoringContext)",
+const Core::Choice LabelScorer::choiceTransitionPreset(
+        "default", TransitionPresetType::DEFAULT,
+        "none", TransitionPresetType::NONE,
+        "ctc", TransitionPresetType::CTC,
+        "transducer", TransitionPresetType::TRANSDUCER,
+        "lm", TransitionPresetType::LM,
+        Core::Choice::endMark());
+
+const Core::ParameterChoice LabelScorer::paramTransitionPreset(
+        "transition-preset",
+        &LabelScorer::choiceTransitionPreset,
+        "Preset for which transition types should be enabled for the label scorer. Disabled transition types get assigned score 0 and do not affect the ScoringContext.",
+        TransitionPresetType::DEFAULT);
+
+const Core::ParameterStringVector LabelScorer::paramExtraTransitionTypes(
+        "extra-transition-types",
+        "Transition types that should be enabled in addition to the ones given by the preset.",
         ",");
 
 LabelScorer::LabelScorer(const Core::Configuration& config)
         : Core::Component(config),
-          ignoredTransitionTypes_() {
-    auto ignoredTransitionTypeStrings = paramIgnoredTransitionTypes(config);
-    for (auto const& transitionTypeString : ignoredTransitionTypeStrings) {
-        auto it = std::find_if(transitionTypeArray_.begin(),
-                               transitionTypeArray_.end(),
-                               [&](auto const& entry) { return entry.first == transitionTypeString; });
-        if (it != transitionTypeArray_.end()) {
-            ignoredTransitionTypes_.insert(it->second);
-        }
-        else {
-            error() << "Ignored transition type name '" << transitionTypeString << "' is not a valid identifier";
-        }
-    }
+          enabledTransitionTypes_() {
+    enableTransitionTypes(config);
 }
 
 void LabelScorer::addInputs(DataView const& input, size_t nTimesteps) {
@@ -53,17 +56,17 @@ void LabelScorer::addInputs(DataView const& input, size_t nTimesteps) {
 }
 
 ScoringContextRef LabelScorer::extendedScoringContext(Request const& request) {
-    if (ignoredTransitionTypes_.contains(request.transitionType)) {
-        return request.context;
+    if (enabledTransitionTypes_.contains(request.transitionType)) {
+        return extendedScoringContextInternal(request);
     }
-    return extendedScoringContextInternal(request);
+    return request.context;
 }
 
 std::optional<LabelScorer::ScoreWithTime> LabelScorer::computeScoreWithTime(Request const& request) {
-    if (ignoredTransitionTypes_.contains(request.transitionType)) {
-        return ScoreWithTime{0.0, 0};
+    if (enabledTransitionTypes_.contains(request.transitionType)) {
+        return computeScoreWithTimeInternal(request);
     }
-    return computeScoreWithTimeInternal(request);
+    return ScoreWithTime{0.0, 0};
 }
 
 std::optional<LabelScorer::ScoresWithTimes> LabelScorer::computeScoresWithTimes(std::vector<LabelScorer::Request> const& requests) {
@@ -76,7 +79,7 @@ std::optional<LabelScorer::ScoresWithTimes> LabelScorer::computeScoresWithTimes(
 
     for (size_t requestIndex = 0ul; requestIndex < requests.size(); ++requestIndex) {
         auto const& request = requests[requestIndex];
-        if (not ignoredTransitionTypes_.contains(request.transitionType)) {
+        if (enabledTransitionTypes_.contains(request.transitionType)) {
             nonIgnoredRequests.push_back(request);
             nonIgnoredRequestIndices.push_back(requestIndex);
         }
@@ -89,12 +92,13 @@ std::optional<LabelScorer::ScoresWithTimes> LabelScorer::computeScoresWithTimes(
     }
 
     // Interleave actual results with 0 scores for requests with ignored transition types
-    ScoresWithTimes result{{requests.size(), 0.0}, {requests.size(), 0}};
+    ScoresWithTimes result{
+            .scores = std::vector<Score>(requests.size(), 0.0),
+            .timeframes{requests.size(), 0}};
     for (size_t i = 0ul; i < nonIgnoredRequestIndices.size(); ++i) {
-        auto requestResult          = nonIgnoredResults[i];
         auto requestIndex           = nonIgnoredRequestIndices[i];
-        result.scores[requestIndex] = requestResult.score;
-        result.timeframes.set(requestIndex, requestResult.timeframe);
+        result.scores[requestIndex] = nonIgnoredResults->scores[i];
+        result.timeframes.set(requestIndex, nonIgnoredResults->timeframes[i]);
     }
 
     return result;
@@ -122,4 +126,66 @@ std::optional<LabelScorer::ScoresWithTimes> LabelScorer::computeScoresWithTimesI
     return result;
 }
 
+void LabelScorer::enableTransitionTypes(Core::Configuration const& config) {
+    auto preset = paramTransitionPreset(config);
+    if (preset == TransitionPresetType::DEFAULT) {
+        preset = defaultPreset();
+    }
+    verify(preset != TransitionPresetType::DEFAULT);
+
+    switch (preset) {
+        case TransitionPresetType::NONE:
+            break;
+        case TransitionPresetType::ALL:
+            for (auto const& [_, transitionType] : transitionTypeArray_) {
+                enabledTransitionTypes_.insert(transitionType);
+            }
+            break;
+        case TransitionPresetType::CTC:
+            enabledTransitionTypes_ = {
+                    LABEL_TO_LABEL,
+                    LABEL_LOOP,
+                    LABEL_TO_BLANK,
+                    BLANK_TO_LABEL,
+                    BLANK_LOOP,
+                    INITIAL_LABEL,
+                    INITIAL_BLANK,
+            };
+            break;
+        case TransitionPresetType::TRANSDUCER:
+            enabledTransitionTypes_ = {
+                    LABEL_TO_LABEL,
+                    LABEL_TO_BLANK,
+                    BLANK_TO_LABEL,
+                    BLANK_LOOP,
+                    INITIAL_LABEL,
+                    INITIAL_BLANK,
+            };
+            break;
+        case TransitionPresetType::LM:
+            enabledTransitionTypes_ = {
+                    LABEL_TO_LABEL,
+                    INITIAL_LABEL,
+            };
+            break;
+    }
+
+    auto extraTransitionTypeStrings = paramExtraTransitionTypes(config);
+    for (auto const& transitionTypeString : extraTransitionTypeStrings) {
+        auto it = std::find_if(transitionTypeArray_.begin(),
+                               transitionTypeArray_.end(),
+                               [&](auto const& entry) { return entry.first == transitionTypeString; });
+        if (it != transitionTypeArray_.end()) {
+            enabledTransitionTypes_.insert(it->second);
+        }
+        else {
+            error() << "Extra transition type name '" << transitionTypeString << "' is not a valid identifier";
+        }
+    }
+
+    if (enabledTransitionTypes_.empty()) {
+        error() << "Label scorer has no enabled transition types. Activate a preset and/or add extra transition types that should be considered.";
+    }
+}
+
 }  // namespace Nn

src/Nn/LabelScorer/LabelScorer.hh

@@ -73,11 +73,14 @@ namespace Nn {
  */
 class LabelScorer : public virtual Core::Component,
                     public Core::ReferenceCounted {
-    static const Core::ParameterStringVector paramIgnoredTransitionTypes;
-
 public:
     typedef Search::Score Score;
 
+    static const Core::Choice          choiceTransitionPreset;
+    static const Core::ParameterChoice paramTransitionPreset;
+
+    static const Core::ParameterStringVector paramExtraTransitionTypes;
+
     enum TransitionType {
         LABEL_TO_LABEL,
         LABEL_LOOP,
@@ -89,6 +92,15 @@ public:
         numTypes,  // must remain at the end
     };
 
+    enum TransitionPresetType {
+        DEFAULT,
+        NONE,
+        ALL,
+        CTC,
+        TRANSDUCER,
+        LM,
+    };
+
     // Request for scoring or context extension
     struct Request {
         ScoringContextRef context;
@@ -166,8 +178,14 @@ protected:
     // By default loops over the single-request version
     virtual std::optional<ScoresWithTimes> computeScoresWithTimesInternal(std::vector<Request> const& requests);
 
+    virtual TransitionPresetType defaultPreset() const {
+        return TransitionPresetType::NONE;
+    }
+
 private:
-    std::unordered_set<TransitionType> ignoredTransitionTypes_;
+    std::unordered_set<TransitionType> enabledTransitionTypes_;
+
+    void enableTransitionTypes(Core::Configuration const& config);
 };
 
 }  // namespace Nn

src/Nn/LabelScorer/NoContextOnnxLabelScorer.hh

@@ -61,6 +61,10 @@ protected:
     // Uses `getScoresWithTimes` internally with some wrapping for vector packing/expansion
     std::optional<LabelScorer::ScoreWithTime> computeScoreWithTimeInternal(LabelScorer::Request const& request) override;
 
+    virtual TransitionPresetType defaultPreset() const override {
+        return TransitionPresetType::CTC;
+    }
+
 private:
     Onnx::Model onnxModel_;
 

src/Nn/LabelScorer/NoOpLabelScorer.hh

@@ -44,6 +44,10 @@ protected:
 
     // Gets the buffered score for the requested token at the requested step
     std::optional<LabelScorer::ScoreWithTime> computeScoreWithTimeInternal(LabelScorer::Request const& request) override;
+
+    virtual TransitionPresetType defaultPreset() const override {
+        return TransitionPresetType::CTC;
+    }
 };
 
 }  // namespace Nn

src/Nn/LabelScorer/StatefulOnnxLabelScorer.hh

@@ -86,6 +86,10 @@ protected:
     std::optional<LabelScorer::ScoreWithTime>   computeScoreWithTimeInternal(LabelScorer::Request const& request) override;
     std::optional<LabelScorer::ScoresWithTimes> computeScoresWithTimesInternal(std::vector<LabelScorer::Request> const& requests) override;
 
+    virtual TransitionPresetType defaultPreset() const override {
+        return TransitionPresetType::LM;
+    }
+
 private:
     // Forward a batch of histories through the ONNX model and put the resulting scores into the score cache
     void forwardBatch(std::vector<OnnxHiddenStateScoringContextRef> const& historyBatch);

src/Nn/LabelScorer/TransitionLabelScorer.hh

@@ -60,6 +60,10 @@ protected:
     // Compute scores of base scorer and add transition scores based on transition types of the requests
     std::optional<ScoresWithTimes> computeScoresWithTimesInternal(std::vector<Request> const& requests) override;
 
+    virtual TransitionPresetType defaultPreset() const override {
+        return TransitionPresetType::ALL;
+    }
+
 private:
     std::unordered_map<TransitionType, Score> transitionScores_;