Source code
Revision control
Copy as Markdown
Other Tools
#pragma once
#include "model_loader.hpp"
#include "streaming_encoder.hpp"
#include "prediction.hpp"
#include "joint.hpp"
#include "rnnt.hpp"
#include "prompt_kernel.hpp"
#include "decode_types.hpp"
#include "transcription.hpp"
#include <functional>
#include <memory>
#include <string>
#include <vector>
#include <cstdint>
namespace pk {
// End-of-utterance / backchannel event emitted by the streaming decoder.
//
// The model nvidia/parakeet_realtime_eou_120m-v1 emits <EOU> (id 1024) and
// <EOB> (id 1025) as regular vocab tokens to mark end-of-utterance /
// backchannel. We surface them as events (stripped from the running text) with
// the encoder frame on which they were emitted and a wall-clock time:
// time_sec = encoder_frame * hop_length * subsampling_factor / sample_rate
// (the subsampled encoder frame stride in seconds), matching NeMo's
// _get_eou_predictions_from_hypotheses timestamp scaling.
struct EouEvent {
int32_t token = 0; // the special token id (<EOU> or <EOB>)
bool is_eob = false; // true for <EOB>, false for <EOU>
int encoder_frame = 0; // encoder-output frame index of the emission
double time_sec = 0.0; // encoder_frame * hop * subsampling / sample_rate
};
// Cache-aware streaming RNN-T session for the pure-RNNT streaming model
// nvidia/parakeet_realtime_eou_120m-v1.
//
// Owns a StreamingEncoder (carries per-layer conv + attention caches across
// chunks) AND the RNN-T greedy decoder state (prediction-net LSTM state, last
// emitted token, accumulated hypothesis). Per chunk:
// 1. StreamingEncoder::step(mel_chunk) -> the chunk's VALID encoder frames
// [valid, d_model] row-major (the cache-aware-equivalent leading frames).
// 2. rnnt_decode_frames over those new frames, carrying the decoder state
// (NOT reset between chunks). Newly emitted token ids are appended to the
// session hypothesis and returned.
//
// Mirrors NeMo's rnnt_decoder_predictions_tensor(..., partial_hypotheses=...):
// the decoder state persists, so feeding the encoder frames in chunks produces
// the identical token sequence as decoding the whole offline encoder output at
// once.
//
// This is the decoder-facing surface only — PCM buffering / the C-API / EOU
// events are Task 7. The caller drives the encoder chunk schedule (mel windows
// already including pre-encode-cache overlap, matching test_streaming_encoder).
class StreamingSession {
public:
// `target_lang` selects the language prompt for multilingual (nemotron)
// prompt-conditioned models (e.g. "en", "de", "auto"); empty -> the model's
// default_lang. It is ignored by non-prompt models (prompt_.present()==false).
// For a prompt model an unknown locale THROWS std::runtime_error (matching the
// offline Model::resolve_prompt_index and the C-API stream_begin_lang
// contract), so a typo fails loudly rather than silently mis-transcribing.
explicit StreamingSession(const ModelLoader& ml, const std::string& target_lang = "");
// Reset the encoder caches AND the decoder state to a fresh stream.
void reset();
// Feed one mel chunk window (row-major [n_mels, n_frames], feat-major
// inner=time — the same orientation StreamingEncoder::step expects, already
// including any pre-encode-cache overlap). Runs the encoder step then the
// RNN-T decode over the new encoder frames. Returns the token ids emitted in
// THIS chunk (and accumulates them into the session hypothesis tokens()).
// Also updates the running stripped text() and the EOU event queue.
//
// is_last selects the encoder keep_all_outputs path (final chunk keeps the
// streaming tail). Defaults to false for mid-stream chunks.
//
// NOTE: this returns token ids (test_streaming_decode relies on that). The
// C-API exposes the newly-finalized *text* via take_new_text() below; the
// text/events surfaces are kept separate so the token-level parity test and
// the text-level streaming API can both be served from one feed call.
std::vector<int32_t> feed_mel_chunk(const std::vector<float>& mel_chunk,
int n_frames, bool is_last = false);
// Flush the end-of-stream tail. Mirrors NeMo's final-chunk keep_all_outputs:
// re-feeds the LAST buffered mel chunk (if the caller did not already mark it
// is_last) so the trailing encoder frames complete, decoding any remaining
// tokens. Returns the newly-finalized text produced by the flush ("" if
// none). It does NOT fabricate an <EOU> NeMo's streaming wouldn't emit: for
// a clip whose final-chunk tail has incomplete right context (parakeet
// streaming drops the trailing <EOU>), finalize emits no extra event — it
// only commits whatever the carried decoder state already produced.
std::string finalize();
// The full accumulated emitted token-id sequence across all fed chunks
// (includes <EOU>/<EOB> ids if any were emitted).
const std::vector<int32_t>& tokens() const { return state_.hyp; }
// Running transcript with <EOU>/<EOB> STRIPPED (specials surface as events,
// not text), detokenized from the non-special emitted tokens.
const std::string& text() const { return text_; }
// The text appended since the previous take_new_text()/feed call — the
// "newly-finalized text". Resets the delta marker. Returns "" if none.
std::string take_new_text();
// True iff the most recent feed_mel_chunk emitted an <EOU>/<EOB> event.
bool last_chunk_had_eou() const { return last_chunk_had_eou_; }
// Move out all EOU/EOB events collected so far (drains the queue).
std::vector<EouEvent> drain_events();
// Peek at the not-yet-drained EOU/EOB events without consuming them (the
// C-API uses the size as a watermark to attribute events to one feed pass
// while leaving the queue intact for a later drain).
const std::vector<EouEvent>& events() const { return events_; }
// Move out the WORDS finalized since the previous drain_words() call, with
// per-word start/end (seconds) and 'min'-aggregate confidence (matching the
// offline pk::group_words / NeMo timestamps=True convention). A word is
// finalized when the NEXT `▁`-token (word-start marker) arrives; the final
// open word is flushed by finalize(). EOU/EOB specials are NOT words (they
// surface via drain_events()). Drains the returned-words queue.
std::vector<Word> drain_words();
// Streaming chunk schedule (delegated from the encoder), so the caller can
// window the mel exactly like test_streaming_encoder.
int chunk_size_first() const { return enc_.chunk_size_first(); }
int chunk_size() const { return enc_.chunk_size(); }
int pre_encode_cache_size() const { return enc_.pre_encode_cache_size(); }
int valid_out_len() const { return enc_.valid_out_len(); }
private:
// Recompute text_ from the non-special tokens in state_.hyp and append any
// brand-new EOU/EOB events found beyond events_seen_tokens_.
void process_emitted(const std::vector<int32_t>& emitted);
const ModelLoader& ml_;
StreamingEncoder enc_;
PredictionNet pred_;
Joint joint_;
PromptKernel prompt_; // post-encoder language conditioning (nemotron)
int prompt_index_ = -1; // resolved language prompt index (-1 if absent)
int d_model_;
int blank_id_;
int max_symbols_;
RnntDecodeState state_;
// EOU/EOB token ids (resolved from the tokenizer pieces; -1 if absent).
int eou_id_ = -1;
int eob_id_ = -1;
double frame_sec_ = 0.0; // hop * subsampling / sample_rate (per enc frame)
int enc_frame_ = 0; // running encoder-output frame counter
// Running stripped transcript + delta marker.
std::vector<int32_t> non_special_; // non-special tokens, for detokenize
std::string text_;
size_t text_taken_ = 0; // byte offset of text_ already returned
bool last_chunk_had_eou_ = false;
std::vector<EouEvent> events_;
// Per-word timestamp accumulation. `word_tokens_` holds the per-token
// TokenInfo (absolute frame, conf, span) for the NON-SPECIAL tokens emitted
// so far, in emission order — the same input pk::group_words consumes
// offline. We regroup the whole accumulated sequence after each chunk and
// surface words that are now FINAL (followed by a later word-start, so their
// text + end-offset can't change). `words_finalized_` counts how many of the
// regrouped words are already considered final; `words_taken_` how many have
// been handed out by drain_words(). finalize() flushes the trailing word.
std::vector<TokenInfo> word_tokens_;
std::vector<Word> words_; // last regrouping of word_tokens_
size_t words_finalized_ = 0; // # of words_ that are final (safe to emit)
size_t words_taken_ = 0; // # of words already returned by drain_words()
// High-water mark of words closed by an <EOU>/<EOB>. regroup_words() runs per
// chunk and recomputes words_finalized_ from scratch, but a caller can feed
// several chunks before draining, so without remembering this the next chunk
// would withhold an already-closed word again.
size_t eou_closed_words_ = 0;
float frame_sec_f_ = 0.0f; // frame_sec as float (group_words uses float)
// Regroup word_tokens_ into words_ and advance words_finalized_ to all but
// the last (still-open) word — flush_all=true (finalize) makes every word
// final, including the trailing one. `eou_word_tokens` is the word_tokens_
// count as of the chunk's last <EOU>/<EOB> (0 if the chunk had none): that
// token closes the utterance, so the words it ends are final as well.
void regroup_words(bool flush_all, size_t eou_word_tokens = 0);
};
// Drive a StreamingSession over a whole 16 kHz mono PCM clip in the model's
// exact cache-aware chunk schedule (the same windowing as test_streaming_decode:
// chunk 0 = chunk_size_first mel frames with no overlap; each later chunk =
// pre_encode_cache_size overlap + chunk_size mel frames; keep_all_outputs only
// on the final chunk). The full-clip mel is computed once (matching NeMo's
// online_normalization=False reference), then fed chunk-by-chunk so the carried
// encoder/decoder caches reproduce NeMo's streaming output EXACTLY.
//
// `on_chunk` (optional) is invoked after each chunk with the newly-finalized
// text (<EOU>/<EOB> stripped), the events emitted in that chunk, and the WORDS
// finalized in that chunk (start/end/conf) — for the CLI's incremental printing.
// The final stripped transcript is sess.text().
//
// This is the authoritative full-clip driver the streaming C-API uses on
// finalize and the CLI uses for --stream. It does not buffer PCM incrementally;
// it processes the supplied clip in one pass (the encoder/decoder are still
// driven chunk-by-chunk with carried state — the streaming numerics, not a
// real-time PCM feeder).
// `target_lang` is currently unused by the driver itself (the session already
// owns its resolved prompt index from construction); it is accepted so callers
// (C-API/CLI, Phase 4) can pass a language through one entry point. Defaults to
// "" (the session's configured language).
void run_stream_over_pcm(
StreamingSession& sess, const ModelLoader& ml,
const std::vector<float>& pcm16k,
const std::function<void(const std::string& new_text,
const std::vector<EouEvent>& chunk_events,
const std::vector<Word>& chunk_words)>& on_chunk
= nullptr,
const std::string& target_lang = "");
} // namespace pk