diff --git a/doc/search_v2.rst b/doc/search_v2.rst index 353a2299f..18156a660 100644 --- a/doc/search_v2.rst +++ b/doc/search_v2.rst @@ -219,9 +219,9 @@ an optional blank label (for CTC/transducer) and optional silence and sentence-e extension candidates need to be compared). Default ``on``. * ``cache-cleanup-interval`` (int): interval (in search steps) after which cached buffered inputs that are no longer needed get freed. Default ``10``. -* ``maximum-stable-delay`` (int): if set, prune away hypotheses that disagree with the current best hypothesis - further back than this many frames. This makes the traceback "stabilize" after at most this many frames, - which is useful for low-latency streaming output. Default: disabled (unbounded). +* ``maximum-stable-delay`` (int): if set, prune hypotheses so that the output older than this many frames can't + change any more, see :ref:`Maximum-stable-delay pruning`. Useful for low-latency streaming output. + Default: disabled (unbounded). * ``maximum-stable-delay-pruning-interval`` (int): how often (in search steps) the above pruning is applied. Default ``10``. Order of operations for one time-synchronous decoding step, assuming two label scorers ``L_1`` and ``L_2`` with @@ -498,9 +498,8 @@ pronunciation) rather than being configurable as a separate parameter, so that i index used for the search tree itself. * ``max-beam-size``, ``max-word-end-beam-size``, ``word-end-score-threshold``, ``num-histogram-bins``, - ``sentence-end-fall-back``, ``recombination-mode``, ``cache-cleanup-interval``, - ``maximum-stable-delay``, ``maximum-stable-delay-pruning-interval``: same meaning and defaults as for - ``tree-timesync-beam-search`` above. + ``sentence-end-fall-back``, ``recombination-mode``, ``cache-cleanup-interval``: same meaning and defaults + as for ``tree-timesync-beam-search`` above. * ``score-threshold``: same meaning as for ``tree-timesync-beam-search`` above. Also interacts with ``length-norm-scale`` the same way as for ``lexiconfree-labelsync-beam-search`` above -- always expressed in un-normalized score units. @@ -609,6 +608,82 @@ listed below. Note that ``minimized-hmm``/``previousBehavior`` is still the code model was trained on fluent speech. Warning: this tree builder is currently experimental and has not been fully tested; some edge cases or input configurations may not be handled correctly yet. +Maximum-stable-delay pruning +^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +In streaming, the stable prefix is the longest partial traceback that is shared by all ongoing hypotheses in search. +This means that regardless of how the search process develops in the future, the stable prefix won't change anymore. + +``maximum-stable-delay`` pruning bounds how much the stable prefix can lag behind the current search step: +The pruning runs every ``maximum-stable-delay-pruning-interval`` steps. After ``t`` decoded frames, the cutoff is +``c = t + 1 - maximum-stable-delay`` and the goal is to make sure that the stable partial traceback ends at ``c`` +or later. It is implemented by ``lexiconfree-timesync-beam-search``, ``tree-timesync-beam-search`` +and the classic ``advanced-tree-search``. + +1. **Reference:** the best hypothesis which has a lemma ending at ``c`` or later in its traceback or is currently + in a pause (i.e. blank/silence segment). The earliest such lemma in the traceback is taken as the pruning root. +2. **Pruning:** every hypothesis that does not also have this pruning root in its traceback is pruned. This guarantees + that the pruning root becomes stable. +3. **Pause splitting:** for every hypothesis that has a pause starting before ``c`` which is ending after ``c`` or still + ongoing: this pause is split into two parts; before and after ``c``. Thus, if the reference contains such a pause, + the partial pause up to ``c`` becomes the pruning root and all hypotheses that share the same partial pause are kept. + The split is only cosmetic; it has no score contributions of its own. The score of the whole pause is attributed + to the remaining part. Only ``lexiconfree-timesync-beam-search`` and ``tree-timesync-beam-search`` split pauses. + +Consequences: + +* A hypothesis in a word that started before the cutoff and still hasn't ended yet doesn't contain a valid pruning root + after the cutoff. Thus, it can't be the reference and is pruned, even if it is the best one. So ``maximum-stable-delay`` + effectively limits the length of words and should be larger than the longest expected word. +* Pauses are not limited this way since a stable partial pause can be split off. +* If no hypothesis exist that fulfills the conditions for being the reference, we take the last finished word of the + current best hypothesis as pruning root and a warning is logged. The delay is exceeded in that case. + +The following examples all use ``maximum-stable-delay = 5`` after 14 decoded frames, i.e. ``c = 10``. Hypotheses +are listed from best to worst. Items are written as ``word[start, end)``, ``…`` marks an item that hasn't ended yet, +``?`` a word that isn't finished yet and ``pause`` a blank or silence segment. + +**Example 1: regular words.** The pruning root is a word ending after the cutoff: + +.. code-block:: text + + hyp traceback decision + A hello[0,6) world[6,12) ?[12,…) reference: best hypothesis with a word ending at c or later + B hello[0,6) world[6,12) pause[12,…) kept: shares the pruning root + C hello[0,6) word[6,11) ?[11,…) pruned: different word + D hello[0,6) world[6,13) ?[13,…) pruned: same word with a different end time is a different traceback item + E yellow[0,6) world[6,12) ?[12,…) pruned: different word earlier in the traceback + +* Pruning root: ``world[6,12)`` +* Common prefix after the pruning: ``hello[0,6) world[6,12)`` + +**Example 2: pause splitting.** The reference is in a pause that started before the cutoff; the pause is split at ``c``: + +.. code-block:: text + + hyp traceback decision + A hello[0,5) pause[5,10) pause[10, …) reference: best hypothesis, in a pause + B hello[0,5) pause[5,10) pause[10,14) kept: same partial pause up to c + C hello[0,5) pause[5,10) pause[10,12) ?[12,…) kept: same partial pause up to c + D hello[0,5) pause[5,8) ?[8,…) pruned: pause ended before c + E yellow[0,5) pause[5,…) pruned: different word earlier in the traceback + F hello[0,5) ?[5,…) pruned: no pause after hello + +* Pruning root: ``pause[5,10)`` +* Common prefix after the pruning: ``hello[0,5) pause[5,10)`` + +**Example 3: no reference.** No hypothesis has a lemma ending at ``c`` or later or is in a pause: + +.. code-block:: text + + hyp traceback decision + A hello[0,5) ?[5,…) best hypothesis, in a word since before c + B hello[0,5) pause[5,7) ?[7,…) kept: has the pruning root + C yellow[0,4) ?[4,…) pruned: doesn't share the pruning root + +* Pruning root: ``hello[0,5)``, the last finished word of the best hypothesis, and a warning is logged +* Common prefix after the pruning: ``hello[0,5)``, which ends before ``c``, so the delay is exceeded + Multiple label scorers and per-stage parameters ------------------------------------------------ @@ -661,8 +736,11 @@ built-in way to do this. Built-in label scorer types include: -* ``no-context-onnx``, ``fixed-context-onnx``, ``stateful-onnx``, ``state-managed-onnx``: forward - features (and, depending on type, label history/hidden state) through an ONNX model. +* ``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. +* ``prefix-speech-lm-onnx``: speech LLM that is prompted with text prompts and the encoder states before + decoding. * ``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 @@ -741,7 +819,7 @@ have no ``transition-preset`` of their own and always report whichever transitio ONNX model configuration ^^^^^^^^^^^^^^^^^^^^^^^^^ -The four ``*-onnx`` label scorers below (and the ``onnx``/``chunked-onnx`` :ref:`encoders ` used by +The ``*-onnx`` label scorers below (and the ``onnx``/``chunked-onnx`` :ref:`encoders ` used by :ref:`encoder-decoder / encoder-only`) all wrap one or more ONNX Runtime sessions and share the same sub-configuration pattern for each wrapped model, under a model-specific selector (e.g. ``onnx-model`` for the single-model scorers): @@ -835,7 +913,8 @@ stateful (recurrent) language model. the respective label types trigger a state update. Default ``false``. * ``max-batch-size`` (int): maximum number of hidden states forwarded through the scorer model at once. Default unbounded. * ``max-cached-score-vectors`` (int): size of an LRU cache mapping scoring contexts to already-computed score - vectors, to avoid recomputation and bound memory use on very long segments. Default ``1000``. + vectors, to avoid recomputation and bound memory use on very long segments. The same size is used for a second + cache that maps scoring contexts to their updated hidden states. Default ``10000``. * Default :ref:`transition-preset `: ``lm``. .. code-block:: ini @@ -846,6 +925,50 @@ stateful (recurrent) language model. state-updater-model.session.file = /path/to/state_update.onnx scorer-model.session.file = /path/to/scorer.onnx +stateful-transducer-onnx +^^^^^^^^^^^^^^^^^^^^^^^^^ + +Transducer counterpart of ``stateful-onnx``: uses the same three ONNX models (state initializer, state updater, +scorer) and the same hidden-state metadata convention, but each scoring context additionally tracks the current +time step, and the scorer receives the input feature of that step alongside the hidden state(s). A typical use +case is a neural transducer whose prediction network has unlimited (recurrent) context, i.e. the counterpart of +``fixed-context-onnx`` without a fixed history window. + +Differences to ``stateful-onnx``: + +* the **state initializer** takes no inputs (in particular no encoder states), it is run once and its output is + reused across segments. +* the **state updater** takes only the next ``token`` (plus the hidden state(s)), no encoder states. +* the **scorer** takes the hidden state(s) and the ``input-feature`` of the context's current time step, with + shape ``[1, F]``. Since all contexts in a scorer batch share the same time step, this single feature row is + passed for the whole batch, so the model has to broadcast it against the ``[B, *]`` hidden states. + +The hidden state is updated lazily. Extending a context with a token only records the token and the state +updater runs (batched) the next time that context is scored. + +* ``state-initializer-model.*`` / ``state-updater-model.*`` / ``scorer-model.*`` : as for ``stateful-onnx``. + Logical I/O names: ``token`` (state updater, required) and ``input-feature`` / ``scores`` (scorer, required). +* ``blank-updates-history`` / ``silence-updates-history`` / ``loop-updates-history`` (bool): whether the + respective label types are fed to the state updater. Default ``false``. +* ``vertical-label-transition`` (bool): as for ``fixed-context-onnx``, whether non-blank label transitions keep + the time step instead of advancing it. Default ``false``. +* ``max-batch-size`` (int): maximum number of scoring contexts forwarded through the state updater and scorer + at once. Default unbounded. +* ``max-cached-score-vectors`` (int): size of the score cache and of the hidden-state cache, as for + ``stateful-onnx``. Default ``10000``. +* Default :ref:`transition-preset `: ``transducer``. + +.. code-block:: ini + + [*.search-algorithm.label-scorer] + type = stateful-transducer-onnx + state-initializer-model.session.file = /path/to/state_init.onnx + state-updater-model.session.file = /path/to/state_update.onnx + state-updater-model.io-map.token = token + scorer-model.session.file = /path/to/scorer.onnx + scorer-model.io-map.input-feature = encoder_frame + scorer-model.io-map.scores = log_probs + state-managed-onnx ^^^^^^^^^^^^^^^^^^^^ @@ -877,6 +1000,57 @@ caches (splitting/merging/rebasing state across beam search steps) without quadr onnx-model.session.file = /path/to/aed_decoder.onnx state-manager.type = transformer +prefix-speech-lm-onnx +^^^^^^^^^^^^^^^^^^^^^ + +Label scorer for speech LLMs (decoder-only LMs conditioned on speech). The LM is first prefilled with the sequence +``[initial prompt tokens, encoder states in LM embedding space, suffix prompt tokens]``, after which the hypothesis +tokens are scored one by one. The encoder states are usually provided by an ``encoder-decoder`` label scorer that +wraps this one as its ``decoder``. Like ``state-managed-onnx``, the KV cache is handled by a ``StateManager``. + +Two ONNX models are required: + +* ``initializer-model.*`` : prefills the prompt. Logical I/O names: ``initial-prompt``, ``initial-prompt-length``, + ``encoder-states``, ``encoder-states-size``, ``suffix-prompt``, ``suffix-prompt-length`` (inputs) and + ``scores`` (output). Everything except ``encoder-states`` and ``scores`` is optional. The state outputs + are the ones declared as state variables in the step model. +* ``step-model.*`` : extends the cached states by one token. Logical I/O names: ``token``, ``token-length``, + ``prefix-length`` and ``scores``, plus the state variables declared in its metadata (see + :ref:`ONNX model configuration`). + +The prompts are given either as text or directly as token ids: + +* ``initial-prompt`` / ``suffix-prompt`` (string): text of the prompt fed before/after the encoder states. + It is tokenized with the tokenizer configured in ``tokenizer.*`` (see :doc:`tokenizer`); the resulting tokens are + logged. Since RASR configuration values cannot contain newlines, lose leading and trailing whitespace and end at + ``#`` or ``"``, the escape sequences ``\n``, ``\t``, ``\r``, ``\\`` and ``\xHH`` (a byte in hex, e.g. ``\x20`` + for a blank) are replaced before tokenization. +* ``initial-prompt-labels`` / ``suffix-prompt-labels`` (int list): token ids of the prompt fed before/after the + encoder states. Can't be combined with the respective text parameter. +* ``tokenizer.*`` : configuration of the tokenizer, only needed if a text prompt is set. Note that special tokens + added by the tokenizer itself (e.g. ``add-special-tokens-on-encode`` of the ``huggingface`` tokenizer) are + inserted into each prompt separately. +* ``state-manager.type`` : as for ``state-managed-onnx``, typically one of the ``transformer*`` variants. +* ``blank-updates-history`` / ``silence-updates-history`` / ``loop-updates-history`` (bool): as above. Default ``false``. +* ``max-batch-size`` (int): maximum number of scoring contexts forwarded through the step model at once. Default unbounded. +* ``max-cached-score-vectors`` (int): size of the score-vector cache, as for ``stateful-onnx``. +* Default :ref:`transition-preset `: ``lm``. + +.. code-block:: ini + + [*.search-algorithm.label-scorer] + type = encoder-decoder + encoder.type = onnx + encoder.onnx-model.session.file = /path/to/speech_encoder.onnx + decoder.type = prefix-speech-lm-onnx + decoder.initializer-model.session.file = /path/to/llm_prefill.onnx + decoder.step-model.session.file = /path/to/llm_step.onnx + decoder.state-manager.type = transformer + decoder.initial-prompt = <|im_start|>user\n + decoder.suffix-prompt = Transcribe the audio.<|im_end|>\n<|im_start|>assistant\n + decoder.tokenizer.type = huggingface + decoder.tokenizer.tokenizer-file = /path/to/tokenizer.json + encoder-decoder / encoder-only ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ @@ -1407,7 +1581,7 @@ Tuning tips * For ``tree-timesync-beam-search``, tune ``max-word-end-beam-size``/``word-end-score-threshold`` independently from the within-word beam if word-end hypotheses are pruned too aggressively (or not enough) relative to within-word hypotheses. -* For low-latency streaming use cases, set ``maximum-stable-delay`` to bound how many frames of output can +* For low-latency streaming use cases, set ``maximum-stable-delay`` (see :ref:`Maximum-stable-delay pruning`) to bound how many frames of output can still change, at the cost of small accuracy loss versus fully offline decoding. * ``recombination-mode = off`` can be used to debug/compare against a search without hypothesis recombination, but is normally left at its default (``on``) since recombination is "free" (it never removes the best diff --git a/src/Nn/LabelScorer/CMakeLists.txt b/src/Nn/LabelScorer/CMakeLists.txt index 4870ce8bb..fe3222a35 100644 --- a/src/Nn/LabelScorer/CMakeLists.txt +++ b/src/Nn/LabelScorer/CMakeLists.txt @@ -21,10 +21,11 @@ target_sources( if(${MODULE_ONNX}) target_sources( RasrNn - PRIVATE FixedContextOnnxLabelScorer.cc NoContextOnnxLabelScorer.cc PRIVATE FixedContextOnnxLabelScorer.cc NoContextOnnxLabelScorer.cc + PrefixSpeechLmOnnxLabelScorer.cc StatefulOnnxLabelScorer.cc + StatefulTransducerOnnxLabelScorer.cc StateManagedOnnxLabelScorer.cc ) endif() @@ -37,5 +38,6 @@ if(${MODULE_PYTHON}) endif() if(${MODULE_ONNX}) - target_link_libraries(RasrNn PUBLIC RasrOnnx) + # RasrTokenizer is only needed by PrefixSpeechLmOnnxLabelScorer + target_link_libraries(RasrNn PUBLIC RasrOnnx RasrTokenizer) endif() diff --git a/src/Nn/LabelScorer/CtcPrefixLabelScorer.hh b/src/Nn/LabelScorer/CtcPrefixLabelScorer.hh index 5ba054cf6..935432ca6 100644 --- a/src/Nn/LabelScorer/CtcPrefixLabelScorer.hh +++ b/src/Nn/LabelScorer/CtcPrefixLabelScorer.hh @@ -16,6 +16,8 @@ #ifndef CTC_PREFIX_LABEL_SCORER_HH #define CTC_PREFIX_LABEL_SCORER_HH +#include + #include "LabelScorer.hh" #include "ModelCache.hh" #include "ScaledLabelScorer.hh" diff --git a/src/Nn/LabelScorer/PrefixSpeechLmOnnxLabelScorer.cc b/src/Nn/LabelScorer/PrefixSpeechLmOnnxLabelScorer.cc new file mode 100644 index 000000000..153d4cdda --- /dev/null +++ b/src/Nn/LabelScorer/PrefixSpeechLmOnnxLabelScorer.cc @@ -0,0 +1,688 @@ +/** 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 "PrefixSpeechLmOnnxLabelScorer.hh" + +#include +#include + +#include +#include +#include +#include + +#include "ScoreAccessor.hh" + +namespace Nn { + +/* + * ==================================== + * === PrefixSpeechLmScoringContext === + * ==================================== + */ +PrefixSpeechLmScoringContext::PrefixSpeechLmScoringContext() + : labelSeq(), + historyLength(0ul), + parent(), + state() { +} + +PrefixSpeechLmScoringContext::PrefixSpeechLmScoringContext(std::vector&& labelSeq, Core::Ref parent) + : labelSeq(std::move(labelSeq)), + historyLength(0ul), + parent(parent), + state() { +} + +size_t PrefixSpeechLmScoringContext::hash() const { + return labelSeqHash(labelSeq); +} + +bool PrefixSpeechLmScoringContext::isEqual(ScoringContextRef const& other) const { + auto* otherPtr = dynamic_cast(other.get()); + if (otherPtr == nullptr) { + return false; + } + + return labelSeqEqual(labelSeq, otherPtr->labelSeq); +} + +/* + * ===================================== + * === PrefixSpeechLmOnnxLabelScorer === + * ===================================== + */ + +static const std::vector initializerIoSpec = { + Onnx::IOSpecification{ + "initial-prompt", + Onnx::IODirection::INPUT, + true, + {Onnx::ValueType::TENSOR}, + {Onnx::ValueDataType::INT32}, + {{-1, -1}, {1, -1}}}, + Onnx::IOSpecification{ + "initial-prompt-length", + Onnx::IODirection::INPUT, + true, + {Onnx::ValueType::TENSOR}, + {Onnx::ValueDataType::INT32}, + {{-1}, {1}}}, + Onnx::IOSpecification{ + "encoder-states", + Onnx::IODirection::INPUT, + false, + {Onnx::ValueType::TENSOR}, + {Onnx::ValueDataType::FLOAT}, + {{-1, -1, -2}, {1, -1, -2}}}, + Onnx::IOSpecification{ + "encoder-states-size", + Onnx::IODirection::INPUT, + true, + {Onnx::ValueType::TENSOR}, + {Onnx::ValueDataType::INT32}, + {{-1}, {1}}}, + Onnx::IOSpecification{ + "suffix-prompt", + Onnx::IODirection::INPUT, + true, + {Onnx::ValueType::TENSOR}, + {Onnx::ValueDataType::INT32}, + {{-1, -1}, {1, -1}}}, + Onnx::IOSpecification{ + "suffix-prompt-length", + Onnx::IODirection::INPUT, + true, + {Onnx::ValueType::TENSOR}, + {Onnx::ValueDataType::INT32}, + {{-1}, {1}}}, + Onnx::IOSpecification{ + "scores", + Onnx::IODirection::OUTPUT, + false, + {Onnx::ValueType::TENSOR}, + {Onnx::ValueDataType::FLOAT}, + {{-1, -2}, {1, -2}}}}; + +static const std::vector stepIoSpec = { + Onnx::IOSpecification{ + "token", + Onnx::IODirection::INPUT, + false, + {Onnx::ValueType::TENSOR}, + {Onnx::ValueDataType::INT32}, + {{-1, -1}, {-1, 1}}}, + Onnx::IOSpecification{ + "token-length", + Onnx::IODirection::INPUT, + true, + {Onnx::ValueType::TENSOR}, + {Onnx::ValueDataType::INT32}, + {{-1}}}, + Onnx::IOSpecification{ + "prefix-length", + Onnx::IODirection::INPUT, + true, + {Onnx::ValueType::TENSOR}, + {Onnx::ValueDataType::INT32}, + {{-1}}}, + Onnx::IOSpecification{ + "scores", + Onnx::IODirection::OUTPUT, + false, + {Onnx::ValueType::TENSOR}, + {Onnx::ValueDataType::FLOAT}, + {{-1, -1, -2}, {-1, 1, -2}}}}; + +const Core::ParameterString PrefixSpeechLmOnnxLabelScorer::paramInitialPrompt( + "initial-prompt", + "Text of the prompt fed into the LM before the input features. Tokenized with the tokenizer configured in the " + "`tokenizer` sub-component. Supports the escape sequences \\n, \\t, \\r, \\\\ and \\xHH. " + "Alternative to `initial-prompt-labels`.", + ""); + +const Core::ParameterString PrefixSpeechLmOnnxLabelScorer::paramSuffixPrompt( + "suffix-prompt", + "Text of the prompt fed into the LM after the input features. Tokenized with the tokenizer configured in the " + "`tokenizer` sub-component. Supports the escape sequences \\n, \\t, \\r, \\\\ and \\xHH. " + "Alternative to `suffix-prompt-labels`.", + ""); + +const Core::ParameterIntVector PrefixSpeechLmOnnxLabelScorer::paramInitialPromptLabels( + "initial-prompt-labels", + "Token ids for the prompt fed into the LM before the input features. Alternative to `initial-prompt`.", + "", + 0, + Core::Type::max, + 0); + +const Core::ParameterIntVector PrefixSpeechLmOnnxLabelScorer::paramSuffixPromptLabels( + "suffix-prompt-labels", + "Token ids for the prompt fed into the LM after the input features. Alternative to `suffix-prompt`.", + "", + 0, + Core::Type::max, + 0); + +const Core::ParameterBool PrefixSpeechLmOnnxLabelScorer::paramBlankUpdatesHistory( + "blank-updates-history", + "Whether previously emitted blank labels should be used to update the recurrent state.", + false); + +const Core::ParameterBool PrefixSpeechLmOnnxLabelScorer::paramSilenceUpdatesHistory( + "silence-updates-history", + "Whether previously emitted silence labels should be used to update the history.", + false); + +const Core::ParameterBool PrefixSpeechLmOnnxLabelScorer::paramLoopUpdatesHistory( + "loop-updates-history", + "Whether loop transitions should update the recurrent state.", + false); + +const Core::ParameterInt PrefixSpeechLmOnnxLabelScorer::paramMaxBatchSize( + "max-batch-size", + "Max number of scoring contexts that can be fed into the step ONNX model at once.", + Core::Type::max); + +const Core::ParameterInt PrefixSpeechLmOnnxLabelScorer::paramMaxCachedScores( + "max-cached-score-vectors", + "Maximum size of cache that maps scoring contexts to scores and state slices.", + 10000); + +PrefixSpeechLmOnnxLabelScorer::PrefixSpeechLmOnnxLabelScorer(Core::Configuration const& config, ModelCache& modelCache) + : Core::Component(config), + Precursor(config, TransitionPresetType::LM), + blankUpdatesHistory_(paramBlankUpdatesHistory(config)), + silenceUpdatesHistory_(paramSilenceUpdatesHistory(config)), + loopUpdatesHistory_(paramLoopUpdatesHistory(config)), + maxBatchSize_(paramMaxBatchSize(config)), + stateManager_(Module::instance().createStateManager(select("state-manager"))), + stateVectorFactory_(Module::instance().createCompressedVectorFactory(select("state-compression"))), + initialContext_(), + initialScores_(), + scoreCache_(paramMaxCachedScores(config)), + stateCache_(scoreCache_.maxSize()) { + Core::Configuration initializerModelConfig(config, "initializer-model"); + Core::Configuration stepModelConfig(config, "step-model"); + + auto initializerKey = initializerModelConfig.getSelection(); + auto stepKey = stepModelConfig.getSelection(); + + initializerOnnxModel_ = modelCache.getOrCreate(initializerKey, initializerModelConfig, initializerIoSpec); + stepOnnxModel_ = modelCache.getOrCreate(stepKey, stepModelConfig, stepIoSpec); + + stateVariables_ = stepOnnxModel_->session.getStateVariablesMetadata(); + + initializerInitialPromptName_ = initializerOnnxModel_->mapping.getOnnxName("initial-prompt"); + initializerInitialPromptLengthName_ = initializerOnnxModel_->mapping.getOnnxName("initial-prompt-length"); + initializerEncoderStatesName_ = initializerOnnxModel_->mapping.getOnnxName("encoder-states"); + initializerEncoderStatesSizeName_ = initializerOnnxModel_->mapping.getOnnxName("encoder-states-size"); + initializerSuffixPromptName_ = initializerOnnxModel_->mapping.getOnnxName("suffix-prompt"); + initializerSuffixPromptLengthName_ = initializerOnnxModel_->mapping.getOnnxName("suffix-prompt-length"); + initializerScoresName_ = initializerOnnxModel_->mapping.getOnnxName("scores"); + stepTokenName_ = stepOnnxModel_->mapping.getOnnxName("token"); + stepTokenLengthName_ = stepOnnxModel_->mapping.getOnnxName("token-length"); + stepPrefixLengthName_ = stepOnnxModel_->mapping.getOnnxName("prefix-length"); + stepScoresName_ = stepOnnxModel_->mapping.getOnnxName("scores"); + + // Only needed to tokenize the prompts, hence not kept beyond the constructor + Tokenizer::TokenizerRef tokenizer; + if (not paramInitialPrompt(config).empty() or not paramSuffixPrompt(config).empty()) { + tokenizer = Tokenizer::Module::instance().tokenizerFactory().createTokenizer(select("tokenizer")); + } + initialPromptLabels_ = getPromptLabels(paramInitialPrompt, paramInitialPromptLabels, tokenizer.get()); + suffixPromptLabels_ = getPromptLabels(paramSuffixPrompt, paramSuffixPromptLabels, tokenizer.get()); + + if (not initialPromptLabels_.empty() and initializerInitialPromptName_.empty()) { + criticalError("Initial prompt is set but the initializer model has no `initial-prompt` input"); + } + if (not suffixPromptLabels_.empty() and initializerSuffixPromptName_.empty()) { + criticalError("Suffix prompt is set but the initializer model has no `suffix-prompt` input"); + } +} + +void PrefixSpeechLmOnnxLabelScorer::addInput(DataView const& input) { + Precursor::addInput(input); + scoreCache_.clear(); + stateCache_.clear(); +} + +ScoringContextRef PrefixSpeechLmOnnxLabelScorer::getInitialScoringContext() { + if (not initialContext_) { + initialContext_ = Core::ref(new PrefixSpeechLmScoringContext()); + } + return initialContext_; +} + +ScoringContextRef PrefixSpeechLmOnnxLabelScorer::extendedScoringContext(ScoringContextRef scoringContext, LabelIndex nextToken, TransitionType transitionType) { + bool updateState = false; + switch (transitionType) { + case TransitionType::BLANK_LOOP: + updateState = blankUpdatesHistory_ and loopUpdatesHistory_; + break; + case TransitionType::SILENCE_LOOP: + updateState = silenceUpdatesHistory_ and loopUpdatesHistory_; + break; + case TransitionType::LABEL_TO_BLANK: + case TransitionType::INITIAL_BLANK: + updateState = blankUpdatesHistory_; + break; + case TransitionType::LABEL_TO_SILENCE: + case TransitionType::INITIAL_SILENCE: + updateState = silenceUpdatesHistory_; + break; + case TransitionType::LABEL_LOOP: + updateState = loopUpdatesHistory_; + break; + case TransitionType::BLANK_TO_LABEL: + case TransitionType::SILENCE_TO_LABEL: + case TransitionType::LABEL_TO_LABEL: + case TransitionType::INITIAL_LABEL: + case TransitionType::SENTENCE_END: + updateState = true; + break; + default: + error() << "Unknown transition type " << transitionType; + } + if (not updateState) { + return scoringContext; + } + + PrefixSpeechLmScoringContextRef context(dynamic_cast(scoringContext.get())); + verify(context); + + std::vector labelSeq(context->labelSeq); + labelSeq.push_back(nextToken); + return Core::ref(new PrefixSpeechLmScoringContext(std::move(labelSeq), context)); +} + +void PrefixSpeechLmOnnxLabelScorer::resetInternal() { + Precursor::resetInternal(); + initializerSessionTime_.reset(); + stepSessionTime_.reset(); + stateMarshallingTime_.reset(); + contextPreparationTime_.reset(); + initialContext_ = PrefixSpeechLmScoringContextRef(); + initialScores_.reset(); + scoreCache_.clear(); + stateCache_.clear(); +} + +std::optional PrefixSpeechLmOnnxLabelScorer::computeScoreAccessor(ScoringContextRef scoringContext) { + return computeScoreAccessors({scoringContext})[0]; +} + +std::vector> PrefixSpeechLmOnnxLabelScorer::computeScoreAccessors(std::vector const& scoringContexts) { + if (scoringContexts.empty()) { + return {}; + } + + std::vector> scoreAccessors(scoringContexts.size(), std::nullopt); + if (not isPrefilled() and (expectMoreFeatures_ or bufferSize() == 0ul)) { + countUncomputedAccessors(scoringContexts.size()); + return scoreAccessors; + } + + setupInitialStates(); + + /* + * Identify unique scoring contexts that still need session runs + */ + contextPreparationTime_.start(); + std::unordered_set uniqueUncachedScoringContexts; + for (auto const& scoringContext : scoringContexts) { + PrefixSpeechLmScoringContextRef context(dynamic_cast(scoringContext.get())); + verify(context); + + if (scoreCache_.contains(context)) { + continue; + } + + verify(not context->labelSeq.empty()); // The empty emitted-label context is scored by setupInitialStates(). + uniqueUncachedScoringContexts.emplace(context); + } + contextPreparationTime_.stop(); + + countUncomputedAccessors(scoringContexts.size() - uniqueUncachedScoringContexts.size()); + + /* + * Fill state and score caches for all uncached scoring contexts + */ + std::vector scoringContextBatch; + scoringContextBatch.reserve(std::min(uniqueUncachedScoringContexts.size(), maxBatchSize_)); + for (auto const& scoringContext : uniqueUncachedScoringContexts) { + scoringContextBatch.push_back(scoringContext); + if (scoringContextBatch.size() == maxBatchSize_) { + cacheStatesAndScores(scoringContextBatch); + scoringContextBatch.clear(); + } + } + cacheStatesAndScores(scoringContextBatch); + + /* + * Assign scores from cache to result vector and states to contexts that are equal to a computed one + */ + contextPreparationTime_.start(); + for (size_t i = 0ul; i < scoringContexts.size(); ++i) { + PrefixSpeechLmScoringContextRef context(dynamic_cast(scoringContexts[i].get())); + if (not context->state) { + auto cachedState = stateCache_.get(context); // Score and state caches evict in sync + verify(cachedState); + verify(context->parent); + context->state = cachedState->get(); + context->historyLength = context->parent->historyLength + 1ul; + } + auto scores = scoreCache_.get(context); + verify(scores); + scoreAccessors[i] = Core::ref(new VectorScoreAccessor(scores->get(), context->labelSeq.size())); + } + contextPreparationTime_.stop(); + + return scoreAccessors; +} + +size_t PrefixSpeechLmOnnxLabelScorer::getMinActiveInputIndex(Core::CollapsedVector const& activeContexts) const { + // All inputs are needed for the prefill and none afterwards + return isPrefilled() ? Core::Type::max : 0ul; +} + +std::vector PrefixSpeechLmOnnxLabelScorer::getPromptLabels(Core::ParameterString const& textParam, + Core::ParameterIntVector const& labelsParam, + Tokenizer::AbstractTokenizer const* tokenizer) const { + std::string text = textParam(config); + auto labels = labelsParam(config); + + if (text.empty()) { + return std::vector(labels.begin(), labels.end()); + } + if (not labels.empty()) { + criticalError("Only one of `%s` and `%s` can be set", textParam.name().c_str(), labelsParam.name().c_str()); + } + + verify(tokenizer); + auto ids = tokenizer->encode(unescapePrompt(text)); + + std::vector result; + result.reserve(ids.size()); + auto logChannel = log("Tokenized %s into %zu tokens:", textParam.name().c_str(), ids.size()); + for (auto id : ids) { + result.push_back(static_cast(id)); + logChannel << " " << id << " (" << tokenizer->tokenIdToStr(id) << ")"; + } + return result; +} + +std::string PrefixSpeechLmOnnxLabelScorer::unescapePrompt(std::string const& text) const { + auto hexValue = [](char c) -> int { + if (c >= '0' and c <= '9') { + return c - '0'; + } + if (c >= 'a' and c <= 'f') { + return c - 'a' + 10; + } + if (c >= 'A' and c <= 'F') { + return c - 'A' + 10; + } + return -1; + }; + + std::string result; + result.reserve(text.size()); + for (size_t i = 0ul; i < text.size(); ++i) { + if (text[i] != '\\') { + result.push_back(text[i]); + continue; + } + if (i + 1ul == text.size()) { + criticalError("Prompt '%s' ends with an incomplete escape sequence", text.c_str()); + break; + } + char c = text[++i]; + switch (c) { + case 'n': result.push_back('\n'); break; + case 't': result.push_back('\t'); break; + case 'r': result.push_back('\r'); break; + case '\\': result.push_back('\\'); break; + case 'x': { + int high = i + 2ul < text.size() ? hexValue(text[i + 1ul]) : -1; + int low = i + 2ul < text.size() ? hexValue(text[i + 2ul]) : -1; + if (high < 0 or low < 0) { + criticalError("Invalid escape sequence '\\x' in prompt '%s', expected two hex digits", text.c_str()); + break; + } + result.push_back(static_cast(high * 16 + low)); + i += 2ul; + break; + } + default: + criticalError("Unknown escape sequence '\\%c' in prompt '%s'", c, text.c_str()); + } + } + return result; +} + +bool PrefixSpeechLmOnnxLabelScorer::isPrefilled() const { + return initialContext_ and initialContext_->state; +} + +void PrefixSpeechLmOnnxLabelScorer::setupInitialStates() { + if (not initialContext_) { + getInitialScoringContext(); + } + if (isPrefilled()) { + if (not scoreCache_.contains(initialContext_)) { + scoreCache_.put(initialContext_, initialScores_); + } + return; + } + + std::vector> inputs; + + if (initializerInitialPromptName_ != "") { + inputs.emplace_back( + initializerInitialPromptName_, + Onnx::Value::create( + initialPromptLabels_.data(), + std::vector{1l, static_cast(initialPromptLabels_.size())})); + if (initializerInitialPromptLengthName_ != "") { + inputs.emplace_back( + initializerInitialPromptLengthName_, + Onnx::Value::create( + std::vector{static_cast(initialPromptLabels_.size())})); + } + } + + inputs.emplace_back(initializerEncoderStatesName_, createEncoderStatesValue()); + if (initializerEncoderStatesSizeName_ != "") { + inputs.emplace_back(initializerEncoderStatesSizeName_, Onnx::Value::create(std::vector{static_cast(bufferSize())})); + } + + if (initializerSuffixPromptName_ != "") { + inputs.emplace_back( + initializerSuffixPromptName_, + Onnx::Value::create( + suffixPromptLabels_.data(), + std::vector{1l, static_cast(suffixPromptLabels_.size())})); + if (initializerSuffixPromptLengthName_ != "") { + inputs.emplace_back( + initializerSuffixPromptLengthName_, + Onnx::Value::create( + std::vector{static_cast(suffixPromptLabels_.size())})); + } + } + + std::vector targets; + targets.reserve(stateVariables_.size() + 1ul); + targets.emplace_back(initializerScoresName_); + for (auto const& stateVariable : stateVariables_) { + targets.emplace_back(stateVariable.output_state_key); + } + + std::vector outputs; + initializerSessionTime_.start(); + initializerOnnxModel_->session.run(std::move(inputs), targets, outputs); + initializerSessionTime_.stop(); + + verify(not stateVariables_.empty()); + verify_ge(outputs.size(), 2ul); + size_t stateLength = 0ul; + // Find time dim and take its size as prompt length + for (size_t d = 1ul; d < stateVariables_.front().shape.size(); ++d) { + if (stateVariables_.front().shape[d] < 0l) { + stateLength = outputs[1].dimSize(d); + break; + } + } + verify_gt(stateLength, 0ul); + + initialScores_ = std::make_shared>(); + outputs.front().get(0, *initialScores_); + + std::vector stateOutputs(std::make_move_iterator(outputs.begin() + 1), std::make_move_iterator(outputs.end())); + std::vector suffixLengths{stateLength}; + stateMarshallingTime_.start(); + auto promptStates = stateManager_->splitStates(stateVariables_, suffixLengths, stateOutputs, *stateVectorFactory_); + stateMarshallingTime_.stop(); + + PrefixSpeechLmScoringContextRef parent; + if (stateManager_->requiresAllParentStates()) { + verify_eq(promptStates.size(), stateLength); + for (size_t i = 0ul; i + 1ul < promptStates.size(); ++i) { + auto context = Core::ref(new PrefixSpeechLmScoringContext(std::vector(), parent)); + context->historyLength = i + 1ul; + context->state = std::make_shared(std::move(promptStates[i])); + parent = context; + } + } + else { + verify_eq(promptStates.size(), 1ul); + } + + initialContext_->historyLength = stateLength; + initialContext_->parent = parent; + initialContext_->state = std::make_shared(std::move(promptStates.back())); + + scoreCache_.put(initialContext_, initialScores_); +} + +Onnx::Value PrefixSpeechLmOnnxLabelScorer::createEncoderStatesValue() const { + u32 T = bufferSize(); + auto inputFeatureDataView = getInput(0); + verify(inputFeatureDataView); + + auto value = Onnx::Value::createEmpty({1l, static_cast(T), static_cast(inputFeatureDataView->size())}); + for (size_t t = 0ul; t < T; ++t) { + inputFeatureDataView = getInput(t); + verify(inputFeatureDataView); + std::copy(inputFeatureDataView->data(), inputFeatureDataView->data() + inputFeatureDataView->size(), value.data(0, t)); + } + return value; +} + +void PrefixSpeechLmOnnxLabelScorer::cacheStatesAndScores(std::vector const& scoringContextBatch) { + if (scoringContextBatch.empty()) { + return; + } + + std::vector prefixLengths; + std::vector prefixStates; + prefixLengths.reserve(scoringContextBatch.size()); + prefixStates.reserve(scoringContextBatch.size()); + + for (auto const& context : scoringContextBatch) { + verify(not context->labelSeq.empty()); + + auto parent = context->parent; + verify(parent); + + if (stateManager_->requiresAllParentStates()) { + size_t prefixLength = parent->historyLength; + prefixLengths.push_back(prefixLength); + prefixStates.resize(prefixStates.size() + prefixLength); + + size_t offset = prefixStates.size() - prefixLength; + for (size_t i = 0ul; i < prefixLength; ++i) { + verify(parent); + verify(parent->state); + prefixStates[offset + prefixLength - i - 1] = parent->state.get(); + parent = parent->parent; + } + } + else { + prefixLengths.push_back(parent->historyLength); + verify(parent->state); + prefixStates.push_back(parent->state.get()); + } + } + + std::vector> inputs; + std::vector targets; + stateMarshallingTime_.start(); + stateManager_->mergeStates(stateVariables_, prefixLengths, prefixStates, inputs, targets); + stateMarshallingTime_.stop(); + + Onnx::Value tokens = Onnx::Value::createEmpty({static_cast(scoringContextBatch.size()), 1l}); + for (size_t i = 0ul; i < scoringContextBatch.size(); ++i) { + verify(not scoringContextBatch[i]->labelSeq.empty()); + tokens.data(i, 0)[0] = static_cast(scoringContextBatch[i]->labelSeq.back()); + } + inputs.emplace_back(stepTokenName_, std::move(tokens)); + + if (stepTokenLengthName_ != "") { + inputs.emplace_back(stepTokenLengthName_, Onnx::Value::create(std::vector(scoringContextBatch.size(), 1))); + } + if (stepPrefixLengthName_ != "") { + std::vector stateLengths(prefixLengths.begin(), prefixLengths.end()); + inputs.emplace_back(stepPrefixLengthName_, Onnx::Value::create(stateLengths)); + } + + targets.emplace(targets.begin(), stepScoresName_); + + std::vector outputs; + stepSessionTime_.start(); + stepOnnxModel_->session.run(std::move(inputs), targets, outputs); + stepSessionTime_.stop(); + + for (size_t b = 0ul; b < scoringContextBatch.size(); ++b) { + auto scores = std::make_shared>(); + outputs.front().get(b, 0, *scores); + scoreCache_.put(scoringContextBatch[b], scores); + } + + std::vector stateOutputs(std::make_move_iterator(outputs.begin() + 1), std::make_move_iterator(outputs.end())); + std::vector suffixLengths(scoringContextBatch.size(), 1ul); + stateMarshallingTime_.start(); + auto splitStates = stateManager_->splitStates(stateVariables_, suffixLengths, stateOutputs, *stateVectorFactory_); + stateMarshallingTime_.stop(); + verify_eq(splitStates.size(), scoringContextBatch.size()); + for (size_t i = 0ul; i < scoringContextBatch.size(); ++i) { + // Not taken from `prefixLengths`, which `mergeStates` may have truncated + scoringContextBatch[i]->historyLength = scoringContextBatch[i]->parent->historyLength + 1ul; + scoringContextBatch[i]->state = std::make_shared(std::move(splitStates[i])); + stateCache_.put(scoringContextBatch[i], scoringContextBatch[i]->state); + if (not stateManager_->requiresAllParentStates()) { + scoringContextBatch[i]->parent.reset(); + } + } +} + +void PrefixSpeechLmOnnxLabelScorer::logScoringBreakdown() const { + statisticsChannel_ << Core::XmlFull("initializer-session-time", initializerSessionTime_.elapsedMilliseconds()); + statisticsChannel_ << Core::XmlFull("step-session-time", stepSessionTime_.elapsedMilliseconds()); + statisticsChannel_ << Core::XmlFull("state-marshalling-time", stateMarshallingTime_.elapsedMilliseconds()); + statisticsChannel_ << Core::XmlFull("context-preparation-time", contextPreparationTime_.elapsedMilliseconds()); +} + +} // namespace Nn diff --git a/src/Nn/LabelScorer/PrefixSpeechLmOnnxLabelScorer.hh b/src/Nn/LabelScorer/PrefixSpeechLmOnnxLabelScorer.hh new file mode 100644 index 000000000..474734d3d --- /dev/null +++ b/src/Nn/LabelScorer/PrefixSpeechLmOnnxLabelScorer.hh @@ -0,0 +1,166 @@ +/** 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 PREFIX_SPEECH_LM_ONNX_LABEL_SCORER_HH +#define PREFIX_SPEECH_LM_ONNX_LABEL_SCORER_HH + +#include +#include + +#include "BufferedLabelScorer.hh" +#include "ModelCache.hh" +#include "ScoringContext.hh" + +namespace Tokenizer { +class AbstractTokenizer; +} // namespace Tokenizer + +namespace Nn { + +/* + * Scoring context for prefix speech LMs. + * `labelSeq` contains only labels emitted by search, not from the prompt and speech-embeddings. + * `historyLength` contains the total LM cache length, including prompt and speech-embedding positions. + * It is only known once the context has been scored and is 0 before. + * Hashing and equality operators are based only on the `labelSeq`. + */ +struct PrefixSpeechLmScoringContext : public ScoringContext { + using StateManager = AbstractStateManager; + using HistoryState = StateManager::HistoryState; + + std::vector labelSeq; + mutable size_t historyLength; + mutable Core::Ref parent; + + mutable std::shared_ptr state; + + PrefixSpeechLmScoringContext(); + PrefixSpeechLmScoringContext(std::vector&& labelSeq, Core::Ref parent); + + bool isEqual(ScoringContextRef const& other) const override; + size_t hash() const override; +}; + +typedef Core::Ref PrefixSpeechLmScoringContextRef; + +/* + * LabelScorer for prefix speech LMs that first initialize with + * [initial prompt tokens, speech encoder states in LM embedding space, suffix prompt tokens] + * before scoring hypothesis tokens one-by-one as usual. + * + * The prompts are given either as text, which is tokenized with the configured tokenizer, or directly as token ids. + * + * The LabelScorer requires two ONNX models, one which takes the prompts and encoder states + * to produce the initial states and score distribution and one which takes a new token and + * existing states to output updated states and scores after incorporating the token. + * The states are managed with a StateManager. + */ +class PrefixSpeechLmOnnxLabelScorer : public BufferedLabelScorer { +public: + PrefixSpeechLmOnnxLabelScorer(Core::Configuration const& config, ModelCache& modelCache); + virtual ~PrefixSpeechLmOnnxLabelScorer() = default; + + void addInput(DataView const& input) override; + + ScoringContextRef getInitialScoringContext() override; + ScoringContextRef extendedScoringContext(ScoringContextRef scoringContext, LabelIndex nextToken, TransitionType transitionType) override; + +protected: + void resetInternal() override; + + std::optional computeScoreAccessor(ScoringContextRef scoringContext) override; + std::vector> computeScoreAccessors(std::vector const& scoringContexts) override; + size_t getMinActiveInputIndex(Core::CollapsedVector const& activeContexts) const override; + +private: + using Precursor = BufferedLabelScorer; + using StateManager = AbstractStateManager; + using HistoryState = StateManager::HistoryState; + + static const Core::ParameterString paramInitialPrompt; + static const Core::ParameterString paramSuffixPrompt; + static const Core::ParameterIntVector paramInitialPromptLabels; + static const Core::ParameterIntVector paramSuffixPromptLabels; + static const Core::ParameterBool paramBlankUpdatesHistory; + static const Core::ParameterBool paramSilenceUpdatesHistory; + static const Core::ParameterBool paramLoopUpdatesHistory; + static const Core::ParameterInt paramMaxBatchSize; + static const Core::ParameterInt paramMaxCachedScores; + + std::vector initialPromptLabels_; + std::vector suffixPromptLabels_; + bool blankUpdatesHistory_; + bool silenceUpdatesHistory_; + bool loopUpdatesHistory_; + size_t maxBatchSize_; + + std::shared_ptr initializerOnnxModel_; + std::shared_ptr stepOnnxModel_; + + std::unique_ptr stateManager_; + std::vector stateVariables_; + CompressedVectorFactoryPtr stateVectorFactory_; + + std::string initializerInitialPromptName_; + std::string initializerInitialPromptLengthName_; + std::string initializerEncoderStatesName_; + std::string initializerEncoderStatesSizeName_; + std::string initializerSuffixPromptName_; + std::string initializerSuffixPromptLengthName_; + std::string initializerScoresName_; + + std::string stepTokenName_; + std::string stepTokenLengthName_; + std::string stepPrefixLengthName_; + std::string stepScoresName_; + + PrefixSpeechLmScoringContextRef initialContext_; + // Kept outside of `scoreCache_` since the prefill can't be recomputed once the inputs are cleaned up + std::shared_ptr> initialScores_; + + // Session runs of the initializer model that prefill the prompt and encoder states + Core::StopWatch initializerSessionTime_; + Core::StopWatch stepSessionTime_; + // Merging the prefix states into the batched session inputs and splitting the outputs back up + Core::StopWatch stateMarshallingTime_; + Core::StopWatch contextPreparationTime_; + + Core::FIFOCache>, ScoringContextHash, ScoringContextEq> scoreCache_; + Core::FIFOCache, ScoringContextHash, ScoringContextEq> stateCache_; + + // Get the token ids of a prompt from either its text parameter (tokenized with `tokenizer`) or its label parameter. + // `tokenizer` may be null if the text parameter is not set. + std::vector getPromptLabels(Core::ParameterString const& textParam, + Core::ParameterIntVector const& labelsParam, + Tokenizer::AbstractTokenizer const* tokenizer) const; + + // Replace the escape sequences \n, \t, \r, \\ and \xHH in a text prompt + std::string unescapePrompt(std::string const& text) const; + + // Whether the prompt and encoder states have been prefilled for the current segment + bool isPrefilled() const; + + void setupInitialStates(); + Onnx::Value createEncoderStatesValue() const; + + // Forward a batch of scoringContexts through the step ONNX model and put the resulting states and scores into the caches + void cacheStatesAndScores(std::vector const& scoringContextBatch); + + void logScoringBreakdown() const override; +}; + +} // namespace Nn + +#endif // PREFIX_SPEECH_LM_ONNX_LABEL_SCORER_HH diff --git a/src/Nn/LabelScorer/StatefulOnnxLabelScorer.cc b/src/Nn/LabelScorer/StatefulOnnxLabelScorer.cc index 309ed522a..87c8a4e6d 100644 --- a/src/Nn/LabelScorer/StatefulOnnxLabelScorer.cc +++ b/src/Nn/LabelScorer/StatefulOnnxLabelScorer.cc @@ -34,21 +34,10 @@ namespace Nn { -OnnxHiddenState::OnnxHiddenState() - : stateValueMap() {} - -OnnxHiddenState::OnnxHiddenState(std::vector&& names, std::vector&& 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])); - } -} - OnnxHiddenStateScoringContext::OnnxHiddenStateScoringContext() : labelSeq(), hiddenState(), requiresFinalize(false) {} -OnnxHiddenStateScoringContext::OnnxHiddenStateScoringContext(std::vector const& labelSeq, OnnxHiddenStateRef state, bool requiresFinalize) +OnnxHiddenStateScoringContext::OnnxHiddenStateScoringContext(std::vector const& labelSeq, Onnx::OnnxNamedValuesRef state, bool requiresFinalize) : labelSeq(labelSeq), hiddenState(state), requiresFinalize(requiresFinalize) {} bool OnnxHiddenStateScoringContext::isEqual(ScoringContextRef const& other) const { @@ -95,19 +84,11 @@ const Core::ParameterInt StatefulOnnxLabelScorer::paramMaxBatchSize( const Core::ParameterInt StatefulOnnxLabelScorer::paramMaxCachedScores( "max-cached-score-vectors", "Maximum size of cache that maps scoring contexts to scores. This prevents memory overflow in case of very long audio segments.", - 1000); - -// Scorer only takes hidden states as input which are not part of the IO spec -const std::vector scorerModelIoSpec = { - Onnx::IOSpecification{ - "scores", - Onnx::IODirection::OUTPUT, - false, - {Onnx::ValueType::TENSOR}, - {Onnx::ValueDataType::FLOAT}, - {{-1, -2}}}}; // [B, V] + 10000); -const std::vector stateInitializerModelIoSpec = { +// The hidden states and the scores output are not part of the IO specs; they are handled by `Onnx::OnnxInitUpdateScoreBundle`. +// The scorer takes nothing but hidden states, so it has no IO spec of its own. +static const std::vector stateInitializerModelIoSpec = { Onnx::IOSpecification{ "encoder-states", Onnx::IODirection::INPUT, @@ -123,7 +104,7 @@ const std::vector stateInitializerModelIoSpec = { {Onnx::ValueDataType::INT32}, {{1}, {-1}}}}; // [1] or [B] -const std::vector stateUpdaterModelIoSpec = { +static const std::vector stateUpdaterModelIoSpec = { Onnx::IOSpecification{ "encoder-states", Onnx::IODirection::INPUT, @@ -153,94 +134,17 @@ StatefulOnnxLabelScorer::StatefulOnnxLabelScorer(Core::Configuration const& conf silenceUpdatesHistory_(paramSilenceUpdatesHistory(config)), loopUpdatesHistory_(paramLoopUpdatesHistory(config)), maxBatchSize_(paramMaxBatchSize(config)), + initUpdateScoreBundle_(config, modelCache, stateInitializerModelIoSpec, stateUpdaterModelIoSpec, {}), initialHiddenState_(), - initializerOutputToStateNameMap_(), - updaterInputToStateNameMap_(), - updaterOutputToStateNameMap_(), - scorerInputToStateNameMap_(), encoderStatesValue_(), encoderStatesSizeValue_(), scoreCache_(paramMaxCachedScores(config)), stateCache_(scoreCache_.maxSize()) { - Core::Configuration initializerModelConfig(config, "state-initializer-model"); - Core::Configuration updaterModelConfig(config, "state-updater-model"); - Core::Configuration scorerModelConfig(config, "scorer-model"); - - auto initializerKey = initializerModelConfig.getSelection(); - auto updaterKey = updaterModelConfig.getSelection(); - auto scorerKey = scorerModelConfig.getSelection(); - - scorerOnnxModel_ = modelCache.getOrCreate(scorerKey, scorerModelConfig, scorerModelIoSpec); - stateInitializerOnnxModel_ = modelCache.getOrCreate(initializerKey, initializerModelConfig, stateInitializerModelIoSpec); - stateUpdaterOnnxModel_ = modelCache.getOrCreate(updaterKey, updaterModelConfig, stateUpdaterModelIoSpec); - - scorerScoresName_ = scorerOnnxModel_->mapping.getOnnxName("scores"); - initializerEncoderStatesName_ = stateInitializerOnnxModel_->mapping.getOnnxName("encoder-states"); - initializerEncoderStatesSizeName_ = stateInitializerOnnxModel_->mapping.getOnnxName("encoder-states-size"); - updaterEncoderStatesName_ = stateUpdaterOnnxModel_->mapping.getOnnxName("encoder-states"); - updaterEncoderStatesSizeName_ = stateUpdaterOnnxModel_->mapping.getOnnxName("encoder-states-size"); - updaterTokenName_ = stateUpdaterOnnxModel_->mapping.getOnnxName("token"); - - auto initializerMetadataKeys = stateInitializerOnnxModel_->session.getCustomMetadataKeys(); - auto updaterMetadataKeys = stateUpdaterOnnxModel_->session.getCustomMetadataKeys(); - auto scorerMetadataKeys = scorerOnnxModel_->session.getCustomMetadataKeys(); - - // Map state initializer outputs to states - std::unordered_set initializerStateNames; - for (auto const& key : initializerMetadataKeys) { - if (stateInitializerOnnxModel_->session.hasOutput(key)) { - auto stateName = stateInitializerOnnxModel_->session.getCustomMetadata(key); - initializerOutputToStateNameMap_.emplace(key, stateName); - initializerStateNames.insert(stateName); - } - } - if (initializerStateNames.empty()) { - error() << "State initializer does not define any hidden states."; - } - - // Map state updater inputs and outputs to states - std::unordered_set updaterStateNames; - for (auto const& key : updaterMetadataKeys) { - if (stateUpdaterOnnxModel_->session.hasInput(key)) { - auto stateName = stateUpdaterOnnxModel_->session.getCustomMetadata(key); - if (initializerStateNames.find(stateName) == initializerStateNames.end()) { - error() << "State updater input " << key << " associated with state " << stateName << " is not present in state initializer"; - } - updaterInputToStateNameMap_.emplace(key, stateName); - } - if (stateUpdaterOnnxModel_->session.hasOutput(key)) { - auto stateName = stateUpdaterOnnxModel_->session.getCustomMetadata(key); - if (initializerStateNames.find(stateName) == initializerStateNames.end()) { - error() << "State updater output " << key << " associated with state " << stateName << " is not present in state initializer"; - } - updaterOutputToStateNameMap_.emplace(key, stateName); - updaterStateNames.insert(stateName); - } - } - if (updaterOutputToStateNameMap_.empty()) { - error() << "State updater does not produce any updated hidden states"; - } - - // In the loop we checked that the updater outputs are a subset of the initializer outputs. - // If they have the same size, they are equal. Otherwise, some initializer outputs - // are not updater outputs. - if (initializerStateNames.size() != updaterStateNames.size()) { - warning() << "State initializer has states that are not updated by the state updater"; - } - - // Map scorer inputs to states - for (auto const& key : scorerMetadataKeys) { - if (scorerOnnxModel_->session.hasInput(key)) { - auto stateName = scorerOnnxModel_->session.getCustomMetadata(key); - if (initializerStateNames.find(stateName) == initializerStateNames.end()) { - error() << "Scorer input " << key << " associated with state " << stateName << " is not present in state initializer"; - } - scorerInputToStateNameMap_.emplace(key, stateName); - } - } - if (scorerInputToStateNameMap_.empty()) { - error() << "Scorer does not take any input"; - } + initializerEncoderStatesName_ = initUpdateScoreBundle_.stateInitializerOnnxName("encoder-states"); + initializerEncoderStatesSizeName_ = initUpdateScoreBundle_.stateInitializerOnnxName("encoder-states-size"); + updaterEncoderStatesName_ = initUpdateScoreBundle_.stateUpdaterOnnxName("encoder-states"); + updaterEncoderStatesSizeName_ = initUpdateScoreBundle_.stateUpdaterOnnxName("encoder-states-size"); + updaterTokenName_ = initUpdateScoreBundle_.stateUpdaterOnnxName("token"); } ScoringContextRef StatefulOnnxLabelScorer::getInitialScoringContext() { @@ -302,7 +206,7 @@ ScoringContextRef StatefulOnnxLabelScorer::extendedScoringContext(ScoringContext void StatefulOnnxLabelScorer::addInput(DataView const& input) { Precursor::addInput(input); - initialHiddenState_ = OnnxHiddenStateRef(); + initialHiddenState_ = Onnx::OnnxNamedValuesRef(); if (not encoderStatesValue_.empty()) { // Any previously computed hidden state values are outdated now so reset them encoderStatesValue_ = Onnx::Value(); @@ -312,8 +216,7 @@ void StatefulOnnxLabelScorer::addInput(DataView const& input) { void StatefulOnnxLabelScorer::resetInternal() { Precursor::resetInternal(); - stateUpdateSessionTime_.reset(); - scorerSessionTime_.reset(); + initUpdateScoreBundle_.resetTimers(); stateMarshallingTime_.reset(); contextPreparationTime_.reset(); stateCache_.clear(); @@ -408,7 +311,7 @@ void StatefulOnnxLabelScorer::cacheScores(std::vectorlabelSeq.empty()) { initialHiddenState = computeInitialHiddenState(); @@ -417,54 +320,33 @@ void StatefulOnnxLabelScorer::cacheScores(std::vector