first
This commit is contained in:
valenti
@@ -0,0 +1,705 @@
|
||||
//! NVIDIA Parakeet TDT 0.6B v3 STT engine via ONNX Runtime (pure Rust).
|
||||
//!
|
||||
//! Implements the full NeMo TDT speech-to-text pipeline in Rust, mirroring the
|
||||
//! reference implementation in `istupakov/onnx-asr`:
|
||||
//!
|
||||
//! ```text
|
||||
//! 16 kHz mono waveform
|
||||
//! └─ nemo128.onnx → mel features [1, 128, T]
|
||||
//! └─ encoder-model → encoder outputs [1, 1024, T']
|
||||
//! └─ decoder_joint (greedy TDT) → token IDs
|
||||
//! └─ vocab.txt → text
|
||||
//! ```
|
||||
//!
|
||||
//! The TDT decoder is transitive: at each encoder frame it predicts a token and
|
||||
//! a "duration" (how many frames to advance). A special blank token (id 8192,
|
||||
//! `<blk>`) means "no token emitted for this frame".
|
||||
//!
|
||||
//! A worker task consumes microphone PCM frames, runs RMS-based Voice Activity
|
||||
//! Detection (VAD), and triggers inference once an end-of-utterance (EOU)
|
||||
//! silence window is detected.
|
||||
|
||||
use crate::config::{EOU_SILENCE_MS, STT_SAMPLE_RATE, VAD_RMS_THRESHOLD};
|
||||
use anyhow::{Context, Result};
|
||||
use ndarray::{Array1, ArrayD, IxDyn};
|
||||
use ort::environment::Environment;
|
||||
use ort::session::builder::GraphOptimizationLevel;
|
||||
use ort::session::{Session, SessionInputValue};
|
||||
use ort::value::{Value, ValueType};
|
||||
use parking_lot::Mutex;
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::mpsc;
|
||||
use tracing::{debug, error, info};
|
||||
|
||||
// --- Model files (relative to the model directory) -------------------------
|
||||
|
||||
const CONFIG_FILE: &str = "config.json";
|
||||
const VOCAB_FILE: &str = "vocab.txt";
|
||||
const PREPROCESSOR_FILE: &str = "nemo128.onnx";
|
||||
/// Quantized encoder (preferred when present: ~2-3x faster on CPU, no external
|
||||
/// data file). Drop this file to fall back to the fp32 encoder.
|
||||
const ENCODER_INT8_FILE: &str = "encoder-model.int8.onnx";
|
||||
const ENCODER_FILE: &str = "encoder-model.onnx";
|
||||
const DECODER_FILE: &str = "decoder_joint-model.onnx";
|
||||
|
||||
// --- ONNX tensor names (fixed by the exported graphs) ----------------------
|
||||
|
||||
const PP_IN_WAVEFORMS: &str = "waveforms";
|
||||
const PP_IN_LENS: &str = "waveforms_lens";
|
||||
const PP_OUT_FEATURES: &str = "features";
|
||||
const PP_OUT_LENS: &str = "features_lens";
|
||||
|
||||
const ENC_IN_SIGNAL: &str = "audio_signal";
|
||||
const ENC_IN_LENGTH: &str = "length";
|
||||
const ENC_OUT_OUTPUTS: &str = "outputs";
|
||||
const ENC_OUT_LENGTHS: &str = "encoded_lengths";
|
||||
|
||||
const DEC_IN_ENCODER: &str = "encoder_outputs";
|
||||
const DEC_IN_TARGETS: &str = "targets";
|
||||
const DEC_IN_TARGET_LENGTH: &str = "target_length";
|
||||
const DEC_IN_STATE1: &str = "input_states_1";
|
||||
const DEC_IN_STATE2: &str = "input_states_2";
|
||||
const DEC_OUT_OUTPUTS: &str = "outputs";
|
||||
const DEC_OUT_STATE1: &str = "output_states_1";
|
||||
const DEC_OUT_STATE2: &str = "output_states_2";
|
||||
|
||||
/// A single decoded turn from the STT engine.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Transcript {
|
||||
/// Token IDs emitted by the TDT decoder.
|
||||
pub token_ids: Vec<i64>,
|
||||
/// Decoded, human-readable transcript.
|
||||
pub text: String,
|
||||
/// Approximate duration of the utterance in milliseconds.
|
||||
pub duration_ms: u64,
|
||||
}
|
||||
|
||||
/// The three ONNX sessions making up the Parakeet TDT pipeline.
|
||||
struct Sessions {
|
||||
/// Keeps the ONNX Runtime environment alive for the lifetime of the sessions.
|
||||
_environment: Environment,
|
||||
preprocessor: Session,
|
||||
encoder: Session,
|
||||
decoder_joint: Session,
|
||||
}
|
||||
|
||||
/// Parakeet TDT ONNX pipeline. Cloning is cheap (Arc + Mutex).
|
||||
#[derive(Clone)]
|
||||
pub struct ParakeetStt {
|
||||
sessions: Arc<Mutex<Sessions>>,
|
||||
/// Token id -> piece, with the SentencePiece word marker (`▁`) replaced by a
|
||||
/// regular space.
|
||||
vocab: Arc<Vec<String>>,
|
||||
/// Index of the blank token (`<blk>`).
|
||||
blank_idx: i64,
|
||||
/// Number of token logits (excludes the trailing duration logits).
|
||||
vocab_size: usize,
|
||||
/// Maximum number of tokens emitted per encoder frame before forcing an advance.
|
||||
max_tokens_per_step: usize,
|
||||
/// Encoder hidden size (1024 for Parakeet TDT 0.6B).
|
||||
d_model: usize,
|
||||
/// Number of decoder prediction-network state layers.
|
||||
state_layers: usize,
|
||||
/// Hidden size of each decoder state layer.
|
||||
state_hidden: usize,
|
||||
}
|
||||
|
||||
impl ParakeetStt {
|
||||
/// Load the Parakeet TDT ONNX pipeline from `model_dir`, which must contain
|
||||
/// `nemo128.onnx`, an `encoder-model*.onnx`, `decoder_joint-model.onnx`,
|
||||
/// `vocab.txt` and `config.json`.
|
||||
pub fn load(model_dir: &Path) -> Result<Self> {
|
||||
info!(dir = %model_dir.display(), "loading Parakeet TDT ONNX models");
|
||||
|
||||
// `config.json` carries `max_tokens_per_step` (default 10).
|
||||
let config_json: serde_json::Value = serde_json::from_slice(
|
||||
&std::fs::read(model_dir.join(CONFIG_FILE))
|
||||
.with_context(|| format!("failed to read {}", model_dir.join(CONFIG_FILE).display()))?,
|
||||
)
|
||||
.context("invalid Parakeet config.json")?;
|
||||
let max_tokens_per_step = config_json
|
||||
.get("max_tokens_per_step")
|
||||
.and_then(|v| v.as_u64())
|
||||
.unwrap_or(10) as usize;
|
||||
|
||||
// Vocabulary: id -> piece, with `▁` (U+2581) mapped to a real space.
|
||||
let vocab = Arc::new(load_vocab(&model_dir.join(VOCAB_FILE))?);
|
||||
let vocab_size = vocab.len();
|
||||
let blank_idx = vocab
|
||||
.iter()
|
||||
.position(|p| p == "<blk>")
|
||||
.map(|i| i as i64)
|
||||
.context("vocab.txt does not contain the <blk> token")?;
|
||||
|
||||
// Prefer the int8 encoder when available; fall back to fp32.
|
||||
let int8_path = model_dir.join(ENCODER_INT8_FILE);
|
||||
let encoder_path = if int8_path.exists() {
|
||||
int8_path
|
||||
} else {
|
||||
model_dir.join(ENCODER_FILE)
|
||||
};
|
||||
|
||||
let environment = crate::onnx::environment()?;
|
||||
|
||||
let preprocessor = build_session(&environment, &model_dir.join(PREPROCESSOR_FILE))?;
|
||||
let encoder = build_session(&environment, &encoder_path)?;
|
||||
let decoder_joint = build_session(&environment, &model_dir.join(DECODER_FILE))?;
|
||||
|
||||
// Introspect the graphs for the dimensions we need at runtime.
|
||||
let d_model = output_dim(&encoder, ENC_OUT_OUTPUTS, 1)?;
|
||||
let decoder_out_dim = output_dim(&decoder_joint, DEC_OUT_OUTPUTS, 3)?;
|
||||
let num_durations = decoder_out_dim
|
||||
.checked_sub(vocab_size)
|
||||
.context("decoder output is smaller than the vocabulary")?;
|
||||
let (state_layers, state_hidden) = input_state_dims(&decoder_joint, DEC_IN_STATE1)?;
|
||||
|
||||
info!(
|
||||
encoder = %encoder_path.display(),
|
||||
vocab_size,
|
||||
blank_idx,
|
||||
num_durations,
|
||||
d_model,
|
||||
state_layers,
|
||||
state_hidden,
|
||||
"Parakeet TDT pipeline loaded"
|
||||
);
|
||||
|
||||
Ok(Self {
|
||||
sessions: Arc::new(Mutex::new(Sessions {
|
||||
_environment: environment,
|
||||
preprocessor,
|
||||
encoder,
|
||||
decoder_joint,
|
||||
})),
|
||||
vocab,
|
||||
blank_idx,
|
||||
vocab_size,
|
||||
max_tokens_per_step,
|
||||
d_model,
|
||||
state_layers,
|
||||
state_hidden,
|
||||
})
|
||||
}
|
||||
|
||||
/// Run the full pipeline on a complete utterance (16 kHz mono float PCM
|
||||
/// normalized to [-1, 1]). Returns the emitted token IDs and decoded text.
|
||||
pub fn infer(&self, audio: &[f32]) -> Result<(Vec<i64>, String)> {
|
||||
let n = audio.len();
|
||||
if n == 0 {
|
||||
return Ok((Vec::new(), String::new()));
|
||||
}
|
||||
|
||||
let mut guard = self.sessions.lock();
|
||||
let sessions = &mut *guard;
|
||||
|
||||
// --- 1. Mel-feature preprocessor (nemo128.onnx) ---------------------
|
||||
let (feature_shape, feature_data, feature_len) = {
|
||||
let waveforms = Value::from_array(ArrayD::from_shape_vec(IxDyn(&[1, n]), audio.to_vec())?)
|
||||
.context("failed to build waveform tensor")?;
|
||||
let waveforms_lens = Value::from_array(Array1::from_vec(vec![n as i64]))
|
||||
.context("failed to build waveform length tensor")?;
|
||||
|
||||
let outputs = sessions
|
||||
.preprocessor
|
||||
.run(vec![
|
||||
(PP_IN_WAVEFORMS.to_string(), SessionInputValue::from(waveforms)),
|
||||
(PP_IN_LENS.to_string(), SessionInputValue::from(waveforms_lens)),
|
||||
])
|
||||
.context("nemo128 preprocessor inference failed")?;
|
||||
|
||||
let (shape, data) = outputs
|
||||
.get(PP_OUT_FEATURES)
|
||||
.context("preprocessor produced no 'features' output")?
|
||||
.try_extract_tensor::<f32>()
|
||||
.context("failed to extract preprocessor features")?;
|
||||
let (_, lens) = outputs
|
||||
.get(PP_OUT_LENS)
|
||||
.context("preprocessor produced no 'features_lens' output")?
|
||||
.try_extract_tensor::<i64>()
|
||||
.context("failed to extract preprocessor feature lengths")?;
|
||||
|
||||
let shape: Vec<usize> = shape.iter().map(|&d| d as usize).collect();
|
||||
(shape, data.to_vec(), lens[0])
|
||||
};
|
||||
|
||||
// --- 2. Encoder -----------------------------------------------------
|
||||
// Encoder output layout is [1, d_model, T'] (already transposed by the
|
||||
// exported graph). We re-pack it to [T', d_model] so each frame is
|
||||
// contiguous for the decoder loop.
|
||||
let (frames, encoded_len) = {
|
||||
let signal = Value::from_array(
|
||||
ArrayD::from_shape_vec(IxDyn(&feature_shape), feature_data)
|
||||
.context("failed to reshape mel features")?,
|
||||
)
|
||||
.context("failed to build encoder input tensor")?;
|
||||
let length = Value::from_array(Array1::from_vec(vec![feature_len]))
|
||||
.context("failed to build encoder length tensor")?;
|
||||
|
||||
let outputs = sessions
|
||||
.encoder
|
||||
.run(vec![
|
||||
(ENC_IN_SIGNAL.to_string(), SessionInputValue::from(signal)),
|
||||
(ENC_IN_LENGTH.to_string(), SessionInputValue::from(length)),
|
||||
])
|
||||
.context("encoder inference failed")?;
|
||||
|
||||
let (shape, data) = outputs
|
||||
.get(ENC_OUT_OUTPUTS)
|
||||
.context("encoder produced no 'outputs'")?
|
||||
.try_extract_tensor::<f32>()
|
||||
.context("failed to extract encoder outputs")?;
|
||||
let (_, lens) = outputs
|
||||
.get(ENC_OUT_LENGTHS)
|
||||
.context("encoder produced no 'encoded_lengths'")?
|
||||
.try_extract_tensor::<i64>()
|
||||
.context("failed to extract encoded lengths")?;
|
||||
|
||||
let t = shape[2] as usize;
|
||||
let mut frames = vec![0f32; t * self.d_model];
|
||||
for tt in 0..t {
|
||||
for d in 0..self.d_model {
|
||||
frames[tt * self.d_model + d] = data[d * t + tt];
|
||||
}
|
||||
}
|
||||
// Clamp: low-precision encoders may report more frames than present.
|
||||
let encoded_len = (lens[0] as usize).min(t);
|
||||
(frames, encoded_len)
|
||||
};
|
||||
|
||||
// --- 3. Greedy TDT decoding ----------------------------------------
|
||||
let mut tokens: Vec<i64> = Vec::new();
|
||||
let mut state1 = vec![0f32; self.state_layers * self.state_hidden];
|
||||
let mut state2 = state1.clone();
|
||||
let mut prev_token = self.blank_idx;
|
||||
let mut emitted = 0usize;
|
||||
let mut t = 0usize;
|
||||
|
||||
while t < encoded_len {
|
||||
let frame = &frames[t * self.d_model..(t + 1) * self.d_model];
|
||||
|
||||
let enc_frame = Value::from_array(
|
||||
ArrayD::from_shape_vec(IxDyn(&[1, self.d_model, 1]), frame.to_vec())?,
|
||||
)
|
||||
.context("failed to build decoder encoder-frame tensor")?;
|
||||
let targets = Value::from_array(ArrayD::from_shape_vec(
|
||||
IxDyn(&[1, 1]),
|
||||
vec![prev_token as i32],
|
||||
)?)
|
||||
.context("failed to build decoder targets tensor")?;
|
||||
let target_length = Value::from_array(Array1::from_vec(vec![1i32]))
|
||||
.context("failed to build decoder target-length tensor")?;
|
||||
let s1 = Value::from_array(ArrayD::from_shape_vec(
|
||||
IxDyn(&[self.state_layers, 1, self.state_hidden]),
|
||||
state1.clone(),
|
||||
)?)
|
||||
.context("failed to build decoder state1 tensor")?;
|
||||
let s2 = Value::from_array(ArrayD::from_shape_vec(
|
||||
IxDyn(&[self.state_layers, 1, self.state_hidden]),
|
||||
state2.clone(),
|
||||
)?)
|
||||
.context("failed to build decoder state2 tensor")?;
|
||||
|
||||
let (token, step, new_state1, new_state2) = {
|
||||
let outputs = sessions
|
||||
.decoder_joint
|
||||
.run(vec![
|
||||
(DEC_IN_ENCODER.to_string(), SessionInputValue::from(enc_frame)),
|
||||
(DEC_IN_TARGETS.to_string(), SessionInputValue::from(targets)),
|
||||
(DEC_IN_TARGET_LENGTH.to_string(), SessionInputValue::from(target_length)),
|
||||
(DEC_IN_STATE1.to_string(), SessionInputValue::from(s1)),
|
||||
(DEC_IN_STATE2.to_string(), SessionInputValue::from(s2)),
|
||||
])
|
||||
.context("decoder_joint inference failed")?;
|
||||
|
||||
let (_, logits) = outputs
|
||||
.get(DEC_OUT_OUTPUTS)
|
||||
.context("decoder produced no 'outputs'")?
|
||||
.try_extract_tensor::<f32>()
|
||||
.context("failed to extract decoder logits")?;
|
||||
if logits.len() < self.vocab_size {
|
||||
anyhow::bail!("decoder logits shorter than vocabulary");
|
||||
}
|
||||
// Token logits come first, then the duration logits (TDT).
|
||||
let token = argmax(&logits[..self.vocab_size]) as i64;
|
||||
let step = argmax(&logits[self.vocab_size..]);
|
||||
|
||||
let (_, next_s1) = outputs
|
||||
.get(DEC_OUT_STATE1)
|
||||
.context("decoder produced no 'output_states_1'")?
|
||||
.try_extract_tensor::<f32>()
|
||||
.context("failed to extract decoder state1")?;
|
||||
let (_, next_s2) = outputs
|
||||
.get(DEC_OUT_STATE2)
|
||||
.context("decoder produced no 'output_states_2'")?
|
||||
.try_extract_tensor::<f32>()
|
||||
.context("failed to extract decoder state2")?;
|
||||
|
||||
(token, step, next_s1.to_vec(), next_s2.to_vec())
|
||||
};
|
||||
|
||||
// Emit the token unless it is the blank symbol.
|
||||
if token != self.blank_idx {
|
||||
state1 = new_state1;
|
||||
state2 = new_state2;
|
||||
tokens.push(token);
|
||||
prev_token = token;
|
||||
emitted += 1;
|
||||
}
|
||||
|
||||
// Advance: a positive duration jumps several frames; otherwise we
|
||||
// only move on when a blank was emitted or we hit the per-step cap.
|
||||
if step > 0 {
|
||||
t += step;
|
||||
emitted = 0;
|
||||
} else if token == self.blank_idx || emitted == self.max_tokens_per_step {
|
||||
t += 1;
|
||||
emitted = 0;
|
||||
}
|
||||
}
|
||||
|
||||
let text = decode_text(&self.vocab, &tokens);
|
||||
Ok((tokens, text))
|
||||
}
|
||||
}
|
||||
|
||||
/// Build a session with full graph optimization.
|
||||
fn build_session(environment: &Environment, path: &Path) -> Result<Session> {
|
||||
let builder = Session::builder(environment)
|
||||
.map_err(|e| anyhow::anyhow!("failed to create ONNX session builder: {e}"))?;
|
||||
builder
|
||||
.with_optimization_level(GraphOptimizationLevel::Level3)
|
||||
.map_err(|e| anyhow::anyhow!("failed to set ONNX optimization level: {e}"))?
|
||||
.commit_from_file(path)
|
||||
.with_context(|| format!("failed to load ONNX model {}", path.display()))
|
||||
}
|
||||
|
||||
/// Read the size of `axis` from the named output of `session`.
|
||||
fn output_dim(session: &Session, name: &str, axis: usize) -> Result<usize> {
|
||||
let outlet = session
|
||||
.outputs()
|
||||
.iter()
|
||||
.find(|o| o.name() == name)
|
||||
.with_context(|| format!("model has no output named '{name}'"))?;
|
||||
match outlet.dtype() {
|
||||
ValueType::Tensor { shape, .. } => {
|
||||
let dim = *shape
|
||||
.get(axis)
|
||||
.with_context(|| format!("output '{name}' has no axis {axis}"))?;
|
||||
usize::try_from(dim).with_context(|| format!("output '{name}' axis {axis} is dynamic"))
|
||||
}
|
||||
_ => anyhow::bail!("output '{name}' is not a tensor"),
|
||||
}
|
||||
}
|
||||
|
||||
/// Read `(layers, hidden)` from a decoder state input shaped `[layers, batch, hidden]`.
|
||||
fn input_state_dims(session: &Session, name: &str) -> Result<(usize, usize)> {
|
||||
let outlet = session
|
||||
.inputs()
|
||||
.iter()
|
||||
.find(|i| i.name() == name)
|
||||
.with_context(|| format!("decoder has no input named '{name}'"))?;
|
||||
match outlet.dtype() {
|
||||
ValueType::Tensor { shape, .. } => {
|
||||
let layers = usize::try_from(*shape.first().context("state input has no axes")?)
|
||||
.context("state input layer count is dynamic")?;
|
||||
let hidden = usize::try_from(*shape.get(2).context("state input has < 3 axes")?)
|
||||
.context("state input hidden size is dynamic")?;
|
||||
Ok((layers, hidden))
|
||||
}
|
||||
_ => anyhow::bail!("input '{name}' is not a tensor"),
|
||||
}
|
||||
}
|
||||
|
||||
/// Index of the maximum value (ties resolved to the lowest index).
|
||||
fn argmax(values: &[f32]) -> usize {
|
||||
let mut best = 0usize;
|
||||
let mut best_val = f32::NEG_INFINITY;
|
||||
for (i, &v) in values.iter().enumerate() {
|
||||
if v > best_val {
|
||||
best_val = v;
|
||||
best = i;
|
||||
}
|
||||
}
|
||||
best
|
||||
}
|
||||
|
||||
/// Load `vocab.txt` (lines of `<piece> <id>`) into an id-indexed table, mapping
|
||||
/// the SentencePiece word marker `▁` (U+2581) to a real space.
|
||||
fn load_vocab(path: &Path) -> Result<Vec<String>> {
|
||||
let text = std::fs::read_to_string(path)
|
||||
.with_context(|| format!("failed to read {}", path.display()))?;
|
||||
|
||||
let mut entries: Vec<(usize, String)> = Vec::new();
|
||||
let mut max_id = 0usize;
|
||||
for line in text.lines() {
|
||||
if line.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let (piece, id) = line
|
||||
.rsplit_once(' ')
|
||||
.with_context(|| format!("malformed vocab line: {line:?}"))?;
|
||||
let id: usize = id
|
||||
.trim()
|
||||
.parse()
|
||||
.with_context(|| format!("invalid token id in vocab line: {line:?}"))?;
|
||||
max_id = max_id.max(id);
|
||||
entries.push((id, piece.replace('\u{2581}', " ")));
|
||||
}
|
||||
|
||||
let mut vocab = vec![String::new(); max_id + 1];
|
||||
for (id, piece) in entries {
|
||||
vocab[id] = piece;
|
||||
}
|
||||
Ok(vocab)
|
||||
}
|
||||
|
||||
/// Decode token IDs to text and normalize SentencePiece spacing.
|
||||
fn decode_text(vocab: &[String], ids: &[i64]) -> String {
|
||||
let mut joined = String::new();
|
||||
for &id in ids {
|
||||
if id >= 0 {
|
||||
if let Some(piece) = vocab.get(id as usize) {
|
||||
joined.push_str(piece);
|
||||
}
|
||||
}
|
||||
}
|
||||
normalize_spaces(&joined)
|
||||
}
|
||||
|
||||
/// Collapse SentencePiece spacing: drop a leading space, keep a space only when
|
||||
/// it precedes a word character, and drop spacing before punctuation. This
|
||||
/// mirrors the `\A\s|\s\B|(\s)\b` substitution used by the reference decoder.
|
||||
fn normalize_spaces(text: &str) -> String {
|
||||
let chars: Vec<char> = text.chars().collect();
|
||||
let mut out = String::with_capacity(text.len());
|
||||
for (i, &c) in chars.iter().enumerate() {
|
||||
if !c.is_whitespace() {
|
||||
out.push(c);
|
||||
continue;
|
||||
}
|
||||
if i == 0 {
|
||||
continue; // leading whitespace
|
||||
}
|
||||
let next_is_word = chars
|
||||
.get(i + 1)
|
||||
.map(|&n| n.is_alphanumeric() || n == '_')
|
||||
.unwrap_or(false);
|
||||
if next_is_word {
|
||||
out.push(' ');
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Input of the STT worker channel.
|
||||
enum SttInput {
|
||||
/// A streaming mic frame: the worker runs its own RMS voice detection.
|
||||
Frame(Vec<i16>),
|
||||
/// A complete utterance already segmented by an external VAD (see
|
||||
/// [`crate::vad`]), which bypasses the internal RMS detection.
|
||||
Utterance(Vec<i16>),
|
||||
}
|
||||
|
||||
/// Handle to the running STT worker.
|
||||
#[derive(Clone)]
|
||||
pub struct SttWorker {
|
||||
tx: mpsc::UnboundedSender<SttInput>,
|
||||
}
|
||||
|
||||
impl SttWorker {
|
||||
/// Spawn the STT worker task. `stt` is the loaded Parakeet pipeline.
|
||||
/// `transcript_tx` receives one [`Transcript`] per detected EOU.
|
||||
pub fn spawn(
|
||||
stt: ParakeetStt,
|
||||
transcript_tx: mpsc::UnboundedSender<Transcript>,
|
||||
) -> Self {
|
||||
let (tx, mut rx) = mpsc::unbounded_channel::<SttInput>();
|
||||
|
||||
tokio::spawn(async move {
|
||||
let mut utterance: Vec<f32> = Vec::new();
|
||||
let mut in_utterance = false;
|
||||
let mut silence_samples: u64 = 0;
|
||||
let eou_samples = (STT_SAMPLE_RATE as u64 * EOU_SILENCE_MS) / 1000;
|
||||
|
||||
while let Some(input) = rx.recv().await {
|
||||
// Frames arrive already resampled to 16 kHz by the session layer.
|
||||
// Externally segmented utterances need no further detection.
|
||||
let frame = match input {
|
||||
SttInput::Utterance(frame) => {
|
||||
let audio: Vec<f32> =
|
||||
frame.iter().map(|&s| s as f32 / 32768.0).collect();
|
||||
debug!(
|
||||
samples = audio.len(),
|
||||
"VAD-segmented utterance received"
|
||||
);
|
||||
let stt = stt.clone();
|
||||
let tx = transcript_tx.clone();
|
||||
tokio::task::spawn_blocking(move || {
|
||||
run_inference_blocking(&stt, &audio, tx);
|
||||
});
|
||||
continue;
|
||||
}
|
||||
SttInput::Frame(frame) => frame,
|
||||
};
|
||||
|
||||
// Convert i16 -> f32 normalized.
|
||||
let normalized: Vec<f32> =
|
||||
frame.iter().map(|&s| s as f32 / 32768.0).collect();
|
||||
let rms = compute_rms(&normalized);
|
||||
|
||||
if rms > VAD_RMS_THRESHOLD {
|
||||
in_utterance = true;
|
||||
silence_samples = 0;
|
||||
utterance.extend_from_slice(&normalized);
|
||||
} else if in_utterance {
|
||||
silence_samples += normalized.len() as u64;
|
||||
utterance.extend_from_slice(&normalized);
|
||||
if silence_samples >= eou_samples {
|
||||
// EOU detected: run inference on a blocking thread.
|
||||
let utterance_clone = std::mem::take(&mut utterance);
|
||||
let stt = stt.clone();
|
||||
let tx = transcript_tx.clone();
|
||||
|
||||
tokio::task::spawn_blocking(move || {
|
||||
run_inference_blocking(&stt, &utterance_clone, tx);
|
||||
});
|
||||
|
||||
in_utterance = false;
|
||||
silence_samples = 0;
|
||||
}
|
||||
}
|
||||
}
|
||||
info!("STT worker channel closed, stopping");
|
||||
});
|
||||
|
||||
Self { tx }
|
||||
}
|
||||
|
||||
/// Push a frame of 16 kHz mono i16 PCM samples into the worker, which will
|
||||
/// segment it with the built-in RMS voice detection.
|
||||
pub fn push_frame(&self, samples: Vec<i16>) {
|
||||
let _ = self.tx.send(SttInput::Frame(samples));
|
||||
}
|
||||
|
||||
/// Push a complete 16 kHz mono i16 utterance segmented by the external VAD.
|
||||
/// Transcribed as-is, without any further voice-activity detection.
|
||||
pub fn push_utterance(&self, samples: Vec<i16>) {
|
||||
if samples.is_empty() {
|
||||
return;
|
||||
}
|
||||
let _ = self.tx.send(SttInput::Utterance(samples));
|
||||
}
|
||||
}
|
||||
|
||||
/// Blocking inference helper used inside `spawn_blocking`.
|
||||
fn run_inference_blocking(
|
||||
stt: &ParakeetStt,
|
||||
audio: &[f32],
|
||||
tx: mpsc::UnboundedSender<Transcript>,
|
||||
) {
|
||||
let (token_ids, text) = match stt.infer(audio) {
|
||||
Ok(result) => result,
|
||||
Err(e) => {
|
||||
error!(error = %e, "Parakeet inference failed");
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
let duration_ms = (audio.len() as u64 * 1000) / STT_SAMPLE_RATE as u64;
|
||||
debug!(tokens = token_ids.len(), duration_ms, text = %text, "EOU transcript ready");
|
||||
|
||||
let _ = tx.send(Transcript {
|
||||
token_ids,
|
||||
text,
|
||||
duration_ms,
|
||||
});
|
||||
}
|
||||
|
||||
/// Compute the Root-Mean-Square energy of a normalized f32 frame.
|
||||
fn compute_rms(samples: &[f32]) -> f32 {
|
||||
if samples.is_empty() {
|
||||
return 0.0;
|
||||
}
|
||||
let sum_sq: f32 = samples.iter().map(|&s| s * s).sum();
|
||||
(sum_sq / samples.len() as f32).sqrt()
|
||||
}
|
||||
|
||||
/// Convert a duration in milliseconds to a sample count at the STT rate.
|
||||
#[allow(dead_code)]
|
||||
pub fn ms_to_samples(ms: u64) -> usize {
|
||||
((STT_SAMPLE_RATE as u64 * ms) / 1000) as usize
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_normalize_spaces() {
|
||||
assert_eq!(normalize_spaces(" Hello world"), "Hello world");
|
||||
assert_eq!(normalize_spaces(" Hello world!"), "Hello world!");
|
||||
assert_eq!(normalize_spaces("a b c"), "a b c");
|
||||
assert_eq!(normalize_spaces(" trailing "), "trailing");
|
||||
assert_eq!(normalize_spaces("a , b"), "a, b");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_load_real_vocab() {
|
||||
let path = Path::new("models/parakeet/vocab.txt");
|
||||
if !path.exists() {
|
||||
eprintln!("{} not found, skipping", path.display());
|
||||
return;
|
||||
}
|
||||
let vocab = load_vocab(path).expect("vocab load failed");
|
||||
assert_eq!(vocab.len(), 8193);
|
||||
assert_eq!(vocab[8192], "<blk>");
|
||||
assert_eq!(vocab[0], "<unk>");
|
||||
}
|
||||
|
||||
/// Read a canonical 16-bit PCM mono WAV file into normalized f32 samples.
|
||||
fn read_wav_mono16(path: &Path) -> Vec<f32> {
|
||||
let bytes = std::fs::read(path).expect("failed to read wav");
|
||||
assert_eq!(&bytes[0..4], b"RIFF", "not a RIFF file");
|
||||
assert_eq!(&bytes[8..12], b"WAVE", "not a WAVE file");
|
||||
// Walk the chunk list to find the `data` chunk.
|
||||
let mut pos = 12usize;
|
||||
while pos + 8 <= bytes.len() {
|
||||
let id = &bytes[pos..pos + 4];
|
||||
let size = u32::from_le_bytes(bytes[pos + 4..pos + 8].try_into().unwrap()) as usize;
|
||||
let body = pos + 8;
|
||||
if id == b"data" {
|
||||
return bytes[body..body + size]
|
||||
.chunks_exact(2)
|
||||
.map(|c| i16::from_le_bytes([c[0], c[1]]) as f32 / 32768.0)
|
||||
.collect();
|
||||
}
|
||||
pos = body + size + (size & 1); // chunks are word-aligned
|
||||
}
|
||||
panic!("no data chunk found");
|
||||
}
|
||||
|
||||
/// End-to-end decode test. Requires the full model directory. Run with:
|
||||
/// `PARAKEET_TEST_WAV=/path/to/16k_mono.wav cargo test -- --ignored`
|
||||
#[test]
|
||||
#[ignore = "requires the full Parakeet ONNX model directory + a 16 kHz WAV"]
|
||||
fn test_infer_real_audio() {
|
||||
let model_dir = std::env::var("PARAKEET_MODEL_DIR")
|
||||
.unwrap_or_else(|_| "models/parakeet".to_string());
|
||||
let wav = std::env::var("PARAKEET_TEST_WAV")
|
||||
.unwrap_or_else(|_| "/tmp/stt_test/solis16k.wav".to_string());
|
||||
|
||||
if !Path::new(&model_dir).join(PREPROCESSOR_FILE).exists() {
|
||||
eprintln!("model dir {model_dir} not found, skipping");
|
||||
return;
|
||||
}
|
||||
let stt = ParakeetStt::load(Path::new(&model_dir)).expect("model load failed");
|
||||
let audio = read_wav_mono16(Path::new(&wav));
|
||||
eprintln!("running inference on {} samples", audio.len());
|
||||
|
||||
let (tokens, text) = stt.infer(&audio).expect("inference failed");
|
||||
eprintln!("tokens: {} | text: {text:?}", tokens.len());
|
||||
assert!(!tokens.is_empty(), "no tokens emitted");
|
||||
assert!(!text.trim().is_empty(), "empty transcript");
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user