client sdk
This commit is contained in:
+35
-33
@@ -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());
|
||||
|
||||
|
||||
Reference in New Issue
Block a user