speechcore 0.1.0

Reusable Rust speech-to-text runtime with audio capture, VAD, backend selection, model provisioning, and transcript streaming.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95

use crate::backend::onnx_utils::{load_session, OnnxSessionOptions};
use anyhow::Result;
use ort::session::Session;
use parking_lot::Mutex;
use std::path::Path;

use super::mel::MelSpectrogram;

pub struct ParakeetModel {
    pub encoder: Mutex<Session>,
    pub decoder: Mutex<Session>,
    pub joiner: Mutex<Session>,
    pub mel: MelSpectrogram,
}

impl ParakeetModel {
    pub fn load(model_dir: impl AsRef<Path>, options: &OnnxSessionOptions) -> Result<Self> {
        let model_dir = model_dir.as_ref();

        let encoder_path =
            resolve_model_file(model_dir, "encoder", &["encoder.int8.onnx", "encoder.onnx"])?;
        let decoder_path =
            resolve_model_file(model_dir, "decoder", &["decoder.int8.onnx", "decoder.onnx"])?;
        let joiner_path =
            resolve_model_file(model_dir, "joiner", &["joiner.int8.onnx", "joiner.onnx"])?;

        tracing::info!("Loading Parakeet encoder from: {}", encoder_path.display());
        let encoder = load_session(&encoder_path, options)?;
        tracing::info!("Loading Parakeet decoder from: {}", decoder_path.display());
        let decoder = load_session(&decoder_path, options)?;
        tracing::info!("Loading Parakeet joiner from: {}", joiner_path.display());
        let joiner = load_session(&joiner_path, options)?;

        let mel = MelSpectrogram::new();

        Ok(Self {
            encoder: Mutex::new(encoder),
            decoder: Mutex::new(decoder),
            joiner: Mutex::new(joiner),
            mel,
        })
    }

    pub fn validate_model_dir(model_dir: impl AsRef<Path>) -> Result<()> {
        let model_dir = model_dir.as_ref();

        let has_encoder =
            model_dir.join("encoder.int8.onnx").exists() || model_dir.join("encoder.onnx").exists();
        if !has_encoder {
            return Err(anyhow::anyhow!(
                "Missing Parakeet encoder model (encoder.int8.onnx or encoder.onnx)"
            ));
        }

        let has_decoder =
            model_dir.join("decoder.int8.onnx").exists() || model_dir.join("decoder.onnx").exists();
        if !has_decoder {
            return Err(anyhow::anyhow!(
                "Missing Parakeet decoder model (decoder.int8.onnx or decoder.onnx)"
            ));
        }

        let has_joiner =
            model_dir.join("joiner.int8.onnx").exists() || model_dir.join("joiner.onnx").exists();
        if !has_joiner {
            return Err(anyhow::anyhow!(
                "Missing Parakeet joiner model (joiner.int8.onnx or joiner.onnx)"
            ));
        }

        if !model_dir.join("tokens.txt").exists() {
            return Err(anyhow::anyhow!("Missing Parakeet tokens.txt"));
        }

        Ok(())
    }
}

fn resolve_model_file(
    model_dir: &Path,
    component: &str,
    candidates: &[&str],
) -> Result<std::path::PathBuf> {
    for candidate in candidates {
        let path = model_dir.join(candidate);
        if path.exists() {
            return Ok(path);
        }
    }
    Err(anyhow::anyhow!(
        "Missing Parakeet {} model. Tried: {}",
        component,
        candidates.join(", ")
    ))
}