StatefulOnnxLabelScorer.hh

/** Copyright 2025 RWTH Aachen University. All rights reserved.
 *
 *  Licensed under the RWTH ASR License (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *      http://www.hltpr.rwth-aachen.de/rwth-asr/rwth-asr-license.html
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */

#ifndef STATEFUL_ONNX_LABEL_SCORER_HH
#define STATEFUL_ONNX_LABEL_SCORER_HH

#include <optional>

#include <Core/Component.hh>
#include <Core/Configuration.hh>
#include <Core/FIFOCache.hh>
#include <Core/ReferenceCounting.hh>
#include <Mm/FeatureScorer.hh>
#include <Onnx/IOSpecification.hh>
#include <Onnx/Model.hh>
#include <Onnx/Session.hh>
#include <Speech/Feature.hh>

#include "BufferedLabelScorer.hh"
#include "ScoringContext.hh"

namespace Nn {

/*
 * Label Scorer that performs scoring by forwarding hidden states through an ONNX model.
 * This Label Scorer requires three ONNX models:
 *  - A State Initializer which produces the hidden states for the first step (optionally based on the input sequence)
 *  - A State Updater which produces updated hidden states based on the previous hidden states and optionally also the input sequence and the next token
 *  - A Scorer which computes scores based on the hidden states
 *
 * The hidden states can be any number of ONNX tensors of any shape and type.
 * Each ONNX model must have metadata that specifies the mapping of its input and output names to the corresponding state names.
 * These state names need to be consistent over all three models.
 *
 * For example:
 *   - The State Initializer has output called "lstm_c" and {"lstm_c": "LSTM_C"} in its metadata
 *   - The State Updater has input "lstm_c_in", output "lstm_c_out" and {"lstm_c_in": "LSTM_C", "lstm_c_out": "LSTM_C"} in its metadata
 *   - The Scorer has input "lstm_c" and {"lstm_c": "LSTM_C"} in its metadata
 * Here, "LSTM_C" is the state name and the same across all three models while the specific input/output names are arbitrary.
 *
 * The State Initializer must have all states as output.
 * The State Updater must have a subset of states as input and all states as output.
 * The Scorer must have a subset of states as input.
 *
 * 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;

    static const Core::ParameterBool paramBlankUpdatesHistory;
    static const Core::ParameterBool paramLoopUpdatesHistory;
    static const Core::ParameterInt  paramMaxBatchSize;
    static const Core::ParameterInt  paramMaxCachedScores;

public:
    StatefulOnnxLabelScorer(Core::Configuration const& config);
    virtual ~StatefulOnnxLabelScorer() = default;

    void reset() override;

    // If startLabelIndex is set, forward that through the state updater to obtain the start ScoringContext
    ScoringContextRef getInitialScoringContext() override;

    // Append the new token to the label sequence; does not update the hidden-state. This is only done once the scoringContext is used for scoring again.
    ScoringContextRef extendedScoringContext(ScoringContextRef scoringContext, LabelIndex nextToken, TransitionType transitionType) override;

    // Add a single input feature to buffer
    void addInput(DataView const& input) override;

    // Update hidden state, run scorer and get an accessor for the output score vector
    std::optional<ScoreAccessorRef> getScoreAccessor(ScoringContextRef scoringContext) override;

    // Update hidden states (batched), run scorers (batched) and get accessor for the output score vectors
    std::vector<std::optional<ScoreAccessorRef>> getScoreAccessors(std::vector<ScoringContextRef> const& scoringContexts) override;

protected:
    size_t getMinActiveInputIndex(Core::CollapsedVector<ScoringContextRef> const& activeContexts) const override;

private:
    // Forward a batch of scoringContexts through the ONNX scorer model and put the resulting scores into the score cache
    void cacheScores(std::vector<OnnxHiddenStateScoringContextRef> const& scoringContextBatch);

    // Computes new hidden state based on previous hidden state and next token with batched state-updater call
    std::vector<OnnxHiddenStateRef> updatedHiddenStates(std::vector<OnnxHiddenStateRef> const& hiddenStatesBatch, std::vector<s32> nextTokensBatch);

    // 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.
    OnnxHiddenStateRef computeInitialHiddenState();

    void setupEncoderStatesValue();
    void setupEncoderStatesSizeValue();

    bool   blankUpdatesHistory_;
    bool   loopUpdatesHistory_;
    size_t maxBatchSize_;

    Onnx::Model scorerOnnxModel_;
    Onnx::Model stateInitializerOnnxModel_;
    Onnx::Model stateUpdaterOnnxModel_;

    OnnxHiddenStateRef initialHiddenState_;

    // Map input/output names of onnx models to hidden state names taken from state initializer model
    std::unordered_map<std::string, std::string> initializerOutputToStateNameMap_;
    std::unordered_map<std::string, std::string> updaterInputToStateNameMap_;
    std::unordered_map<std::string, std::string> updaterOutputToStateNameMap_;
    std::unordered_map<std::string, std::string> scorerInputToStateNameMap_;

    std::string scorerScoresName_;

    std::string initializerEncoderStatesName_;
    std::string initializerEncoderStatesSizeName_;

    std::string updaterEncoderStatesName_;
    std::string updaterEncoderStatesSizeName_;
    std::string updaterTokenName_;

    // Store the onnx values with all encoder states and lengths inside so that it doesn't have to be recomputed every time
    Onnx::Value encoderStatesValue_;
    Onnx::Value encoderStatesSizeValue_;

    Core::FIFOCache<OnnxHiddenStateScoringContextRef, std::shared_ptr<std::vector<Score>>, ScoringContextHash, ScoringContextEq> scoreCache_;
    Core::FIFOCache<OnnxHiddenStateScoringContextRef, OnnxHiddenStateRef, ScoringContextHash, ScoringContextEq>                  stateCache_;
};

}  // namespace Nn

#endif  // STATEFUL_ONNX_LABEL_SCORER_HH