#![allow(clippy::all, clippy::pedantic, clippy::restriction, clippy::nursery)]
#[cfg(feature = "realizar-gpu")]
use crate::cuda;
use crate::{
audio, detection, error, format, inference, model, progress, timestamps, tokenizer, vad,
};
pub use error::{WhisperError, WhisperResult};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ModelType {
Tiny,
TinyEn,
Base,
BaseEn,
Small,
SmallEn,
Medium,
MediumEn,
Large,
LargeV1,
LargeV2,
LargeV3,
LargeV3Turbo,
}
#[derive(Debug, Clone, Copy, Default)]
pub enum DecodingStrategy {
#[default]
Greedy,
BeamSearch {
beam_size: usize,
temperature: f32,
patience: f32,
},
Sampling {
temperature: f32,
top_k: Option<usize>,
top_p: Option<f32>,
},
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub enum Task {
#[default]
Transcribe,
Translate,
}
#[derive(Debug, Clone, Default)]
pub struct TranscribeOptions {
pub language: Option<String>,
pub task: Task,
pub strategy: DecodingStrategy,
pub word_timestamps: bool,
pub profile: bool,
pub prompt: Option<String>,
pub hotwords: Vec<String>,
}
#[derive(Debug, Clone, Default)]
#[cfg_attr(feature = "cli", derive(serde::Serialize, serde::Deserialize))]
pub struct Segment {
pub start: f32,
pub end: f32,
pub text: String,
pub tokens: Vec<u32>,
}
#[derive(Debug, Clone, Default)]
#[cfg_attr(feature = "cli", derive(serde::Serialize, serde::Deserialize))]
pub struct ProfilingStats {
pub total_ms: f64,
pub breakdown: std::collections::HashMap<String, f64>,
#[cfg_attr(feature = "cli", serde(skip_serializing_if = "Option::is_none"))]
pub trace_json: Option<String>,
}
#[derive(Debug, Clone, Default)]
#[cfg_attr(feature = "cli", derive(serde::Serialize, serde::Deserialize))]
pub struct TranscriptionResult {
pub text: String,
pub language: String,
pub segments: Vec<Segment>,
#[cfg_attr(feature = "cli", serde(skip_serializing_if = "Option::is_none"))]
pub profiling: Option<ProfilingStats>,
}
#[derive(Debug, Clone)]
pub struct BatchTranscriptionResult {
pub results: Vec<TranscriptionResult>,
pub total_duration_secs: f32,
}
impl BatchTranscriptionResult {
#[must_use]
pub fn len(&self) -> usize {
self.results.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.results.is_empty()
}
#[must_use]
pub fn get(&self, index: usize) -> Option<&TranscriptionResult> {
self.results.get(index)
}
pub fn iter(&self) -> impl Iterator<Item = &TranscriptionResult> {
self.results.iter()
}
#[must_use]
pub fn texts(&self) -> Vec<&str> {
self.results.iter().map(|r| r.text.as_str()).collect()
}
}
pub struct SummarizeOptions<'a> {
pub model: &'a model::lfm2::Lfm2,
pub tokenizer: &'a model::lfm2::Lfm2Tokenizer,
pub max_tokens: usize,
pub temperature: f32,
}
impl<'a> SummarizeOptions<'a> {
#[must_use]
pub fn new(model: &'a model::lfm2::Lfm2, tokenizer: &'a model::lfm2::Lfm2Tokenizer) -> Self {
Self {
model,
tokenizer,
max_tokens: 256,
temperature: 0.3,
}
}
#[must_use]
pub const fn with_max_tokens(mut self, max_tokens: usize) -> Self {
self.max_tokens = max_tokens;
self
}
#[must_use]
pub const fn with_temperature(mut self, temperature: f32) -> Self {
self.temperature = temperature;
self
}
}
#[derive(Debug, Clone)]
pub struct TranscribeSummaryResult {
pub transcription: TranscriptionResult,
pub summary: String,
pub generation_stats: Option<model::lfm2::GenerationStats>,
}
impl TranscribeSummaryResult {
#[must_use]
pub fn transcript(&self) -> &str {
&self.transcription.text
}
#[must_use]
pub fn summary(&self) -> &str {
&self.summary
}
#[must_use]
pub fn has_summary(&self) -> bool {
!self.summary.is_empty()
}
}
#[derive(Debug, Clone)]
pub struct WhisperApr {
config: model::ModelConfig,
encoder: model::Encoder,
decoder: model::Decoder,
tokenizer: tokenizer::Tokenizer,
mel_filters: Option<audio::MelFilterbank>,
conv_stem: Option<audio::ConvStem>,
resampler: Option<audio::SincResampler>,
weights_loaded: bool,
}
impl WhisperApr {
#[must_use]
pub fn from_config(config: model::ModelConfig) -> Self {
let encoder = model::Encoder::new(&config);
let decoder = model::Decoder::new(&config);
let tokenizer = if config.model_family == format::ModelFamily::Moonshine {
tokenizer::Tokenizer::SentencePiece(
tokenizer::SentencePieceTokenizer::moonshine_default(),
)
} else {
tokenizer::Tokenizer::Bpe(tokenizer::BpeTokenizer::with_base_tokens())
};
let (mel_filters, conv_stem) = match config.audio_frontend {
model::AudioFrontend::MelFilterbank => (
Some(audio::MelFilterbank::new(&audio::MelConfig {
n_mels: config.n_mels as usize,
..audio::MelConfig::whisper()
})),
None,
),
model::AudioFrontend::LearnedConv => (
None,
Some(audio::ConvStem::new(config.n_audio_state as usize)),
),
};
Self {
config,
encoder,
decoder,
tokenizer,
mel_filters,
conv_stem,
resampler: None,
weights_loaded: false,
}
}
#[must_use]
pub fn tiny() -> Self {
Self::from_config(model::ModelConfig::tiny())
}
#[must_use]
pub fn base() -> Self {
Self::from_config(model::ModelConfig::base())
}
#[must_use]
pub fn small() -> Self {
Self::from_config(model::ModelConfig::small())
}
#[must_use]
pub fn medium() -> Self {
Self::from_config(model::ModelConfig::medium())
}
#[must_use]
pub fn large() -> Self {
Self::from_config(model::ModelConfig::large())
}
#[must_use]
pub fn moonshine_tiny() -> Self {
Self::from_config(model::ModelConfig::moonshine_tiny())
}
#[must_use]
pub fn moonshine_base() -> Self {
Self::from_config(model::ModelConfig::moonshine_base())
}
#[must_use]
pub const fn config(&self) -> &model::ModelConfig {
&self.config
}
#[must_use]
pub const fn model_type(&self) -> ModelType {
self.config.model_type
}
const CHUNK_SAMPLES: usize = 30 * audio::SAMPLE_RATE as usize;
const OVERLAP_SAMPLES: usize = 5 * audio::SAMPLE_RATE as usize;
const MAX_SUBTITLE_SECS: f32 = 10.0;
pub fn transcribe(
&self,
audio: &[f32],
options: TranscribeOptions,
) -> WhisperResult<TranscriptionResult> {
if audio.len() > Self::CHUNK_SAMPLES {
return self.transcribe_chunked(audio, options);
}
self.transcribe_single_chunk(audio, options)
}
fn transcribe_single_chunk(
&self,
audio: &[f32],
options: TranscribeOptions,
) -> WhisperResult<TranscriptionResult> {
#[cfg(feature = "std")]
let start_total = if options.profile {
Some(std::time::Instant::now())
} else {
None
};
#[cfg(feature = "std")]
let start_audio = start_total.map(|_| std::time::Instant::now());
#[cfg(feature = "std")]
let mut mel_ms: Option<f64> = None;
#[cfg(feature = "std")]
let mut brick_profiler: Option<trueno::BrickProfiler> = if options.profile {
Some(trueno::BrickProfiler::enabled())
} else {
None
};
#[cfg(feature = "std")]
let mut page_faults: Option<(u64, u64)> = None;
#[cfg(feature = "std")]
let mut blis_profiler_stats: Option<trueno::blis::BlisProfiler> = None;
#[cfg(all(feature = "std", feature = "realizar-inference"))]
let mut inference_tracer: Option<realizar::InferenceTracer> = if options.profile {
let config = realizar::TraceConfig::enabled();
let mut tracer = realizar::InferenceTracer::new(config);
tracer.set_model_info(realizar::ModelInfo {
name: format!("{:?}", self.config.model_type),
num_layers: self.config.n_audio_layer as usize,
hidden_dim: self.config.n_audio_state as usize,
vocab_size: self.config.n_vocab as usize,
num_heads: self.config.n_audio_head as usize,
quant_type: Some("f32".into()),
});
Some(tracer)
} else {
None
};
let audio_features = match self.config.audio_frontend {
model::AudioFrontend::MelFilterbank => {
#[cfg(feature = "std")]
let mel_start = std::time::Instant::now();
let mel = self.compute_mel(audio)?;
#[cfg(feature = "std")]
{
mel_ms = Some(mel_start.elapsed().as_secs_f64() * 1000.0);
}
#[cfg(feature = "std")]
if let Some(ref mut profiler) = brick_profiler {
crate::simd::enable_blis_profiling();
let (pf_minor_before, pf_major_before) = trueno::brick::get_page_faults();
#[cfg(feature = "realizar-inference")]
let features =
self.encode_profiled(&mel, profiler, inference_tracer.as_mut())?;
#[cfg(not(feature = "realizar-inference"))]
let features = self.encode_profiled(&mel, profiler)?;
blis_profiler_stats = crate::simd::take_blis_profiler();
let (pf_minor_after, pf_major_after) = trueno::brick::get_page_faults();
page_faults = Some((
pf_minor_after.saturating_sub(pf_minor_before),
pf_major_after.saturating_sub(pf_major_before),
));
features
} else {
self.encode(&mel)?
}
#[cfg(not(feature = "std"))]
self.encode(&mel)?
}
model::AudioFrontend::LearnedConv => {
let stem = self
.conv_stem
.as_ref()
.ok_or_else(|| WhisperError::Model("Moonshine requires ConvStem".into()))?;
let stem_out = stem.forward(audio)?;
self.encoder.forward(&stem_out)?
}
};
#[cfg(feature = "std")]
let enc_ms = start_audio.map(|s| s.elapsed().as_secs_f64() * 1000.0);
let language = options.language.clone().unwrap_or_else(|| "en".to_string());
let initial_tokens =
self.build_initial_tokens(&language, options.task, options.prompt.as_deref());
#[cfg(feature = "std")]
let start_dec = start_total.map(|_| std::time::Instant::now());
let tokens = self.decode(&audio_features, &initial_tokens, &options)?;
#[cfg(feature = "std")]
let dec_ms = start_dec.map(|s| s.elapsed().as_secs_f64() * 1000.0);
let text = self.tokenizer.decode(&tokens)?;
let segments = if timestamps::has_timestamps(&tokens) {
timestamps::extract_segments(&tokens, |ts| self.tokenizer.decode(ts).ok())
} else if !text.trim().is_empty() {
let duration = audio.len() as f32 / audio::SAMPLE_RATE as f32;
let single = vec![Segment {
start: 0.0,
end: duration,
text: text.clone(),
tokens: tokens.clone(),
}];
timestamps::split_long_segments(&single, Self::MAX_SUBTITLE_SECS)
} else {
Vec::new()
};
#[cfg(feature = "std")]
let profiling = start_total.map(|st| {
let mut breakdown = std::collections::HashMap::new();
let pairs: &[(&str, Option<f64>)] = &[
("mel_ms", mel_ms),
("audio_ms", enc_ms),
("decoder_ms", dec_ms),
];
for &(key, val) in pairs {
if let Some(ms) = val {
breakdown.insert(key.to_string(), ms);
}
}
if let Some(audio) = enc_ms {
let enc_total = audio - mel_ms.unwrap_or(0.0);
breakdown.insert("encoder_ms".to_string(), enc_total);
}
if let Some(ref profiler) = brick_profiler {
let cats = profiler.category_stats();
let total_prof_ns = profiler.total_ns();
let norm_ns = cats[trueno::BrickCategory::Norm as usize].total_ns;
let attn_ns = cats[trueno::BrickCategory::Attention as usize].total_ns;
let ffn_ns = cats[trueno::BrickCategory::Ffn as usize].total_ns;
let other_ns = cats[trueno::BrickCategory::Other as usize].total_ns;
breakdown.insert("brick_norm_ms".to_string(), norm_ns as f64 / 1_000_000.0);
breakdown.insert("brick_attn_ms".to_string(), attn_ns as f64 / 1_000_000.0);
breakdown.insert("brick_ffn_ms".to_string(), ffn_ns as f64 / 1_000_000.0);
breakdown.insert("brick_other_ms".to_string(), other_ns as f64 / 1_000_000.0);
breakdown.insert("brick_total_ns".to_string(), total_prof_ns as f64);
for &brick_id in &[
trueno::BrickId::LayerNorm,
trueno::BrickId::AttentionScore,
trueno::BrickId::GateProjection,
trueno::BrickId::Embedding,
] {
let stats = profiler.brick_stats(brick_id);
if stats.count > 0 {
let name = brick_id.name();
breakdown.insert(
format!("brick_{name}_ms"),
stats.total_ns as f64 / 1_000_000.0,
);
breakdown.insert(format!("brick_{name}_count"), stats.count as f64);
breakdown.insert(
format!("brick_{name}_cycles_per_elem"),
stats.cycles_per_element(),
);
let diagnosis = stats.diagnose_from_cycles();
let diag_code = match diagnosis {
"memory-bound (low IPC, likely cache misses)" => 1.0,
"compute-bound (efficient)" => 2.0,
"throttled or context-switched" => 3.0,
"balanced" => 4.0,
_ => 0.0,
};
breakdown.insert(format!("brick_{name}_bottleneck"), diag_code);
}
}
if let Some((minor, major)) = page_faults {
breakdown.insert("page_faults_minor".to_string(), minor as f64);
breakdown.insert("page_faults_major".to_string(), major as f64);
}
}
if let Some(ref blis) = blis_profiler_stats {
breakdown.insert("blis_macro_gflops".to_string(), blis.macro_stats.gflops());
breakdown.insert(
"blis_macro_calls".to_string(),
blis.macro_stats.count as f64,
);
breakdown.insert(
"blis_macro_ns".to_string(),
blis.macro_stats.total_ns as f64,
);
breakdown.insert("blis_midi_gflops".to_string(), blis.midi_stats.gflops());
breakdown.insert("blis_midi_calls".to_string(), blis.midi_stats.count as f64);
breakdown.insert("blis_micro_gflops".to_string(), blis.micro_stats.gflops());
breakdown.insert(
"blis_micro_calls".to_string(),
blis.micro_stats.count as f64,
);
breakdown.insert("blis_pack_ns".to_string(), blis.pack_stats.total_ns as f64);
breakdown.insert("blis_pack_calls".to_string(), blis.pack_stats.count as f64);
breakdown.insert("blis_total_gflops".to_string(), blis.total_gflops());
if blis.macro_stats.total_ns > 0 {
let pack_pct =
blis.pack_stats.total_ns as f64 / blis.macro_stats.total_ns as f64 * 100.0;
breakdown.insert("blis_pack_pct".to_string(), pack_pct);
}
}
#[cfg(feature = "realizar-inference")]
let trace_json = inference_tracer.as_ref().map(|t| t.to_json());
#[cfg(not(feature = "realizar-inference"))]
let trace_json: Option<String> = None;
ProfilingStats {
total_ms: st.elapsed().as_secs_f64() * 1000.0,
breakdown,
trace_json,
}
});
#[cfg(not(feature = "std"))]
let profiling = None;
Ok(TranscriptionResult {
text,
language,
segments,
profiling,
})
}
fn transcribe_chunked(
&self,
audio: &[f32],
options: TranscribeOptions,
) -> WhisperResult<TranscriptionResult> {
let chunk_size = Self::CHUNK_SAMPLES;
let overlap = Self::OVERLAP_SAMPLES;
let step = chunk_size;
let language = options.language.clone().unwrap_or_else(|| "en".to_string());
let mut all_segments: Vec<Segment> = Vec::new();
let mut all_text = String::new();
let mut chunk_idx = 0;
let mut offset = 0;
while offset < audio.len() {
let chunk_end = (offset + chunk_size + overlap).min(audio.len());
let chunk = &audio[offset..chunk_end];
if chunk.len() < audio::SAMPLE_RATE as usize / 2 {
break;
}
let chunk_options = TranscribeOptions {
language: Some(language.clone()),
task: options.task,
strategy: options.strategy,
word_timestamps: options.word_timestamps,
profile: options.profile,
prompt: options.prompt.clone(),
hotwords: options.hotwords.clone(),
};
let chunk_result = self.transcribe_single_chunk(chunk, chunk_options)?;
let time_offset = offset as f32 / audio::SAMPLE_RATE as f32;
let chunk_text = if chunk_idx == 0 {
chunk_result.text.clone()
} else {
let overlap_ratio = overlap as f32 / chunk.len() as f32;
let words: Vec<&str> = chunk_result.text.split_whitespace().collect();
let skip_words = ((words.len() as f32) * overlap_ratio * 0.8) as usize;
words
.into_iter()
.skip(skip_words)
.collect::<Vec<_>>()
.join(" ")
};
if !chunk_text.is_empty() {
if !all_text.is_empty() {
all_text.push(' ');
}
all_text.push_str(&chunk_text);
}
for mut seg in chunk_result.segments {
seg.start += time_offset;
seg.end += time_offset;
all_segments.push(seg);
}
offset += step;
chunk_idx += 1;
}
let merged_segments = self.merge_overlapping_segments(all_segments);
let final_segments =
timestamps::split_long_segments(&merged_segments, Self::MAX_SUBTITLE_SECS);
Ok(TranscriptionResult {
text: all_text,
language,
segments: final_segments,
profiling: None,
})
}
fn merge_overlapping_segments(&self, segments: Vec<Segment>) -> Vec<Segment> {
if segments.is_empty() {
return segments;
}
let mut merged: Vec<Segment> = Vec::with_capacity(segments.len());
let mut current = segments[0].clone();
for seg in segments.into_iter().skip(1) {
if seg.start < current.end + 0.1 {
current.end = current.end.max(seg.end);
if !seg.text.is_empty() {
if !current.text.is_empty() {
current.text.push(' ');
}
current.text.push_str(&seg.text);
}
current.tokens.extend(seg.tokens);
} else {
merged.push(current);
current = seg;
}
}
merged.push(current);
merged
}
pub fn forward_probed(
&self,
audio: &[f32],
tokens: &[u32],
probe: &mut crate::probe::ActivationProbe,
) -> WhisperResult<Vec<f32>> {
let audio_features = match self.config.audio_frontend {
model::AudioFrontend::MelFilterbank => {
let mel = self.compute_mel(audio)?;
self.encoder.forward_probed(&mel, probe)?
}
model::AudioFrontend::LearnedConv => {
let stem = self
.conv_stem
.as_ref()
.ok_or_else(|| WhisperError::Model("Moonshine requires ConvStem".into()))?;
let stem_out = stem.forward_probed(audio, probe)?;
self.encoder.forward_probed(&stem_out, probe)?
}
};
self.decoder.forward_probed(tokens, &audio_features, probe)
}
fn eot_token(&self) -> u32 {
if self.config.model_family == format::ModelFamily::Moonshine {
2 } else {
tokenizer::special_tokens::EOT
}
}
pub fn compute_mel(&self, audio: &[f32]) -> WhisperResult<Vec<f32>> {
const N_SAMPLES_30S: usize = 480_000; const N_FRAMES: usize = 3000; const N_MELS: usize = 80;
let padded_audio = match audio.len().cmp(&N_SAMPLES_30S) {
std::cmp::Ordering::Equal => audio.to_vec(),
std::cmp::Ordering::Less => {
let mut padded = vec![0.0_f32; N_SAMPLES_30S];
padded[..audio.len()].copy_from_slice(audio);
padded
}
std::cmp::Ordering::Greater => {
audio[..N_SAMPLES_30S].to_vec()
}
};
let mel_fb = self.mel_filters.as_ref().ok_or_else(|| {
WhisperError::Audio("mel filterbank not available (Moonshine model?)".into())
})?;
let mut mel = mel_fb
.compute(&padded_audio)
.map_err(|e| WhisperError::Audio(e.to_string()))?;
let actual_frames = mel.len() / N_MELS;
if actual_frames < N_FRAMES {
let pad_value = -1.0_f32;
let mut padded_mel = vec![pad_value; N_FRAMES * N_MELS];
padded_mel[..mel.len()].copy_from_slice(&mel);
mel = padded_mel;
} else if actual_frames > N_FRAMES {
mel.truncate(N_FRAMES * N_MELS);
}
Ok(mel)
}
pub fn encode(&self, mel: &[f32]) -> WhisperResult<Vec<f32>> {
self.encoder.forward_mel(mel)
}
#[cfg(feature = "realizar-inference")]
pub fn encode_profiled(
&self,
mel: &[f32],
profiler: &mut trueno::BrickProfiler,
tracer: Option<&mut realizar::InferenceTracer>,
) -> WhisperResult<Vec<f32>> {
self.encoder.forward_mel_profiled(mel, profiler, tracer)
}
#[cfg(not(feature = "realizar-inference"))]
pub fn encode_profiled(
&self,
mel: &[f32],
profiler: &mut trueno::BrickProfiler,
) -> WhisperResult<Vec<f32>> {
self.encoder.forward_mel_profiled(mel, profiler)
}
fn get_initial_tokens(&self, language: &str, task: Task) -> Vec<u32> {
if self.config.model_family == format::ModelFamily::Moonshine {
return vec![1]; }
use tokenizer::special_tokens::{self, SpecialTokens};
let specials = SpecialTokens::for_vocab_size(self.config.n_vocab as usize);
let mut tokens = vec![specials.sot];
if specials.is_multilingual {
let lang_offset = special_tokens::language_offset(language).unwrap_or(0);
tokens.push(specials.lang_base + lang_offset);
}
match task {
Task::Transcribe => tokens.push(specials.transcribe),
Task::Translate => tokens.push(special_tokens::TRANSLATE),
}
tokens
}
fn build_initial_tokens(&self, language: &str, task: Task, prompt: Option<&str>) -> Vec<u32> {
let mut tokens = self.get_initial_tokens(language, task);
let prompt_text = match prompt {
Some(p) if !p.is_empty() => p,
_ => return tokens,
};
let prompt_tokens = match self.tokenizer.encode(prompt_text) {
Ok(t) if !t.is_empty() => t,
_ => return tokens,
};
let max_prompt = (self.config.n_text_ctx as usize / 2).min(224);
let truncated = if prompt_tokens.len() > max_prompt {
&prompt_tokens[prompt_tokens.len() - max_prompt..]
} else {
&prompt_tokens
};
let mut prefix = Vec::with_capacity(1 + truncated.len() + tokens.len());
prefix.push(tokenizer::special_tokens::PREV);
prefix.extend_from_slice(truncated);
prefix.append(&mut tokens);
prefix
}
pub fn detect_language(&self, audio: &[f32]) -> WhisperResult<detection::LanguageProbs> {
let mel = self.compute_mel(audio)?;
let audio_features = self.encode(&mel)?;
let n_vocab = self.config.n_vocab as usize;
let logits_fn = |tokens: &[u32]| -> WhisperResult<Vec<f32>> {
let all_logits = self.decoder.forward(tokens, &audio_features)?;
let seq_len = tokens.len();
let last_start = (seq_len - 1) * n_vocab;
if all_logits.len() >= last_start + n_vocab {
Ok(all_logits[last_start..last_start + n_vocab].to_vec())
} else {
let mut padded = vec![f32::NEG_INFINITY; n_vocab];
let available = all_logits.len().saturating_sub(last_start);
if available > 0 {
padded[..available].copy_from_slice(&all_logits[last_start..]);
}
Ok(padded)
}
};
let detector = detection::LanguageDetector::new();
detector.detect(logits_fn)
}
fn decode(
&self,
audio_features: &[f32],
initial_tokens: &[u32],
options: &TranscribeOptions,
) -> WhisperResult<Vec<u32>> {
use std::cell::RefCell;
let n_vocab = self.config.n_vocab as usize;
let max_tokens = self.config.n_text_ctx as usize;
let cache = RefCell::new(self.decoder.create_kv_cache());
let processed_count = RefCell::new(0usize);
let suppressor = inference::WhisperTokenSuppressor::new()
.with_timestamp_suppression(false)
.with_vocab_size(n_vocab);
let hotword_booster = if !options.hotwords.is_empty() {
let mut booster = crate::vocabulary::HotwordBooster::new();
for word in &options.hotwords {
if let Ok(tokens) = self.tokenizer.encode(word) {
if !tokens.is_empty() {
booster.add_hotword_with_tokens_default(word, tokens);
}
}
}
if booster.is_empty() {
None
} else {
Some(booster)
}
} else {
None
};
let scratch = std::cell::RefCell::new(self.decoder.create_decoder_scratch());
let logits_fn = |tokens: &[u32]| -> WhisperResult<Vec<f32>> {
let seq_len = tokens.len();
let already_processed = *processed_count.borrow();
let mut logits = vec![f32::NEG_INFINITY; n_vocab];
for &token in tokens.iter().take(seq_len).skip(already_processed) {
logits = self.decoder.forward_one_with_scratch(
token,
audio_features,
&mut cache.borrow_mut(),
&mut scratch.borrow_mut(),
)?;
}
*processed_count.borrow_mut() = seq_len;
suppressor.apply(&mut logits);
if let Some(ref booster) = hotword_booster {
booster.apply_bias(&mut logits, tokens);
}
let eot_id = self.eot_token();
Self::suppress_repetitions(&mut logits, tokens, eot_id, n_vocab);
Ok(logits)
};
let eot = self.eot_token();
match options.strategy {
DecodingStrategy::Greedy => {
let decoder = inference::GreedyDecoder::new(max_tokens);
decoder.decode(logits_fn, initial_tokens, eot)
}
DecodingStrategy::BeamSearch {
beam_size,
temperature,
patience,
} => {
let decoder = inference::BeamSearchDecoder::new(beam_size, max_tokens)
.with_temperature(temperature)
.with_patience(patience);
decoder.decode(logits_fn, initial_tokens, eot)
}
DecodingStrategy::Sampling { temperature, .. } => {
let decoder =
inference::GreedyDecoder::new(max_tokens).with_temperature(temperature);
decoder.decode(logits_fn, initial_tokens, eot)
}
}
}
fn suppress_repetitions(logits: &mut [f32], tokens: &[u32], eot_id: u32, n_vocab: usize) {
let text_tokens: Vec<u32> = if eot_id <= 3 {
tokens.iter().copied().filter(|&t| t > eot_id).collect()
} else {
tokens.iter().copied().filter(|&t| t < eot_id).collect()
};
let window_start = text_tokens.len().saturating_sub(50);
for &prev_tok in &text_tokens[window_start..] {
if (prev_tok as usize) < n_vocab {
logits[prev_tok as usize] -= 2.0;
}
}
if text_tokens.len() >= 3 {
let prev2 = text_tokens[text_tokens.len() - 2];
let prev1 = text_tokens[text_tokens.len() - 1];
for w in text_tokens[..text_tokens.len() - 2].windows(3) {
if w[0] == prev2 && w[1] == prev1 && (w[2] as usize) < n_vocab {
logits[w[2] as usize] -= 10.0;
}
}
}
Self::suppress_degenerate_loops(logits, &text_tokens, eot_id, n_vocab);
}
fn suppress_degenerate_loops(
logits: &mut [f32],
text_tokens: &[u32],
eot_id: u32,
n_vocab: usize,
) {
if text_tokens.len() >= 12 {
let mut fourgram_counts = std::collections::HashMap::<[u32; 4], u32>::new();
for w in text_tokens.windows(4) {
*fourgram_counts.entry([w[0], w[1], w[2], w[3]]).or_insert(0) += 1;
}
if fourgram_counts.values().any(|&c| c >= 3) && (eot_id as usize) < n_vocab {
logits[eot_id as usize] = 100.0;
return;
}
}
if text_tokens.len() >= 80 {
let window = &text_tokens[text_tokens.len() - 80..];
let mut freq = std::collections::HashMap::<u32, u32>::new();
for &t in window {
*freq.entry(t).or_insert(0) += 1;
}
for (&tok, &count) in &freq {
if count >= 8 && (tok as usize) < n_vocab {
logits[tok as usize] -= (count as f32 - 7.0) * 3.0;
}
}
if freq.len() as f32 / 80.0 < 0.35 && (eot_id as usize) < n_vocab {
logits[eot_id as usize] = 100.0;
}
}
}
pub fn set_resampler(&mut self, input_rate: u32) -> WhisperResult<()> {
if input_rate == audio::SAMPLE_RATE {
self.resampler = None;
} else {
self.resampler = Some(audio::SincResampler::new(input_rate, audio::SAMPLE_RATE)?);
}
Ok(())
}
pub fn resample(&self, audio: &[f32]) -> WhisperResult<Vec<f32>> {
self.resampler
.as_ref()
.map_or_else(|| Ok(audio.to_vec()), |resampler| resampler.resample(audio))
}
#[must_use]
pub const fn tokenizer(&self) -> &tokenizer::Tokenizer {
&self.tokenizer
}
#[must_use]
pub fn memory_size(&self) -> usize {
let params = match self.config.model_type {
ModelType::Tiny | ModelType::TinyEn => 39_000_000,
ModelType::Base | ModelType::BaseEn => 74_000_000,
ModelType::Small | ModelType::SmallEn => 244_000_000,
ModelType::Medium | ModelType::MediumEn => 769_000_000,
ModelType::Large | ModelType::LargeV1 | ModelType::LargeV2 | ModelType::LargeV3 => {
1_550_000_000
}
ModelType::LargeV3Turbo => 809_000_000,
};
params * 4 }
#[must_use]
pub fn has_weights(&self) -> bool {
self.weights_loaded
}
pub fn load_from_apr(data: &[u8]) -> WhisperResult<Self> {
Self::load_from_apr_with_progress(data, &mut progress::null_callback)
}
pub fn load_from_apr_with_progress(
data: &[u8],
callback: progress::ProgressCallback<'_>,
) -> WhisperResult<Self> {
let mut tracker = progress::ProgressTracker::model_loading();
callback(&tracker.to_progress());
let reader = format::AprV2ReaderRef::from_bytes(data)
.map_err(|e| error::WhisperError::Format(e.to_string()))?;
let config = format::metadata_to_model_config(reader.metadata());
tracker.next_phase();
let is_f16 = reader
.tensor_names()
.iter()
.find(|n| n.ends_with(".weight") && !n.starts_with("__"))
.and_then(|n| reader.get_tensor(n))
.map_or(false, |t| t.dtype == format::TensorDType::F16);
callback(&tracker.to_progress());
let mut encoder = model::Encoder::new(&config);
Self::load_encoder_weights(&reader, &mut encoder, &mut tracker, callback);
encoder.finalize_weights();
tracker.next_phase();
callback(&tracker.to_progress());
let mut decoder = model::Decoder::new(&config);
if is_f16 {
Self::load_decoder_weights_f16(&reader, &mut decoder, &mut tracker, callback);
} else {
Self::load_decoder_weights(&reader, &mut decoder, &mut tracker, callback);
decoder.convert_to_f16();
}
decoder.finalize_weights();
tracker.next_phase();
callback(&tracker.to_progress());
let tokenizer = Self::build_tokenizer(&config, &reader);
tracker.next_phase();
callback(&tracker.to_progress());
let (mel_filters, conv_stem) = match config.audio_frontend {
model::AudioFrontend::MelFilterbank => {
let mel_config = audio::MelConfig {
n_mels: config.n_mels as usize,
..audio::MelConfig::whisper()
};
let mf = Self::read_mel_filterbank(&reader).map_or_else(
|| audio::MelFilterbank::new(&mel_config),
|fb| audio::MelFilterbank::from_filters(fb.data, &mel_config),
);
(Some(mf), None)
}
model::AudioFrontend::LearnedConv => {
let d_model = config.n_audio_state as usize;
let mut stem = audio::ConvStem::new(d_model);
Self::load_conv_stem_weights(&reader, &mut stem);
(None, Some(stem))
}
};
tracker.complete();
callback(&tracker.to_progress());
Ok(Self {
config,
encoder,
decoder,
tokenizer,
mel_filters,
conv_stem,
resampler: None,
weights_loaded: true,
})
}
fn build_tokenizer(
config: &model::ModelConfig,
reader: &format::AprV2ReaderRef<'_>,
) -> tokenizer::Tokenizer {
if config.model_family == format::ModelFamily::Moonshine {
let mut sp = tokenizer::SentencePieceTokenizer::moonshine_default();
if let Some(vocab) = Self::read_vocabulary(reader) {
Self::populate_sentencepiece(&mut sp, &vocab);
}
tokenizer::Tokenizer::SentencePiece(sp)
} else {
tokenizer::Tokenizer::Bpe(Self::read_vocabulary(reader).map_or_else(
tokenizer::BpeTokenizer::with_base_tokens,
tokenizer::BpeTokenizer::from_vocabulary,
))
}
}
fn read_vocabulary(reader: &format::AprV2ReaderRef<'_>) -> Option<tokenizer::Vocabulary> {
let raw = reader.get_tensor_data("__vocab__")?;
tokenizer::Vocabulary::from_bytes(raw)
}
fn read_mel_filterbank(
reader: &format::AprV2ReaderRef<'_>,
) -> Option<format::MelFilterbankData> {
let data = reader.get_tensor_as_f32("__mel_filters__")?;
let entry = reader.get_tensor("__mel_filters__")?;
let shape = &entry.shape;
if shape.len() == 2 {
Some(format::MelFilterbankData::new(
shape[0] as u32,
shape[1] as u32,
data,
))
} else {
None
}
}
fn load_f16_raw(reader: &format::AprV2ReaderRef<'_>, name: &str) -> Option<Vec<u16>> {
let raw = reader.get_tensor_data(name)?;
Some(
raw.chunks_exact(2)
.map(|b| u16::from_le_bytes([b[0], b[1]]))
.collect(),
)
}
fn populate_sentencepiece(
sp: &mut tokenizer::SentencePieceTokenizer,
vocab: &tokenizer::Vocabulary,
) {
for id in 0..vocab.len() as u32 {
if let Some(bytes) = vocab.get_bytes(id) {
if let Ok(piece) = core::str::from_utf8(bytes) {
if !piece.is_empty() {
sp.add_piece(id, piece);
}
}
}
}
}
fn load_encoder_weights(
reader: &format::AprV2ReaderRef<'_>,
encoder: &mut model::Encoder,
tracker: &mut progress::ProgressTracker,
callback: progress::ProgressCallback<'_>,
) {
let n_layers = encoder.n_layers();
if let Some(conv_frontend) = encoder.conv_frontend_mut() {
if let Some(weight) = reader.get_tensor_as_f32("encoder.conv1.weight") {
let target = conv_frontend.conv1.weight_mut();
let len = weight.len().min(target.len());
target[..len].copy_from_slice(&weight[..len]);
}
if let Some(bias) = reader.get_tensor_as_f32("encoder.conv1.bias") {
let target = conv_frontend.conv1.bias_mut();
let len = bias.len().min(target.len());
target[..len].copy_from_slice(&bias[..len]);
}
if let Some(weight) = reader.get_tensor_as_f32("encoder.conv2.weight") {
let target = conv_frontend.conv2.weight_mut();
let len = weight.len().min(target.len());
target[..len].copy_from_slice(&weight[..len]);
}
if let Some(bias) = reader.get_tensor_as_f32("encoder.conv2.bias") {
let target = conv_frontend.conv2.bias_mut();
let len = bias.len().min(target.len());
target[..len].copy_from_slice(&bias[..len]);
}
}
let pe_result = reader
.get_tensor_as_f32("encoder.embed_positions.weight")
.or_else(|| reader.get_tensor_as_f32("encoder.positional_embedding"));
if let Some(pe) = pe_result {
let target = encoder.positional_embedding_mut();
let len = pe.len().min(target.len());
target[..len].copy_from_slice(&pe[..len]);
}
if !encoder.moonshine_blocks().is_empty() {
for layer_idx in 0..n_layers {
let progress = layer_idx as f32 / n_layers as f32;
tracker.update_phase_progress(progress);
callback(&tracker.to_progress());
let block = &mut encoder.moonshine_blocks_mut()[layer_idx];
Self::load_layernorm_nobias_weights(
reader,
&format!("encoder.blocks.{layer_idx}.ln1"),
&mut block.ln1,
);
Self::load_gqa_weights(
reader,
&format!("encoder.blocks.{layer_idx}.attn"),
&mut block.self_attn,
);
Self::load_layernorm_nobias_weights(
reader,
&format!("encoder.blocks.{layer_idx}.ln2"),
&mut block.ln2,
);
Self::load_mlp_weights(
reader,
&format!("encoder.blocks.{layer_idx}.ffn"),
&mut block.ffn,
);
}
} else {
for layer_idx in 0..n_layers {
let progress = layer_idx as f32 / n_layers as f32;
tracker.update_phase_progress(progress);
callback(&tracker.to_progress());
let block = &mut encoder.blocks_mut()[layer_idx];
Self::load_layer_norm_weights(
reader,
&format!("encoder.layers.{layer_idx}.self_attn_layer_norm"),
&mut block.ln1,
);
Self::load_attention_weights(
reader,
&format!("encoder.layers.{layer_idx}.self_attn"),
&mut block.self_attn,
);
Self::load_layer_norm_weights(
reader,
&format!("encoder.layers.{layer_idx}.final_layer_norm"),
&mut block.ln2,
);
Self::load_ffn_weights(
reader,
&format!("encoder.layers.{layer_idx}"),
&mut block.ffn,
);
}
}
if let Some(ln) = encoder.ln_post_rms_mut() {
Self::load_layernorm_nobias_weights(reader, "encoder.layer_norm", ln);
} else {
Self::load_layer_norm_weights(reader, "encoder.layer_norm", encoder.ln_post_mut());
}
}
fn load_decoder_weights(
reader: &format::AprV2ReaderRef<'_>,
decoder: &mut model::Decoder,
tracker: &mut progress::ProgressTracker,
callback: progress::ProgressCallback<'_>,
) {
let n_layers = decoder.n_layers();
let te_result = reader
.get_tensor_as_f32("decoder.embed_tokens.weight")
.or_else(|| reader.get_tensor_as_f32("decoder.token_embedding.weight"))
.or_else(|| reader.get_tensor_as_f32("decoder.token_embedding"));
if let Some(te) = te_result {
let target = decoder.token_embedding_mut();
let len = te.len().min(target.len());
target[..len].copy_from_slice(&te[..len]);
}
let pe_result = reader
.get_tensor_as_f32("decoder.embed_positions.weight")
.or_else(|| reader.get_tensor_as_f32("decoder.positional_embedding"));
if let Some(pe) = pe_result {
let target = decoder.positional_embedding_mut();
let len = pe.len().min(target.len());
target[..len].copy_from_slice(&pe[..len]);
}
if !decoder.moonshine_blocks().is_empty() {
for layer_idx in 0..n_layers {
let progress = layer_idx as f32 / n_layers as f32;
tracker.update_phase_progress(progress);
callback(&tracker.to_progress());
let block = &mut decoder.moonshine_blocks_mut()[layer_idx];
Self::load_layernorm_nobias_weights(
reader,
&format!("decoder.blocks.{layer_idx}.ln1"),
&mut block.ln1,
);
Self::load_gqa_weights(
reader,
&format!("decoder.blocks.{layer_idx}.attn"),
&mut block.self_attn,
);
Self::load_layernorm_nobias_weights(
reader,
&format!("decoder.blocks.{layer_idx}.ln_cross"),
&mut block.ln_cross,
);
Self::load_gqa_weights(
reader,
&format!("decoder.blocks.{layer_idx}.cross_attn"),
&mut block.cross_attn,
);
Self::load_layernorm_nobias_weights(
reader,
&format!("decoder.blocks.{layer_idx}.ln2"),
&mut block.ln2,
);
Self::load_gated_mlp_weights(
reader,
&format!("decoder.blocks.{layer_idx}.ffn"),
&mut block.ffn,
);
}
if let Some(ln) = decoder.ln_post_rms_mut() {
Self::load_layernorm_nobias_weights(reader, "decoder.ln_post", ln);
}
} else {
for layer_idx in 0..n_layers {
let progress = layer_idx as f32 / n_layers as f32;
tracker.update_phase_progress(progress);
callback(&tracker.to_progress());
let block = &mut decoder.blocks_mut()[layer_idx];
Self::load_layer_norm_weights(
reader,
&format!("decoder.layers.{layer_idx}.self_attn_layer_norm"),
&mut block.ln1,
);
Self::load_attention_weights(
reader,
&format!("decoder.layers.{layer_idx}.self_attn"),
&mut block.self_attn,
);
Self::load_layer_norm_weights(
reader,
&format!("decoder.layers.{layer_idx}.encoder_attn_layer_norm"),
&mut block.ln2,
);
Self::load_attention_weights(
reader,
&format!("decoder.layers.{layer_idx}.encoder_attn"),
&mut block.cross_attn,
);
Self::load_layer_norm_weights(
reader,
&format!("decoder.layers.{layer_idx}.final_layer_norm"),
&mut block.ln3,
);
Self::load_ffn_weights(
reader,
&format!("decoder.layers.{layer_idx}"),
&mut block.ffn,
);
}
Self::load_layer_norm_weights(reader, "decoder.layer_norm", decoder.ln_post_mut());
}
decoder.finalize_weights();
}
fn load_decoder_weights_f16(
reader: &format::AprV2ReaderRef<'_>,
decoder: &mut model::Decoder,
tracker: &mut progress::ProgressTracker,
callback: progress::ProgressCallback<'_>,
) {
let n_layers = decoder.n_layers();
let te_result = reader
.get_tensor_as_f32("decoder.embed_tokens.weight")
.or_else(|| reader.get_tensor_as_f32("decoder.token_embedding.weight"))
.or_else(|| reader.get_tensor_as_f32("decoder.token_embedding"));
if let Some(te) = te_result {
let target = decoder.token_embedding_mut();
let len = te.len().min(target.len());
target[..len].copy_from_slice(&te[..len]);
}
let pe_result = reader
.get_tensor_as_f32("decoder.embed_positions.weight")
.or_else(|| reader.get_tensor_as_f32("decoder.positional_embedding"));
if let Some(pe) = pe_result {
let target = decoder.positional_embedding_mut();
let len = pe.len().min(target.len());
target[..len].copy_from_slice(&pe[..len]);
}
for layer_idx in 0..n_layers {
let progress = layer_idx as f32 / n_layers as f32;
tracker.update_phase_progress(progress);
callback(&tracker.to_progress());
let block = &mut decoder.blocks_mut()[layer_idx];
Self::load_layer_norm_weights(
reader,
&format!("decoder.layers.{layer_idx}.self_attn_layer_norm"),
&mut block.ln1,
);
Self::load_attention_weights_f16(
reader,
&format!("decoder.layers.{layer_idx}.self_attn"),
&mut block.self_attn,
);
Self::load_layer_norm_weights(
reader,
&format!("decoder.layers.{layer_idx}.encoder_attn_layer_norm"),
&mut block.ln2,
);
Self::load_attention_weights_f16(
reader,
&format!("decoder.layers.{layer_idx}.encoder_attn"),
&mut block.cross_attn,
);
Self::load_layer_norm_weights(
reader,
&format!("decoder.layers.{layer_idx}.final_layer_norm"),
&mut block.ln3,
);
Self::load_ffn_weights_f16(
reader,
&format!("decoder.layers.{layer_idx}"),
&mut block.ffn,
);
}
Self::load_layer_norm_weights(reader, "decoder.layer_norm", decoder.ln_post_mut());
decoder.finalize_weights();
decoder.convert_embeddings_to_f16();
}
fn load_layer_norm_weights(
reader: &format::AprV2ReaderRef<'_>,
prefix: &str,
ln: &mut model::LayerNorm,
) {
if let Some(weight) = reader.get_tensor_as_f32(&format!("{prefix}.weight")) {
let len = weight.len().min(ln.weight.len());
ln.weight[..len].copy_from_slice(&weight[..len]);
}
if let Some(bias) = reader.get_tensor_as_f32(&format!("{prefix}.bias")) {
let len = bias.len().min(ln.bias.len());
ln.bias[..len].copy_from_slice(&bias[..len]);
}
}
fn load_attention_weights(
reader: &format::AprV2ReaderRef<'_>,
prefix: &str,
attn: &mut model::MultiHeadAttention,
) {
if let Some(q_weight) = reader.get_tensor_as_f32(&format!("{prefix}.q_proj.weight")) {
attn.set_query_weight(&q_weight);
}
if let Some(q_bias) = reader.get_tensor_as_f32(&format!("{prefix}.q_proj.bias")) {
attn.set_query_bias(&q_bias);
}
if let Some(k_weight) = reader.get_tensor_as_f32(&format!("{prefix}.k_proj.weight")) {
attn.set_key_weight(&k_weight);
}
if let Some(k_bias) = reader.get_tensor_as_f32(&format!("{prefix}.k_proj.bias")) {
attn.set_key_bias(&k_bias);
}
if let Some(v_weight) = reader.get_tensor_as_f32(&format!("{prefix}.v_proj.weight")) {
attn.set_value_weight(&v_weight);
}
if let Some(v_bias) = reader.get_tensor_as_f32(&format!("{prefix}.v_proj.bias")) {
attn.set_value_bias(&v_bias);
}
if let Some(out_weight) = reader.get_tensor_as_f32(&format!("{prefix}.out_proj.weight")) {
attn.set_out_weight(&out_weight);
}
if let Some(out_bias) = reader.get_tensor_as_f32(&format!("{prefix}.out_proj.bias")) {
attn.set_out_bias(&out_bias);
}
}
fn load_ffn_weights(
reader: &format::AprV2ReaderRef<'_>,
prefix: &str,
ffn: &mut model::FeedForward,
) {
if let Some(fc1_weight) = reader.get_tensor_as_f32(&format!("{prefix}.fc1.weight")) {
ffn.fc1.set_weight(&fc1_weight);
}
if let Some(fc1_bias) = reader.get_tensor_as_f32(&format!("{prefix}.fc1.bias")) {
ffn.fc1.set_bias(&fc1_bias);
}
if let Some(fc2_weight) = reader.get_tensor_as_f32(&format!("{prefix}.fc2.weight")) {
ffn.fc2.set_weight(&fc2_weight);
}
if let Some(fc2_bias) = reader.get_tensor_as_f32(&format!("{prefix}.fc2.bias")) {
ffn.fc2.set_bias(&fc2_bias);
}
}
fn load_attention_weights_f16(
reader: &format::AprV2ReaderRef<'_>,
prefix: &str,
attn: &mut model::MultiHeadAttention,
) {
if let Some(q_weight) = Self::load_f16_raw(reader, &format!("{prefix}.q_proj.weight")) {
attn.set_query_weight_f16(&q_weight);
}
if let Some(k_weight) = Self::load_f16_raw(reader, &format!("{prefix}.k_proj.weight")) {
attn.set_key_weight_f16(&k_weight);
}
if let Some(v_weight) = Self::load_f16_raw(reader, &format!("{prefix}.v_proj.weight")) {
attn.set_value_weight_f16(&v_weight);
}
if let Some(out_weight) = Self::load_f16_raw(reader, &format!("{prefix}.out_proj.weight")) {
attn.set_out_weight_f16(&out_weight);
}
if let Some(q_bias) = reader.get_tensor_as_f32(&format!("{prefix}.q_proj.bias")) {
attn.set_query_bias(&q_bias);
}
if let Some(k_bias) = reader.get_tensor_as_f32(&format!("{prefix}.k_proj.bias")) {
attn.set_key_bias(&k_bias);
}
if let Some(v_bias) = reader.get_tensor_as_f32(&format!("{prefix}.v_proj.bias")) {
attn.set_value_bias(&v_bias);
}
if let Some(out_bias) = reader.get_tensor_as_f32(&format!("{prefix}.out_proj.bias")) {
attn.set_out_bias(&out_bias);
}
}
fn load_ffn_weights_f16(
reader: &format::AprV2ReaderRef<'_>,
prefix: &str,
ffn: &mut model::FeedForward,
) {
if let Some(fc1_weight) = Self::load_f16_raw(reader, &format!("{prefix}.fc1.weight")) {
ffn.fc1.set_weight_f16(&fc1_weight);
}
if let Some(fc1_bias) = reader.get_tensor_as_f32(&format!("{prefix}.fc1.bias")) {
ffn.fc1.set_bias(&fc1_bias);
}
if let Some(fc2_weight) = Self::load_f16_raw(reader, &format!("{prefix}.fc2.weight")) {
ffn.fc2.set_weight_f16(&fc2_weight);
}
if let Some(fc2_bias) = reader.get_tensor_as_f32(&format!("{prefix}.fc2.bias")) {
ffn.fc2.set_bias(&fc2_bias);
}
}
#[allow(dead_code)]
fn load_rms_norm_weights(
reader: &format::AprV2ReaderRef<'_>,
prefix: &str,
rms: &mut model::lfm2::layer::RmsNorm,
) {
if let Some(weight) = reader.get_tensor_as_f32(&format!("{prefix}.weight")) {
let len = weight.len().min(rms.weight.len());
rms.weight[..len].copy_from_slice(&weight[..len]);
}
}
fn load_layernorm_nobias_weights(
reader: &format::AprV2ReaderRef<'_>,
prefix: &str,
ln: &mut model::lfm2::layer::LayerNormNoBias,
) {
if let Some(weight) = reader.get_tensor_as_f32(&format!("{prefix}.weight")) {
let len = weight.len().min(ln.weight.len());
ln.weight[..len].copy_from_slice(&weight[..len]);
}
}
fn load_gqa_weights(
reader: &format::AprV2ReaderRef<'_>,
prefix: &str,
gqa: &mut model::lfm2::gqa::GroupedQueryAttention,
) {
if let Some(w) = reader.get_tensor_as_f32(&format!("{prefix}.q.weight")) {
let len = w.len().min(gqa.w_q.len());
gqa.w_q[..len].copy_from_slice(&w[..len]);
}
if let Some(w) = reader.get_tensor_as_f32(&format!("{prefix}.k.weight")) {
let len = w.len().min(gqa.w_k.len());
gqa.w_k[..len].copy_from_slice(&w[..len]);
}
if let Some(w) = reader.get_tensor_as_f32(&format!("{prefix}.v.weight")) {
let len = w.len().min(gqa.w_v.len());
gqa.w_v[..len].copy_from_slice(&w[..len]);
}
if let Some(w) = reader.get_tensor_as_f32(&format!("{prefix}.o.weight")) {
let len = w.len().min(gqa.w_o.len());
gqa.w_o[..len].copy_from_slice(&w[..len]);
}
}
fn load_mlp_weights(
reader: &format::AprV2ReaderRef<'_>,
prefix: &str,
ffn: &mut model::lfm2::mlp::MlpFfn,
) {
if let Some(w) = reader.get_tensor_as_f32(&format!("{prefix}.fc1.weight")) {
let len = w.len().min(ffn.fc1.len());
ffn.fc1[..len].copy_from_slice(&w[..len]);
}
if let Some(b) = reader.get_tensor_as_f32(&format!("{prefix}.fc1.bias")) {
ffn.b1 = Some(b);
}
if let Some(w) = reader.get_tensor_as_f32(&format!("{prefix}.fc2.weight")) {
let len = w.len().min(ffn.fc2.len());
ffn.fc2[..len].copy_from_slice(&w[..len]);
}
if let Some(b) = reader.get_tensor_as_f32(&format!("{prefix}.fc2.bias")) {
ffn.b2 = Some(b);
}
}
fn load_gated_mlp_weights(
reader: &format::AprV2ReaderRef<'_>,
prefix: &str,
ffn: &mut model::lfm2::mlp::GatedMlpFfn,
) {
if let Some(w) = reader.get_tensor_as_f32(&format!("{prefix}.fc1.weight")) {
let len = w.len().min(ffn.fc1.len());
ffn.fc1[..len].copy_from_slice(&w[..len]);
}
if let Some(b) = reader.get_tensor_as_f32(&format!("{prefix}.fc1.bias")) {
ffn.b1 = Some(b);
}
if let Some(w) = reader.get_tensor_as_f32(&format!("{prefix}.fc2.weight")) {
let len = w.len().min(ffn.fc2.len());
ffn.fc2[..len].copy_from_slice(&w[..len]);
}
if let Some(b) = reader.get_tensor_as_f32(&format!("{prefix}.fc2.bias")) {
ffn.b2 = Some(b);
}
}
fn load_conv_stem_weights(reader: &format::AprV2ReaderRef<'_>, stem: &mut audio::ConvStem) {
if let Some(w) = reader.get_tensor_as_f32("encoder.conv1.weight") {
let target = stem.conv1.weight_mut();
let len = w.len().min(target.len());
target[..len].copy_from_slice(&w[..len]);
}
if let Some(w) = reader.get_tensor_as_f32("encoder.conv2.weight") {
let target = stem.conv2.weight_mut();
let len = w.len().min(target.len());
target[..len].copy_from_slice(&w[..len]);
}
if let Some(b) = reader.get_tensor_as_f32("encoder.conv2.bias") {
let target = stem.conv2.bias_mut();
let len = b.len().min(target.len());
target[..len].copy_from_slice(&b[..len]);
}
if let Some(w) = reader.get_tensor_as_f32("encoder.conv3.weight") {
let target = stem.conv3.weight_mut();
let len = w.len().min(target.len());
target[..len].copy_from_slice(&w[..len]);
}
if let Some(b) = reader.get_tensor_as_f32("encoder.conv3.bias") {
let target = stem.conv3.bias_mut();
let len = b.len().min(target.len());
target[..len].copy_from_slice(&b[..len]);
}
if let Some(w) = reader.get_tensor_as_f32("encoder.groupnorm.weight") {
let len = w.len().min(stem.groupnorm.weight.len());
stem.groupnorm.weight[..len].copy_from_slice(&w[..len]);
}
if let Some(b) = reader.get_tensor_as_f32("encoder.groupnorm.bias") {
let len = b.len().min(stem.groupnorm.bias.len());
stem.groupnorm.bias[..len].copy_from_slice(&b[..len]);
}
}
pub fn encoder_mut(&mut self) -> &mut model::Encoder {
&mut self.encoder
}
pub fn decoder_mut(&mut self) -> &mut model::Decoder {
&mut self.decoder
}
pub fn encoder(&self) -> &model::Encoder {
&self.encoder
}
pub fn decoder(&self) -> &model::Decoder {
&self.decoder
}
#[cfg(feature = "realizar-gpu")]
pub fn into_cuda(self, device_ordinal: i32) -> WhisperResult<cuda::WhisperCuda> {
let mel_filters = self.mel_filters.ok_or_else(|| {
WhisperError::Audio("CUDA requires mel filterbank (Whisper model)".into())
})?;
let bpe_tokenizer = match self.tokenizer {
tokenizer::Tokenizer::Bpe(bpe) => bpe,
tokenizer::Tokenizer::SentencePiece(_) => {
return Err(WhisperError::Model(
"CUDA backend requires Whisper BPE tokenizer".into(),
));
}
};
cuda::WhisperCuda::new_with_components(
self.encoder,
self.decoder,
self.config,
bpe_tokenizer,
mel_filters,
device_ordinal,
)
}
#[must_use]
pub const fn conv_stem(&self) -> Option<&audio::ConvStem> {
self.conv_stem.as_ref()
}
#[must_use]
pub const fn mel_filters(&self) -> Option<&audio::MelFilterbank> {
self.mel_filters.as_ref()
}
pub fn transcribe_batch(
&self,
audio_batch: &[Vec<f32>],
options: TranscribeOptions,
) -> WhisperResult<BatchTranscriptionResult> {
if audio_batch.is_empty() {
return Err(WhisperError::Audio("empty batch".into()));
}
let start_time = std::time::Instant::now();
let mut results = Vec::with_capacity(audio_batch.len());
for audio in audio_batch {
let result = self.transcribe(audio, options.clone())?;
results.push(result);
}
let total_duration_secs = start_time.elapsed().as_secs_f32();
Ok(BatchTranscriptionResult {
results,
total_duration_secs,
})
}
pub fn transcribe_audio_batch(
&self,
batch: &audio::AudioBatch,
options: TranscribeOptions,
) -> WhisperResult<BatchTranscriptionResult> {
if batch.is_empty() {
return Err(WhisperError::Audio("empty batch".into()));
}
let start_time = std::time::Instant::now();
let preprocessor = audio::BatchPreprocessor::new(audio::MelConfig::default());
let mel_result = preprocessor.process_batch(batch)?;
let mut results = Vec::with_capacity(batch.len());
let language = options.language.clone().unwrap_or_else(|| "en".to_string());
for mel in &mel_result.mels {
let audio_features = self.encode(mel)?;
let initial_tokens = self.get_initial_tokens(&language, options.task);
let tokens = self.decode(&audio_features, &initial_tokens, &options)?;
let segments = if timestamps::has_timestamps(&tokens) {
timestamps::extract_segments(&tokens, |ts| self.tokenizer.decode(ts).ok())
} else {
Vec::new()
};
let text = self.tokenizer.decode(&tokens)?;
results.push(TranscriptionResult {
text,
language: language.clone(),
segments,
profiling: None,
});
}
let total_duration_secs = start_time.elapsed().as_secs_f32();
Ok(BatchTranscriptionResult {
results,
total_duration_secs,
})
}
#[must_use]
pub fn create_audio_batch(audio_segments: &[Vec<f32>]) -> audio::AudioBatch {
let mut batch = audio::AudioBatch::with_default_config();
for segment in audio_segments {
batch.add_segment(segment.clone());
}
batch
}
pub fn transcribe_batch_optimized(
&self,
audio_batch: &[Vec<f32>],
options: TranscribeOptions,
) -> WhisperResult<BatchTranscriptionResult> {
if audio_batch.is_empty() {
return Err(WhisperError::Audio("empty batch".into()));
}
let start_time = std::time::Instant::now();
let mut mels = Vec::with_capacity(audio_batch.len());
for audio in audio_batch {
let mel = self.compute_mel(audio)?;
mels.push(mel);
}
let encoder_outputs = self.encoder.forward_batch(&mels)?;
let mut results = Vec::with_capacity(audio_batch.len());
let language = options.language.clone().unwrap_or_else(|| "en".to_string());
for features in &encoder_outputs {
let initial_tokens = self.get_initial_tokens(&language, options.task);
let tokens = self.decode(features, &initial_tokens, &options)?;
let segments = if timestamps::has_timestamps(&tokens) {
timestamps::extract_segments(&tokens, |ts| self.tokenizer.decode(ts).ok())
} else {
Vec::new()
};
let text = self.tokenizer.decode(&tokens)?;
results.push(TranscriptionResult {
text,
language: language.clone(),
segments,
profiling: None,
});
}
let total_duration_secs = start_time.elapsed().as_secs_f32();
Ok(BatchTranscriptionResult {
results,
total_duration_secs,
})
}
pub fn transcribe_with_vad(
&self,
audio: &[f32],
options: TranscribeOptions,
vad_config: Option<vad::VadConfig>,
) -> WhisperResult<VadTranscriptionResult> {
let start_time = std::time::Instant::now();
let config = vad_config.unwrap_or_default();
let mut vad = vad::VoiceActivityDetector::new(config);
let speech_segments = vad.detect(audio);
if speech_segments.is_empty() {
return Ok(VadTranscriptionResult {
text: String::new(),
language: options.language.unwrap_or_else(|| "en".to_string()),
segments: Vec::new(),
speech_segments: Vec::new(),
total_duration_secs: start_time.elapsed().as_secs_f32(),
speech_duration_secs: 0.0,
});
}
let sample_rate = audio::SAMPLE_RATE as f32;
let mut speech_audios = Vec::with_capacity(speech_segments.len());
let mut speech_duration = 0.0f32;
for segment in &speech_segments {
let start_sample = (segment.start * sample_rate) as usize;
let end_sample = ((segment.end * sample_rate) as usize).min(audio.len());
if end_sample > start_sample {
speech_audios.push((
segment.start,
segment.end,
audio[start_sample..end_sample].to_vec(),
));
speech_duration += segment.duration();
}
}
let mut all_segments = Vec::new();
let mut full_text = String::new();
let language = options.language.clone().unwrap_or_else(|| "en".to_string());
for (seg_start, seg_end, speech_audio) in &speech_audios {
let result = self.transcribe(speech_audio, options.clone())?;
let segment = VadSpeechSegment {
start: *seg_start,
end: *seg_end,
text: result.text.clone(),
tokens: result
.segments
.first()
.map(|s| s.tokens.clone())
.unwrap_or_default(),
};
if !full_text.is_empty() {
full_text.push(' ');
}
full_text.push_str(&result.text);
all_segments.push(segment);
}
let total_duration_secs = start_time.elapsed().as_secs_f32();
Ok(VadTranscriptionResult {
text: full_text,
language,
segments: all_segments,
speech_segments: speech_segments
.into_iter()
.map(|s| (s.start, s.end))
.collect(),
total_duration_secs,
speech_duration_secs: speech_duration,
})
}
pub fn transcribe_with_silence_detection(
&self,
audio: &[f32],
options: TranscribeOptions,
silence_config: Option<vad::SilenceConfig>,
) -> WhisperResult<VadTranscriptionResult> {
let start_time = std::time::Instant::now();
let config = silence_config.unwrap_or_default();
let mut detector = vad::SilenceDetector::new(config, audio::SAMPLE_RATE);
let frame_size = 480; let silence_segments = detector.detect(audio, frame_size);
let speech_segments = self.invert_silence_segments(&silence_segments, audio.len());
if speech_segments.is_empty() {
return Ok(VadTranscriptionResult {
text: String::new(),
language: options.language.unwrap_or_else(|| "en".to_string()),
segments: Vec::new(),
speech_segments: Vec::new(),
total_duration_secs: start_time.elapsed().as_secs_f32(),
speech_duration_secs: 0.0,
});
}
let sample_rate = audio::SAMPLE_RATE as f32;
let mut all_segments = Vec::new();
let mut full_text = String::new();
let mut speech_duration = 0.0f32;
let language = options.language.clone().unwrap_or_else(|| "en".to_string());
for (start, end) in &speech_segments {
let start_sample = (start * sample_rate) as usize;
let end_sample = ((end * sample_rate) as usize).min(audio.len());
if end_sample > start_sample {
let speech_audio = &audio[start_sample..end_sample];
let result = self.transcribe(speech_audio, options.clone())?;
let segment = VadSpeechSegment {
start: *start,
end: *end,
text: result.text.clone(),
tokens: result
.segments
.first()
.map(|s| s.tokens.clone())
.unwrap_or_default(),
};
if !full_text.is_empty() {
full_text.push(' ');
}
full_text.push_str(&result.text);
speech_duration += end - start;
all_segments.push(segment);
}
}
let total_duration_secs = start_time.elapsed().as_secs_f32();
Ok(VadTranscriptionResult {
text: full_text,
language,
segments: all_segments,
speech_segments,
total_duration_secs,
speech_duration_secs: speech_duration,
})
}
fn invert_silence_segments(
&self,
silence_segments: &[vad::SilenceSegment],
audio_len: usize,
) -> Vec<(f32, f32)> {
let _ = self; let sample_rate = audio::SAMPLE_RATE as f32;
let total_duration = audio_len as f32 / sample_rate;
let mut speech_segments = Vec::new();
let mut current_pos = 0.0f32;
for silence in silence_segments {
if silence.start > current_pos {
speech_segments.push((current_pos, silence.start));
}
current_pos = silence.end;
}
if current_pos < total_duration {
speech_segments.push((current_pos, total_duration));
}
speech_segments
}
pub fn transcribe_partial(
&self,
partial_audio: &[f32],
options: TranscribeOptions,
is_final: bool,
) -> WhisperResult<PartialTranscriptionResult> {
let start_time = std::time::Instant::now();
let min_samples = (audio::SAMPLE_RATE as f32 * 0.5) as usize;
if partial_audio.len() < min_samples {
return Ok(PartialTranscriptionResult {
text: String::new(),
language: options.language.unwrap_or_else(|| "en".to_string()),
is_final,
confidence: 0.0,
duration_secs: partial_audio.len() as f32 / audio::SAMPLE_RATE as f32,
processing_time_secs: start_time.elapsed().as_secs_f32(),
});
}
let result = self.transcribe(partial_audio, options)?;
let processing_time = start_time.elapsed().as_secs_f32();
Ok(PartialTranscriptionResult {
text: result.text,
language: result.language,
is_final,
confidence: 1.0, duration_secs: partial_audio.len() as f32 / audio::SAMPLE_RATE as f32,
processing_time_secs: processing_time,
})
}
#[must_use]
pub fn create_streaming_session(
&self,
options: TranscribeOptions,
input_sample_rate: u32,
) -> StreamingSession<'_> {
let streaming_config = audio::StreamingConfig::with_sample_rate(input_sample_rate);
let processor = audio::StreamingProcessor::new(streaming_config);
StreamingSession {
whisper: self,
processor,
options,
last_partial_text: String::new(),
}
}
pub fn transcribe_and_summarize(
&self,
audio: &[f32],
transcribe_options: TranscribeOptions,
summarize_options: SummarizeOptions<'_>,
) -> WhisperResult<TranscribeSummaryResult> {
let transcription = self.transcribe(audio, transcribe_options)?;
if transcription.text.trim().is_empty() {
return Ok(TranscribeSummaryResult {
transcription,
summary: String::new(),
generation_stats: None,
});
}
let input_tokens = summarize_options.tokenizer.encode(&transcription.text);
let (output_tokens, stats) = summarize_options.model.generate_with_stats(
&input_tokens,
summarize_options.max_tokens,
summarize_options.temperature,
Some(|_token: u32, _idx: usize| true), )?;
let summary = summarize_options.tokenizer.decode(&output_tokens);
Ok(TranscribeSummaryResult {
transcription,
summary,
generation_stats: Some(stats),
})
}
}
#[derive(Debug, Clone)]
pub struct PartialTranscriptionResult {
pub text: String,
pub language: String,
pub is_final: bool,
pub confidence: f32,
pub duration_secs: f32,
pub processing_time_secs: f32,
}
impl PartialTranscriptionResult {
#[must_use]
pub fn has_text(&self) -> bool {
!self.text.is_empty()
}
#[must_use]
pub fn is_empty_interim(&self) -> bool {
self.text.is_empty() && !self.is_final
}
#[must_use]
pub fn real_time_factor(&self) -> f32 {
if self.duration_secs <= 0.0 {
0.0
} else {
self.processing_time_secs / self.duration_secs
}
}
}
#[derive(Debug)]
pub struct StreamingSession<'a> {
whisper: &'a WhisperApr,
processor: audio::StreamingProcessor,
options: TranscribeOptions,
last_partial_text: String,
}
impl StreamingSession<'_> {
pub fn push(&mut self, audio: &[f32]) -> WhisperResult<Option<PartialTranscriptionResult>> {
self.processor.push_audio(audio);
self.processor.process();
if self.processor.has_partial() {
if let Some(partial_audio) = self.processor.get_partial() {
let result =
self.whisper
.transcribe_partial(&partial_audio, self.options.clone(), false)?;
if result.text != self.last_partial_text {
result.text.clone_into(&mut self.last_partial_text);
return Ok(Some(result));
}
}
}
Ok(None)
}
#[must_use]
pub fn has_chunk(&self) -> bool {
self.processor.has_chunk()
}
#[must_use]
pub fn has_events(&self) -> bool {
self.processor.has_events()
}
pub fn drain_events(&mut self) -> Vec<audio::StreamingEvent> {
self.processor.drain_events()
}
pub fn finalize(&mut self) -> WhisperResult<PartialTranscriptionResult> {
let chunk = self
.processor
.get_chunk()
.ok_or_else(|| WhisperError::Audio("no chunk ready for finalization".into()))?;
let result = self
.whisper
.transcribe_partial(&chunk, self.options.clone(), true)?;
self.last_partial_text.clear();
Ok(result)
}
pub fn flush(&mut self) -> WhisperResult<Option<PartialTranscriptionResult>> {
if let Some(chunk) = self.processor.flush() {
let result = self
.whisper
.transcribe_partial(&chunk, self.options.clone(), true)?;
self.last_partial_text.clear();
Ok(Some(result))
} else {
Ok(None)
}
}
pub fn reset(&mut self) {
self.processor.reset();
self.last_partial_text.clear();
}
#[must_use]
pub fn state(&self) -> audio::ProcessorState {
self.processor.state()
}
#[must_use]
pub fn chunk_progress(&self) -> f32 {
self.processor.chunk_progress()
}
#[must_use]
pub fn partial_duration(&self) -> f32 {
self.processor.partial_duration()
}
pub fn set_partial_threshold(&mut self, seconds: f32) {
self.processor.set_partial_threshold(seconds);
}
}
#[derive(Debug, Clone)]
pub struct VadTranscriptionResult {
pub text: String,
pub language: String,
pub segments: Vec<VadSpeechSegment>,
pub speech_segments: Vec<(f32, f32)>,
pub total_duration_secs: f32,
pub speech_duration_secs: f32,
}
impl VadTranscriptionResult {
#[must_use]
pub fn num_segments(&self) -> usize {
self.segments.len()
}
#[must_use]
pub fn has_speech(&self) -> bool {
!self.segments.is_empty()
}
#[must_use]
pub fn silence_ratio(&self, audio_duration: f32) -> f32 {
if audio_duration <= 0.0 {
return 1.0;
}
1.0 - (self.speech_duration_secs / audio_duration)
}
#[must_use]
pub fn first_segment(&self) -> Option<&VadSpeechSegment> {
self.segments.first()
}
#[must_use]
pub fn last_segment(&self) -> Option<&VadSpeechSegment> {
self.segments.last()
}
pub fn iter(&self) -> impl Iterator<Item = &VadSpeechSegment> {
self.segments.iter()
}
}
#[derive(Debug, Clone)]
pub struct VadSpeechSegment {
pub start: f32,
pub end: f32,
pub text: String,
pub tokens: Vec<u32>,
}
impl VadSpeechSegment {
#[must_use]
pub fn duration(&self) -> f32 {
self.end - self.start
}
#[must_use]
pub fn has_text(&self) -> bool {
!self.text.is_empty()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_default_options() {
let options = TranscribeOptions::default();
assert!(options.language.is_none());
assert_eq!(options.task, Task::Transcribe);
assert!(!options.word_timestamps);
assert!(options.prompt.is_none());
assert!(options.hotwords.is_empty());
}
#[test]
fn test_options_with_prompt() {
let options = TranscribeOptions {
prompt: Some("This lecture covers AWS, YAML, and Rust programming.".into()),
..TranscribeOptions::default()
};
assert_eq!(
options.prompt.as_deref(),
Some("This lecture covers AWS, YAML, and Rust programming.")
);
}
#[test]
fn test_options_with_hotwords() {
let options = TranscribeOptions {
hotwords: vec!["AWS".into(), "YAML".into(), "SIMD".into(), "Rust".into()],
..TranscribeOptions::default()
};
assert_eq!(options.hotwords.len(), 4);
assert_eq!(options.hotwords[0], "AWS");
}
#[test]
fn test_options_with_prompt_and_hotwords() {
let options = TranscribeOptions {
language: Some("en".into()),
prompt: Some("Technical lecture on cloud computing".into()),
hotwords: vec!["Kubernetes".into(), "Docker".into()],
..TranscribeOptions::default()
};
assert!(options.prompt.is_some());
assert_eq!(options.hotwords.len(), 2);
assert_eq!(options.language, Some("en".into()));
}
#[test]
fn test_decoding_strategy_default() {
let strategy = DecodingStrategy::default();
assert!(matches!(strategy, DecodingStrategy::Greedy));
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_whisper_tiny() {
let whisper = WhisperApr::tiny();
assert_eq!(whisper.model_type(), ModelType::Tiny);
assert_eq!(whisper.config().n_audio_layer, 4);
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_whisper_base() {
let whisper = WhisperApr::base();
assert_eq!(whisper.model_type(), ModelType::Base);
assert_eq!(whisper.config().n_audio_layer, 6);
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_whisper_memory_size() {
let tiny = WhisperApr::tiny();
let base = WhisperApr::base();
assert!(tiny.memory_size() < base.memory_size());
assert!(tiny.memory_size() > 100_000_000); }
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_whisper_initial_tokens() {
let whisper = WhisperApr::tiny();
let tokens = whisper.get_initial_tokens("en", Task::Transcribe);
assert_eq!(tokens[0], tokenizer::special_tokens::SOT);
assert!(tokens.len() >= 3);
let translate_tokens = whisper.get_initial_tokens("es", Task::Translate);
assert!(translate_tokens.contains(&tokenizer::special_tokens::TRANSLATE));
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_whisper_set_resampler() {
let mut whisper = WhisperApr::tiny();
assert!(whisper.resampler.is_none());
whisper.set_resampler(44100).expect("should succeed");
assert!(whisper.resampler.is_some());
whisper.set_resampler(16000).expect("should succeed");
assert!(whisper.resampler.is_none());
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_whisper_resample_passthrough() {
let whisper = WhisperApr::tiny();
let audio = vec![0.1, 0.2, 0.3, 0.4];
let resampled = whisper.resample(&audio).expect("should succeed");
assert_eq!(resampled, audio);
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_whisper_resample_with_resampler() {
let mut whisper = WhisperApr::tiny();
whisper.set_resampler(32000).expect("should succeed");
let n_samples = 16000;
let audio: Vec<f32> = (0..n_samples)
.map(|i| (2.0 * std::f32::consts::PI * 440.0 * i as f32 / 32000.0).sin())
.collect();
let resampled = whisper.resample(&audio).expect("should succeed");
assert!(resampled.len() > n_samples / 3);
assert!(resampled.len() < n_samples);
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_whisper_tokenizer() {
let whisper = WhisperApr::tiny();
let tokenizer = whisper.tokenizer();
let vocab_size = tokenizer.vocab_size();
assert!(
vocab_size >= 256,
"vocab_size should include base byte tokens"
);
}
#[test]
fn test_transcribe_options_with_language() {
let options = TranscribeOptions {
language: Some("es".to_string()),
task: Task::Transcribe,
strategy: DecodingStrategy::Greedy,
word_timestamps: false,
profile: false,
..Default::default()
};
assert_eq!(options.language, Some("es".to_string()));
}
#[test]
fn test_transcribe_options_beam_search() {
let options = TranscribeOptions {
language: None,
task: Task::Transcribe,
strategy: DecodingStrategy::BeamSearch {
beam_size: 5,
temperature: 0.0,
patience: 1.0,
},
word_timestamps: false,
profile: false,
..Default::default()
};
assert!(matches!(
options.strategy,
DecodingStrategy::BeamSearch { .. }
));
}
#[test]
fn test_segment_struct() {
let segment = Segment {
start: 0.0,
end: 2.5,
text: "Hello world".to_string(),
tokens: vec![1, 2, 3],
};
assert!((segment.start - 0.0).abs() < f32::EPSILON);
assert!((segment.end - 2.5).abs() < f32::EPSILON);
assert_eq!(segment.text, "Hello world");
assert_eq!(segment.tokens.len(), 3);
}
#[test]
fn test_transcription_result_struct() {
let result = TranscriptionResult {
text: "Test transcription".to_string(),
language: "en".to_string(),
segments: vec![],
profiling: None,
};
assert_eq!(result.text, "Test transcription");
assert_eq!(result.language, "en");
assert!(result.segments.is_empty());
}
#[test]
fn test_load_from_apr_basic() {
let data = format::create_test_apr();
let result = WhisperApr::load_from_apr(&data);
assert!(result.is_ok());
let whisper = result.expect("should load");
assert_eq!(whisper.model_type(), ModelType::Tiny);
}
#[test]
fn test_load_from_apr_with_progress_callback() {
let data = format::create_test_apr();
let mut progress_updates = Vec::new();
let mut callback = |p: &progress::Progress| {
progress_updates.push(p.percent());
};
let result = WhisperApr::load_from_apr_with_progress(&data, &mut callback);
assert!(result.is_ok());
assert!(!progress_updates.is_empty());
}
#[test]
fn test_load_from_apr_invalid_magic() {
let mut data = format::create_test_apr();
data[0] = b'X';
let result = WhisperApr::load_from_apr(&data);
assert!(result.is_err());
}
#[test]
fn test_load_from_apr_too_short() {
let data = vec![b'A', b'P', b'R', b'1'];
let result = WhisperApr::load_from_apr(&data);
assert!(result.is_err());
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_encoder_mut_accessor() {
let mut whisper = WhisperApr::tiny();
let encoder = whisper.encoder_mut();
assert_eq!(encoder.n_layers(), 4);
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_decoder_mut_accessor() {
let mut whisper = WhisperApr::tiny();
let decoder = whisper.decoder_mut();
assert_eq!(decoder.n_layers(), 4);
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_full_pipeline_tiny_model() {
let whisper = WhisperApr::tiny();
assert_eq!(whisper.model_type(), ModelType::Tiny);
assert_eq!(whisper.config().n_audio_layer, 4);
assert_eq!(whisper.config().n_text_layer, 4);
assert_eq!(whisper.config().n_audio_state, 384);
assert_eq!(whisper.config().n_text_state, 384);
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_full_pipeline_base_model() {
let whisper = WhisperApr::base();
assert_eq!(whisper.model_type(), ModelType::Base);
assert_eq!(whisper.config().n_audio_layer, 6);
assert_eq!(whisper.config().n_text_layer, 6);
assert_eq!(whisper.config().n_audio_state, 512);
assert_eq!(whisper.config().n_text_state, 512);
}
#[test]
fn test_transcribe_options_all_strategies() {
let opts_greedy = TranscribeOptions::default();
assert!(matches!(opts_greedy.strategy, DecodingStrategy::Greedy));
let opts_beam = TranscribeOptions {
language: Some("en".to_string()),
task: Task::Transcribe,
strategy: DecodingStrategy::BeamSearch {
beam_size: 5,
temperature: 0.0,
patience: 1.0,
},
word_timestamps: false,
profile: false,
..Default::default()
};
assert!(matches!(
opts_beam.strategy,
DecodingStrategy::BeamSearch { .. }
));
let opts_sampling = TranscribeOptions {
language: None,
task: Task::Translate,
strategy: DecodingStrategy::Sampling {
temperature: 0.7,
top_k: Some(50),
top_p: Some(0.9),
},
word_timestamps: true,
profile: false,
..Default::default()
};
assert!(matches!(
opts_sampling.strategy,
DecodingStrategy::Sampling { .. }
));
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_memory_estimation_consistency() {
let tiny = WhisperApr::tiny();
let base = WhisperApr::base();
assert!(base.config().weights_memory_mb() > tiny.config().weights_memory_mb());
assert!(base.config().peak_memory_mb() > tiny.config().peak_memory_mb());
assert!(base.config().parameter_count() > tiny.config().parameter_count());
}
#[test]
fn test_simd_operations_integration() {
use crate::simd;
let a = vec![1.0, 2.0, 3.0, 4.0];
let b = vec![5.0, 6.0, 7.0, 8.0];
let sum = simd::add(&a, &b);
assert_eq!(sum.len(), 4);
assert!((sum[0] - 6.0).abs() < 1e-5);
assert!((sum[3] - 12.0).abs() < 1e-5);
let softmax = simd::softmax(&a);
let sum_softmax: f32 = softmax.iter().sum();
assert!((sum_softmax - 1.0).abs() < 1e-5);
let mat_a = vec![1.0, 2.0, 3.0, 4.0]; let mat_b = vec![5.0, 6.0, 7.0, 8.0]; let result = simd::matmul(&mat_a, &mat_b, 2, 2, 2);
assert_eq!(result.len(), 4);
}
#[test]
fn test_memory_pool_integration() {
use crate::memory::{get_buffer, pool_stats, return_buffer, MemoryPool};
let pool = MemoryPool::new();
let buf1 = pool.get(1024);
assert_eq!(buf1.len(), 1024);
pool.return_buffer(buf1);
let buf2 = pool.get(1024);
assert_eq!(buf2.len(), 1024);
let stats = pool.stats();
assert_eq!(stats.hits, 1);
let tlbuf = get_buffer(512);
assert_eq!(tlbuf.len(), 512);
return_buffer(tlbuf);
let tl_stats = pool_stats();
assert!(tl_stats.allocations > 0);
}
#[test]
fn test_audio_resampling_integration() {
use audio::SincResampler;
let resampler = SincResampler::new(44100, 16000).expect("resampler should work");
let duration_ms = 100;
let samples_at_44100 = (44100 * duration_ms) / 1000;
let input: Vec<f32> = (0..samples_at_44100)
.map(|i| {
let t = i as f32 / 44100.0;
(2.0 * std::f32::consts::PI * 440.0 * t).sin()
})
.collect();
let output = resampler.resample(&input).expect("resample should work");
let expected_len = (input.len() as f32 * 16000.0 / 44100.0) as usize;
assert!(
(output.len() as i32 - expected_len as i32).abs() < 10,
"output len {} vs expected ~{}",
output.len(),
expected_len
);
}
#[test]
fn test_vad_integration() {
use vad::VadConfig;
let config = VadConfig::default();
assert!(config.energy_threshold > 0.0);
assert!(config.zcr_threshold > 0.0);
assert!(config.min_speech_frames > 0);
}
#[test]
fn test_vad_transcription_result_new() {
let result = VadTranscriptionResult {
text: "hello world".to_string(),
language: "en".to_string(),
segments: vec![],
speech_segments: vec![],
total_duration_secs: 1.0,
speech_duration_secs: 0.5,
};
assert_eq!(result.text, "hello world");
assert_eq!(result.language, "en");
assert!(!result.has_speech());
assert_eq!(result.num_segments(), 0);
}
#[test]
fn test_vad_transcription_result_with_segments() {
let result = VadTranscriptionResult {
text: "hello world".to_string(),
language: "en".to_string(),
segments: vec![
VadSpeechSegment {
start: 0.0,
end: 1.0,
text: "hello".to_string(),
tokens: vec![1, 2],
},
VadSpeechSegment {
start: 1.5,
end: 2.5,
text: "world".to_string(),
tokens: vec![3, 4],
},
],
speech_segments: vec![(0.0, 1.0), (1.5, 2.5)],
total_duration_secs: 3.0,
speech_duration_secs: 2.0,
};
assert!(result.has_speech());
assert_eq!(result.num_segments(), 2);
assert!(result.first_segment().is_some());
assert!(result.last_segment().is_some());
assert_eq!(
result.first_segment().map(|s| &s.text),
Some(&"hello".to_string())
);
assert_eq!(
result.last_segment().map(|s| &s.text),
Some(&"world".to_string())
);
}
#[test]
fn test_vad_transcription_result_silence_ratio() {
let result = VadTranscriptionResult {
text: String::new(),
language: "en".to_string(),
segments: vec![],
speech_segments: vec![],
total_duration_secs: 1.0,
speech_duration_secs: 0.5,
};
let ratio = result.silence_ratio(2.0);
assert!((ratio - 0.75).abs() < 0.01); }
#[test]
fn test_vad_transcription_result_silence_ratio_zero_duration() {
let result = VadTranscriptionResult {
text: String::new(),
language: "en".to_string(),
segments: vec![],
speech_segments: vec![],
total_duration_secs: 0.0,
speech_duration_secs: 0.0,
};
let ratio = result.silence_ratio(0.0);
assert!((ratio - 1.0).abs() < 0.01);
}
#[test]
fn test_vad_transcription_result_iter() {
let result = VadTranscriptionResult {
text: "a b".to_string(),
language: "en".to_string(),
segments: vec![
VadSpeechSegment {
start: 0.0,
end: 1.0,
text: "a".to_string(),
tokens: vec![1],
},
VadSpeechSegment {
start: 1.0,
end: 2.0,
text: "b".to_string(),
tokens: vec![2],
},
],
speech_segments: vec![(0.0, 1.0), (1.0, 2.0)],
total_duration_secs: 2.0,
speech_duration_secs: 2.0,
};
let texts: Vec<_> = result.iter().map(|s| s.text.as_str()).collect();
assert_eq!(texts, vec!["a", "b"]);
}
#[test]
fn test_vad_speech_segment_duration() {
let segment = VadSpeechSegment {
start: 1.5,
end: 3.0,
text: "test".to_string(),
tokens: vec![1, 2, 3],
};
assert!((segment.duration() - 1.5).abs() < 0.01);
}
#[test]
fn test_vad_speech_segment_has_text() {
let with_text = VadSpeechSegment {
start: 0.0,
end: 1.0,
text: "hello".to_string(),
tokens: vec![1],
};
let empty = VadSpeechSegment {
start: 0.0,
end: 1.0,
text: String::new(),
tokens: vec![],
};
assert!(with_text.has_text());
assert!(!empty.has_text());
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_transcribe_with_vad_silence_only() {
let whisper = WhisperApr::tiny();
let silence = vec![0.0; 16000];
let result = whisper
.transcribe_with_vad(&silence, TranscribeOptions::default(), None)
.expect("should succeed");
assert!(!result.has_speech());
assert_eq!(result.num_segments(), 0);
assert!(result.text.is_empty());
}
#[test]
fn test_transcribe_with_silence_detection_config() {
let config = vad::SilenceConfig::new()
.with_min_silence_duration(0.3)
.with_max_silence_duration(2.0)
.with_silence_threshold(0.001);
assert!((config.min_silence_duration - 0.3).abs() < 0.01);
assert!((config.max_silence_duration - 2.0).abs() < 0.01);
assert!((config.silence_threshold - 0.001).abs() < 0.001);
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_invert_silence_segments_empty() {
let whisper = WhisperApr::tiny();
let audio_len = 16000;
let silence_segments: Vec<vad::SilenceSegment> = vec![];
let speech = whisper.invert_silence_segments(&silence_segments, audio_len);
assert_eq!(speech.len(), 1);
assert!((speech[0].0 - 0.0).abs() < 0.01);
assert!((speech[0].1 - 1.0).abs() < 0.01);
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_invert_silence_segments_single() {
let whisper = WhisperApr::tiny();
let audio_len = 32000;
let silence_segments = vec![vad::SilenceSegment {
start: 0.5,
end: 1.5,
noise_floor: 0.001,
}];
let speech = whisper.invert_silence_segments(&silence_segments, audio_len);
assert_eq!(speech.len(), 2);
assert!((speech[0].0 - 0.0).abs() < 0.01); assert!((speech[0].1 - 0.5).abs() < 0.01); assert!((speech[1].0 - 1.5).abs() < 0.01); assert!((speech[1].1 - 2.0).abs() < 0.01); }
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_invert_silence_segments_multiple() {
let whisper = WhisperApr::tiny();
let audio_len = 48000;
let silence_segments = vec![
vad::SilenceSegment {
start: 0.5,
end: 1.0,
noise_floor: 0.001,
},
vad::SilenceSegment {
start: 2.0,
end: 2.5,
noise_floor: 0.001,
},
];
let speech = whisper.invert_silence_segments(&silence_segments, audio_len);
assert_eq!(speech.len(), 3);
}
#[test]
fn test_tokenizer_integration() {
use tokenizer::special_tokens;
assert!(special_tokens::SOT > 0);
assert!(special_tokens::EOT > 0);
assert!(special_tokens::TRANSCRIBE > 0);
assert!(special_tokens::TRANSLATE > 0);
assert!(special_tokens::NO_TIMESTAMPS > 0);
}
#[test]
fn test_inference_beam_search_integration() {
use inference::BeamSearchDecoder;
let decoder = BeamSearchDecoder::new(5, 448);
assert_eq!(decoder.beam_size(), 5);
}
#[test]
fn test_format_decompression_integration() {
use format::Decompressor;
let mut decompressor = Decompressor::new();
assert!(decompressor.is_empty());
decompressor.reset();
assert!(decompressor.is_empty());
}
#[test]
fn test_timestamps_generation() {
let segment = Segment {
start: 1.5,
end: 3.25,
text: "This is a test".to_string(),
tokens: vec![1, 2, 3, 4],
};
let duration = segment.end - segment.start;
assert!((duration - 1.75).abs() < 1e-5);
}
#[test]
fn test_language_detection_integration() {
use detection::{is_supported, language_name};
assert!(is_supported("en"));
assert!(is_supported("es"));
assert!(is_supported("ja"));
assert!(!is_supported("invalid_lang"));
assert_eq!(language_name("en"), Some("English"));
assert_eq!(language_name("es"), Some("Spanish"));
assert_eq!(language_name("zh"), Some("Chinese"));
for lang in &["en", "es", "fr", "de", "it", "ja", "zh", "ko", "pt", "ru"] {
assert!(
language_name(lang).is_some(),
"language {} should have a name",
lang
);
}
}
#[test]
fn test_model_kv_cache_integration() {
use model::{Decoder, ModelConfig};
let config = ModelConfig::tiny();
let decoder = Decoder::new(&config);
let cache = decoder.create_kv_cache();
assert!(cache.self_attn_cache.iter().all(|c| c.is_empty()));
assert!(cache.cross_attn_cache.iter().all(|c| c.is_empty()));
}
#[test]
fn test_progress_tracking_integration() {
use progress::{format_bytes, Progress};
let progress = Progress::new(50, 100);
assert_eq!(progress.percent(), 50.0);
assert_eq!(progress.current, 50);
assert_eq!(progress.total, 100);
let kb_str = format_bytes(1024);
assert!(kb_str.contains("KB"), "should contain KB: {}", kb_str);
let mb_str = format_bytes(1024 * 1024);
assert!(mb_str.contains("MB"), "should contain MB: {}", mb_str);
}
#[test]
fn test_model_type_variants() {
let tiny = ModelType::Tiny;
let tiny_en = ModelType::TinyEn;
let base = ModelType::Base;
let base_en = ModelType::BaseEn;
let small = ModelType::Small;
assert!(format!("{tiny:?}").contains("Tiny"));
assert!(format!("{tiny_en:?}").contains("TinyEn"));
assert!(format!("{base:?}").contains("Base"));
assert!(format!("{base_en:?}").contains("BaseEn"));
assert!(format!("{small:?}").contains("Small"));
let tiny_clone = tiny;
assert_eq!(tiny_clone, ModelType::Tiny);
assert_eq!(tiny, ModelType::Tiny);
assert_ne!(tiny, base);
}
#[test]
fn test_task_variants() {
let transcribe = Task::Transcribe;
let translate = Task::Translate;
assert!(format!("{transcribe:?}").contains("Transcribe"));
assert!(format!("{translate:?}").contains("Translate"));
let transcribe_clone = transcribe;
assert_eq!(transcribe_clone, Task::Transcribe);
assert_eq!(transcribe, Task::Transcribe);
assert_ne!(transcribe, translate);
assert_eq!(Task::default(), Task::Transcribe);
}
#[test]
fn test_decoding_strategy_sampling() {
let sampling = DecodingStrategy::Sampling {
temperature: 0.8,
top_k: Some(50),
top_p: Some(0.9),
};
let debug_str = format!("{sampling:?}");
assert!(debug_str.contains("Sampling"));
let cloned = sampling.clone();
assert!(matches!(cloned, DecodingStrategy::Sampling { .. }));
}
#[test]
fn test_decoding_strategy_beam_search() {
let beam = DecodingStrategy::BeamSearch {
beam_size: 5,
temperature: 0.0,
patience: 1.0,
};
let debug_str = format!("{beam:?}");
assert!(debug_str.contains("BeamSearch"));
}
#[test]
fn test_transcribe_options_clone() {
let options = TranscribeOptions {
language: Some("fr".to_string()),
task: Task::Translate,
strategy: DecodingStrategy::Greedy,
word_timestamps: true,
profile: false,
..Default::default()
};
let cloned = options.clone();
assert_eq!(cloned.language, Some("fr".to_string()));
assert_eq!(cloned.task, Task::Translate);
assert!(cloned.word_timestamps);
}
#[test]
fn test_transcribe_options_debug() {
let options = TranscribeOptions::default();
let debug_str = format!("{options:?}");
assert!(debug_str.contains("TranscribeOptions"));
}
#[test]
fn test_segment_clone() {
let segment = Segment {
start: 1.0,
end: 2.0,
text: "test".to_string(),
tokens: vec![1, 2],
};
let cloned = segment.clone();
assert_eq!(cloned.text, "test");
assert_eq!(cloned.tokens, vec![1, 2]);
}
#[test]
fn test_segment_debug() {
let segment = Segment {
start: 0.0,
end: 1.0,
text: "hello".to_string(),
tokens: vec![],
};
let debug_str = format!("{segment:?}");
assert!(debug_str.contains("hello"));
}
#[test]
fn test_transcription_result_clone() {
let result = TranscriptionResult {
text: "hello world".to_string(),
language: "en".to_string(),
segments: vec![Segment {
start: 0.0,
end: 1.0,
text: "hello".to_string(),
tokens: vec![1],
}],
profiling: None,
};
let cloned = result.clone();
assert_eq!(cloned.text, "hello world");
assert_eq!(cloned.segments.len(), 1);
}
#[test]
fn test_transcription_result_debug() {
let result = TranscriptionResult {
text: "test".to_string(),
language: "en".to_string(),
segments: vec![],
profiling: None,
};
let debug_str = format!("{result:?}");
assert!(debug_str.contains("TranscriptionResult"));
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_whisper_debug() {
let whisper = WhisperApr::tiny();
let debug_str = format!("{whisper:?}");
assert!(debug_str.contains("WhisperApr"));
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_whisper_memory_size_all_models() {
let tiny = WhisperApr::tiny();
let base = WhisperApr::base();
assert!(tiny.memory_size() < base.memory_size());
assert!(tiny.memory_size() > 0);
assert!(base.memory_size() > 0);
}
#[test]
fn test_transcribe_options_sampling_strategy() {
let options = TranscribeOptions {
language: None,
task: Task::Transcribe,
strategy: DecodingStrategy::Sampling {
temperature: 1.0,
top_k: None,
top_p: None,
},
word_timestamps: false,
profile: false,
..Default::default()
};
assert!(matches!(
options.strategy,
DecodingStrategy::Sampling { .. }
));
}
#[test]
fn test_transcribe_options_word_timestamps_enabled() {
let options = TranscribeOptions {
language: Some("en".to_string()),
task: Task::Transcribe,
strategy: DecodingStrategy::default(),
word_timestamps: true,
profile: false,
..Default::default()
};
assert!(options.word_timestamps);
}
#[test]
fn test_model_type_small() {
let small = ModelType::Small;
assert!(format!("{small:?}").contains("Small"));
}
#[test]
fn test_model_type_medium() {
let medium = ModelType::Medium;
assert!(format!("{medium:?}").contains("Medium"));
}
#[test]
fn test_model_type_medium_en() {
let medium_en = ModelType::MediumEn;
assert!(format!("{medium_en:?}").contains("MediumEn"));
}
#[test]
fn test_model_type_large() {
let large = ModelType::Large;
assert!(format!("{large:?}").contains("Large"));
}
#[test]
fn test_model_type_large_v1() {
let large_v1 = ModelType::LargeV1;
assert!(format!("{large_v1:?}").contains("LargeV1"));
}
#[test]
fn test_model_type_large_v2() {
let large_v2 = ModelType::LargeV2;
assert!(format!("{large_v2:?}").contains("LargeV2"));
}
#[test]
fn test_model_type_large_v3() {
let large_v3 = ModelType::LargeV3;
assert!(format!("{large_v3:?}").contains("LargeV3"));
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_whisper_small() {
let whisper = WhisperApr::small();
assert_eq!(whisper.model_type(), ModelType::Small);
assert_eq!(whisper.config().n_audio_layer, 12);
assert_eq!(whisper.config().n_audio_state, 768);
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_whisper_medium() {
let whisper = WhisperApr::medium();
assert_eq!(whisper.model_type(), ModelType::Medium);
assert_eq!(whisper.config().n_audio_layer, 24);
assert_eq!(whisper.config().n_audio_state, 1024);
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_whisper_large() {
let whisper = WhisperApr::large();
assert_eq!(whisper.model_type(), ModelType::Large);
assert_eq!(whisper.config().n_audio_layer, 32);
assert_eq!(whisper.config().n_audio_state, 1280);
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_whisper_memory_size_all_extended_models() {
let small = WhisperApr::small();
let medium = WhisperApr::medium();
let large = WhisperApr::large();
assert!(small.memory_size() < medium.memory_size());
assert!(medium.memory_size() < large.memory_size());
assert!(small.memory_size() > 0);
assert!(medium.memory_size() > 0);
assert!(large.memory_size() > 0);
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_extended_model_memory_size_estimates() {
let small = WhisperApr::small();
let medium = WhisperApr::medium();
let large = WhisperApr::large();
assert!(small.memory_size() > 900_000_000);
assert!(small.memory_size() < 1_200_000_000);
assert!(medium.memory_size() > 2_500_000_000);
assert!(medium.memory_size() < 3_500_000_000);
assert!(large.memory_size() > 5_000_000_000);
assert!(large.memory_size() < 7_000_000_000);
}
#[test]
fn test_batch_transcription_result_len() {
let result = BatchTranscriptionResult {
results: vec![
TranscriptionResult {
text: "Hello".to_string(),
language: "en".to_string(),
segments: vec![],
profiling: None,
},
TranscriptionResult {
text: "World".to_string(),
language: "en".to_string(),
segments: vec![],
profiling: None,
},
],
total_duration_secs: 1.5,
};
assert_eq!(result.len(), 2);
assert!(!result.is_empty());
}
#[test]
fn test_batch_transcription_result_empty() {
let result = BatchTranscriptionResult {
results: vec![],
total_duration_secs: 0.0,
};
assert!(result.is_empty());
assert_eq!(result.len(), 0);
}
#[test]
fn test_batch_transcription_result_get() {
let result = BatchTranscriptionResult {
results: vec![
TranscriptionResult {
text: "First".to_string(),
language: "en".to_string(),
segments: vec![],
profiling: None,
},
TranscriptionResult {
text: "Second".to_string(),
language: "es".to_string(),
segments: vec![],
profiling: None,
},
],
total_duration_secs: 2.0,
};
assert!(result.get(0).is_some());
assert_eq!(result.get(0).map(|r| r.text.as_str()), Some("First"));
assert_eq!(result.get(1).map(|r| r.language.as_str()), Some("es"));
assert!(result.get(2).is_none());
}
#[test]
fn test_batch_transcription_result_texts() {
let result = BatchTranscriptionResult {
results: vec![
TranscriptionResult {
text: "One".to_string(),
language: "en".to_string(),
segments: vec![],
profiling: None,
},
TranscriptionResult {
text: "Two".to_string(),
language: "en".to_string(),
segments: vec![],
profiling: None,
},
TranscriptionResult {
text: "Three".to_string(),
language: "en".to_string(),
segments: vec![],
profiling: None,
},
],
total_duration_secs: 3.0,
};
let texts = result.texts();
assert_eq!(texts, vec!["One", "Two", "Three"]);
}
#[test]
fn test_batch_transcription_result_iter() {
let result = BatchTranscriptionResult {
results: vec![
TranscriptionResult {
text: "A".to_string(),
language: "en".to_string(),
segments: vec![],
profiling: None,
},
TranscriptionResult {
text: "B".to_string(),
language: "en".to_string(),
segments: vec![],
profiling: None,
},
],
total_duration_secs: 1.0,
};
let collected: Vec<&str> = result.iter().map(|r| r.text.as_str()).collect();
assert_eq!(collected, vec!["A", "B"]);
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_transcribe_batch_empty() {
let whisper = WhisperApr::tiny();
let result = whisper.transcribe_batch(&[], TranscribeOptions::default());
assert!(result.is_err());
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_transcribe_audio_batch_empty() {
let whisper = WhisperApr::tiny();
let batch = audio::AudioBatch::with_default_config();
let result = whisper.transcribe_audio_batch(&batch, TranscribeOptions::default());
assert!(result.is_err());
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_transcribe_batch_optimized_empty() {
let whisper = WhisperApr::tiny();
let result = whisper.transcribe_batch_optimized(&[], TranscribeOptions::default());
assert!(result.is_err());
}
#[test]
fn test_create_audio_batch() {
let segments = vec![vec![0.1_f32, 0.2, 0.3], vec![0.4_f32, 0.5]];
let batch = WhisperApr::create_audio_batch(&segments);
assert_eq!(batch.len(), 2);
assert!(!batch.is_empty());
}
#[test]
fn test_create_audio_batch_empty() {
let segments: Vec<Vec<f32>> = vec![];
let batch = WhisperApr::create_audio_batch(&segments);
assert!(batch.is_empty());
}
#[test]
fn test_batch_transcription_result_duration() {
let result = BatchTranscriptionResult {
results: vec![],
total_duration_secs: 5.25,
};
assert!((result.total_duration_secs - 5.25).abs() < f32::EPSILON);
}
#[test]
fn test_partial_transcription_result_new() {
let result = PartialTranscriptionResult {
text: "hello".to_string(),
language: "en".to_string(),
is_final: false,
confidence: 0.95,
duration_secs: 1.5,
processing_time_secs: 0.3,
};
assert_eq!(result.text, "hello");
assert_eq!(result.language, "en");
assert!(!result.is_final);
assert!((result.confidence - 0.95).abs() < 0.01);
}
#[test]
fn test_partial_transcription_result_has_text() {
let with_text = PartialTranscriptionResult {
text: "hello".to_string(),
language: "en".to_string(),
is_final: false,
confidence: 1.0,
duration_secs: 1.0,
processing_time_secs: 0.1,
};
let empty = PartialTranscriptionResult {
text: String::new(),
language: "en".to_string(),
is_final: false,
confidence: 0.0,
duration_secs: 0.5,
processing_time_secs: 0.05,
};
assert!(with_text.has_text());
assert!(!empty.has_text());
}
#[test]
fn test_partial_transcription_result_is_empty_interim() {
let empty_interim = PartialTranscriptionResult {
text: String::new(),
language: "en".to_string(),
is_final: false,
confidence: 0.0,
duration_secs: 0.5,
processing_time_secs: 0.05,
};
let empty_final = PartialTranscriptionResult {
text: String::new(),
language: "en".to_string(),
is_final: true,
confidence: 0.0,
duration_secs: 0.5,
processing_time_secs: 0.05,
};
let with_text = PartialTranscriptionResult {
text: "hello".to_string(),
language: "en".to_string(),
is_final: false,
confidence: 1.0,
duration_secs: 1.0,
processing_time_secs: 0.1,
};
assert!(empty_interim.is_empty_interim());
assert!(!empty_final.is_empty_interim()); assert!(!with_text.is_empty_interim()); }
#[test]
fn test_partial_transcription_result_real_time_factor() {
let result = PartialTranscriptionResult {
text: "hello".to_string(),
language: "en".to_string(),
is_final: false,
confidence: 1.0,
duration_secs: 2.0,
processing_time_secs: 0.5,
};
assert!((result.real_time_factor() - 0.25).abs() < 0.01);
}
#[test]
fn test_partial_transcription_result_real_time_factor_zero_duration() {
let result = PartialTranscriptionResult {
text: String::new(),
language: "en".to_string(),
is_final: false,
confidence: 0.0,
duration_secs: 0.0,
processing_time_secs: 0.0,
};
assert!((result.real_time_factor() - 0.0).abs() < 0.01);
}
#[test]
fn test_partial_transcription_result_debug_clone() {
let result = PartialTranscriptionResult {
text: "test".to_string(),
language: "en".to_string(),
is_final: true,
confidence: 0.9,
duration_secs: 1.0,
processing_time_secs: 0.1,
};
let debug_str = format!("{result:?}");
assert!(debug_str.contains("PartialTranscriptionResult"));
let cloned = result.clone();
assert_eq!(cloned.text, "test");
assert!(cloned.is_final);
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_transcribe_partial_too_short() {
let whisper = WhisperApr::tiny();
let short_audio = vec![0.0; 4000];
let result = whisper
.transcribe_partial(&short_audio, TranscribeOptions::default(), false)
.expect("should succeed with empty result");
assert!(result.text.is_empty());
assert!(!result.is_final);
assert!((result.confidence - 0.0).abs() < 0.01);
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_encode_3_second_chunk() {
let whisper = WhisperApr::tiny();
let audio = vec![0.0; 48000];
let mel = whisper
.compute_mel(&audio)
.expect("mel computation should succeed");
let result = whisper.encode(&mel);
assert!(
result.is_ok(),
"encode should succeed for 3s audio mel: {:?}",
result.err()
);
let encoded = result.expect("encode should succeed");
let d_model = whisper.config().n_text_state as usize; assert_eq!(
encoded.len() % d_model,
0,
"encoded output should be multiple of d_model"
);
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_create_streaming_session() {
let whisper = WhisperApr::tiny();
let session = whisper.create_streaming_session(TranscribeOptions::default(), 44100);
assert_eq!(session.state(), audio::ProcessorState::WaitingForSpeech);
assert!((session.chunk_progress() - 0.0).abs() < 0.01);
assert!(!session.has_chunk());
assert!(!session.has_events());
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_streaming_session_reset() {
let whisper = WhisperApr::tiny();
let mut session = whisper.create_streaming_session(TranscribeOptions::default(), 16000);
session.push(&vec![0.1; 1000]).expect("push should work");
session.reset();
assert_eq!(session.state(), audio::ProcessorState::WaitingForSpeech);
assert!((session.partial_duration() - 0.0).abs() < 0.01);
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_streaming_session_set_partial_threshold() {
let whisper = WhisperApr::tiny();
let mut session = whisper.create_streaming_session(TranscribeOptions::default(), 16000);
session.set_partial_threshold(5.0);
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_streaming_session_finalize_no_chunk() {
let whisper = WhisperApr::tiny();
let mut session = whisper.create_streaming_session(TranscribeOptions::default(), 16000);
let result = session.finalize();
assert!(result.is_err());
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_streaming_session_flush_empty() {
let whisper = WhisperApr::tiny();
let mut session = whisper.create_streaming_session(TranscribeOptions::default(), 16000);
let result = session.flush().expect("flush should work");
assert!(result.is_none());
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_streaming_session_drain_events() {
let whisper = WhisperApr::tiny();
let mut session = whisper.create_streaming_session(TranscribeOptions::default(), 16000);
session.reset();
let events = session.drain_events();
assert!(!events.is_empty());
assert!(events
.iter()
.any(|e| matches!(e, audio::StreamingEvent::Reset)));
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_streaming_session_push_silence() {
let whisper = WhisperApr::tiny();
let mut session = whisper.create_streaming_session(TranscribeOptions::default(), 16000);
let result = session.push(&vec![0.0; 16000]).expect("push should work");
assert!(result.is_none()); }
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_streaming_session_debug() {
let whisper = WhisperApr::tiny();
let session = whisper.create_streaming_session(TranscribeOptions::default(), 16000);
let debug_str = format!("{session:?}");
assert!(debug_str.contains("StreamingSession"));
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_streaming_session_state() {
let whisper = WhisperApr::tiny();
let session = whisper.create_streaming_session(TranscribeOptions::default(), 16000);
let state = session.state();
assert_eq!(state, audio::ProcessorState::WaitingForSpeech);
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_streaming_session_chunk_progress() {
let whisper = WhisperApr::tiny();
let session = whisper.create_streaming_session(TranscribeOptions::default(), 16000);
let progress = session.chunk_progress();
assert!(progress >= 0.0 && progress <= 1.0);
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_streaming_session_partial_duration() {
let whisper = WhisperApr::tiny();
let session = whisper.create_streaming_session(TranscribeOptions::default(), 16000);
let duration = session.partial_duration();
assert!(duration >= 0.0);
}
#[test]
fn test_partial_transcription_result_rtf_with_zero_processing() {
let result = PartialTranscriptionResult {
text: "test".to_string(),
language: "en".to_string(),
is_final: true,
confidence: 0.9,
duration_secs: 5.0,
processing_time_secs: 0.0,
};
let rtf = result.real_time_factor();
assert!(rtf >= 0.0);
}
#[test]
fn test_vad_transcription_result_methods_empty() {
let result = VadTranscriptionResult {
text: "Hello world".to_string(),
language: "en".to_string(),
segments: vec![],
speech_segments: vec![],
total_duration_secs: 5.0,
speech_duration_secs: 0.0,
};
assert_eq!(result.num_segments(), 0);
assert!(!result.has_speech());
}
#[test]
fn test_batch_transcription_result_defaults() {
let result = BatchTranscriptionResult {
results: vec![],
total_duration_secs: 0.0,
};
assert_eq!(result.len(), 0);
assert!(result.is_empty());
assert!(result.get(0).is_none());
assert!(result.texts().is_empty());
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_whisper_config_accessors() {
let whisper = WhisperApr::tiny();
let config = whisper.config();
assert!(config.n_vocab > 0);
assert!(config.n_audio_ctx > 0);
assert!(config.n_text_ctx > 0);
}
#[test]
fn test_transcribe_options_all_fields() {
let options = TranscribeOptions {
language: Some("fr".to_string()),
task: Task::Translate,
strategy: DecodingStrategy::BeamSearch {
beam_size: 3,
temperature: 0.2,
patience: 1.5,
},
word_timestamps: true,
profile: false,
prompt: Some("domain-specific prompt".into()),
hotwords: vec!["hotword1".into(), "hotword2".into()],
};
assert_eq!(options.language, Some("fr".to_string()));
assert_eq!(options.task, Task::Translate);
assert!(options.word_timestamps);
}
#[test]
fn test_segment_with_tokens() {
let segment = Segment {
text: "Hello".to_string(),
start: 0.0,
end: 1.0,
tokens: vec![1, 2, 3, 4, 5],
};
assert_eq!(segment.tokens.len(), 5);
assert!((segment.end - segment.start - 1.0).abs() < f32::EPSILON);
}
#[test]
fn test_model_from_config() {
let config = model::ModelConfig::tiny();
let whisper = WhisperApr::from_config(config);
assert_eq!(whisper.model_type(), ModelType::Tiny);
}
#[test]
fn test_decoding_strategy_variants() {
let greedy = DecodingStrategy::Greedy;
assert!(matches!(greedy, DecodingStrategy::Greedy));
let sampling = DecodingStrategy::Sampling {
temperature: 0.5,
top_k: Some(40),
top_p: Some(0.9),
};
if let DecodingStrategy::Sampling {
temperature,
top_k,
top_p,
} = sampling
{
assert!((temperature - 0.5).abs() < f32::EPSILON);
assert_eq!(top_k, Some(40));
assert_eq!(top_p, Some(0.9));
}
}
#[test]
fn test_task_variants_eq() {
assert_eq!(Task::Transcribe, Task::Transcribe);
assert_ne!(Task::Transcribe, Task::Translate);
}
#[test]
fn test_model_type_all_variants() {
let variants = vec![
ModelType::Tiny,
ModelType::TinyEn,
ModelType::Base,
ModelType::BaseEn,
ModelType::Small,
ModelType::SmallEn,
ModelType::Medium,
ModelType::MediumEn,
ModelType::Large,
ModelType::LargeV1,
ModelType::LargeV2,
ModelType::LargeV3,
];
for variant in variants {
let debug_str = format!("{variant:?}");
assert!(!debug_str.is_empty());
}
}
#[test]
fn test_vad_speech_segment_empty_text() {
let segment = VadSpeechSegment {
start: 0.0,
end: 1.0,
text: String::new(),
tokens: vec![],
};
assert!(!segment.has_text());
assert!((segment.duration() - 1.0).abs() < f32::EPSILON);
}
#[test]
fn test_transcription_result_empty() {
let result = TranscriptionResult {
text: String::new(),
language: "en".to_string(),
segments: vec![],
profiling: None,
};
assert!(result.text.is_empty());
assert!(result.segments.is_empty());
}
#[test]
fn test_batch_transcription_result_iter_coverage() {
let results = vec![
TranscriptionResult {
text: "First".to_string(),
language: "en".to_string(),
segments: vec![],
profiling: None,
},
TranscriptionResult {
text: "Second".to_string(),
language: "en".to_string(),
segments: vec![],
profiling: None,
},
];
let batch = BatchTranscriptionResult {
results,
total_duration_secs: 1.0,
};
let mut count = 0;
for result in batch.iter() {
count += 1;
assert!(!result.text.is_empty());
}
assert_eq!(count, 2);
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_whisper_clone() {
let whisper = WhisperApr::tiny();
let cloned = whisper.clone();
assert_eq!(whisper.model_type(), cloned.model_type());
assert_eq!(whisper.memory_size(), cloned.memory_size());
}
#[test]
fn test_vad_transcription_result_first_last_segment() {
let result = VadTranscriptionResult {
text: "hello world".to_string(),
language: "en".to_string(),
segments: vec![
VadSpeechSegment {
start: 0.0,
end: 1.0,
text: "hello".to_string(),
tokens: vec![1],
},
VadSpeechSegment {
start: 1.5,
end: 2.5,
text: "world".to_string(),
tokens: vec![2],
},
],
speech_segments: vec![(0.0, 1.0), (1.5, 2.5)],
total_duration_secs: 3.0,
speech_duration_secs: 2.0,
};
let first = result.first_segment().expect("first segment");
assert_eq!(first.text, "hello");
let last = result.last_segment().expect("last segment");
assert_eq!(last.text, "world");
}
#[test]
fn test_vad_transcription_result_iter_segments() {
let result = VadTranscriptionResult {
text: "test".to_string(),
language: "en".to_string(),
segments: vec![
VadSpeechSegment {
start: 0.0,
end: 1.0,
text: "a".to_string(),
tokens: vec![1],
},
VadSpeechSegment {
start: 1.0,
end: 2.0,
text: "b".to_string(),
tokens: vec![2],
},
],
speech_segments: vec![(0.0, 1.0), (1.0, 2.0)],
total_duration_secs: 2.0,
speech_duration_secs: 2.0,
};
let mut count = 0;
for segment in result.iter() {
count += 1;
assert!(segment.has_text());
}
assert_eq!(count, 2);
}
#[test]
fn test_partial_transcription_result_methods() {
let result = PartialTranscriptionResult {
text: "hello".to_string(),
language: "en".to_string(),
is_final: false,
confidence: 0.85,
duration_secs: 2.0,
processing_time_secs: 0.5,
};
assert!(result.has_text());
assert!((result.real_time_factor() - 0.25).abs() < f32::EPSILON);
}
#[test]
fn test_partial_transcription_result_zero_duration() {
let result = PartialTranscriptionResult {
text: "".to_string(),
language: "en".to_string(),
is_final: true,
confidence: 0.0,
duration_secs: 0.0,
processing_time_secs: 0.1,
};
assert!(!result.has_text());
assert!((result.real_time_factor() - 0.0).abs() < f32::EPSILON);
}
#[test]
#[ignore = "Requires model file - run with --ignored"]
fn test_e2e_transcribe_with_int8_model() {
let model_path = std::path::Path::new("models/whisper-tiny-int8.apr");
if !model_path.exists() {
eprintln!(
"Skipping E2E test: model file not found at {:?}",
model_path
);
return;
}
let model_data = std::fs::read(model_path).expect("Failed to read model file");
eprintln!("Loaded model: {} bytes", model_data.len());
let whisper = WhisperApr::load_from_apr(&model_data).expect("Failed to load model");
eprintln!("Model loaded: {:?}", whisper.model_type());
let sample_rate = 16000;
let duration_secs = 3.0;
let num_samples = (sample_rate as f32 * duration_secs) as usize;
let audio: Vec<f32> = (0..num_samples)
.map(|i| {
let noise = ((i as f32 * 0.1).sin() * 0.001) + ((i as f32 * 0.37).cos() * 0.001);
noise
})
.collect();
eprintln!("Generated {} samples of test audio", audio.len());
let start = std::time::Instant::now();
let result = whisper.transcribe(&audio, TranscribeOptions::default());
let elapsed = start.elapsed();
eprintln!("Transcription completed in {:?}", elapsed);
match result {
Ok(transcription) => {
eprintln!("Result: '{}'", transcription.text);
eprintln!("Language: {}", transcription.language);
eprintln!("Segments: {}", transcription.segments.len());
assert!(
transcription.text.len() < 100,
"Unexpected long transcription for silence"
);
}
Err(e) => {
panic!("Transcription failed: {e:?}");
}
}
let rtf = elapsed.as_secs_f32() / duration_secs;
eprintln!("Real-time factor: {rtf:.2}x");
assert!(rtf < 50.0, "RTF {rtf} is too slow, SIMD may not be working");
}
#[test]
fn test_summarize_options_default_params() {
use model::lfm2::{Lfm2, Lfm2Tokenizer};
let config = format::apr2::Lfm2Config {
hidden_size: 64,
num_layers: 2,
num_q_heads: 4,
num_kv_heads: 2,
intermediate_size: 128,
vocab_size: 1000,
max_seq_len: 512,
rope_theta: 10000.0,
conv_dimension: 32,
layer_types: vec![
format::apr2::LayerType::Convolution {
kernel_size: 4,
cache_len: 3,
},
format::apr2::LayerType::Attention { use_gqa: true },
],
};
let model = Lfm2::new(config).expect("Model creation should succeed");
let tokenizer = Lfm2Tokenizer::new();
let options = SummarizeOptions::new(&model, &tokenizer);
assert_eq!(options.max_tokens, 256);
assert!((options.temperature - 0.3).abs() < f32::EPSILON);
}
#[test]
fn test_summarize_options_builder() {
use model::lfm2::{Lfm2, Lfm2Tokenizer};
let config = format::apr2::Lfm2Config {
hidden_size: 64,
num_layers: 2,
num_q_heads: 4,
num_kv_heads: 2,
intermediate_size: 128,
vocab_size: 1000,
max_seq_len: 512,
rope_theta: 10000.0,
conv_dimension: 32,
layer_types: vec![
format::apr2::LayerType::Convolution {
kernel_size: 4,
cache_len: 3,
},
format::apr2::LayerType::Attention { use_gqa: true },
],
};
let model = Lfm2::new(config).expect("Model creation should succeed");
let tokenizer = Lfm2Tokenizer::new();
let options = SummarizeOptions::new(&model, &tokenizer)
.with_max_tokens(512)
.with_temperature(0.7);
assert_eq!(options.max_tokens, 512);
assert!((options.temperature - 0.7).abs() < f32::EPSILON);
}
#[test]
fn test_transcribe_summary_result_accessors() {
let transcription = TranscriptionResult {
text: "Hello world".to_string(),
language: "en".to_string(),
segments: vec![],
profiling: None,
};
let result = TranscribeSummaryResult {
transcription,
summary: "Summary text".to_string(),
generation_stats: None,
};
assert_eq!(result.transcript(), "Hello world");
assert_eq!(result.summary(), "Summary text");
assert!(result.has_summary());
}
#[test]
fn test_transcribe_summary_result_empty_summary() {
let transcription = TranscriptionResult {
text: "Hello world".to_string(),
language: "en".to_string(),
segments: vec![],
profiling: None,
};
let result = TranscribeSummaryResult {
transcription,
summary: String::new(),
generation_stats: None,
};
assert!(!result.has_summary());
}
#[test]
#[ignore = "Slow test - run with --ignored"]
fn test_transcribe_and_summarize_empty_audio() {
use model::lfm2::{Lfm2, Lfm2Tokenizer};
let whisper = WhisperApr::tiny();
let config = format::apr2::Lfm2Config {
hidden_size: 64,
num_layers: 2,
num_q_heads: 4,
num_kv_heads: 2,
intermediate_size: 128,
vocab_size: 1000,
max_seq_len: 2048, rope_theta: 10000.0,
conv_dimension: 32,
layer_types: vec![
format::apr2::LayerType::Convolution {
kernel_size: 4,
cache_len: 3,
},
format::apr2::LayerType::Attention { use_gqa: true },
],
};
let lfm2 = Lfm2::new(config).expect("Model creation should succeed");
let tokenizer = Lfm2Tokenizer::new();
let audio = vec![0.0f32; 16000]; let transcribe_options = TranscribeOptions::default();
let summarize_options = SummarizeOptions::new(&lfm2, &tokenizer).with_max_tokens(8);
let result = whisper
.transcribe_and_summarize(&audio, transcribe_options, summarize_options)
.expect("Should not fail");
if result.transcription.text.trim().is_empty() {
assert!(!result.has_summary());
}
}
#[test]
#[ignore = "Slow test - run with --ignored"]
fn test_transcribe_and_summarize_integration() {
use model::lfm2::{Lfm2, Lfm2Tokenizer};
let whisper = WhisperApr::tiny();
let config = format::apr2::Lfm2Config {
hidden_size: 64,
num_layers: 2,
num_q_heads: 4,
num_kv_heads: 2,
intermediate_size: 128,
vocab_size: 1000,
max_seq_len: 2048, rope_theta: 10000.0,
conv_dimension: 32,
layer_types: vec![
format::apr2::LayerType::Convolution {
kernel_size: 4,
cache_len: 3,
},
format::apr2::LayerType::Attention { use_gqa: true },
],
};
let lfm2 = Lfm2::new(config).expect("Model creation should succeed");
let tokenizer = Lfm2Tokenizer::new();
let sample_rate = 16000;
let duration_secs = 1.0;
let num_samples = (sample_rate as f32 * duration_secs) as usize;
let audio: Vec<f32> = (0..num_samples)
.map(|i| (2.0 * std::f32::consts::PI * 440.0 * i as f32 / sample_rate as f32).sin())
.collect();
let transcribe_options = TranscribeOptions::default();
let summarize_options = SummarizeOptions::new(&lfm2, &tokenizer).with_max_tokens(8);
let result =
whisper.transcribe_and_summarize(&audio, transcribe_options, summarize_options);
assert!(result.is_ok());
let result = result.expect("Should succeed");
assert_eq!(result.transcription.language, "en");
}
#[test]
fn test_chunk_constants() {
assert_eq!(WhisperApr::CHUNK_SAMPLES, 30 * 16000);
assert_eq!(WhisperApr::OVERLAP_SAMPLES, 5 * 16000);
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_short_audio_uses_single_chunk() {
let _whisper = WhisperApr::tiny();
let short_audio = vec![0.0_f32; 10 * 16000]; assert!(short_audio.len() <= WhisperApr::CHUNK_SAMPLES);
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_long_audio_uses_chunking() {
let _whisper = WhisperApr::tiny();
let long_audio = vec![0.0_f32; 60 * 16000]; assert!(long_audio.len() > WhisperApr::CHUNK_SAMPLES);
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_merge_overlapping_segments_empty() {
let whisper = WhisperApr::tiny();
let segments: Vec<Segment> = vec![];
let merged = whisper.merge_overlapping_segments(segments);
assert!(merged.is_empty());
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_merge_overlapping_segments_single() {
let whisper = WhisperApr::tiny();
let segments = vec![Segment {
start: 0.0,
end: 5.0,
text: "Hello".to_string(),
tokens: vec![1, 2, 3],
}];
let merged = whisper.merge_overlapping_segments(segments);
assert_eq!(merged.len(), 1);
assert_eq!(merged[0].text, "Hello");
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_merge_overlapping_segments_no_overlap() {
let whisper = WhisperApr::tiny();
let segments = vec![
Segment {
start: 0.0,
end: 5.0,
text: "Hello".to_string(),
tokens: vec![1],
},
Segment {
start: 10.0,
end: 15.0,
text: "World".to_string(),
tokens: vec![2],
},
];
let merged = whisper.merge_overlapping_segments(segments);
assert_eq!(merged.len(), 2);
assert_eq!(merged[0].text, "Hello");
assert_eq!(merged[1].text, "World");
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_merge_overlapping_segments_with_overlap() {
let whisper = WhisperApr::tiny();
let segments = vec![
Segment {
start: 0.0,
end: 5.0,
text: "Hello".to_string(),
tokens: vec![1],
},
Segment {
start: 4.9, end: 10.0,
text: "World".to_string(),
tokens: vec![2],
},
];
let merged = whisper.merge_overlapping_segments(segments);
assert_eq!(merged.len(), 1);
assert_eq!(merged[0].text, "Hello World");
assert_eq!(merged[0].start, 0.0);
assert_eq!(merged[0].end, 10.0);
assert_eq!(merged[0].tokens, vec![1, 2]);
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_merge_overlapping_segments_multiple_overlaps() {
let whisper = WhisperApr::tiny();
let segments = vec![
Segment {
start: 0.0,
end: 5.0,
text: "A".to_string(),
tokens: vec![1],
},
Segment {
start: 4.95,
end: 10.0,
text: "B".to_string(),
tokens: vec![2],
},
Segment {
start: 9.95,
end: 15.0,
text: "C".to_string(),
tokens: vec![3],
},
Segment {
start: 20.0, end: 25.0,
text: "D".to_string(),
tokens: vec![4],
},
];
let merged = whisper.merge_overlapping_segments(segments);
assert_eq!(merged.len(), 2);
assert_eq!(merged[0].text, "A B C");
assert_eq!(merged[0].end, 15.0);
assert_eq!(merged[1].text, "D");
assert_eq!(merged[1].start, 20.0);
}
#[test]
fn test_chunk_boundary_calculation() {
let total_samples = 90 * 16000; let chunk_size = WhisperApr::CHUNK_SAMPLES; let overlap = WhisperApr::OVERLAP_SAMPLES;
let chunk1_end = (0 + chunk_size + overlap).min(total_samples);
assert_eq!(chunk1_end, 35 * 16000);
let chunk2_start = chunk_size; let chunk2_end = (chunk2_start + chunk_size + overlap).min(total_samples);
assert_eq!(chunk2_end, 65 * 16000);
let chunk3_start = 2 * chunk_size; let chunk3_end = (chunk3_start + chunk_size + overlap).min(total_samples);
assert_eq!(chunk3_end, total_samples);
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_falsification_point_25_long_audio_full_transcription() {
let _whisper = WhisperApr::tiny();
let ten_minutes_samples = 10 * 60 * 16000;
assert!(ten_minutes_samples > WhisperApr::CHUNK_SAMPLES);
let expected_chunks =
(ten_minutes_samples + WhisperApr::CHUNK_SAMPLES - 1) / WhisperApr::CHUNK_SAMPLES;
assert_eq!(expected_chunks, 20);
let mut offset = 0;
let mut chunk_count = 0;
while offset < ten_minutes_samples {
let chunk_end = (offset + WhisperApr::CHUNK_SAMPLES + WhisperApr::OVERLAP_SAMPLES)
.min(ten_minutes_samples);
let _chunk_len = chunk_end - offset;
offset += WhisperApr::CHUNK_SAMPLES;
chunk_count += 1;
}
assert_eq!(chunk_count, 20);
}
#[test]
fn test_falsification_point_30_streaming_consistency() {
let overlap_seconds = WhisperApr::OVERLAP_SAMPLES as f32 / 16000.0;
assert!((overlap_seconds - 5.0).abs() < 0.01);
let expected_overlap_words = overlap_seconds * 2.5;
assert!(expected_overlap_words >= 10.0);
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_moonshine_tiny_constructs() {
let model = WhisperApr::moonshine_tiny();
assert_eq!(model.config.n_audio_state, 288);
assert_eq!(model.config.n_text_state, 288);
assert_eq!(model.config.n_audio_layer, 6);
assert_eq!(model.config.n_text_layer, 6);
assert_eq!(model.config.n_audio_head, 8);
assert_eq!(model.config.n_text_head, 8);
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_moonshine_tiny_encoder_has_moonshine_blocks() {
let model = WhisperApr::moonshine_tiny();
assert!(model.encoder.moonshine_blocks().len() > 0);
assert!(model.encoder.blocks().is_empty());
assert!(model.encoder.rope().is_some());
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_moonshine_tiny_decoder_has_moonshine_blocks() {
let model = WhisperApr::moonshine_tiny();
assert!(model.decoder.moonshine_blocks().len() > 0);
assert!(model.decoder.blocks().is_empty());
assert!(model.decoder.rope().is_some());
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_moonshine_tiny_has_sentencepiece_tokenizer() {
let model = WhisperApr::moonshine_tiny();
match model.tokenizer() {
tokenizer::Tokenizer::SentencePiece(_) => {} tokenizer::Tokenizer::Bpe(_) => panic!("Moonshine should use SentencePiece tokenizer"),
}
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_moonshine_tiny_has_conv_stem() {
let model = WhisperApr::moonshine_tiny();
assert!(model.conv_stem.is_some());
assert!(model.mel_filters.is_none());
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_moonshine_tiny_initial_tokens() {
let model = WhisperApr::moonshine_tiny();
let tokens = model.get_initial_tokens("en", crate::Task::Transcribe);
assert_eq!(tokens, vec![1]);
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_moonshine_tiny_eot_token() {
let model = WhisperApr::moonshine_tiny();
assert_eq!(model.eot_token(), 2); }
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_moonshine_encoder_forward_shape() {
let model = WhisperApr::moonshine_tiny();
let d_model = 288;
let seq_len = 7;
let features = vec![0.1_f32; seq_len * d_model];
let output = model.encoder.forward(&features).expect("encoder forward");
assert_eq!(output.len(), seq_len * d_model);
assert!(output.iter().all(|v| v.is_finite()));
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_moonshine_decoder_forward_shape() {
let model = WhisperApr::moonshine_tiny();
let d_model = 288;
let n_vocab = model.config.n_vocab as usize;
let enc_seq_len = 7;
let encoder_output = vec![0.1_f32; enc_seq_len * d_model];
let tokens = vec![1_u32, 100];
let logits = model
.decoder
.forward(&tokens, &encoder_output)
.expect("decoder forward");
assert_eq!(logits.len(), tokens.len() * n_vocab);
assert!(logits.iter().all(|v| v.is_finite()));
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_moonshine_variable_length_input() {
let model = WhisperApr::moonshine_tiny();
let d_model = 288;
for seq_len in [3, 7, 15, 31] {
let features = vec![0.1_f32; seq_len * d_model];
let output = model.encoder.forward(&features).expect("encoder forward");
assert_eq!(
output.len(),
seq_len * d_model,
"seq_len={seq_len} output mismatch"
);
}
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_whisper_tiny_unaffected() {
let model = WhisperApr::tiny();
assert!(model.encoder.blocks().len() > 0);
assert!(model.encoder.moonshine_blocks().is_empty());
assert!(model.encoder.rope().is_none());
assert!(model.decoder.blocks().len() > 0);
assert!(model.decoder.moonshine_blocks().is_empty());
assert!(model.decoder.rope().is_none());
assert!(model.mel_filters.is_some());
assert!(model.conv_stem.is_none());
match model.tokenizer() {
tokenizer::Tokenizer::Bpe(_) => {} tokenizer::Tokenizer::SentencePiece(_) => {
panic!("Whisper should use BPE tokenizer")
}
}
}
#[test]
fn test_moonshine_encoder_uses_rmsnorm_final() {
let config = model::ModelConfig::moonshine_tiny();
let encoder = model::Encoder::new(&config);
assert!(
encoder.ln_post_rms().is_some(),
"Moonshine encoder must use RmsNorm for final layer norm"
);
}
#[test]
fn test_whisper_encoder_uses_layernorm_final() {
let config = model::ModelConfig::tiny();
let encoder = model::Encoder::new(&config);
assert!(
encoder.ln_post_rms().is_none(),
"Whisper encoder must use LayerNorm, not RmsNorm"
);
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_moonshine_decoder_forward_one_sets_cross_attn_cached() {
let model = WhisperApr::moonshine_tiny();
let d_model = 288;
let n_layers = model.decoder.n_layers();
let max_len = model.decoder.max_len();
let enc_seq_len = 7;
let encoder_output = vec![0.1_f32; enc_seq_len * d_model];
let mut cache = model::DecoderKVCache::new(n_layers, d_model, max_len);
assert!(!cache.cross_attn_cached);
let _logits = model
.decoder
.forward_one(1, &encoder_output, &mut cache)
.expect("forward_one");
assert!(
cache.cross_attn_cached,
"cross_attn_cached must be true after forward_one for Moonshine"
);
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_moonshine_decoder_forward_one_cache_has_layers() {
let config = model::ModelConfig::moonshine_tiny();
let decoder = model::Decoder::new(&config);
let cache =
model::DecoderKVCache::new(decoder.n_layers(), decoder.d_model(), decoder.max_len());
assert!(
!cache.self_attn_cache.is_empty(),
"Moonshine decoder cache must have at least 1 layer"
);
}
#[test]
fn test_moonshine_decoder_is_finalized() {
let config = model::ModelConfig::moonshine_tiny();
let decoder = model::Decoder::new(&config);
assert!(
decoder.is_finalized(),
"Moonshine decoder should be considered finalized"
);
}
#[test]
fn test_whisper_decoder_is_finalized_initially_false() {
let config = model::ModelConfig::tiny();
let decoder = model::Decoder::new(&config);
assert!(
!decoder.is_finalized(),
"Whisper decoder should NOT be finalized before finalize_weights()"
);
}
#[test]
#[ignore = "Allocates large model - run with --ignored"]
fn test_moonshine_decoder_forward_traced() {
let model = WhisperApr::moonshine_tiny();
let d_model = 288;
let enc_seq_len = 7;
let encoder_output = vec![0.1_f32; enc_seq_len * d_model];
let (logits, trace) = model
.decoder
.forward_traced(&[1, 2], &encoder_output)
.expect("forward_traced");
let layer_traces: Vec<_> = trace
.iter()
.filter(|(name, _)| name.starts_with("layer_"))
.collect();
assert!(
!layer_traces.is_empty(),
"forward_traced must produce layer traces for Moonshine"
);
assert!(!logits.is_empty());
assert!(logits.iter().all(|v| v.is_finite()));
}
#[test]
fn test_moonshine_encoder_forward_batch_uses_forward() {
let config = model::ModelConfig::moonshine_tiny();
let encoder = model::Encoder::new(&config);
let d_model = 288;
let batch = vec![vec![0.1_f32; 7 * d_model], vec![0.1_f32; 5 * d_model]];
let results = encoder.forward_batch(&batch).expect("forward_batch");
assert_eq!(results.len(), 2);
assert_eq!(results[0].len(), 7 * d_model);
assert_eq!(results[1].len(), 5 * d_model);
}
}