Source code
Revision control
Copy as Markdown
Other Tools
#include "mel.hpp"
#include "fft.hpp"
#include "backend.hpp"
#include "ggml.h"
#include <cassert>
#include <cmath>
#include <cstring>
#include <algorithm>
#include <vector>
namespace pk {
// CONSTANT epsilon added to std in NeMo normalize_batch (features.py: CONSTANT = 1e-5).
static constexpr double kNormEps = 1e-5;
MelKernel::MelKernel(const ModelLoader& ml) {
const ParakeetConfig& c = ml.config();
n_fft_ = (int)c.n_fft;
hop_ = (int)c.hop_length;
assert((n_fft_ & (n_fft_ - 1)) == 0 && "n_fft must be a power of two for rfft");
n_mels_ = (int)c.n_mels;
n_bins_ = n_fft_ / 2 + 1;
preemph_ = c.preemph;
mag_power_ = c.mag_power;
log_guard_ = c.log_zero_guard;
per_feature_ = (c.normalize == "per_feature");
// --- Window: lift the Hann window from the GGUF and center-pad to n_fft. ---
// NeMo builds torch.hann_window(win_length, periodic=False) and passes it to
// torch.stft with win_length < n_fft, so torch zero-pads the window to n_fft,
// centered: pad (n_fft - win_length) // 2 zeros on each side.
window_.assign(n_fft_, 0.0f);
ggml_tensor* w = ml.tensor("preprocessor.featurizer.window");
if (w) {
std::vector<float> wbuf;
pk::weight_to_host_f32(ml, "preprocessor.featurizer.window", wbuf);
const int wlen = (int)w->ne[0];
const float* wd = wbuf.data();
if (wlen == n_fft_) {
std::memcpy(window_.data(), wd, sizeof(float) * n_fft_);
} else {
assert(wlen <= n_fft_ && "window tensor wider than n_fft — wrong model?");
const int left = (n_fft_ - wlen) / 2;
for (int i = 0; i < wlen && (left + i) < n_fft_; ++i) {
if (left + i < 0) continue;
window_[left + i] = wd[i];
}
}
}
// --- Filterbank: preprocessor.featurizer.fb has ggml ne=[n_bins, n_mels, 1]
// i.e. numpy [1, n_mels, n_bins] row-major -> fb_[m*n_bins + b]. ---
fb_.assign((size_t)n_mels_ * n_bins_, 0.0f);
ggml_tensor* fb = ml.tensor("preprocessor.featurizer.fb");
if (fb) {
std::vector<float> fbuf;
pk::weight_to_host_f32(ml, "preprocessor.featurizer.fb", fbuf);
const float* fd = fbuf.data();
std::memcpy(fb_.data(), fd, sizeof(float) * (size_t)n_mels_ * n_bins_);
}
}
void MelKernel::frame_logmel(const double* frame_in, float* out_col, int out_stride,
std::vector<float>& frame,
std::vector<float>& re, std::vector<float>& im,
std::vector<double>& power) const {
// Window the (already-preemphasized, frame-major, double) samples exactly
// like MelFrontend: wframe[i] = (float)(padded[i] * (double)window[i]).
frame.resize((size_t)n_fft_);
for (int i = 0; i < n_fft_; ++i)
frame[i] = (float)(frame_in[i] * (double)window_[i]);
pk::rfft(frame, re, im);
// ----- power spectrum -----
// NeMo: x = sqrt(re^2 + im^2) (guard=0 at inference), then x = x^mag_power.
power.resize((size_t)n_bins_);
for (int b = 0; b < n_bins_; ++b) {
const double mag = std::sqrt((double)re[b] * re[b] + (double)im[b] * im[b]);
power[b] = (mag_power_ == 1.0f) ? mag : std::pow(mag, (double)mag_power_);
}
// ----- mel projection + log(mel + guard) -----
for (int m = 0; m < n_mels_; ++m) {
const float* fbm = &fb_[(size_t)m * n_bins_];
double acc = 0.0;
for (int b = 0; b < n_bins_; ++b)
acc += (double)fbm[b] * power[b];
out_col[(size_t)m * out_stride] = (float)std::log(acc + (double)log_guard_);
}
}
MelFrontend::MelFrontend(const ModelLoader& ml) : k_(ml) {}
void MelFrontend::compute(const std::vector<float>& samples,
std::vector<float>& feats, int& n_mels, int& T) const {
if (samples.empty()) { n_mels = k_.n_mels_; T = 0; feats.clear(); return; }
const int S = (int)samples.size();
const int n_fft_ = k_.n_fft_, hop_ = k_.hop_, n_mels_ = k_.n_mels_;
const float preemph_ = k_.preemph_;
// ----- Step 2: Preemphasis -----
// NeMo: x = cat((x[:,0], x[:,1:] - preemph * x[:,:-1])); first sample unchanged.
std::vector<double> x((size_t)S);
if (preemph_ > 0.0f && S > 0) {
x[0] = samples[0];
for (int t = 1; t < S; ++t)
x[t] = (double)samples[t] - (double)preemph_ * (double)samples[t - 1];
} else {
for (int t = 0; t < S; ++t) x[t] = samples[t];
}
// seq_len = floor(S / hop) (NeMo get_seq_len for center=True:
// floor((S + n_fft//2*2 - n_fft) / hop) = floor(S / hop)).
const int seq_len = (S > 0) ? (S / hop_) : 0;
// ----- Step 3: STFT with center=True -----
// Reflect-pad n_fft//2 on each side, then frame with hop, apply the window,
// FFT -> keep bins 0..n_fft/2. Frame count = 1 + floor((S + 2*pad - n_fft)/hop).
const int pad = n_fft_ / 2;
const int padded_len = S + 2 * pad;
const int n_frames = (padded_len >= n_fft_) ? (1 + (padded_len - n_fft_) / hop_) : 0;
T = n_frames;
n_mels = n_mels_;
// Zero-padded signal. NeMo's stft uses pad_mode="constant" (features.py),
// so center=True pads n_fft//2 ZEROS on each side (not reflection).
// The original signal sits at [pad, pad+S); everything else is 0.
std::vector<double> padded((size_t)padded_len, 0.0);
for (int j = 0; j < S; ++j) padded[pad + j] = x[j];
// Power spectrum per frame: power[b][t] = (re^2 + im^2) (mag then ^mag_power).
// We store frame-major scratch; mel is computed directly into feats.
feats.assign((size_t)n_mels_ * T, 0.0f);
std::vector<float> frame, re, im;
std::vector<double> power;
// Each frame's mel column shares the exact per-frame math with the streaming
// front end via MelKernel::frame_logmel (one source of truth). out stride T
// writes feats[m*T + t] (feat-major, inner = time).
for (int t = 0; t < T; ++t) {
const int start = t * hop_;
k_.frame_logmel(&padded[start], &feats[t], T, frame, re, im, power);
}
// ----- Step 7: per-feature normalization (normalize_batch, per_feature) -----
// Per mel-bin (row): mean over the first `seq_len` frames, unbiased std
// (ddof=1), std += CONSTANT, then (x - mean) / std. Frames >= seq_len are
// zeroed (valid_mask). B=1.
if (k_.per_feature_ && T > 0) {
const int valid = std::min(seq_len, T);
for (int m = 0; m < n_mels_; ++m) {
float* row = &feats[(size_t)m * T];
double mean = 0.0;
if (valid > 0) {
for (int t = 0; t < valid; ++t) mean += row[t];
mean /= (double)valid;
}
double var = 0.0;
if (valid > 1) {
for (int t = 0; t < valid; ++t) {
const double d = (double)row[t] - mean;
var += d * d;
}
var /= (double)(valid - 1); // unbiased (ddof=1)
}
double sd = std::sqrt(var); // NaN-guard not needed: valid>1 here
sd += kNormEps;
for (int t = 0; t < T; ++t) {
if (t < valid) row[t] = (float)(((double)row[t] - mean) / sd);
else row[t] = 0.0f; // masked beyond seq_len
}
}
}
}
// ===========================================================================
// StreamingMel — incremental, frame-local log-mel for the streaming front end.
// ===========================================================================
StreamingMel::StreamingMel(const ModelLoader& ml) : k_(ml), pad_(k_.n_fft_ / 2) {
// Incremental causal mel is only equivalent to the full-buffer compute when
// there is no whole-utterance normalization. The only streaming model
// (nvidia/parakeet_realtime_eou_120m-v1) uses normalize="NA"; per_feature
// would need cross-frame stats and must use MelFrontend::compute instead.
assert(!k_.per_feature_ &&
"StreamingMel is frame-local; per_feature needs whole-utterance "
"stats — use MelFrontend::compute for per_feature models");
reset();
}
void StreamingMel::reset() {
samples_seen_ = 0;
emitted_ = 0;
have_prev_ = false;
prev_raw_ = 0.0f;
buf_.clear();
buf_origin_ = 0;
framebuf_.assign((size_t)k_.n_fft_, 0.0);
}
// Build every frame whose right edge index (t*hop + pad) is <= `avail_x` real
// preemph samples, appending each frame's mel column to `out`. `avail_x` is the
// number of real preemph samples conceptually available for THIS pass: during a
// feed it is samples_seen_; during finalize the padding lets all T frames go.
static inline void build_frames(const MelKernel& k, int pad,
const std::vector<double>& buf, long long buf_origin,
long long S /* total real preemph samples */,
int& emitted, int last_frame_excl,
std::vector<double>& framebuf,
std::vector<float>& wframe,
std::vector<float>& re, std::vector<float>& im,
std::vector<double>& power,
std::vector<float>& out, int& n_new) {
const int hop = k.hop_, n_fft = k.n_fft_, n_mels = k.n_mels_;
// Highest frame index whose right edge is real audio: t*hop + pad <= S.
// (For finalize we pass last_frame_excl = T directly.)
int start_frame = emitted;
for (int t = start_frame; t < last_frame_excl; ++t) {
const long long lo = (long long)t * hop - pad; // preemph index of frame start
// Gather n_fft preemph samples [lo, lo+n_fft), zero where out of [0, S).
for (int i = 0; i < n_fft; ++i) {
const long long idx = lo + i;
double v = 0.0;
if (idx >= 0 && idx < S) {
const long long k_off = idx - buf_origin;
// idx must be retained (we never drop a sample a pending frame needs)
v = buf[(size_t)k_off];
}
framebuf[(size_t)i] = v;
}
// Append one mel column (feat-major: we grow `out` by n_mels, one frame
// at a time; the caller transposes/concatenates frame-major below).
const size_t base = out.size();
out.resize(base + (size_t)n_mels);
// Write column with stride 1 here (per-frame contiguous), the wrapper
// reorders into [n_mels, n_new] feat-major after all frames are built.
k.frame_logmel(framebuf.data(), out.data() + base, 1, wframe, re, im, power);
++emitted;
++n_new;
}
}
std::vector<float> StreamingMel::feed(const float* pcm, int n, int& n_new_frames) {
n_new_frames = 0;
// ----- Preemphasis (carried across feeds) -----
// y[0] (very first sample of the stream) = x[0]; otherwise
// y[t] = x[t] - preemph * x[t-1], where x[t-1] may live in the prior feed.
const float preemph = k_.preemph_;
for (int i = 0; i < n; ++i) {
const float raw = pcm[i];
double y;
if (preemph > 0.0f) {
if (!have_prev_) y = raw; // first ever sample
else y = (double)raw - (double)preemph * (double)prev_raw_;
} else {
y = raw;
}
buf_.push_back(y);
prev_raw_ = raw;
have_prev_ = true;
}
samples_seen_ += n;
const long long S = samples_seen_; // real preemph samples available now
// Frames ready mid-stream: t*hop + pad <= S => t <= (S - pad)/hop.
int last_frame_excl = 0;
if (S >= pad_) last_frame_excl = (int)((S - pad_) / k_.hop_) + 1;
std::vector<float> col_major; // [n_new frames] x [n_mels], frame-contiguous
build_frames(k_, pad_, buf_, buf_origin_, S, emitted_, last_frame_excl,
framebuf_, wframe_, re_, im_, power_, col_major, n_new_frames);
// Drop retained preemph samples no longer needed by any future frame. The
// next frame to emit is `emitted_`; its lowest needed preemph index is
// emitted_*hop - pad. Keep from there on. This bounds the buffer to
// ~n_fft + hop samples (NOT the stream length).
const long long keep_from = std::max<long long>(0, (long long)emitted_ * k_.hop_ - pad_);
if (keep_from > buf_origin_) {
const long long drop = keep_from - buf_origin_;
if (drop >= (long long)buf_.size()) { buf_.clear(); }
else buf_.erase(buf_.begin(), buf_.begin() + (size_t)drop);
buf_origin_ = keep_from;
}
return to_feat_major(col_major, n_new_frames);
}
std::vector<float> StreamingMel::finalize(int& n_tail_frames) {
n_tail_frames = 0;
const long long S = samples_seen_;
// Total frames for the full padded clip: T = 1 + floor(S / hop) (even n_fft),
// 0 if no audio at all.
const int T = (S > 0) ? (int)(S / k_.hop_) + 1 : 0;
std::vector<float> col_major;
build_frames(k_, pad_, buf_, buf_origin_, S, emitted_, T,
framebuf_, wframe_, re_, im_, power_, col_major, n_tail_frames);
// All frames emitted; the retained buffer is no longer needed.
buf_.clear();
buf_origin_ = (long long)emitted_ * k_.hop_; // (consistent; not reused)
return to_feat_major(col_major, n_tail_frames);
}
// Reorder the frame-contiguous columns produced by build_frames (n frames each
// of length n_mels, frame-major) into the [n_mels, n] feat-major layout the
// encoder/test expect (out[m*n + t]).
std::vector<float> StreamingMel::to_feat_major(const std::vector<float>& col_major,
int n) const {
const int n_mels = k_.n_mels_;
std::vector<float> out((size_t)n_mels * n);
for (int t = 0; t < n; ++t)
for (int m = 0; m < n_mels; ++m)
out[(size_t)m * n + t] = col_major[(size_t)t * n_mels + m];
return out;
}
} // namespace pk