708 lines
27 KiB
Rust
708 lines
27 KiB
Rust
//! 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");
|
|
}
|
|
}
|