Source code
Revision control
Copy as Markdown
Other Tools
#include "streaming.hpp"
#include "tokenizer.hpp"
#include "mel.hpp"
#include <algorithm>
#include <cassert>
namespace pk {
StreamingSession::StreamingSession(const ModelLoader& ml, const std::string& target_lang)
: ml_(ml), enc_(ml), pred_(ml), joint_(ml), prompt_(ml) {
const ParakeetConfig& cfg = ml.config();
d_model_ = (int)cfg.d_model;
blank_id_ = (int)cfg.blank_id;
// Resolve the language prompt index for multilingual (nemotron) models. The
// one-hot is constant over time (one language per utterance), so the prompt
// projection is applied per chunk in feed_mel_chunk using this fixed index.
// Non-prompt models leave prompt_index_ = -1 and skip prompt_.apply().
if (cfg.prompt.present) {
// Empty target_lang -> the model default; an unknown locale THROWS
// std::runtime_error (same message as Model::resolve_prompt_index), so a
// typo (e.g. --lang xx) fails loudly instead of silently mis-transcribing.
// Matches the offline path and the parakeet_capi_stream_begin_lang
// contract (NULL + ctx last_error on an unknown locale).
prompt_index_ = cfg.prompt.resolve_index_or_throw(target_lang);
}
// Greedy max symbols per frame, from model metadata (NeMo default 10);
// matches the offline pk::transcribe path in model.cpp.
max_symbols_ = (int)cfg.max_symbols;
assert(joint_.num_durations() == 0 && "StreamingSession is RNN-T only (no TDT durations)");
// Resolve the <EOU>/<EOB> special token ids from the tokenizer pieces (do
// NOT hardcode 1024/1025; read them from the loaded vocab).
const auto& pieces = cfg.tokenizer_pieces;
for (int i = 0; i < (int)pieces.size(); ++i) {
if (pieces[i] == "<EOU>") eou_id_ = i;
else if (pieces[i] == "<EOB>") eob_id_ = i;
}
// Per-encoder-frame stride in seconds: each encoder output frame spans
// hop_length * subsampling_factor input samples (the subsampling downsamples
// the mel-frame time axis by subsampling_factor). time = frame * stride.
const double hop = (double)cfg.hop_length;
const double sub = (double)(cfg.subsampling_factor ? cfg.subsampling_factor : 1);
const double sr = (double)(cfg.sample_rate ? cfg.sample_rate : 16000);
frame_sec_ = (hop * sub) / sr;
frame_sec_f_ = (float)frame_sec_;
reset();
}
void StreamingSession::reset() {
enc_.reset();
state_ = rnnt_decode_init(pred_);
// eou_id_/eob_id_/frame_sec_ are resolved once in the constructor.
enc_frame_ = 0;
non_special_.clear();
text_.clear();
text_taken_ = 0;
last_chunk_had_eou_ = false;
events_.clear();
word_tokens_.clear();
words_.clear();
words_finalized_ = 0;
words_taken_ = 0;
eou_closed_words_ = 0;
}
void StreamingSession::process_emitted(const std::vector<int32_t>& emitted) {
last_chunk_had_eou_ = false;
bool text_changed = false;
for (int32_t tok : emitted) {
if (tok == eou_id_ || tok == eob_id_) {
// EOU/EOB: surface as an event, do NOT add to the text.
EouEvent ev;
ev.token = tok;
ev.is_eob = (tok == eob_id_);
// The frame index is recorded into events_ by feed/finalize (which
// knows the per-token frame); we patch it there. Here we only know
// the token, so push a placeholder the caller will overwrite.
ev.encoder_frame = enc_frame_; // refined by caller via per-token frame
ev.time_sec = ev.encoder_frame * frame_sec_;
events_.push_back(ev);
last_chunk_had_eou_ = true;
} else {
non_special_.push_back(tok);
text_changed = true;
}
}
if (text_changed) {
text_ = detokenize(ml_.config().tokenizer_pieces, non_special_);
}
}
std::vector<int32_t> StreamingSession::feed_mel_chunk(const std::vector<float>& mel_chunk,
int n_frames, bool is_last) {
// 1. Encoder step: the chunk's valid encoder frames, row-major [valid, d_model]
// (d_model fastest) — exactly the orientation rnnt_decode_frames expects.
int n_valid = 0;
std::vector<float> enc_frames = enc_.step(mel_chunk, n_frames, is_last, n_valid);
if (n_valid <= 0) {
last_chunk_had_eou_ = false;
return {};
}
// 1b. Prompt conditioning (nemotron multilingual): project the chunk's
// encoder frames through prompt_kernel for the resolved language before
// the RNN-T decode. The one-hot is constant over time, so applying it
// per chunk is exact (== the offline forward's single application).
// prompt_.apply() wants channels-first [d_model, valid]; enc_.step gives
// time-major [valid, d_model], so transpose in, apply, transpose back.
// No-op for non-prompt models (prompt_.present()==false): enc_frames is
// left byte-identical.
if (prompt_.present()) {
std::vector<float> chunk_cf((size_t)d_model_ * n_valid); // [d_model, valid]
for (int t = 0; t < n_valid; ++t)
for (int c = 0; c < d_model_; ++c)
chunk_cf[(size_t)c * n_valid + t] = enc_frames[(size_t)t * d_model_ + c];
std::vector<float> projected;
prompt_.apply(chunk_cf, d_model_, n_valid, prompt_index_, projected); // [d_model, valid]
for (int t = 0; t < n_valid; ++t)
for (int c = 0; c < d_model_; ++c)
enc_frames[(size_t)t * d_model_ + c] = projected[(size_t)c * n_valid + t];
}
// 2. RNN-T greedy over the new encoder frames, carrying the decoder state
// across chunks (do NOT reset). Appends to state_.hyp and returns the ids
// emitted in this chunk, with their LOCAL frame index in [0, n_valid) and
// per-token TokenInfo (LOCAL frame, max_prob conf, span==1).
const int base_frame = enc_frame_;
std::vector<int32_t> local_frames;
std::vector<TokenInfo> chunk_tokens;
std::vector<int32_t> emitted =
rnnt_decode_frames(pred_, joint_, enc_frames, n_valid, d_model_,
state_, blank_id_, max_symbols_, &local_frames,
&chunk_tokens);
enc_frame_ += n_valid;
// 3. Update text + EOU events; refine each new event's absolute frame index
// from the per-token local frame the decoder reported.
const size_t prev_events = events_.size();
process_emitted(emitted);
// Re-walk emitted to assign the correct absolute frame to each new event,
// and accumulate NON-special tokens (absolute frame) for word grouping.
size_t evi = prev_events;
size_t eou_word_tokens = 0;
for (size_t i = 0; i < emitted.size(); ++i) {
if (emitted[i] == eou_id_ || emitted[i] == eob_id_) {
const int abs_frame = base_frame + (int)local_frames[i];
events_[evi].encoder_frame = abs_frame;
events_[evi].time_sec = abs_frame * frame_sec_;
++evi;
// Tokens are walked in emission order, so this is the word-token
// count the <EOU> closes.
eou_word_tokens = word_tokens_.size();
} else {
TokenInfo ti = chunk_tokens[i];
ti.frame += base_frame; // local -> absolute encoder frame
word_tokens_.push_back(ti);
}
}
// Regroup the accumulated tokens; words before the last (still-open) one are
// final and become available to drain_words(), as are the words closed by an
// <EOU> in this chunk.
regroup_words(/*flush_all=*/false, eou_word_tokens);
// 4. End-of-utterance reset. The realtime EOU model is trained to emit <EOU>
// (end of utterance) / <EOB> (backchannel) and have the decoder START THE
// NEXT UTTERANCE FROM A FRESH STATE — exactly what NeMo's reference
// streaming driver does (examples/voice_agent .../nemo/streaming_asr.py
// NemoStreamingASRService.transcribe -> reset_state() when <EOU>/<EOB>
// appears in the chunk text). Without this the prediction net stays
// conditioned on the just-emitted <EOU> and the joint scores blank on every
// subsequent frame, so the stream goes silent after the first utterance
// (issue #13). We reset the carried RNN-T decoder state — LSTM h/c to zero
// and last_token back to SOS — for the next chunk. We deliberately reset
// ONLY the decoder, not the StreamingEncoder cache: NeMo's reset_state also
// drops the encoder cache, but that was verified to be a no-op for the
// decoded tokens (decoder-only reset == NeMo's full reset_state byte-for-
// byte on multi-utterance clips), so the validated streaming-encoder path is
// left untouched. enc_frame_ keeps running so <EOU> timestamps stay absolute
// in the clip, and state_.hyp keeps the full token record across utterances.
if (last_chunk_had_eou_) {
state_.state = pred_.zero_state();
state_.last_token = -1; // SOS sentinel (nothing emitted yet)
state_.have_token = false;
}
return emitted;
}
std::string StreamingSession::finalize() {
// The end-of-stream tail is flushed by the caller feeding the final buffered
// mel chunk with is_last=true (which keeps the streaming tail frames). At the
// session level there is no further audio to process here, so finalize just
// returns whatever newly-finalized text remains since the last take. NeMo's
// cache-aware streaming behaves identically: the final chunk's incomplete
// right context means a trailing <EOU> is NOT recovered, so we never
// fabricate one — finalize emits only the carried-state tokens already
// decoded.
//
// Word side: the end-of-stream has no further word-start markers, so the
// trailing open word is now final too — regroup with flush_all so it becomes
// available to drain_words().
regroup_words(/*flush_all=*/true);
return take_new_text();
}
void StreamingSession::regroup_words(bool flush_all, size_t eou_word_tokens) {
// Re-run the validated offline grouping over the whole accumulated
// non-special token sequence (it does the punctuation lookahead / refinement
// exactly like the offline transcribe_with_timestamps path). The last word is
// still "open" (its text/end can change when more tokens arrive), so only
// words BEFORE it are considered final mid-stream; flush_all makes every word
// final at end-of-stream.
words_ = group_words(word_tokens_, ml_.config().tokenizer_pieces, frame_sec_f_);
if (words_.empty()) {
words_finalized_ = 0;
} else if (flush_all) {
words_finalized_ = words_.size();
} else {
words_finalized_ = words_.size() - 1;
// An <EOU> ends the utterance: the decoder restarts from a fresh state
// (see feed_mel_chunk), so no later token can extend a word decoded
// before it and those words are final now. Without this the trailing
// word of every utterance is withheld until the NEXT utterance emits a
// token, so each utterance's last word lags by a whole utterance.
// Regroup the closed prefix to count them: the next utterance opens on a
// word-start token, so its words are exactly words_[closed.size()..].
if (eou_word_tokens > 0) {
const std::vector<Word> closed = group_words(
std::vector<TokenInfo>(word_tokens_.begin(),
word_tokens_.begin() + eou_word_tokens),
ml_.config().tokenizer_pieces, frame_sec_f_);
if (closed.size() > eou_closed_words_) {
eou_closed_words_ = closed.size();
}
}
// Applied on every chunk, not just the one that carried the <EOU>: a
// caller feeding several chunks per call only drains once at the end, so
// recomputing the line above would un-finalize the closed word again and
// hold it until the NEXT utterance emits a token.
if (eou_closed_words_ > words_finalized_) {
words_finalized_ = eou_closed_words_;
}
}
// Never "un-finalize" a word we've already handed out.
if (words_finalized_ < words_taken_) words_finalized_ = words_taken_;
}
std::vector<Word> StreamingSession::drain_words() {
std::vector<Word> out;
for (size_t i = words_taken_; i < words_finalized_ && i < words_.size(); ++i)
out.push_back(words_[i]);
words_taken_ = words_finalized_;
return out;
}
std::string StreamingSession::take_new_text() {
if (text_taken_ >= text_.size()) return std::string();
std::string delta = text_.substr(text_taken_);
text_taken_ = text_.size();
return delta;
}
std::vector<EouEvent> StreamingSession::drain_events() {
std::vector<EouEvent> out;
out.swap(events_);
return out;
}
void run_stream_over_pcm(
StreamingSession& sess, const ModelLoader& ml,
const std::vector<float>& pcm16k,
const std::function<void(const std::string&,
const std::vector<EouEvent>&,
const std::vector<Word>&)>& on_chunk,
const std::string& target_lang) {
// target_lang is intentionally unused here: the session already carries its
// resolved prompt index from construction. The parameter exists so callers
// can route a language through a single entry point (Phase 4 C-API/CLI).
(void)target_lang;
// 1. Full-clip mel [n_mels, T] (feat-major inner=T), matching the offline /
// NeMo online_normalization=False reference (normalization over the whole
// clip). The streaming numerics come from the carried encoder/decoder
// caches, not from chunking the front end.
MelFrontend mel_fe(ml);
std::vector<float> mel;
int n_mels = 0, T = 0;
mel_fe.compute(pcm16k, mel, n_mels, T);
if (T <= 0) return;
const int chunk0 = sess.chunk_size_first(); // 9
const int chunk_main = sess.chunk_size(); // 16
const int pre_cache = sess.pre_encode_cache_size(); // 9
// mel[:, lo:hi] in feat-major layout.
auto window = [&](int lo, int hi) {
const int len = hi - lo;
std::vector<float> w((size_t)n_mels * len);
for (int m = 0; m < n_mels; ++m)
for (int t = 0; t < len; ++t)
w[(size_t)m * len + t] = mel[(size_t)m * T + (lo + t)];
return w;
};
int buffer_idx = 0;
bool first = true;
while (buffer_idx < T) {
const int chunk_size = first ? chunk0 : chunk_main;
const int shift = chunk_size; // shift_size == chunk_size here
const int chunk_hi = std::min(buffer_idx + chunk_size, T);
if (chunk_hi - buffer_idx <= 0) break;
const int lo = first ? buffer_idx : std::max(0, buffer_idx - pre_cache);
std::vector<float> win = window(lo, chunk_hi);
const int win_frames = chunk_hi - lo;
const bool is_last = (chunk_hi >= T);
sess.feed_mel_chunk(win, win_frames, is_last);
if (on_chunk) {
std::string nt = sess.take_new_text();
std::vector<EouEvent> ev = sess.drain_events();
std::vector<Word> wd = sess.drain_words();
if (!nt.empty() || !ev.empty() || !wd.empty()) on_chunk(nt, ev, wd);
}
buffer_idx += shift;
first = false;
}
}
} // namespace pk