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
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
use crate::audio::{self, FeatureCache};
use crate::config::PreprocessorConfig;
use crate::decoder::{ParakeetDecoder, TranscriptionResult};
use crate::error::{Error, Result};
use crate::execution::ModelConfig as ExecutionConfig;
use crate::model::ParakeetModel;
use crate::timestamps::{process_timestamps, rebuild_text, TimestampMode};
use crate::transcriber::Transcriber;
use std::path::{Path, PathBuf};
pub struct Parakeet {
model: ParakeetModel,
decoder: ParakeetDecoder,
preprocessor_config: PreprocessorConfig,
feature_cache: FeatureCache,
model_dir: PathBuf,
}
impl Parakeet {
/// Load Parakeet model from path with optional configuration.
///
/// # Arguments
/// * `path` - Directory containing model files, or path to specific model file
/// * `config` - Execution configuration; ONNX Runtime needs one with an environment
///
/// # Examples
/// ```no_run
/// # #[cfg(feature = "ort")]
/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
/// use parakeet_rs::{ExecutionConfig, Parakeet};
///
/// // One ONNX Runtime environment per process, shared by every model
/// let env = parakeet_rs::ort::init().build()?;
/// let config = ExecutionConfig::new().with_environment(env);
///
/// // Load from directory on the CPU
/// let parakeet = Parakeet::from_pretrained(".", Some(config.clone()))?;
///
/// // Or load from specific model file
/// let parakeet = Parakeet::from_pretrained("model_q4.onnx", Some(config))?;
/// # Ok(())
/// # }
/// # #[cfg(not(feature = "ort"))]
/// # fn main() {}
/// ```
///
/// For GPU acceleration, enable the corresponding feature (cuda, tensorrt, webgpu, etc.)
/// and pass an `ExecutionConfig` with the desired execution provider.
pub fn from_pretrained<P: AsRef<Path>>(
path: P,
config: Option<ExecutionConfig>,
) -> Result<Self> {
let path = path.as_ref();
// Determine if path is a directory or file
let (model_path, tokenizer_path, model_dir) = if path.is_dir() {
// Directory mode: auto-detect model file
let model_path = Self::find_model_file(path)?;
let tokenizer_path = path.join("tokenizer.json");
(model_path, tokenizer_path, path.to_path_buf())
} else if path.is_file() {
// File mode: path points directly to model file
let model_dir = path
.parent()
.ok_or_else(|| Error::Config("Invalid model path".to_string()))?;
let tokenizer_path = model_dir.join("tokenizer.json");
(path.to_path_buf(), tokenizer_path, model_dir.to_path_buf())
} else {
return Err(Error::Config(format!(
"Path does not exist: {}",
path.display()
)));
};
// Check tokenizer exists
if !tokenizer_path.exists() {
return Err(Error::Config(format!(
"Required file 'tokenizer.json' not found in {}",
model_dir.display()
)));
}
let preprocessor_config = PreprocessorConfig::default();
let exec_config = config.unwrap_or_default();
let model = ParakeetModel::from_pretrained_with_config(&model_path, exec_config)?;
let decoder = ParakeetDecoder::from_pretrained(&tokenizer_path)?;
let feature_cache = FeatureCache::from_config(&preprocessor_config);
Ok(Self {
model,
decoder,
preprocessor_config,
feature_cache,
model_dir,
})
}
fn find_model_file(dir: &Path) -> Result<PathBuf> {
// Priority order: model.onnx > model_fp16.onnx > model_int8.onnx > model_q4.onnx
let candidates = [
"model.onnx",
"model_fp16.onnx",
"model_int8.onnx",
"model_q4.onnx",
];
for candidate in &candidates {
let path = dir.join(candidate);
if path.exists() {
return Ok(path);
}
}
// If none of the standard names found, search for any .onnx file
if let Ok(entries) = std::fs::read_dir(dir) {
for entry in entries.flatten() {
let path = entry.path();
if path.extension().and_then(|s| s.to_str()) == Some("onnx") {
return Ok(path);
}
}
}
Err(Error::Config(format!(
"No model file (*.onnx) found in directory: {}",
dir.display()
)))
}
pub fn model_dir(&self) -> &Path {
&self.model_dir
}
pub fn preprocessor_config(&self) -> &PreprocessorConfig {
&self.preprocessor_config
}
}
impl Transcriber for Parakeet {
fn transcribe_samples(
&mut self,
audio: Vec<f32>,
sample_rate: u32,
channels: u16,
mode: Option<TimestampMode>,
) -> Result<TranscriptionResult> {
let features = audio::extract_features_with_cache(
audio,
sample_rate,
channels,
&self.preprocessor_config,
&self.feature_cache,
)?;
let logits = self.model.forward(features)?;
let mut result = self.decoder.decode_with_timestamps(
&logits,
self.preprocessor_config.hop_length,
self.preprocessor_config.sampling_rate,
)?;
// Process timestamps to requested output mode
let mode = mode.unwrap_or(TimestampMode::Tokens);
result.tokens = process_timestamps(&result.tokens, mode);
// Rebuild full text from processed tokens to ensure consistency
result.text = rebuild_text(&result.tokens, mode);
Ok(result)
}
}