diff --git a/doc/search_v2.rst b/doc/search_v2.rst index 035ad5cc5..dcd564f8b 100644 --- a/doc/search_v2.rst +++ b/doc/search_v2.rst @@ -736,9 +736,9 @@ built-in way to do this. Built-in label scorer types include: -* ``no-context-onnx``, ``fixed-context-onnx``, ``stateful-onnx``, ``stateful-transducer-onnx``, - ``state-managed-onnx``: forward features (and, depending on type, label history/hidden state) through an ONNX - model. +* ``no-context-onnx``, ``fixed-context-onnx``, ``full-context-onnx``, ``stateful-onnx``, + ``stateful-transducer-onnx``, ``state-managed-onnx``: forward features (and, depending on type, label + history/hidden state) through an ONNX model. * ``encoder-decoder`` / ``encoder-only``: wrap a separate encoder (see :ref:`Encoders` below) that pre-processes features, combined with a decoder label scorer (or no decoder, for encoder-only models). * ``ctc-prefix``: wraps a time-synchronous CTC scorer and derives label-synchronous prefix scores from it @@ -887,6 +887,47 @@ history instead of a recurrent state. loop-updates-history = true onnx-model.session.file = /path/to/transducer_predictor.onnx +full-context-onnx +^^^^^^^^^^^^^^^^^^ + +Like ``fixed-context-onnx``, but forwards the *entire* label history through the ONNX model instead of a +fixed-size window -- the history is never truncated. The acoustic input can be given to the model in one of two +ways, depending on which inputs are mapped in the model's ONNX I/O spec (both can also be mapped at once, in +which case the model receives both): + +* ``input-feature``: only the input feature at the hypothesis' current timestep is fed in, re-selected on every + scoring call. This works incrementally as features arrive. +* ``encoder-states`` / ``encoder-states-size``: the complete input sequence collected so far is fed in as one + tensor, together with its length. Since this requires the whole sequence, scoring only starts once all + features of the segment have been passed (i.e. after ``signalNoMoreFeatures``/``finishSegment``), and the + input buffer is never trimmed. Hypotheses can still only be scored for timesteps that exist in the input sequence. + +Since histories in a batch generally differ in length, the ``history`` input is always padded to the longest history +in the batch, and the mandatory ``history-size`` input tells the model the true length of each entry, so the model +itself is responsible for handling the padding correctly (e.g. via masking). + +* ``onnx-model.io-map.input-feature`` / ``encoder-states`` / ``encoder-states-size`` / ``history`` / + ``history-size`` / ``scores``: ONNX tensor names for the respective inputs/outputs. ``history``, + ``history-size`` and ``scores`` are always required; at least one of ``input-feature`` or ``encoder-states`` + must be mapped as well. +* ``start-label-index`` (int): label index used to seed the very first scoring context, whose history then + consists of this single label. Default ``0``. +* ``blank-updates-history`` / ``silence-updates-history`` / ``loop-updates-history`` / ``vertical-label-transition`` + / ``max-batch-size``: same meaning and defaults as for ``fixed-context-onnx`` above. +* Default :ref:`transition-preset `: ``aed``. + +.. code-block:: ini + + [*.search-algorithm.label-scorer] + type = full-context-onnx + start-label-index = 0 + onnx-model.session.file = /path/to/model.onnx + onnx-model.io-map.encoder-states = encoder_states + onnx-model.io-map.encoder-states-size = encoder_states_size + onnx-model.io-map.history = history + onnx-model.io-map.history-size = history_size + onnx-model.io-map.scores = scores + stateful-onnx ^^^^^^^^^^^^^^ diff --git a/src/Nn/LabelScorer/CMakeLists.txt b/src/Nn/LabelScorer/CMakeLists.txt index 3a39bb84c..d83c77755 100644 --- a/src/Nn/LabelScorer/CMakeLists.txt +++ b/src/Nn/LabelScorer/CMakeLists.txt @@ -22,6 +22,7 @@ if(${MODULE_ONNX}) target_sources( RasrNn PRIVATE FixedContextOnnxLabelScorer.cc + FullContextOnnxLabelScorer.cc NoContextOnnxLabelScorer.cc StatefulOnnxLabelScorer.cc StatefulTransducerOnnxLabelScorer.cc diff --git a/src/Nn/LabelScorer/FixedContextOnnxLabelScorer.cc b/src/Nn/LabelScorer/FixedContextOnnxLabelScorer.cc index 78c4dca56..08d58dcec 100644 --- a/src/Nn/LabelScorer/FixedContextOnnxLabelScorer.cc +++ b/src/Nn/LabelScorer/FixedContextOnnxLabelScorer.cc @@ -21,34 +21,6 @@ namespace Nn { -SeqStepScoringContext::SeqStepScoringContext() - : labelSeq(), currentStep(0ul) {} - -SeqStepScoringContext::SeqStepScoringContext(std::vector const& seq, Speech::TimeframeIndex step) - : labelSeq(seq), currentStep(step) {} - -SeqStepScoringContext::SeqStepScoringContext(std::vector&& seq, Speech::TimeframeIndex step) - : labelSeq(std::move(seq)), currentStep(step) {} - -bool SeqStepScoringContext::isEqual(ScoringContextRef const& other) const { - auto* otherPtr = dynamic_cast(other.get()); - if (otherPtr == nullptr) { - return false; - } - - if (currentStep != otherPtr->currentStep) { - return false; - } - - return labelSeqEqual(labelSeq, otherPtr->labelSeq); -} - -size_t SeqStepScoringContext::hash() const { - return Core::combineHashes(currentStep, labelSeqHash(labelSeq)); -} - -typedef Core::Ref SeqStepScoringContextRef; - const Core::ParameterInt FixedContextOnnxLabelScorer::paramStartLabelIndex( "start-label-index", "Initial history in the first step is filled with this label index.", diff --git a/src/Nn/LabelScorer/FixedContextOnnxLabelScorer.hh b/src/Nn/LabelScorer/FixedContextOnnxLabelScorer.hh index bca09f7ed..916a8265d 100644 --- a/src/Nn/LabelScorer/FixedContextOnnxLabelScorer.hh +++ b/src/Nn/LabelScorer/FixedContextOnnxLabelScorer.hh @@ -23,23 +23,6 @@ namespace Nn { -/* - * Scoring context that describes a sequence of previously observed labels as well as the current decoding step - */ -struct SeqStepScoringContext : public ScoringContext { - std::vector labelSeq; - Speech::TimeframeIndex currentStep; - - SeqStepScoringContext(); - SeqStepScoringContext(std::vector const& seq, Speech::TimeframeIndex step); - SeqStepScoringContext(std::vector&& seq, Speech::TimeframeIndex step); - - bool isEqual(ScoringContextRef const& other) const override; - size_t hash() const override; -}; - -typedef Core::Ref SeqStepScoringContextRef; - /* * Label Scorer that performs scoring by forwarding the input feature at the current timestep together * with a fixed-size sequence of history tokens through an ONNX model. diff --git a/src/Nn/LabelScorer/FullContextOnnxLabelScorer.cc b/src/Nn/LabelScorer/FullContextOnnxLabelScorer.cc new file mode 100644 index 000000000..e4c339e55 --- /dev/null +++ b/src/Nn/LabelScorer/FullContextOnnxLabelScorer.cc @@ -0,0 +1,441 @@ +/** Copyright 2026 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. + */ +#include "FullContextOnnxLabelScorer.hh" + +#include +#include + +#include +#include + +#include "ScoreAccessor.hh" + +namespace Nn { + +const Core::ParameterInt FullContextOnnxLabelScorer::paramStartLabelIndex( + "start-label-index", + "Initial history in the first step is filled with this label index.", + 0); + +const Core::ParameterBool FullContextOnnxLabelScorer::paramBlankUpdatesHistory( + "blank-updates-history", + "Whether previously emitted blank labels should be included in the history.", + false); + +const Core::ParameterBool FullContextOnnxLabelScorer::paramSilenceUpdatesHistory( + "silence-updates-history", + "Whether previously emitted silence labels should be included in the history.", + false); + +const Core::ParameterBool FullContextOnnxLabelScorer::paramLoopUpdatesHistory( + "loop-updates-history", + "Whether in the case of loop transitions every repeated emission should be separately included in the history.", + false); + +const Core::ParameterInt FullContextOnnxLabelScorer::paramMaxBatchSize( + "max-batch-size", + "Max number of histories that can be fed into the ONNX model at once.", + Core::Type::max); + +const Core::ParameterBool FullContextOnnxLabelScorer::paramVerticalLabelTransition( + "vertical-label-transition", + "Whether (non-blank) label transitions should be vertical, i.e. not increase the time step.", + false); + +static const std::vector ioSpec = { + Onnx::IOSpecification{ + "input-feature", + Onnx::IODirection::INPUT, + true, + {Onnx::ValueType::TENSOR}, + {Onnx::ValueDataType::FLOAT}, + {{-1, -2}, {1, -2}}}, // [B, F] or [1, F] + Onnx::IOSpecification{ + "encoder-states", + Onnx::IODirection::INPUT, + true, + {Onnx::ValueType::TENSOR}, + {Onnx::ValueDataType::FLOAT}, + {{1, -1, -2}, {-1, -1, -2}}}, // [1, T, E] or [B, T, E] + Onnx::IOSpecification{ + "encoder-states-size", + Onnx::IODirection::INPUT, + true, + {Onnx::ValueType::TENSOR}, + {Onnx::ValueDataType::INT32}, + {{1}, {-1}}}, // [1] or [B] + Onnx::IOSpecification{ + "history", + Onnx::IODirection::INPUT, + false, + {Onnx::ValueType::TENSOR}, + {Onnx::ValueDataType::INT32}, + {{-1, -2}}}, // [B, historyLength] + Onnx::IOSpecification{ + "history-size", + Onnx::IODirection::INPUT, + false, + {Onnx::ValueType::TENSOR}, + {Onnx::ValueDataType::INT32}, + {{1}, {-1}}}, // [1] or [B] + Onnx::IOSpecification{ + "scores", + Onnx::IODirection::OUTPUT, + false, + {Onnx::ValueType::TENSOR}, + {Onnx::ValueDataType::FLOAT}, + {{-1, -2}}}}; // [B, numLabels] + +FullContextOnnxLabelScorer::FullContextOnnxLabelScorer(Core::Configuration const& config, ModelCache& modelCache) + : Core::Component(config), + Precursor(config, TransitionPresetType::AED), + startLabelIndex_(paramStartLabelIndex(config)), + blankUpdatesHistory_(paramBlankUpdatesHistory(config)), + silenceUpdatesHistory_(paramSilenceUpdatesHistory(config)), + loopUpdatesHistory_(paramLoopUpdatesHistory(config)), + verticalLabelTransition_(paramVerticalLabelTransition(config)), + maxBatchSize_(paramMaxBatchSize(config)), + encoderStatesValue_(), + encoderStatesSizeValue_(), + scoreCache_() { + Core::Configuration modelConfig(config, "onnx-model"); + auto key = modelConfig.getSelection(); + onnxModel_ = modelCache.getOrCreate(key, modelConfig, ioSpec); + + inputFeatureName_ = onnxModel_->mapping.getOnnxName("input-feature"); + encoderStatesName_ = onnxModel_->mapping.getOnnxName("encoder-states"); + encoderStatesSizeName_ = onnxModel_->mapping.getOnnxName("encoder-states-size"); + historyName_ = onnxModel_->mapping.getOnnxName("history"); + historySizeName_ = onnxModel_->mapping.getOnnxName("history-size"); + scoresName_ = onnxModel_->mapping.getOnnxName("scores"); +} + +void FullContextOnnxLabelScorer::resetInternal() { + Precursor::resetInternal(); + onnxSessionTime_.reset(); + tensorMarshallingTime_.reset(); + contextPreparationTime_.reset(); + scoreCache_.clear(); + encoderStatesValue_ = Onnx::Value(); + encoderStatesSizeValue_ = Onnx::Value(); +} + +void FullContextOnnxLabelScorer::addInput(DataView const& input) { + Precursor::addInput(input); + + if (not encoderStatesValue_.empty()) { // Any previously computed encoder-states values are outdated now, so reset them + encoderStatesValue_ = Onnx::Value(); + encoderStatesSizeValue_ = Onnx::Value(); + } +} + +ScoringContextRef FullContextOnnxLabelScorer::getInitialScoringContext() { + auto hist = Core::ref(new SeqStepScoringContext()); + hist->labelSeq.push_back(startLabelIndex_); + return hist; +} + +size_t FullContextOnnxLabelScorer::getMinActiveInputIndex(Core::CollapsedVector const& activeContexts) const { + if (encoderStatesName_ != "" or encoderStatesSizeName_ != "") { + // Full encoder-states sequence is needed for every scoring call, so the input buffer must not be trimmed + return 0u; + } + + auto minTimeIndex = Core::Type::max; + for (auto const& context : activeContexts.internalData()) { + SeqStepScoringContextRef stepHistory(dynamic_cast(context.get())); + minTimeIndex = std::min(minTimeIndex, stepHistory->currentStep); + } + + return minTimeIndex; +} + +void FullContextOnnxLabelScorer::cleanupCaches(Core::CollapsedVector const& activeContexts) { + Precursor::cleanupCaches(activeContexts); + + std::unordered_set 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 FullContextOnnxLabelScorer::extendedScoringContext(ScoringContextRef scoringContext, LabelIndex nextToken, TransitionType transitionType) { + bool pushToken = false; + size_t timeIncrement = 0ul; + switch (transitionType) { + case TransitionType::BLANK_LOOP: + pushToken = blankUpdatesHistory_ and loopUpdatesHistory_; + timeIncrement = 1ul; + break; + case TransitionType::SILENCE_LOOP: + pushToken = silenceUpdatesHistory_ and loopUpdatesHistory_; + timeIncrement = 1ul; + break; + case TransitionType::LABEL_TO_BLANK: + case TransitionType::INITIAL_BLANK: + pushToken = blankUpdatesHistory_; + timeIncrement = 1ul; + break; + case TransitionType::LABEL_TO_SILENCE: + case TransitionType::INITIAL_SILENCE: + pushToken = silenceUpdatesHistory_; + timeIncrement = 1ul; + break; + case TransitionType::LABEL_LOOP: + pushToken = loopUpdatesHistory_; + timeIncrement = not verticalLabelTransition_; + break; + case TransitionType::BLANK_TO_LABEL: + case TransitionType::SILENCE_TO_LABEL: + case TransitionType::LABEL_TO_LABEL: + case TransitionType::INITIAL_LABEL: + case TransitionType::SENTENCE_END: + pushToken = true; + timeIncrement = not verticalLabelTransition_; + break; + default: + error() << "Unknown transition type " << transitionType; + } + + // If context is not going to be modified, return the original one to avoid copying + if (not pushToken and timeIncrement == 0ul) { + return scoringContext; + } + + SeqStepScoringContextRef seqStepScoringContext(dynamic_cast(scoringContext.get())); + + std::vector newLabelSeq; + if (pushToken) { + // Copy the complete old history and append the new token + newLabelSeq.reserve(seqStepScoringContext->labelSeq.size() + 1); + newLabelSeq.insert(newLabelSeq.end(), seqStepScoringContext->labelSeq.begin(), seqStepScoringContext->labelSeq.end()); + newLabelSeq.push_back(nextToken); + } + else { + newLabelSeq = seqStepScoringContext->labelSeq; + } + + return Core::ref(new SeqStepScoringContext(std::move(newLabelSeq), seqStepScoringContext->currentStep + timeIncrement)); +} + +std::vector> FullContextOnnxLabelScorer::computeScoreAccessors(std::vector const& scoringContexts) { + if (scoringContexts.empty()) { + return {}; + } + + if ((encoderStatesName_ != "" or encoderStatesSizeName_ != "") and (expectMoreFeatures_ or bufferSize() == 0)) { + // Only allow scoring once all encoder states have been passed + countUncomputedAccessors(scoringContexts.size()); + return std::vector>(scoringContexts.size(), std::nullopt); + } + + contextPreparationTime_.start(); + + // Cast scoring contexts to concrete types + std::vector seqStepScoringContexts; + seqStepScoringContexts.reserve(scoringContexts.size()); + for (auto const& scoringContext : scoringContexts) { + seqStepScoringContexts.push_back(Core::ref(dynamic_cast(scoringContext.get()))); + } + + /* + * Collect all requests. If a per-frame input-feature is used, requests are grouped by their + * timestep since the fed-in feature differs by step. The batched call otherwise uses the same + * full encoder-states input for every step, so no such grouping is needed in that case. + */ + std::unordered_map> contextsByTimestep; + + for (size_t contextIndex = 0ul; contextIndex < scoringContexts.size(); ++contextIndex) { + auto step = seqStepScoringContexts[contextIndex]->currentStep; + + // If no input exists for this step (e.g. past the end of the segment), this context can't be forwarded + if (not getInput(step)) { + continue; + } + + size_t groupKey = inputFeatureName_ != "" ? step : 0ul; + + contextsByTimestep[groupKey].push_back(contextIndex); + } + + contextPreparationTime_.stop(); + + std::vector> scoreAccessors(scoringContexts.size(), std::nullopt); + + size_t numComputed = 0ul; + + /* + * Iterate over distinct groups + */ + for (auto const& [groupKey, contextIndices] : contextsByTimestep) { + /* + * Identify unique histories that still need session runs + */ + contextPreparationTime_.start(); + std::unordered_set uniqueUncachedContexts; + + for (auto contextIndex : contextIndices) { + if (scoreCache_.find(seqStepScoringContexts[contextIndex]) == scoreCache_.end()) { + // Group by unique context + uniqueUncachedContexts.emplace(seqStepScoringContexts[contextIndex]); + } + } + contextPreparationTime_.stop(); + + numComputed += uniqueUncachedContexts.size(); + + std::vector contextBatch; + contextBatch.reserve(std::min(uniqueUncachedContexts.size(), maxBatchSize_)); + for (auto context : uniqueUncachedContexts) { + contextBatch.push_back(context); + if (contextBatch.size() == maxBatchSize_) { // Batch is full -> forward now + forwardBatch(contextBatch); + contextBatch.clear(); + } + } + forwardBatch(contextBatch); // Forward remaining histories + + // Create score accessors from cache + contextPreparationTime_.start(); + for (auto const& contextIndex : contextIndices) { + auto const& scoreVec = scoreCache_.at(seqStepScoringContexts[contextIndex]); + scoreAccessors[contextIndex] = Core::ref(new VectorScoreAccessor(scoreVec, seqStepScoringContexts[contextIndex]->currentStep)); + } + contextPreparationTime_.stop(); + } + + // Includes contexts served from cache as well as contexts whose input is not available yet + countUncomputedAccessors(scoringContexts.size() - numComputed); + + return scoreAccessors; +} + +std::optional FullContextOnnxLabelScorer::computeScoreAccessor(ScoringContextRef scoringContext) { + return computeScoreAccessors({scoringContext})[0]; +} + +void FullContextOnnxLabelScorer::forwardBatch(std::vector const& scoringContextBatch) { + if (scoringContextBatch.empty()) { + return; + } + + /* + * Create session inputs + */ + tensorMarshallingTime_.start(); + + std::vector> sessionInputs; + + if (inputFeatureName_ != "") { + // All requests in this batch share the same input feature which is set up here + auto inputFeatureDataView = getInput(scoringContextBatch.front()->currentStep); + f32 const* inputFeatureData = inputFeatureDataView->data(); + std::vector inputFeatureShape = {1ul, static_cast(inputFeatureDataView->size())}; + sessionInputs.emplace_back(inputFeatureName_, Onnx::Value::create(inputFeatureData, inputFeatureShape)); + } + + if (encoderStatesName_ != "") { + setupEncoderStatesValue(); + sessionInputs.emplace_back(encoderStatesName_, encoderStatesValue_); + } + if (encoderStatesSizeName_ != "") { + setupEncoderStatesSizeValue(); + sessionInputs.emplace_back(encoderStatesSizeName_, encoderStatesSizeValue_); + } + + // Requests in this batch may have different history lengths, therefore the history tensor is + // padded to the longest one and history-size tells the model the true length of each entry + size_t maxHistoryLength = 0ul; + for (auto const& context : scoringContextBatch) { + maxHistoryLength = std::max(maxHistoryLength, context->labelSeq.size()); + } + + Math::FastMatrix historyMat(maxHistoryLength, scoringContextBatch.size()); + historyMat.fill(0); + + std::vector historySizes(scoringContextBatch.size()); + + // Create batched context input + for (size_t b = 0ul; b < scoringContextBatch.size(); ++b) { + auto const& context = scoringContextBatch[b]; + historySizes[b] = static_cast(context->labelSeq.size()); + if (not context->labelSeq.empty()) { + std::copy(context->labelSeq.begin(), context->labelSeq.end(), &(historyMat.at(0, b))); // Pointer to first element in column b + } + } + + sessionInputs.emplace_back(historyName_, Onnx::Value::create(historyMat, true)); + sessionInputs.emplace_back(historySizeName_, Onnx::Value::create(historySizes)); + + tensorMarshallingTime_.stop(); + + /* + * Run session + */ + std::vector sessionOutputs; + onnxSessionTime_.start(); + onnxModel_->session.run(std::move(sessionInputs), {scoresName_}, sessionOutputs); + onnxSessionTime_.stop(); + + /* + * Put resulting scores into cache map + */ + Core::StopWatch::Scope marshallingTimer(tensorMarshallingTime_); + for (size_t b = 0ul; b < scoringContextBatch.size(); ++b) { + auto scoreVec = std::make_shared>(); + sessionOutputs.front().get(b, *scoreVec); + scoreCache_.emplace(scoringContextBatch[b], scoreVec); + } +} + +void FullContextOnnxLabelScorer::setupEncoderStatesValue() { + if (not encoderStatesValue_.empty()) { + return; + } + + u32 T = bufferSize(); + auto inputFeatureDataView = getInput(0); + + encoderStatesValue_ = Onnx::Value::createEmpty({1l, static_cast(T), static_cast(inputFeatureDataView->size())}); + + for (size_t t = 0ul; t < T; ++t) { + inputFeatureDataView = getInput(t); + std::copy(inputFeatureDataView->data(), inputFeatureDataView->data() + inputFeatureDataView->size(), encoderStatesValue_.data(0, t)); + } +} + +void FullContextOnnxLabelScorer::setupEncoderStatesSizeValue() { + if (not encoderStatesSizeValue_.empty()) { + return; + } + + u32 T = bufferSize(); + + encoderStatesSizeValue_ = Onnx::Value::create(std::vector{static_cast(T)}); +} + +void FullContextOnnxLabelScorer::logScoringBreakdown() const { + statisticsChannel_ << Core::XmlFull("onnx-session-time", onnxSessionTime_.elapsedMilliseconds()); + statisticsChannel_ << Core::XmlFull("tensor-marshalling-time", tensorMarshallingTime_.elapsedMilliseconds()); + statisticsChannel_ << Core::XmlFull("context-preparation-time", contextPreparationTime_.elapsedMilliseconds()); +} + +} // namespace Nn diff --git a/src/Nn/LabelScorer/FullContextOnnxLabelScorer.hh b/src/Nn/LabelScorer/FullContextOnnxLabelScorer.hh new file mode 100644 index 000000000..fbd8ca6dc --- /dev/null +++ b/src/Nn/LabelScorer/FullContextOnnxLabelScorer.hh @@ -0,0 +1,138 @@ +/** Copyright 2026 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 FULL_CONTEXT_ONNX_LABEL_SCORER_HH +#define FULL_CONTEXT_ONNX_LABEL_SCORER_HH + +#include + +#include +#include +#include +#include +#include + +#include "BufferedLabelScorer.hh" +#include "ModelCache.hh" + +namespace Nn { + +/* + * Label scorer that forwards the complete label history together with some form of + * acoustic input through a single ONNX model. + * + * This is stateless with respect to the ONNX model: the full history is supplied + * as input on every scoring call. Unlike FixedContextOnnxLabelScorer, the history + * is never truncated. + * + * Since histories in a batch can have different lengths, the "history" input is padded + * to the longest one in the batch and the "history-size" input tells the model the true + * length of each entry, so the model itself is responsible for handling the padding + * correctly (e.g. via masking). + * + * The acoustic input can be given to the ONNX model in one of two ways, depending on + * which inputs are mapped in the model's ONNX I/O spec: + * - "input-feature": only the single input feature at the hypothesis' current timestep + * is fed in, re-selected on every scoring call. This works incrementally as features + * arrive. + * - "encoder-states"/"encoder-states-size": the complete input sequence collected so far + * is fed in as one tensor, together with its length. Since this requires the whole + * sequence, scoring only starts once all features of the segment have been passed + * (i.e. after `signalNoMoreFeatures`), and the input buffer is never trimmed. + * Both can also be mapped at once, in which case the model receives both the current + * frame and the full sequence. + */ +class FullContextOnnxLabelScorer : public BufferedLabelScorer { +public: + FullContextOnnxLabelScorer(Core::Configuration const& config, ModelCache& modelCache); + virtual ~FullContextOnnxLabelScorer() = default; + + // Add a single input feature to the buffer + void addInput(DataView const& input) override; + + // Initial scoring context contains step 0 and a single start-label-index as history + ScoringContextRef getInitialScoringContext() 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(ScoringContextRef scoringContext, LabelIndex nextToken, TransitionType transitionType) override; + + // Clean up input buffer as well as cached score vectors that are no longer needed + void cleanupCaches(Core::CollapsedVector const& activeContexts) override; + +protected: + // Clear feature buffer and cached scores + void resetInternal() 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 + std::vector> computeScoreAccessors(std::vector const& scoringContexts) override; + + // Uses `computeScoreAccessors` internally with some wrapping for vector packing/expansion + std::optional computeScoreAccessor(ScoringContextRef scoringContext) override; + size_t getMinActiveInputIndex(Core::CollapsedVector const& activeContexts) const override; + +private: + using Precursor = BufferedLabelScorer; + + static const Core::ParameterInt paramStartLabelIndex; + static const Core::ParameterBool paramBlankUpdatesHistory; + static const Core::ParameterBool paramSilenceUpdatesHistory; + static const Core::ParameterBool paramLoopUpdatesHistory; + static const Core::ParameterBool paramVerticalLabelTransition; + static const Core::ParameterInt paramMaxBatchSize; + + size_t startLabelIndex_; + bool blankUpdatesHistory_; + bool silenceUpdatesHistory_; + bool loopUpdatesHistory_; + bool verticalLabelTransition_; + size_t maxBatchSize_; + + std::shared_ptr onnxModel_; + + std::string inputFeatureName_; + std::string encoderStatesName_; + std::string encoderStatesSizeName_; + std::string historyName_; + std::string historySizeName_; + std::string scoresName_; + + // 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::StopWatch onnxSessionTime_; + // Building the batched session inputs (incl. encoder states and padded histories) and copying the resulting scores into the cache + Core::StopWatch tensorMarshallingTime_; + Core::StopWatch contextPreparationTime_; + + std::unordered_map>, ScoringContextHash, ScoringContextEq> scoreCache_; + + // Forward a batch of histories through the ONNX model and put the resulting scores into the score cache. + // The history tensor is padded to the longest one and history-size contains the true length of each entry. + // If a per-frame input-feature is used, all histories in the batch must be based on the same timestep. + void forwardBatch(std::vector const& scoringContextBatch); + + // Set up encoderStatesValue_/encoderStatesSizeValue_ from the full input buffer, unless already cached + void setupEncoderStatesValue(); + void setupEncoderStatesSizeValue(); + + void logScoringBreakdown() const override; +}; + +} // namespace Nn + +#endif // FULL_CONTEXT_ONNX_LABEL_SCORER_HH diff --git a/src/Nn/LabelScorer/ScoringContext.cc b/src/Nn/LabelScorer/ScoringContext.cc index 37b9b17e2..0c1c37105 100644 --- a/src/Nn/LabelScorer/ScoringContext.cc +++ b/src/Nn/LabelScorer/ScoringContext.cc @@ -55,4 +55,29 @@ bool StepScoringContext::isEqual(ScoringContextRef const& other) const { return o != nullptr and currentStep == o->currentStep; } +/* + * ================================ + * ==== SeqStepScoringContext ===== + * ================================ + */ +size_t SeqStepScoringContext::hash() const { + return Core::combineHashes(currentStep, labelSeqHash(labelSeq)); +} + +bool SeqStepScoringContext::isEqual(ScoringContextRef const& other) const { + auto* otherPtr = dynamic_cast(other.get()); + if (otherPtr == nullptr) { + return false; + } + if (currentStep != otherPtr->currentStep) { + return false; + } + if (labelSeq.size() != otherPtr->labelSeq.size()) { + return false; + } + return labelSeqEqual(labelSeq, otherPtr->labelSeq); +} + +typedef Core::Ref FullSeqStepScoringContextRef; + } // namespace Nn diff --git a/src/Nn/LabelScorer/ScoringContext.hh b/src/Nn/LabelScorer/ScoringContext.hh index e4e11b589..8bc77421e 100644 --- a/src/Nn/LabelScorer/ScoringContext.hh +++ b/src/Nn/LabelScorer/ScoringContext.hh @@ -33,8 +33,8 @@ bool labelSeqEqual(std::vector const& lhs, std::vector struct ScoringContext : public Core::ReferenceCounted { virtual ~ScoringContext() = default; - virtual bool isEqual(Core::Ref const& other) const; virtual size_t hash() const; + virtual bool isEqual(Core::Ref const& other) const; }; typedef Core::Ref ScoringContextRef; @@ -63,12 +63,32 @@ struct StepScoringContext : public ScoringContext { StepScoringContext(Speech::TimeframeIndex step) : currentStep(step) {} - bool isEqual(ScoringContextRef const& other) const; size_t hash() const; + bool isEqual(ScoringContextRef const& other) const; }; typedef Core::Ref StepScoringContextRef; +/* + * Scoring context that describes the complete sequence of previously observed labels as well as the current decoding step + */ +struct SeqStepScoringContext : public ScoringContext { + std::vector labelSeq; + Speech::TimeframeIndex currentStep; + + SeqStepScoringContext() + : labelSeq(), currentStep(0ul) {} + SeqStepScoringContext(std::vector const& seq, Speech::TimeframeIndex step) + : labelSeq(seq), currentStep(step) {} + SeqStepScoringContext(std::vector&& seq, Speech::TimeframeIndex step) + : labelSeq(std::move(seq)), currentStep(step) {} + + size_t hash() const override; + bool isEqual(ScoringContextRef const& other) const override; +}; + +typedef Core::Ref SeqStepScoringContextRef; + } // namespace Nn #endif // SCORING_CONTEXT_HH diff --git a/src/Nn/Module.cc b/src/Nn/Module.cc index 190d506e5..7343a73c2 100644 --- a/src/Nn/Module.cc +++ b/src/Nn/Module.cc @@ -35,6 +35,7 @@ #include #include "LabelScorer/FixedContextOnnxLabelScorer.hh" +#include "LabelScorer/FullContextOnnxLabelScorer.hh" #include "LabelScorer/NoContextOnnxLabelScorer.hh" #include "LabelScorer/StateManagedOnnxLabelScorer.hh" #include "LabelScorer/StatefulOnnxLabelScorer.hh" @@ -204,6 +205,13 @@ Module_::Module_() return Core::ref(new FixedContextOnnxLabelScorer(config, modelCache)); }); + // Compute scores by forwarding a single input feature vector together with the full history through an ONNX model + labelScorerFactory_.registerLabelScorer( + "full-context-onnx", + [](Core::Configuration const& config, ModelCache& modelCache) { + return Core::ref(new FullContextOnnxLabelScorer(config, modelCache)); + }); + // Compute scores based on hidden state tensors. labelScorerFactory_.registerLabelScorer( "stateful-onnx",