client sdk

This commit is contained in:
valenti
2026-10-21 17:56:04 +01:00
parent 1cafe3dd3c
commit d6bc2437c9
237 changed files with 5227 additions and 1019 deletions
+35 -33
View File
@@ -113,11 +113,11 @@ impl ParakeetStt {
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 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())
@@ -195,16 +195,23 @@ impl ParakeetStt {
// --- 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 =
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)),
(
PP_IN_WAVEFORMS.to_string(),
SessionInputValue::from(waveforms),
),
(
PP_IN_LENS.to_string(),
SessionInputValue::from(waveforms_lens),
),
])
.context("nemo128 preprocessor inference failed")?;
@@ -278,9 +285,10 @@ impl ParakeetStt {
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())?,
)
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]),
@@ -304,9 +312,15 @@ impl ParakeetStt {
let outputs = sessions
.decoder_joint
.run(vec![
(DEC_IN_ENCODER.to_string(), SessionInputValue::from(enc_frame)),
(
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_TARGET_LENGTH.to_string(),
SessionInputValue::from(target_length),
),
(DEC_IN_STATE1.to_string(), SessionInputValue::from(s1)),
(DEC_IN_STATE2.to_string(), SessionInputValue::from(s2)),
])
@@ -510,10 +524,7 @@ pub struct SttWorker {
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 {
pub fn spawn(stt: ParakeetStt, transcript_tx: mpsc::UnboundedSender<Transcript>) -> Self {
let (tx, mut rx) = mpsc::unbounded_channel::<SttInput>();
tokio::spawn(async move {
@@ -527,12 +538,8 @@ impl SttWorker {
// 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 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 || {
@@ -544,8 +551,7 @@ impl SttWorker {
};
// Convert i16 -> f32 normalized.
let normalized: Vec<f32> =
frame.iter().map(|&s| s as f32 / 32768.0).collect();
let normalized: Vec<f32> = frame.iter().map(|&s| s as f32 / 32768.0).collect();
let rms = compute_rms(&normalized);
if rms > VAD_RMS_THRESHOLD {
@@ -593,11 +599,7 @@ impl SttWorker {
}
/// Blocking inference helper used inside `spawn_blocking`.
fn run_inference_blocking(
stt: &ParakeetStt,
audio: &[f32],
tx: mpsc::UnboundedSender<Transcript>,
) {
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) => {
@@ -684,8 +686,8 @@ mod tests {
#[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 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());