1#![allow(clippy::all, clippy::pedantic, clippy::restriction, clippy::nursery)]
3
4#[cfg(feature = "realizar-gpu")]
6use crate::cuda;
7use crate::{
8 audio, detection, error, format, inference, model, progress, timestamps, tokenizer, vad,
9};
10pub use error::{WhisperError, WhisperResult};
11
12#[derive(Debug, Clone, Copy, PartialEq, Eq)]
14pub enum ModelType {
15 Tiny,
17 TinyEn,
19 Base,
21 BaseEn,
23 Small,
25 SmallEn,
27 Medium,
29 MediumEn,
31 Large,
33 LargeV1,
35 LargeV2,
37 LargeV3,
39 LargeV3Turbo,
41}
42
43#[derive(Debug, Clone, Copy, Default)]
45pub enum DecodingStrategy {
46 #[default]
48 Greedy,
49 BeamSearch {
51 beam_size: usize,
53 temperature: f32,
55 patience: f32,
57 },
58 Sampling {
60 temperature: f32,
62 top_k: Option<usize>,
64 top_p: Option<f32>,
66 },
67}
68
69#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
71pub enum Task {
72 #[default]
74 Transcribe,
75 Translate,
77}
78
79#[derive(Debug, Clone, Default)]
81pub struct TranscribeOptions {
82 pub language: Option<String>,
84 pub task: Task,
86 pub strategy: DecodingStrategy,
88 pub word_timestamps: bool,
90 pub profile: bool,
92 pub prompt: Option<String>,
95 pub hotwords: Vec<String>,
98}
99
100#[derive(Debug, Clone, Default)]
102#[cfg_attr(feature = "cli", derive(serde::Serialize, serde::Deserialize))]
103pub struct Segment {
104 pub start: f32,
106 pub end: f32,
108 pub text: String,
110 pub tokens: Vec<u32>,
112}
113
114#[derive(Debug, Clone, Default)]
116#[cfg_attr(feature = "cli", derive(serde::Serialize, serde::Deserialize))]
117pub struct ProfilingStats {
118 pub total_ms: f64,
120 pub breakdown: std::collections::HashMap<String, f64>,
122 #[cfg_attr(feature = "cli", serde(skip_serializing_if = "Option::is_none"))]
124 pub trace_json: Option<String>,
125}
126
127#[derive(Debug, Clone, Default)]
129#[cfg_attr(feature = "cli", derive(serde::Serialize, serde::Deserialize))]
130pub struct TranscriptionResult {
131 pub text: String,
133 pub language: String,
135 pub segments: Vec<Segment>,
137 #[cfg_attr(feature = "cli", serde(skip_serializing_if = "Option::is_none"))]
139 pub profiling: Option<ProfilingStats>,
140}
141
142#[derive(Debug, Clone)]
144pub struct BatchTranscriptionResult {
145 pub results: Vec<TranscriptionResult>,
147 pub total_duration_secs: f32,
149}
150
151impl BatchTranscriptionResult {
152 #[must_use]
154 pub fn len(&self) -> usize {
155 self.results.len()
156 }
157
158 #[must_use]
160 pub fn is_empty(&self) -> bool {
161 self.results.is_empty()
162 }
163
164 #[must_use]
166 pub fn get(&self, index: usize) -> Option<&TranscriptionResult> {
167 self.results.get(index)
168 }
169
170 pub fn iter(&self) -> impl Iterator<Item = &TranscriptionResult> {
172 self.results.iter()
173 }
174
175 #[must_use]
177 pub fn texts(&self) -> Vec<&str> {
178 self.results.iter().map(|r| r.text.as_str()).collect()
179 }
180}
181
182pub struct SummarizeOptions<'a> {
186 pub model: &'a model::lfm2::Lfm2,
188 pub tokenizer: &'a model::lfm2::Lfm2Tokenizer,
190 pub max_tokens: usize,
192 pub temperature: f32,
194}
195
196impl<'a> SummarizeOptions<'a> {
197 #[must_use]
203 pub fn new(model: &'a model::lfm2::Lfm2, tokenizer: &'a model::lfm2::Lfm2Tokenizer) -> Self {
204 Self {
205 model,
206 tokenizer,
207 max_tokens: 256,
208 temperature: 0.3,
209 }
210 }
211
212 #[must_use]
214 pub const fn with_max_tokens(mut self, max_tokens: usize) -> Self {
215 self.max_tokens = max_tokens;
216 self
217 }
218
219 #[must_use]
221 pub const fn with_temperature(mut self, temperature: f32) -> Self {
222 self.temperature = temperature;
223 self
224 }
225}
226
227#[derive(Debug, Clone)]
231pub struct TranscribeSummaryResult {
232 pub transcription: TranscriptionResult,
234 pub summary: String,
236 pub generation_stats: Option<model::lfm2::GenerationStats>,
238}
239
240impl TranscribeSummaryResult {
241 #[must_use]
243 pub fn transcript(&self) -> &str {
244 &self.transcription.text
245 }
246
247 #[must_use]
249 pub fn summary(&self) -> &str {
250 &self.summary
251 }
252
253 #[must_use]
255 pub fn has_summary(&self) -> bool {
256 !self.summary.is_empty()
257 }
258}
259
260#[derive(Debug, Clone)]
281pub struct WhisperApr {
282 config: model::ModelConfig,
284 encoder: model::Encoder,
286 decoder: model::Decoder,
288 tokenizer: tokenizer::Tokenizer,
290 mel_filters: Option<audio::MelFilterbank>,
292 conv_stem: Option<audio::ConvStem>,
294 resampler: Option<audio::SincResampler>,
296 weights_loaded: bool,
298}
299
300impl WhisperApr {
301 #[must_use]
306 pub fn from_config(config: model::ModelConfig) -> Self {
307 let encoder = model::Encoder::new(&config);
308 let decoder = model::Decoder::new(&config);
309 let tokenizer = if config.model_family == format::ModelFamily::Moonshine {
310 tokenizer::Tokenizer::SentencePiece(
311 tokenizer::SentencePieceTokenizer::moonshine_default(),
312 )
313 } else {
314 tokenizer::Tokenizer::Bpe(tokenizer::BpeTokenizer::with_base_tokens())
315 };
316
317 let (mel_filters, conv_stem) = match config.audio_frontend {
319 model::AudioFrontend::MelFilterbank => (
320 Some(audio::MelFilterbank::new(&audio::MelConfig {
321 n_mels: config.n_mels as usize,
322 ..audio::MelConfig::whisper()
323 })),
324 None,
325 ),
326 model::AudioFrontend::LearnedConv => (
327 None,
328 Some(audio::ConvStem::new(config.n_audio_state as usize)),
329 ),
330 };
331
332 Self {
333 config,
334 encoder,
335 decoder,
336 tokenizer,
337 mel_filters,
338 conv_stem,
339 resampler: None,
340 weights_loaded: false,
341 }
342 }
343
344 #[must_use]
346 pub fn tiny() -> Self {
347 Self::from_config(model::ModelConfig::tiny())
348 }
349
350 #[must_use]
352 pub fn base() -> Self {
353 Self::from_config(model::ModelConfig::base())
354 }
355
356 #[must_use]
358 pub fn small() -> Self {
359 Self::from_config(model::ModelConfig::small())
360 }
361
362 #[must_use]
364 pub fn medium() -> Self {
365 Self::from_config(model::ModelConfig::medium())
366 }
367
368 #[must_use]
370 pub fn large() -> Self {
371 Self::from_config(model::ModelConfig::large())
372 }
373
374 #[must_use]
376 pub fn moonshine_tiny() -> Self {
377 Self::from_config(model::ModelConfig::moonshine_tiny())
378 }
379
380 #[must_use]
382 pub fn moonshine_base() -> Self {
383 Self::from_config(model::ModelConfig::moonshine_base())
384 }
385
386 #[must_use]
388 pub const fn config(&self) -> &model::ModelConfig {
389 &self.config
390 }
391
392 #[must_use]
394 pub const fn model_type(&self) -> ModelType {
395 self.config.model_type
396 }
397
398 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;
406
407 pub fn transcribe(
423 &self,
424 audio: &[f32],
425 options: TranscribeOptions,
426 ) -> WhisperResult<TranscriptionResult> {
427 if audio.len() > Self::CHUNK_SAMPLES {
429 return self.transcribe_chunked(audio, options);
430 }
431
432 self.transcribe_single_chunk(audio, options)
434 }
435
436 fn transcribe_single_chunk(
440 &self,
441 audio: &[f32],
442 options: TranscribeOptions,
443 ) -> WhisperResult<TranscriptionResult> {
444 #[cfg(feature = "std")]
445 let start_total = if options.profile {
446 Some(std::time::Instant::now())
447 } else {
448 None
449 };
450
451 #[cfg(feature = "std")]
453 let start_audio = start_total.map(|_| std::time::Instant::now());
454
455 #[cfg(feature = "std")]
456 let mut mel_ms: Option<f64> = None;
457
458 #[cfg(feature = "std")]
460 let mut brick_profiler: Option<trueno::BrickProfiler> = if options.profile {
461 Some(trueno::BrickProfiler::enabled())
462 } else {
463 None
464 };
465
466 #[cfg(feature = "std")]
468 let mut page_faults: Option<(u64, u64)> = None;
469
470 #[cfg(feature = "std")]
472 let mut blis_profiler_stats: Option<trueno::blis::BlisProfiler> = None;
473
474 #[cfg(all(feature = "std", feature = "realizar-inference"))]
476 let mut inference_tracer: Option<realizar::InferenceTracer> = if options.profile {
477 let config = realizar::TraceConfig::enabled();
478 let mut tracer = realizar::InferenceTracer::new(config);
479 tracer.set_model_info(realizar::ModelInfo {
480 name: format!("{:?}", self.config.model_type),
481 num_layers: self.config.n_audio_layer as usize,
482 hidden_dim: self.config.n_audio_state as usize,
483 vocab_size: self.config.n_vocab as usize,
484 num_heads: self.config.n_audio_head as usize,
485 quant_type: Some("f32".into()),
486 });
487 Some(tracer)
488 } else {
489 None
490 };
491
492 let audio_features = match self.config.audio_frontend {
493 model::AudioFrontend::MelFilterbank => {
494 #[cfg(feature = "std")]
495 let mel_start = std::time::Instant::now();
496
497 let mel = self.compute_mel(audio)?;
498
499 #[cfg(feature = "std")]
500 {
501 mel_ms = Some(mel_start.elapsed().as_secs_f64() * 1000.0);
502 }
503
504 #[cfg(feature = "std")]
506 if let Some(ref mut profiler) = brick_profiler {
507 crate::simd::enable_blis_profiling();
509 let (pf_minor_before, pf_major_before) = trueno::brick::get_page_faults();
510 #[cfg(feature = "realizar-inference")]
511 let features =
512 self.encode_profiled(&mel, profiler, inference_tracer.as_mut())?;
513 #[cfg(not(feature = "realizar-inference"))]
514 let features = self.encode_profiled(&mel, profiler)?;
515 blis_profiler_stats = crate::simd::take_blis_profiler();
517 let (pf_minor_after, pf_major_after) = trueno::brick::get_page_faults();
518 page_faults = Some((
519 pf_minor_after.saturating_sub(pf_minor_before),
520 pf_major_after.saturating_sub(pf_major_before),
521 ));
522 features
523 } else {
524 self.encode(&mel)?
525 }
526
527 #[cfg(not(feature = "std"))]
528 self.encode(&mel)?
529 }
530 model::AudioFrontend::LearnedConv => {
531 let stem = self
532 .conv_stem
533 .as_ref()
534 .ok_or_else(|| WhisperError::Model("Moonshine requires ConvStem".into()))?;
535 let stem_out = stem.forward(audio)?;
536 self.encoder.forward(&stem_out)?
537 }
538 };
539
540 #[cfg(feature = "std")]
541 let enc_ms = start_audio.map(|s| s.elapsed().as_secs_f64() * 1000.0);
542
543 let language = options.language.clone().unwrap_or_else(|| "en".to_string());
545
546 let initial_tokens =
548 self.build_initial_tokens(&language, options.task, options.prompt.as_deref());
549
550 #[cfg(feature = "std")]
552 let start_dec = start_total.map(|_| std::time::Instant::now());
553 let tokens = self.decode(&audio_features, &initial_tokens, &options)?;
554 #[cfg(feature = "std")]
555 let dec_ms = start_dec.map(|s| s.elapsed().as_secs_f64() * 1000.0);
556
557 let text = self.tokenizer.decode(&tokens)?;
559
560 let segments = if timestamps::has_timestamps(&tokens) {
562 timestamps::extract_segments(&tokens, |ts| self.tokenizer.decode(ts).ok())
563 } else if !text.trim().is_empty() {
564 let duration = audio.len() as f32 / audio::SAMPLE_RATE as f32;
565 let single = vec![Segment {
566 start: 0.0,
567 end: duration,
568 text: text.clone(),
569 tokens: tokens.clone(),
570 }];
571 timestamps::split_long_segments(&single, Self::MAX_SUBTITLE_SECS)
572 } else {
573 Vec::new()
574 };
575
576 #[cfg(feature = "std")]
578 let profiling = start_total.map(|st| {
579 let mut breakdown = std::collections::HashMap::new();
580 let pairs: &[(&str, Option<f64>)] = &[
581 ("mel_ms", mel_ms),
582 ("audio_ms", enc_ms),
583 ("decoder_ms", dec_ms),
584 ];
585 for &(key, val) in pairs {
586 if let Some(ms) = val {
587 breakdown.insert(key.to_string(), ms);
588 }
589 }
590 if let Some(audio) = enc_ms {
592 let enc_total = audio - mel_ms.unwrap_or(0.0);
593 breakdown.insert("encoder_ms".to_string(), enc_total);
594 }
595 if let Some(ref profiler) = brick_profiler {
597 let cats = profiler.category_stats();
598 let total_prof_ns = profiler.total_ns();
599 let norm_ns = cats[trueno::BrickCategory::Norm as usize].total_ns;
601 let attn_ns = cats[trueno::BrickCategory::Attention as usize].total_ns;
602 let ffn_ns = cats[trueno::BrickCategory::Ffn as usize].total_ns;
603 let other_ns = cats[trueno::BrickCategory::Other as usize].total_ns;
604 breakdown.insert("brick_norm_ms".to_string(), norm_ns as f64 / 1_000_000.0);
605 breakdown.insert("brick_attn_ms".to_string(), attn_ns as f64 / 1_000_000.0);
606 breakdown.insert("brick_ffn_ms".to_string(), ffn_ns as f64 / 1_000_000.0);
607 breakdown.insert("brick_other_ms".to_string(), other_ns as f64 / 1_000_000.0);
608 breakdown.insert("brick_total_ns".to_string(), total_prof_ns as f64);
609 for &brick_id in &[
611 trueno::BrickId::LayerNorm,
612 trueno::BrickId::AttentionScore,
613 trueno::BrickId::GateProjection,
614 trueno::BrickId::Embedding,
615 ] {
616 let stats = profiler.brick_stats(brick_id);
617 if stats.count > 0 {
618 let name = brick_id.name();
619 breakdown.insert(
620 format!("brick_{name}_ms"),
621 stats.total_ns as f64 / 1_000_000.0,
622 );
623 breakdown.insert(format!("brick_{name}_count"), stats.count as f64);
624 breakdown.insert(
625 format!("brick_{name}_cycles_per_elem"),
626 stats.cycles_per_element(),
627 );
628 let diagnosis = stats.diagnose_from_cycles();
629 let diag_code = match diagnosis {
631 "memory-bound (low IPC, likely cache misses)" => 1.0,
632 "compute-bound (efficient)" => 2.0,
633 "throttled or context-switched" => 3.0,
634 "balanced" => 4.0,
635 _ => 0.0,
636 };
637 breakdown.insert(format!("brick_{name}_bottleneck"), diag_code);
638 }
639 }
640 if let Some((minor, major)) = page_faults {
642 breakdown.insert("page_faults_minor".to_string(), minor as f64);
643 breakdown.insert("page_faults_major".to_string(), major as f64);
644 }
645 }
646 if let Some(ref blis) = blis_profiler_stats {
648 breakdown.insert("blis_macro_gflops".to_string(), blis.macro_stats.gflops());
649 breakdown.insert(
650 "blis_macro_calls".to_string(),
651 blis.macro_stats.count as f64,
652 );
653 breakdown.insert(
654 "blis_macro_ns".to_string(),
655 blis.macro_stats.total_ns as f64,
656 );
657 breakdown.insert("blis_midi_gflops".to_string(), blis.midi_stats.gflops());
658 breakdown.insert("blis_midi_calls".to_string(), blis.midi_stats.count as f64);
659 breakdown.insert("blis_micro_gflops".to_string(), blis.micro_stats.gflops());
660 breakdown.insert(
661 "blis_micro_calls".to_string(),
662 blis.micro_stats.count as f64,
663 );
664 breakdown.insert("blis_pack_ns".to_string(), blis.pack_stats.total_ns as f64);
665 breakdown.insert("blis_pack_calls".to_string(), blis.pack_stats.count as f64);
666 breakdown.insert("blis_total_gflops".to_string(), blis.total_gflops());
667 if blis.macro_stats.total_ns > 0 {
669 let pack_pct =
670 blis.pack_stats.total_ns as f64 / blis.macro_stats.total_ns as f64 * 100.0;
671 breakdown.insert("blis_pack_pct".to_string(), pack_pct);
672 }
673 }
674 #[cfg(feature = "realizar-inference")]
676 let trace_json = inference_tracer.as_ref().map(|t| t.to_json());
677 #[cfg(not(feature = "realizar-inference"))]
678 let trace_json: Option<String> = None;
679
680 ProfilingStats {
681 total_ms: st.elapsed().as_secs_f64() * 1000.0,
682 breakdown,
683 trace_json,
684 }
685 });
686
687 #[cfg(not(feature = "std"))]
688 let profiling = None;
689
690 Ok(TranscriptionResult {
691 text,
692 language,
693 segments,
694 profiling,
695 })
696 }
697
698 fn transcribe_chunked(
718 &self,
719 audio: &[f32],
720 options: TranscribeOptions,
721 ) -> WhisperResult<TranscriptionResult> {
722 let chunk_size = Self::CHUNK_SAMPLES;
723 let overlap = Self::OVERLAP_SAMPLES;
724 let step = chunk_size; let language = options.language.clone().unwrap_or_else(|| "en".to_string());
728
729 let mut all_segments: Vec<Segment> = Vec::new();
730 let mut all_text = String::new();
731 let mut chunk_idx = 0;
732
733 let mut offset = 0;
735 while offset < audio.len() {
736 let chunk_end = (offset + chunk_size + overlap).min(audio.len());
738 let chunk = &audio[offset..chunk_end];
739
740 if chunk.len() < audio::SAMPLE_RATE as usize / 2 {
742 break;
743 }
744
745 let chunk_options = TranscribeOptions {
747 language: Some(language.clone()),
748 task: options.task,
749 strategy: options.strategy,
750 word_timestamps: options.word_timestamps,
751 profile: options.profile,
752 prompt: options.prompt.clone(),
753 hotwords: options.hotwords.clone(),
754 };
755
756 let chunk_result = self.transcribe_single_chunk(chunk, chunk_options)?;
757
758 let time_offset = offset as f32 / audio::SAMPLE_RATE as f32;
760
761 let chunk_text = if chunk_idx == 0 {
763 chunk_result.text.clone()
765 } else {
766 let overlap_ratio = overlap as f32 / chunk.len() as f32;
770 let words: Vec<&str> = chunk_result.text.split_whitespace().collect();
771 let skip_words = ((words.len() as f32) * overlap_ratio * 0.8) as usize;
772 words
773 .into_iter()
774 .skip(skip_words)
775 .collect::<Vec<_>>()
776 .join(" ")
777 };
778
779 if !chunk_text.is_empty() {
781 if !all_text.is_empty() {
782 all_text.push(' ');
783 }
784 all_text.push_str(&chunk_text);
785 }
786
787 for mut seg in chunk_result.segments {
789 seg.start += time_offset;
790 seg.end += time_offset;
791 all_segments.push(seg);
792 }
793
794 offset += step;
796 chunk_idx += 1;
797 }
798
799 let merged_segments = self.merge_overlapping_segments(all_segments);
801
802 let final_segments =
804 timestamps::split_long_segments(&merged_segments, Self::MAX_SUBTITLE_SECS);
805
806 Ok(TranscriptionResult {
807 text: all_text,
808 language,
809 segments: final_segments,
810 profiling: None,
811 })
812 }
813
814 fn merge_overlapping_segments(&self, segments: Vec<Segment>) -> Vec<Segment> {
819 if segments.is_empty() {
820 return segments;
821 }
822
823 let mut merged: Vec<Segment> = Vec::with_capacity(segments.len());
824 let mut current = segments[0].clone();
825
826 for seg in segments.into_iter().skip(1) {
827 if seg.start < current.end + 0.1 {
829 current.end = current.end.max(seg.end);
831 if !seg.text.is_empty() {
832 if !current.text.is_empty() {
833 current.text.push(' ');
834 }
835 current.text.push_str(&seg.text);
836 }
837 current.tokens.extend(seg.tokens);
838 } else {
839 merged.push(current);
841 current = seg;
842 }
843 }
844
845 merged.push(current);
847
848 merged
849 }
850
851 pub fn forward_probed(
869 &self,
870 audio: &[f32],
871 tokens: &[u32],
872 probe: &mut crate::probe::ActivationProbe,
873 ) -> WhisperResult<Vec<f32>> {
874 let audio_features = match self.config.audio_frontend {
876 model::AudioFrontend::MelFilterbank => {
877 let mel = self.compute_mel(audio)?;
878 self.encoder.forward_probed(&mel, probe)?
879 }
880 model::AudioFrontend::LearnedConv => {
881 let stem = self
882 .conv_stem
883 .as_ref()
884 .ok_or_else(|| WhisperError::Model("Moonshine requires ConvStem".into()))?;
885 let stem_out = stem.forward_probed(audio, probe)?;
886 self.encoder.forward_probed(&stem_out, probe)?
887 }
888 };
889
890 self.decoder.forward_probed(tokens, &audio_features, probe)
892 }
893
894 fn eot_token(&self) -> u32 {
896 if self.config.model_family == format::ModelFamily::Moonshine {
897 2 } else {
899 tokenizer::special_tokens::EOT
900 }
901 }
902
903 pub fn compute_mel(&self, audio: &[f32]) -> WhisperResult<Vec<f32>> {
909 let mel_fb = self.mel_filters.as_ref().ok_or_else(|| {
910 WhisperError::Audio("mel filterbank not available (Moonshine model?)".into())
911 })?;
912 let mel = mel_fb
913 .compute(&audio)
914 .map_err(|e| WhisperError::Audio(e.to_string()))?;
915
916 Ok(mel)
920 }
921
922 pub fn encode(&self, mel: &[f32]) -> WhisperResult<Vec<f32>> {
924 self.encoder.forward_mel(mel)
926 }
927
928 #[cfg(feature = "realizar-inference")]
933 pub fn encode_profiled(
934 &self,
935 mel: &[f32],
936 profiler: &mut trueno::BrickProfiler,
937 tracer: Option<&mut realizar::InferenceTracer>,
938 ) -> WhisperResult<Vec<f32>> {
939 self.encoder.forward_mel_profiled(mel, profiler, tracer)
940 }
941
942 #[cfg(not(feature = "realizar-inference"))]
946 pub fn encode_profiled(
947 &self,
948 mel: &[f32],
949 profiler: &mut trueno::BrickProfiler,
950 ) -> WhisperResult<Vec<f32>> {
951 self.encoder.forward_mel_profiled(mel, profiler)
952 }
953
954 fn get_initial_tokens(&self, language: &str, task: Task) -> Vec<u32> {
959 if self.config.model_family == format::ModelFamily::Moonshine {
961 return vec![1]; }
963
964 use tokenizer::special_tokens::{self, SpecialTokens};
966
967 let specials = SpecialTokens::for_vocab_size(self.config.n_vocab as usize);
969
970 let mut tokens = vec![specials.sot];
971
972 if specials.is_multilingual {
976 let lang_offset = special_tokens::language_offset(language).unwrap_or(0);
977 tokens.push(specials.lang_base + lang_offset);
978 }
979
980 match task {
982 Task::Transcribe => tokens.push(specials.transcribe),
983 Task::Translate => tokens.push(special_tokens::TRANSLATE),
984 }
985
986 tokens
991 }
992
993 fn build_initial_tokens(&self, language: &str, task: Task, prompt: Option<&str>) -> Vec<u32> {
998 let mut tokens = self.get_initial_tokens(language, task);
999
1000 let prompt_text = match prompt {
1001 Some(p) if !p.is_empty() => p,
1002 _ => return tokens,
1003 };
1004
1005 let prompt_tokens = match self.tokenizer.encode(prompt_text) {
1006 Ok(t) if !t.is_empty() => t,
1007 _ => return tokens,
1008 };
1009
1010 let max_prompt = (self.config.n_text_ctx as usize / 2).min(224);
1012 let truncated = if prompt_tokens.len() > max_prompt {
1013 &prompt_tokens[prompt_tokens.len() - max_prompt..]
1014 } else {
1015 &prompt_tokens
1016 };
1017
1018 let mut prefix = Vec::with_capacity(1 + truncated.len() + tokens.len());
1020 prefix.push(tokenizer::special_tokens::PREV);
1021 prefix.extend_from_slice(truncated);
1022 prefix.append(&mut tokens);
1023 prefix
1024 }
1025
1026 pub fn detect_language(&self, audio: &[f32]) -> WhisperResult<detection::LanguageProbs> {
1039 let mel = self.compute_mel(audio)?;
1041
1042 let audio_features = self.encode(&mel)?;
1044
1045 let n_vocab = self.config.n_vocab as usize;
1047 let logits_fn = |tokens: &[u32]| -> WhisperResult<Vec<f32>> {
1048 let all_logits = self.decoder.forward(tokens, &audio_features)?;
1049 let seq_len = tokens.len();
1050 let last_start = (seq_len - 1) * n_vocab;
1051
1052 if all_logits.len() >= last_start + n_vocab {
1053 Ok(all_logits[last_start..last_start + n_vocab].to_vec())
1054 } else {
1055 let mut padded = vec![f32::NEG_INFINITY; n_vocab];
1056 let available = all_logits.len().saturating_sub(last_start);
1057 if available > 0 {
1058 padded[..available].copy_from_slice(&all_logits[last_start..]);
1059 }
1060 Ok(padded)
1061 }
1062 };
1063
1064 let detector = detection::LanguageDetector::new();
1066 detector.detect(logits_fn)
1067 }
1068
1069 fn decode(
1073 &self,
1074 audio_features: &[f32],
1075 initial_tokens: &[u32],
1076 options: &TranscribeOptions,
1077 ) -> WhisperResult<Vec<u32>> {
1078 use std::cell::RefCell;
1079
1080 let n_vocab = self.config.n_vocab as usize;
1081 let max_tokens = self.config.n_text_ctx as usize;
1082
1083 let cache = RefCell::new(self.decoder.create_kv_cache());
1086 let processed_count = RefCell::new(0usize);
1087
1088 let suppressor = inference::WhisperTokenSuppressor::new()
1091 .with_timestamp_suppression(false)
1092 .with_vocab_size(n_vocab);
1093
1094 let hotword_booster = if !options.hotwords.is_empty() {
1096 let mut booster = crate::vocabulary::HotwordBooster::new();
1097 for word in &options.hotwords {
1098 if let Ok(tokens) = self.tokenizer.encode(word) {
1099 if !tokens.is_empty() {
1100 booster.add_hotword_with_tokens_default(word, tokens);
1101 }
1102 }
1103 }
1104 if booster.is_empty() {
1105 None
1106 } else {
1107 Some(booster)
1108 }
1109 } else {
1110 None
1111 };
1112
1113 let scratch = std::cell::RefCell::new(self.decoder.create_decoder_scratch());
1115
1116 let logits_fn = |tokens: &[u32]| -> WhisperResult<Vec<f32>> {
1117 let seq_len = tokens.len();
1118 let already_processed = *processed_count.borrow();
1119
1120 let mut logits = vec![f32::NEG_INFINITY; n_vocab];
1122
1123 for &token in tokens.iter().take(seq_len).skip(already_processed) {
1124 logits = self.decoder.forward_one_with_scratch(
1125 token,
1126 audio_features,
1127 &mut cache.borrow_mut(),
1128 &mut scratch.borrow_mut(),
1129 )?;
1130 }
1131
1132 *processed_count.borrow_mut() = seq_len;
1133
1134 suppressor.apply(&mut logits);
1137
1138 if let Some(ref booster) = hotword_booster {
1140 booster.apply_bias(&mut logits, tokens);
1141 }
1142
1143 let eot_id = self.eot_token();
1145 Self::suppress_repetitions(&mut logits, tokens, eot_id, n_vocab);
1146
1147 Ok(logits)
1148 };
1149
1150 let eot = self.eot_token();
1152 match options.strategy {
1153 DecodingStrategy::Greedy => {
1154 let decoder = inference::GreedyDecoder::new(max_tokens);
1155 decoder.decode(logits_fn, initial_tokens, eot)
1156 }
1157 DecodingStrategy::BeamSearch {
1158 beam_size,
1159 temperature,
1160 patience,
1161 } => {
1162 let decoder = inference::BeamSearchDecoder::new(beam_size, max_tokens)
1163 .with_temperature(temperature)
1164 .with_patience(patience);
1165 decoder.decode(logits_fn, initial_tokens, eot)
1166 }
1167 DecodingStrategy::Sampling { temperature, .. } => {
1168 let decoder =
1170 inference::GreedyDecoder::new(max_tokens).with_temperature(temperature);
1171 decoder.decode(logits_fn, initial_tokens, eot)
1172 }
1173 }
1174 }
1175
1176 fn suppress_repetitions(logits: &mut [f32], tokens: &[u32], eot_id: u32, n_vocab: usize) {
1182 let text_tokens: Vec<u32> = if eot_id <= 3 {
1186 tokens.iter().copied().filter(|&t| t > eot_id).collect()
1187 } else {
1188 tokens.iter().copied().filter(|&t| t < eot_id).collect()
1189 };
1190
1191 let window_start = text_tokens.len().saturating_sub(50);
1193 for &prev_tok in &text_tokens[window_start..] {
1194 if (prev_tok as usize) < n_vocab {
1195 logits[prev_tok as usize] -= 2.0;
1196 }
1197 }
1198
1199 if text_tokens.len() >= 3 {
1201 let prev2 = text_tokens[text_tokens.len() - 2];
1202 let prev1 = text_tokens[text_tokens.len() - 1];
1203 for w in text_tokens[..text_tokens.len() - 2].windows(3) {
1204 if w[0] == prev2 && w[1] == prev1 && (w[2] as usize) < n_vocab {
1205 logits[w[2] as usize] -= 10.0;
1206 }
1207 }
1208 }
1209
1210 Self::suppress_degenerate_loops(logits, &text_tokens, eot_id, n_vocab);
1211 }
1212
1213 fn suppress_degenerate_loops(
1215 logits: &mut [f32],
1216 text_tokens: &[u32],
1217 eot_id: u32,
1218 n_vocab: usize,
1219 ) {
1220 if text_tokens.len() >= 12 {
1222 let mut fourgram_counts = std::collections::HashMap::<[u32; 4], u32>::new();
1223 for w in text_tokens.windows(4) {
1224 *fourgram_counts.entry([w[0], w[1], w[2], w[3]]).or_insert(0) += 1;
1225 }
1226 if fourgram_counts.values().any(|&c| c >= 3) && (eot_id as usize) < n_vocab {
1227 logits[eot_id as usize] = 100.0;
1228 return;
1229 }
1230 }
1231
1232 if text_tokens.len() >= 80 {
1234 let window = &text_tokens[text_tokens.len() - 80..];
1235 let mut freq = std::collections::HashMap::<u32, u32>::new();
1236 for &t in window {
1237 *freq.entry(t).or_insert(0) += 1;
1238 }
1239 for (&tok, &count) in &freq {
1240 if count >= 8 && (tok as usize) < n_vocab {
1241 logits[tok as usize] -= (count as f32 - 7.0) * 3.0;
1242 }
1243 }
1244 if freq.len() as f32 / 80.0 < 0.35 && (eot_id as usize) < n_vocab {
1245 logits[eot_id as usize] = 100.0;
1246 }
1247 }
1248 }
1249
1250 pub fn set_resampler(&mut self, input_rate: u32) -> WhisperResult<()> {
1255 if input_rate == audio::SAMPLE_RATE {
1256 self.resampler = None;
1257 } else {
1258 self.resampler = Some(audio::SincResampler::new(input_rate, audio::SAMPLE_RATE)?);
1259 }
1260 Ok(())
1261 }
1262
1263 pub fn resample(&self, audio: &[f32]) -> WhisperResult<Vec<f32>> {
1268 self.resampler
1269 .as_ref()
1270 .map_or_else(|| Ok(audio.to_vec()), |resampler| resampler.resample(audio))
1271 }
1272
1273 #[must_use]
1275 pub const fn tokenizer(&self) -> &tokenizer::Tokenizer {
1276 &self.tokenizer
1277 }
1278
1279 #[must_use]
1281 pub fn memory_size(&self) -> usize {
1282 let params = match self.config.model_type {
1284 ModelType::Tiny | ModelType::TinyEn => 39_000_000,
1285 ModelType::Base | ModelType::BaseEn => 74_000_000,
1286 ModelType::Small | ModelType::SmallEn => 244_000_000,
1287 ModelType::Medium | ModelType::MediumEn => 769_000_000,
1288 ModelType::Large | ModelType::LargeV1 | ModelType::LargeV2 | ModelType::LargeV3 => {
1289 1_550_000_000
1290 }
1291 ModelType::LargeV3Turbo => 809_000_000,
1292 };
1293 params * 4 }
1295
1296 #[must_use]
1301 pub fn has_weights(&self) -> bool {
1302 self.weights_loaded
1303 }
1304
1305 pub fn load_from_apr(data: &[u8]) -> WhisperResult<Self> {
1325 Self::load_from_apr_with_progress(data, &mut progress::null_callback)
1326 }
1327
1328 pub fn load_from_apr_with_progress(
1352 data: &[u8],
1353 callback: progress::ProgressCallback<'_>,
1354 ) -> WhisperResult<Self> {
1355 let mut tracker = progress::ProgressTracker::model_loading();
1357
1358 callback(&tracker.to_progress());
1360 let reader = format::AprV2ReaderRef::from_bytes(data)
1361 .map_err(|e| error::WhisperError::Format(e.to_string()))?;
1362 let config = format::metadata_to_model_config(reader.metadata());
1363 tracker.next_phase();
1364
1365 let is_f16 = reader
1366 .tensor_names()
1367 .iter()
1368 .find(|n| n.ends_with(".weight") && !n.starts_with("__"))
1369 .and_then(|n| reader.get_tensor(n))
1370 .map_or(false, |t| t.dtype == format::TensorDType::F16);
1371
1372 callback(&tracker.to_progress());
1374 let mut encoder = model::Encoder::new(&config);
1375 if is_f16 {
1376 Self::load_encoder_weights_f16(&reader, &mut encoder, &mut tracker, callback);
1377 } else {
1378 Self::load_encoder_weights(&reader, &mut encoder, &mut tracker, callback);
1379 }
1380 encoder.finalize_weights();
1383 tracker.next_phase();
1384
1385 callback(&tracker.to_progress());
1387 let mut decoder = model::Decoder::new(&config);
1388 if is_f16 {
1389 Self::load_decoder_weights_f16(&reader, &mut decoder, &mut tracker, callback);
1391 } else {
1392 Self::load_decoder_weights(&reader, &mut decoder, &mut tracker, callback);
1393 decoder.convert_to_f16();
1396 }
1397 decoder.finalize_weights();
1399 tracker.next_phase();
1400
1401 callback(&tracker.to_progress());
1403 let tokenizer = Self::build_tokenizer(&config, &reader);
1405 tracker.next_phase();
1406
1407 callback(&tracker.to_progress());
1409 let (mel_filters, conv_stem) = match config.audio_frontend {
1410 model::AudioFrontend::MelFilterbank => {
1411 let mel_config = audio::MelConfig {
1412 n_mels: config.n_mels as usize,
1413 ..audio::MelConfig::whisper()
1414 };
1415 let mf = Self::read_mel_filterbank(&reader).map_or_else(
1416 || audio::MelFilterbank::new(&mel_config),
1417 |fb| audio::MelFilterbank::from_filters(fb.data, &mel_config),
1418 );
1419 (Some(mf), None)
1420 }
1421 model::AudioFrontend::LearnedConv => {
1422 let d_model = config.n_audio_state as usize;
1423 let mut stem = audio::ConvStem::new(d_model);
1424 Self::load_conv_stem_weights(&reader, &mut stem);
1425 (None, Some(stem))
1426 }
1427 };
1428 tracker.complete();
1429 callback(&tracker.to_progress());
1430
1431 Ok(Self {
1432 config,
1433 encoder,
1434 decoder,
1435 tokenizer,
1436 mel_filters,
1437 conv_stem,
1438 resampler: None,
1439 weights_loaded: true,
1440 })
1441 }
1442
1443 fn build_tokenizer(
1445 config: &model::ModelConfig,
1446 reader: &format::AprV2ReaderRef<'_>,
1447 ) -> tokenizer::Tokenizer {
1448 if config.model_family == format::ModelFamily::Moonshine {
1449 let mut sp = tokenizer::SentencePieceTokenizer::moonshine_default();
1450 if let Some(vocab) = Self::read_vocabulary(reader) {
1451 Self::populate_sentencepiece(&mut sp, &vocab);
1452 }
1453 tokenizer::Tokenizer::SentencePiece(sp)
1454 } else {
1455 tokenizer::Tokenizer::Bpe(Self::read_vocabulary(reader).map_or_else(
1456 tokenizer::BpeTokenizer::with_base_tokens,
1457 tokenizer::BpeTokenizer::from_vocabulary,
1458 ))
1459 }
1460 }
1461
1462 fn read_vocabulary(reader: &format::AprV2ReaderRef<'_>) -> Option<tokenizer::Vocabulary> {
1464 let raw = reader.get_tensor_data("__vocab__")?;
1465 tokenizer::Vocabulary::from_bytes(raw)
1466 }
1467
1468 fn read_mel_filterbank(
1470 reader: &format::AprV2ReaderRef<'_>,
1471 ) -> Option<format::MelFilterbankData> {
1472 let data = reader.get_tensor_as_f32("__mel_filters__")?;
1473 let entry = reader.get_tensor("__mel_filters__")?;
1474 let shape = &entry.shape;
1475 if shape.len() == 2 {
1476 Some(format::MelFilterbankData::new(
1477 shape[0] as u32,
1478 shape[1] as u32,
1479 data,
1480 ))
1481 } else {
1482 None
1483 }
1484 }
1485
1486 fn load_f16_raw(reader: &format::AprV2ReaderRef<'_>, name: &str) -> Option<Vec<u16>> {
1488 let raw = reader.get_tensor_data(name)?;
1489 Some(
1490 raw.chunks_exact(2)
1491 .map(|b| u16::from_le_bytes([b[0], b[1]]))
1492 .collect(),
1493 )
1494 }
1495
1496 fn populate_sentencepiece(
1498 sp: &mut tokenizer::SentencePieceTokenizer,
1499 vocab: &tokenizer::Vocabulary,
1500 ) {
1501 for id in 0..vocab.len() as u32 {
1502 if let Some(bytes) = vocab.get_bytes(id) {
1503 if let Ok(piece) = core::str::from_utf8(bytes) {
1504 if !piece.is_empty() {
1505 sp.add_piece(id, piece);
1506 }
1507 }
1508 }
1509 }
1510 }
1511
1512 fn load_encoder_weights(
1514 reader: &format::AprV2ReaderRef<'_>,
1515 encoder: &mut model::Encoder,
1516 tracker: &mut progress::ProgressTracker,
1517 callback: progress::ProgressCallback<'_>,
1518 ) {
1519 let n_layers = encoder.n_layers();
1520
1521 if let Some(conv_frontend) = encoder.conv_frontend_mut() {
1523 if let Some(weight) = reader.get_tensor_as_f32("encoder.conv1.weight") {
1525 let target = conv_frontend.conv1.weight_mut();
1526 let len = weight.len().min(target.len());
1527 target[..len].copy_from_slice(&weight[..len]);
1528 }
1529 if let Some(bias) = reader.get_tensor_as_f32("encoder.conv1.bias") {
1530 let target = conv_frontend.conv1.bias_mut();
1531 let len = bias.len().min(target.len());
1532 target[..len].copy_from_slice(&bias[..len]);
1533 }
1534
1535 if let Some(weight) = reader.get_tensor_as_f32("encoder.conv2.weight") {
1537 let target = conv_frontend.conv2.weight_mut();
1538 let len = weight.len().min(target.len());
1539 target[..len].copy_from_slice(&weight[..len]);
1540 }
1541 if let Some(bias) = reader.get_tensor_as_f32("encoder.conv2.bias") {
1542 let target = conv_frontend.conv2.bias_mut();
1543 let len = bias.len().min(target.len());
1544 target[..len].copy_from_slice(&bias[..len]);
1545 }
1546 }
1547
1548 let pe_result = reader
1550 .get_tensor_as_f32("encoder.embed_positions.weight")
1551 .or_else(|| reader.get_tensor_as_f32("encoder.positional_embedding"));
1552 if let Some(pe) = pe_result {
1553 let target = encoder.positional_embedding_mut();
1554 let len = pe.len().min(target.len());
1555 target[..len].copy_from_slice(&pe[..len]);
1556 }
1557
1558 if !encoder.moonshine_blocks().is_empty() {
1560 for layer_idx in 0..n_layers {
1562 let progress = layer_idx as f32 / n_layers as f32;
1563 tracker.update_phase_progress(progress);
1564 callback(&tracker.to_progress());
1565
1566 let block = &mut encoder.moonshine_blocks_mut()[layer_idx];
1567
1568 Self::load_layernorm_nobias_weights(
1570 reader,
1571 &format!("encoder.blocks.{layer_idx}.ln1"),
1572 &mut block.ln1,
1573 );
1574
1575 Self::load_gqa_weights(
1577 reader,
1578 &format!("encoder.blocks.{layer_idx}.attn"),
1579 &mut block.self_attn,
1580 );
1581
1582 Self::load_layernorm_nobias_weights(
1584 reader,
1585 &format!("encoder.blocks.{layer_idx}.ln2"),
1586 &mut block.ln2,
1587 );
1588
1589 Self::load_mlp_weights(
1591 reader,
1592 &format!("encoder.blocks.{layer_idx}.ffn"),
1593 &mut block.ffn,
1594 );
1595 }
1596 } else {
1597 for layer_idx in 0..n_layers {
1599 let progress = layer_idx as f32 / n_layers as f32;
1600 tracker.update_phase_progress(progress);
1601 callback(&tracker.to_progress());
1602
1603 let block = &mut encoder.blocks_mut()[layer_idx];
1604
1605 Self::load_layer_norm_weights(
1607 reader,
1608 &format!("encoder.layers.{layer_idx}.self_attn_layer_norm"),
1609 &mut block.ln1,
1610 );
1611
1612 Self::load_attention_weights(
1614 reader,
1615 &format!("encoder.layers.{layer_idx}.self_attn"),
1616 &mut block.self_attn,
1617 );
1618
1619 Self::load_layer_norm_weights(
1621 reader,
1622 &format!("encoder.layers.{layer_idx}.final_layer_norm"),
1623 &mut block.ln2,
1624 );
1625
1626 Self::load_ffn_weights(
1628 reader,
1629 &format!("encoder.layers.{layer_idx}"),
1630 &mut block.ffn,
1631 );
1632 }
1633 }
1634
1635 if let Some(ln) = encoder.ln_post_rms_mut() {
1637 Self::load_layernorm_nobias_weights(reader, "encoder.layer_norm", ln);
1639 } else {
1640 Self::load_layer_norm_weights(reader, "encoder.layer_norm", encoder.ln_post_mut());
1642 }
1643 }
1644
1645 fn load_encoder_weights_f16(
1646 reader: &crate::format::AprV2ReaderRef<'_>,
1647 encoder: &mut crate::model::Encoder,
1648 tracker: &mut crate::progress::ProgressTracker,
1649 callback: &mut dyn FnMut(&crate::progress::Progress),
1650 ) {
1651 let n_layers = encoder.n_layers();
1652
1653 if let Some(conv_frontend) = encoder.conv_frontend_mut() {
1654 if let Some(weight) = reader.get_tensor_as_f32("encoder.conv1.weight") {
1655 let target = conv_frontend.conv1.weight_mut();
1656 let len = weight.len().min(target.len());
1657 target[..len].copy_from_slice(&weight[..len]);
1658 }
1659 if let Some(bias) = reader.get_tensor_as_f32("encoder.conv1.bias") {
1660 let target = conv_frontend.conv1.bias_mut();
1661 let len = bias.len().min(target.len());
1662 target[..len].copy_from_slice(&bias[..len]);
1663 }
1664 if let Some(weight) = reader.get_tensor_as_f32("encoder.conv2.weight") {
1665 let target = conv_frontend.conv2.weight_mut();
1666 let len = weight.len().min(target.len());
1667 target[..len].copy_from_slice(&weight[..len]);
1668 }
1669 if let Some(bias) = reader.get_tensor_as_f32("encoder.conv2.bias") {
1670 let target = conv_frontend.conv2.bias_mut();
1671 let len = bias.len().min(target.len());
1672 target[..len].copy_from_slice(&bias[..len]);
1673 }
1674 }
1675
1676 let pe_result = reader
1677 .get_tensor_as_f32("encoder.embed_positions.weight")
1678 .or_else(|| reader.get_tensor_as_f32("encoder.positional_embedding"));
1679 if let Some(pe) = pe_result {
1680 let target = encoder.positional_embedding_mut();
1681 let len = pe.len().min(target.len());
1682 target[..len].copy_from_slice(&pe[..len]);
1683 }
1684
1685 if encoder.moonshine_blocks_mut().len() > 0 {
1686 for layer_idx in 0..n_layers {
1687 let progress = layer_idx as f32 / n_layers as f32;
1688 tracker.update_phase_progress(progress);
1689 callback(&tracker.to_progress());
1690 let block = &mut encoder.moonshine_blocks_mut()[layer_idx];
1691 Self::load_layernorm_nobias_weights(
1692 reader,
1693 &format!("encoder.layers.{layer_idx}.pre_attn_layer_norm"),
1694 &mut block.ln1,
1695 );
1696 Self::load_gqa_weights(
1697 reader,
1698 &format!("encoder.layers.{layer_idx}.attn"),
1699 &mut block.self_attn,
1700 );
1701 Self::load_layernorm_nobias_weights(
1702 reader,
1703 &format!("encoder.layers.{layer_idx}.pre_mlp_layer_norm"),
1704 &mut block.ln2,
1705 );
1706 Self::load_mlp_weights(
1707 reader,
1708 &format!("encoder.blocks.{layer_idx}.ffn"),
1709 &mut block.ffn,
1710 );
1711 }
1712 } else {
1713 for layer_idx in 0..n_layers {
1714 let progress = layer_idx as f32 / n_layers as f32;
1715 tracker.update_phase_progress(progress);
1716 callback(&tracker.to_progress());
1717 let block = &mut encoder.blocks_mut()[layer_idx];
1718 Self::load_layer_norm_weights(
1719 reader,
1720 &format!("encoder.layers.{layer_idx}.self_attn_layer_norm"),
1721 &mut block.ln1,
1722 );
1723 Self::load_attention_weights_f16(
1724 reader,
1725 &format!("encoder.layers.{layer_idx}.self_attn"),
1726 &mut block.self_attn,
1727 );
1728 Self::load_layer_norm_weights(
1729 reader,
1730 &format!("encoder.layers.{layer_idx}.final_layer_norm"),
1731 &mut block.ln2,
1732 );
1733 Self::load_ffn_weights_f16(
1734 reader,
1735 &format!("encoder.layers.{layer_idx}"),
1736 &mut block.ffn,
1737 );
1738 }
1739 }
1740
1741 if let Some(ln) = encoder.ln_post_rms_mut() {
1742 Self::load_layernorm_nobias_weights(reader, "encoder.layer_norm", ln);
1743 } else {
1744 Self::load_layer_norm_weights(reader, "encoder.layer_norm", encoder.ln_post_mut());
1745 }
1746 }
1747
1748 fn load_decoder_weights(
1750 reader: &format::AprV2ReaderRef<'_>,
1751 decoder: &mut model::Decoder,
1752 tracker: &mut progress::ProgressTracker,
1753 callback: progress::ProgressCallback<'_>,
1754 ) {
1755 let n_layers = decoder.n_layers();
1756
1757 let te_result = reader
1762 .get_tensor_as_f32("decoder.embed_tokens.weight")
1763 .or_else(|| reader.get_tensor_as_f32("decoder.token_embedding.weight"))
1764 .or_else(|| reader.get_tensor_as_f32("decoder.token_embedding"));
1765 if let Some(te) = te_result {
1766 let target = decoder.token_embedding_mut();
1767 let len = te.len().min(target.len());
1768 target[..len].copy_from_slice(&te[..len]);
1769 }
1770
1771 let pe_result = reader
1773 .get_tensor_as_f32("decoder.embed_positions.weight")
1774 .or_else(|| reader.get_tensor_as_f32("decoder.positional_embedding"));
1775 if let Some(pe) = pe_result {
1776 let target = decoder.positional_embedding_mut();
1777 let len = pe.len().min(target.len());
1778 target[..len].copy_from_slice(&pe[..len]);
1779 }
1780
1781 if !decoder.moonshine_blocks().is_empty() {
1783 for layer_idx in 0..n_layers {
1785 let progress = layer_idx as f32 / n_layers as f32;
1786 tracker.update_phase_progress(progress);
1787 callback(&tracker.to_progress());
1788
1789 let block = &mut decoder.moonshine_blocks_mut()[layer_idx];
1790
1791 Self::load_layernorm_nobias_weights(
1793 reader,
1794 &format!("decoder.blocks.{layer_idx}.ln1"),
1795 &mut block.ln1,
1796 );
1797
1798 Self::load_gqa_weights(
1800 reader,
1801 &format!("decoder.blocks.{layer_idx}.attn"),
1802 &mut block.self_attn,
1803 );
1804
1805 Self::load_layernorm_nobias_weights(
1807 reader,
1808 &format!("decoder.blocks.{layer_idx}.ln_cross"),
1809 &mut block.ln_cross,
1810 );
1811
1812 Self::load_gqa_weights(
1814 reader,
1815 &format!("decoder.blocks.{layer_idx}.cross_attn"),
1816 &mut block.cross_attn,
1817 );
1818
1819 Self::load_layernorm_nobias_weights(
1821 reader,
1822 &format!("decoder.blocks.{layer_idx}.ln2"),
1823 &mut block.ln2,
1824 );
1825
1826 Self::load_gated_mlp_weights(
1828 reader,
1829 &format!("decoder.blocks.{layer_idx}.ffn"),
1830 &mut block.ffn,
1831 );
1832 }
1833
1834 if let Some(ln) = decoder.ln_post_rms_mut() {
1836 Self::load_layernorm_nobias_weights(reader, "decoder.ln_post", ln);
1837 }
1838 } else {
1839 for layer_idx in 0..n_layers {
1841 let progress = layer_idx as f32 / n_layers as f32;
1842 tracker.update_phase_progress(progress);
1843 callback(&tracker.to_progress());
1844
1845 let block = &mut decoder.blocks_mut()[layer_idx];
1846
1847 Self::load_layer_norm_weights(
1849 reader,
1850 &format!("decoder.layers.{layer_idx}.self_attn_layer_norm"),
1851 &mut block.ln1,
1852 );
1853
1854 Self::load_attention_weights(
1856 reader,
1857 &format!("decoder.layers.{layer_idx}.self_attn"),
1858 &mut block.self_attn,
1859 );
1860
1861 Self::load_layer_norm_weights(
1863 reader,
1864 &format!("decoder.layers.{layer_idx}.encoder_attn_layer_norm"),
1865 &mut block.ln2,
1866 );
1867
1868 Self::load_attention_weights(
1870 reader,
1871 &format!("decoder.layers.{layer_idx}.encoder_attn"),
1872 &mut block.cross_attn,
1873 );
1874
1875 Self::load_layer_norm_weights(
1877 reader,
1878 &format!("decoder.layers.{layer_idx}.final_layer_norm"),
1879 &mut block.ln3,
1880 );
1881
1882 Self::load_ffn_weights(
1884 reader,
1885 &format!("decoder.layers.{layer_idx}"),
1886 &mut block.ffn,
1887 );
1888 }
1889
1890 Self::load_layer_norm_weights(reader, "decoder.layer_norm", decoder.ln_post_mut());
1892 }
1893
1894 decoder.finalize_weights();
1896 }
1897
1898 fn load_decoder_weights_f16(
1906 reader: &format::AprV2ReaderRef<'_>,
1907 decoder: &mut model::Decoder,
1908 tracker: &mut progress::ProgressTracker,
1909 callback: progress::ProgressCallback<'_>,
1910 ) {
1911 let n_layers = decoder.n_layers();
1912
1913 let te_result = reader
1915 .get_tensor_as_f32("decoder.embed_tokens.weight")
1916 .or_else(|| reader.get_tensor_as_f32("decoder.token_embedding.weight"))
1917 .or_else(|| reader.get_tensor_as_f32("decoder.token_embedding"));
1918 if let Some(te) = te_result {
1919 let target = decoder.token_embedding_mut();
1920 let len = te.len().min(target.len());
1921 target[..len].copy_from_slice(&te[..len]);
1922 }
1923
1924 let pe_result = reader
1926 .get_tensor_as_f32("decoder.embed_positions.weight")
1927 .or_else(|| reader.get_tensor_as_f32("decoder.positional_embedding"));
1928 if let Some(pe) = pe_result {
1929 let target = decoder.positional_embedding_mut();
1930 let len = pe.len().min(target.len());
1931 target[..len].copy_from_slice(&pe[..len]);
1932 }
1933
1934 for layer_idx in 0..n_layers {
1936 let progress = layer_idx as f32 / n_layers as f32;
1937 tracker.update_phase_progress(progress);
1938 callback(&tracker.to_progress());
1939
1940 let block = &mut decoder.blocks_mut()[layer_idx];
1941
1942 Self::load_layer_norm_weights(
1944 reader,
1945 &format!("decoder.layers.{layer_idx}.self_attn_layer_norm"),
1946 &mut block.ln1,
1947 );
1948
1949 Self::load_attention_weights_f16(
1951 reader,
1952 &format!("decoder.layers.{layer_idx}.self_attn"),
1953 &mut block.self_attn,
1954 );
1955
1956 Self::load_layer_norm_weights(
1957 reader,
1958 &format!("decoder.layers.{layer_idx}.encoder_attn_layer_norm"),
1959 &mut block.ln2,
1960 );
1961
1962 Self::load_attention_weights_f16(
1964 reader,
1965 &format!("decoder.layers.{layer_idx}.encoder_attn"),
1966 &mut block.cross_attn,
1967 );
1968
1969 Self::load_layer_norm_weights(
1970 reader,
1971 &format!("decoder.layers.{layer_idx}.final_layer_norm"),
1972 &mut block.ln3,
1973 );
1974
1975 Self::load_ffn_weights_f16(
1977 reader,
1978 &format!("decoder.layers.{layer_idx}"),
1979 &mut block.ffn,
1980 );
1981 }
1982
1983 Self::load_layer_norm_weights(reader, "decoder.layer_norm", decoder.ln_post_mut());
1985
1986 decoder.finalize_weights();
1988 decoder.convert_embeddings_to_f16();
1989 }
1990
1991 fn load_layer_norm_weights(
1993 reader: &format::AprV2ReaderRef<'_>,
1994 prefix: &str,
1995 ln: &mut model::LayerNorm,
1996 ) {
1997 if let Some(weight) = reader.get_tensor_as_f32(&format!("{prefix}.weight")) {
1998 let len = weight.len().min(ln.weight.len());
1999 ln.weight[..len].copy_from_slice(&weight[..len]);
2000 }
2001 if let Some(bias) = reader.get_tensor_as_f32(&format!("{prefix}.bias")) {
2002 let len = bias.len().min(ln.bias.len());
2003 ln.bias[..len].copy_from_slice(&bias[..len]);
2004 }
2005 }
2006
2007 fn load_attention_weights(
2009 reader: &format::AprV2ReaderRef<'_>,
2010 prefix: &str,
2011 attn: &mut model::MultiHeadAttention,
2012 ) {
2013 if let Some(q_weight) = reader.get_tensor_as_f32(&format!("{prefix}.q_proj.weight")) {
2015 attn.set_query_weight(&q_weight);
2016 }
2017 if let Some(q_bias) = reader.get_tensor_as_f32(&format!("{prefix}.q_proj.bias")) {
2018 attn.set_query_bias(&q_bias);
2019 }
2020 if let Some(k_weight) = reader.get_tensor_as_f32(&format!("{prefix}.k_proj.weight")) {
2021 attn.set_key_weight(&k_weight);
2022 }
2023 if let Some(k_bias) = reader.get_tensor_as_f32(&format!("{prefix}.k_proj.bias")) {
2024 attn.set_key_bias(&k_bias);
2025 }
2026 if let Some(v_weight) = reader.get_tensor_as_f32(&format!("{prefix}.v_proj.weight")) {
2027 attn.set_value_weight(&v_weight);
2028 }
2029 if let Some(v_bias) = reader.get_tensor_as_f32(&format!("{prefix}.v_proj.bias")) {
2030 attn.set_value_bias(&v_bias);
2031 }
2032 if let Some(out_weight) = reader.get_tensor_as_f32(&format!("{prefix}.out_proj.weight")) {
2033 attn.set_out_weight(&out_weight);
2034 }
2035 if let Some(out_bias) = reader.get_tensor_as_f32(&format!("{prefix}.out_proj.bias")) {
2036 attn.set_out_bias(&out_bias);
2037 }
2038 }
2039
2040 fn load_ffn_weights(
2042 reader: &format::AprV2ReaderRef<'_>,
2043 prefix: &str,
2044 ffn: &mut model::FeedForward,
2045 ) {
2046 if let Some(fc1_weight) = reader.get_tensor_as_f32(&format!("{prefix}.fc1.weight")) {
2047 ffn.fc1.set_weight(&fc1_weight);
2048 }
2049 if let Some(fc1_bias) = reader.get_tensor_as_f32(&format!("{prefix}.fc1.bias")) {
2050 ffn.fc1.set_bias(&fc1_bias);
2051 }
2052 if let Some(fc2_weight) = reader.get_tensor_as_f32(&format!("{prefix}.fc2.weight")) {
2053 ffn.fc2.set_weight(&fc2_weight);
2054 }
2055 if let Some(fc2_bias) = reader.get_tensor_as_f32(&format!("{prefix}.fc2.bias")) {
2056 ffn.fc2.set_bias(&fc2_bias);
2057 }
2058 }
2059
2060 fn load_attention_weights_f16(
2065 reader: &format::AprV2ReaderRef<'_>,
2066 prefix: &str,
2067 attn: &mut model::MultiHeadAttention,
2068 ) {
2069 if let Some(q_weight) = Self::load_f16_raw(reader, &format!("{prefix}.q_proj.weight")) {
2071 attn.set_query_weight_f16(&q_weight);
2072 }
2073 if let Some(k_weight) = Self::load_f16_raw(reader, &format!("{prefix}.k_proj.weight")) {
2074 attn.set_key_weight_f16(&k_weight);
2075 }
2076 if let Some(v_weight) = Self::load_f16_raw(reader, &format!("{prefix}.v_proj.weight")) {
2077 attn.set_value_weight_f16(&v_weight);
2078 }
2079 if let Some(out_weight) = Self::load_f16_raw(reader, &format!("{prefix}.out_proj.weight")) {
2080 attn.set_out_weight_f16(&out_weight);
2081 }
2082 if let Some(q_bias) = reader.get_tensor_as_f32(&format!("{prefix}.q_proj.bias")) {
2084 attn.set_query_bias(&q_bias);
2085 }
2086 if let Some(k_bias) = reader.get_tensor_as_f32(&format!("{prefix}.k_proj.bias")) {
2087 attn.set_key_bias(&k_bias);
2088 }
2089 if let Some(v_bias) = reader.get_tensor_as_f32(&format!("{prefix}.v_proj.bias")) {
2090 attn.set_value_bias(&v_bias);
2091 }
2092 if let Some(out_bias) = reader.get_tensor_as_f32(&format!("{prefix}.out_proj.bias")) {
2093 attn.set_out_bias(&out_bias);
2094 }
2095 }
2096
2097 fn load_ffn_weights_f16(
2101 reader: &format::AprV2ReaderRef<'_>,
2102 prefix: &str,
2103 ffn: &mut model::FeedForward,
2104 ) {
2105 if let Some(fc1_weight) = Self::load_f16_raw(reader, &format!("{prefix}.fc1.weight")) {
2106 ffn.fc1.set_weight_f16(&fc1_weight);
2107 }
2108 if let Some(fc1_bias) = reader.get_tensor_as_f32(&format!("{prefix}.fc1.bias")) {
2109 ffn.fc1.set_bias(&fc1_bias);
2110 }
2111 if let Some(fc2_weight) = Self::load_f16_raw(reader, &format!("{prefix}.fc2.weight")) {
2112 ffn.fc2.set_weight_f16(&fc2_weight);
2113 }
2114 if let Some(fc2_bias) = reader.get_tensor_as_f32(&format!("{prefix}.fc2.bias")) {
2115 ffn.fc2.set_bias(&fc2_bias);
2116 }
2117 }
2118
2119 #[allow(dead_code)]
2121 fn load_rms_norm_weights(
2122 reader: &format::AprV2ReaderRef<'_>,
2123 prefix: &str,
2124 rms: &mut model::lfm2::layer::RmsNorm,
2125 ) {
2126 if let Some(weight) = reader.get_tensor_as_f32(&format!("{prefix}.weight")) {
2127 let len = weight.len().min(rms.weight.len());
2128 rms.weight[..len].copy_from_slice(&weight[..len]);
2129 }
2130 }
2131
2132 fn load_layernorm_nobias_weights(
2134 reader: &format::AprV2ReaderRef<'_>,
2135 prefix: &str,
2136 ln: &mut model::lfm2::layer::LayerNormNoBias,
2137 ) {
2138 if let Some(weight) = reader.get_tensor_as_f32(&format!("{prefix}.weight")) {
2139 let len = weight.len().min(ln.weight.len());
2140 ln.weight[..len].copy_from_slice(&weight[..len]);
2141 }
2142 }
2143
2144 fn load_gqa_weights(
2146 reader: &format::AprV2ReaderRef<'_>,
2147 prefix: &str,
2148 gqa: &mut model::lfm2::gqa::GroupedQueryAttention,
2149 ) {
2150 if let Some(w) = reader.get_tensor_as_f32(&format!("{prefix}.q.weight")) {
2151 let len = w.len().min(gqa.w_q.len());
2152 gqa.w_q[..len].copy_from_slice(&w[..len]);
2153 }
2154 if let Some(w) = reader.get_tensor_as_f32(&format!("{prefix}.k.weight")) {
2155 let len = w.len().min(gqa.w_k.len());
2156 gqa.w_k[..len].copy_from_slice(&w[..len]);
2157 }
2158 if let Some(w) = reader.get_tensor_as_f32(&format!("{prefix}.v.weight")) {
2159 let len = w.len().min(gqa.w_v.len());
2160 gqa.w_v[..len].copy_from_slice(&w[..len]);
2161 }
2162 if let Some(w) = reader.get_tensor_as_f32(&format!("{prefix}.o.weight")) {
2163 let len = w.len().min(gqa.w_o.len());
2164 gqa.w_o[..len].copy_from_slice(&w[..len]);
2165 }
2166 }
2167
2168 fn load_mlp_weights(
2170 reader: &format::AprV2ReaderRef<'_>,
2171 prefix: &str,
2172 ffn: &mut model::lfm2::mlp::MlpFfn,
2173 ) {
2174 if let Some(w) = reader.get_tensor_as_f32(&format!("{prefix}.fc1.weight")) {
2175 let len = w.len().min(ffn.fc1.len());
2176 ffn.fc1[..len].copy_from_slice(&w[..len]);
2177 }
2178 if let Some(b) = reader.get_tensor_as_f32(&format!("{prefix}.fc1.bias")) {
2179 ffn.b1 = Some(b);
2180 }
2181 if let Some(w) = reader.get_tensor_as_f32(&format!("{prefix}.fc2.weight")) {
2182 let len = w.len().min(ffn.fc2.len());
2183 ffn.fc2[..len].copy_from_slice(&w[..len]);
2184 }
2185 if let Some(b) = reader.get_tensor_as_f32(&format!("{prefix}.fc2.bias")) {
2186 ffn.b2 = Some(b);
2187 }
2188 }
2189
2190 fn load_gated_mlp_weights(
2192 reader: &format::AprV2ReaderRef<'_>,
2193 prefix: &str,
2194 ffn: &mut model::lfm2::mlp::GatedMlpFfn,
2195 ) {
2196 if let Some(w) = reader.get_tensor_as_f32(&format!("{prefix}.fc1.weight")) {
2197 let len = w.len().min(ffn.fc1.len());
2198 ffn.fc1[..len].copy_from_slice(&w[..len]);
2199 }
2200 if let Some(b) = reader.get_tensor_as_f32(&format!("{prefix}.fc1.bias")) {
2201 ffn.b1 = Some(b);
2202 }
2203 if let Some(w) = reader.get_tensor_as_f32(&format!("{prefix}.fc2.weight")) {
2204 let len = w.len().min(ffn.fc2.len());
2205 ffn.fc2[..len].copy_from_slice(&w[..len]);
2206 }
2207 if let Some(b) = reader.get_tensor_as_f32(&format!("{prefix}.fc2.bias")) {
2208 ffn.b2 = Some(b);
2209 }
2210 }
2211
2212 fn load_conv_stem_weights(reader: &format::AprV2ReaderRef<'_>, stem: &mut audio::ConvStem) {
2214 if let Some(w) = reader.get_tensor_as_f32("encoder.conv1.weight") {
2216 let target = stem.conv1.weight_mut();
2217 let len = w.len().min(target.len());
2218 target[..len].copy_from_slice(&w[..len]);
2219 }
2220
2221 if let Some(w) = reader.get_tensor_as_f32("encoder.conv2.weight") {
2223 let target = stem.conv2.weight_mut();
2224 let len = w.len().min(target.len());
2225 target[..len].copy_from_slice(&w[..len]);
2226 }
2227 if let Some(b) = reader.get_tensor_as_f32("encoder.conv2.bias") {
2228 let target = stem.conv2.bias_mut();
2229 let len = b.len().min(target.len());
2230 target[..len].copy_from_slice(&b[..len]);
2231 }
2232
2233 if let Some(w) = reader.get_tensor_as_f32("encoder.conv3.weight") {
2235 let target = stem.conv3.weight_mut();
2236 let len = w.len().min(target.len());
2237 target[..len].copy_from_slice(&w[..len]);
2238 }
2239 if let Some(b) = reader.get_tensor_as_f32("encoder.conv3.bias") {
2240 let target = stem.conv3.bias_mut();
2241 let len = b.len().min(target.len());
2242 target[..len].copy_from_slice(&b[..len]);
2243 }
2244
2245 if let Some(w) = reader.get_tensor_as_f32("encoder.groupnorm.weight") {
2247 let len = w.len().min(stem.groupnorm.weight.len());
2248 stem.groupnorm.weight[..len].copy_from_slice(&w[..len]);
2249 }
2250 if let Some(b) = reader.get_tensor_as_f32("encoder.groupnorm.bias") {
2251 let len = b.len().min(stem.groupnorm.bias.len());
2252 stem.groupnorm.bias[..len].copy_from_slice(&b[..len]);
2253 }
2254
2255 }
2258
2259 pub fn encoder_mut(&mut self) -> &mut model::Encoder {
2261 &mut self.encoder
2262 }
2263
2264 pub fn decoder_mut(&mut self) -> &mut model::Decoder {
2266 &mut self.decoder
2267 }
2268
2269 pub fn encoder(&self) -> &model::Encoder {
2271 &self.encoder
2272 }
2273
2274 pub fn decoder(&self) -> &model::Decoder {
2276 &self.decoder
2277 }
2278
2279 #[cfg(feature = "realizar-gpu")]
2289 pub fn into_cuda(self, device_ordinal: i32) -> WhisperResult<cuda::WhisperCuda> {
2290 let mel_filters = self.mel_filters.ok_or_else(|| {
2291 WhisperError::Audio("CUDA requires mel filterbank (Whisper model)".into())
2292 })?;
2293 let bpe_tokenizer = match self.tokenizer {
2294 tokenizer::Tokenizer::Bpe(bpe) => bpe,
2295 tokenizer::Tokenizer::SentencePiece(_) => {
2296 return Err(WhisperError::Model(
2297 "CUDA backend requires Whisper BPE tokenizer".into(),
2298 ));
2299 }
2300 };
2301 cuda::WhisperCuda::new_with_components(
2302 self.encoder,
2303 self.decoder,
2304 self.config,
2305 bpe_tokenizer,
2306 mel_filters,
2307 device_ordinal,
2308 )
2309 }
2310
2311 #[must_use]
2313 pub const fn conv_stem(&self) -> Option<&audio::ConvStem> {
2314 self.conv_stem.as_ref()
2315 }
2316
2317 #[must_use]
2319 pub const fn mel_filters(&self) -> Option<&audio::MelFilterbank> {
2320 self.mel_filters.as_ref()
2321 }
2322
2323 pub fn transcribe_batch(
2352 &self,
2353 audio_batch: &[Vec<f32>],
2354 options: TranscribeOptions,
2355 ) -> WhisperResult<BatchTranscriptionResult> {
2356 if audio_batch.is_empty() {
2357 return Err(WhisperError::Audio("empty batch".into()));
2358 }
2359
2360 let start_time = std::time::Instant::now();
2361 let mut results = Vec::with_capacity(audio_batch.len());
2362
2363 for audio in audio_batch {
2365 let result = self.transcribe(audio, options.clone())?;
2366 results.push(result);
2367 }
2368
2369 let total_duration_secs = start_time.elapsed().as_secs_f32();
2370
2371 Ok(BatchTranscriptionResult {
2372 results,
2373 total_duration_secs,
2374 })
2375 }
2376
2377 pub fn transcribe_audio_batch(
2392 &self,
2393 batch: &audio::AudioBatch,
2394 options: TranscribeOptions,
2395 ) -> WhisperResult<BatchTranscriptionResult> {
2396 if batch.is_empty() {
2397 return Err(WhisperError::Audio("empty batch".into()));
2398 }
2399
2400 let start_time = std::time::Instant::now();
2401
2402 let preprocessor = audio::BatchPreprocessor::new(audio::MelConfig::default());
2404 let mel_result = preprocessor.process_batch(batch)?;
2405
2406 let mut results = Vec::with_capacity(batch.len());
2407 let language = options.language.clone().unwrap_or_else(|| "en".to_string());
2408
2409 for mel in &mel_result.mels {
2411 let audio_features = self.encode(mel)?;
2413
2414 let initial_tokens = self.get_initial_tokens(&language, options.task);
2416
2417 let tokens = self.decode(&audio_features, &initial_tokens, &options)?;
2419
2420 let segments = if timestamps::has_timestamps(&tokens) {
2422 timestamps::extract_segments(&tokens, |ts| self.tokenizer.decode(ts).ok())
2423 } else {
2424 Vec::new()
2425 };
2426
2427 let text = self.tokenizer.decode(&tokens)?;
2429
2430 results.push(TranscriptionResult {
2431 text,
2432 language: language.clone(),
2433 segments,
2434 profiling: None,
2435 });
2436 }
2437
2438 let total_duration_secs = start_time.elapsed().as_secs_f32();
2439
2440 Ok(BatchTranscriptionResult {
2441 results,
2442 total_duration_secs,
2443 })
2444 }
2445
2446 #[must_use]
2448 pub fn create_audio_batch(audio_segments: &[Vec<f32>]) -> audio::AudioBatch {
2449 let mut batch = audio::AudioBatch::with_default_config();
2450 for segment in audio_segments {
2451 batch.add_segment(segment.clone());
2452 }
2453 batch
2454 }
2455
2456 pub fn transcribe_batch_optimized(
2471 &self,
2472 audio_batch: &[Vec<f32>],
2473 options: TranscribeOptions,
2474 ) -> WhisperResult<BatchTranscriptionResult> {
2475 if audio_batch.is_empty() {
2476 return Err(WhisperError::Audio("empty batch".into()));
2477 }
2478
2479 let start_time = std::time::Instant::now();
2480
2481 let mut mels = Vec::with_capacity(audio_batch.len());
2483 for audio in audio_batch {
2484 let mel = self.compute_mel(audio)?;
2485 mels.push(mel);
2486 }
2487
2488 let encoder_outputs = self.encoder.forward_batch(&mels)?;
2490
2491 let mut results = Vec::with_capacity(audio_batch.len());
2492 let language = options.language.clone().unwrap_or_else(|| "en".to_string());
2493
2494 for features in &encoder_outputs {
2496 let initial_tokens = self.get_initial_tokens(&language, options.task);
2497 let tokens = self.decode(features, &initial_tokens, &options)?;
2498
2499 let segments = if timestamps::has_timestamps(&tokens) {
2500 timestamps::extract_segments(&tokens, |ts| self.tokenizer.decode(ts).ok())
2501 } else {
2502 Vec::new()
2503 };
2504
2505 let text = self.tokenizer.decode(&tokens)?;
2506
2507 results.push(TranscriptionResult {
2508 text,
2509 language: language.clone(),
2510 segments,
2511 profiling: None,
2512 });
2513 }
2514
2515 let total_duration_secs = start_time.elapsed().as_secs_f32();
2516
2517 Ok(BatchTranscriptionResult {
2518 results,
2519 total_duration_secs,
2520 })
2521 }
2522
2523 pub fn transcribe_with_vad(
2558 &self,
2559 audio: &[f32],
2560 options: TranscribeOptions,
2561 vad_config: Option<vad::VadConfig>,
2562 ) -> WhisperResult<VadTranscriptionResult> {
2563 let start_time = std::time::Instant::now();
2564
2565 let config = vad_config.unwrap_or_default();
2567 let mut vad = vad::VoiceActivityDetector::new(config);
2568
2569 let speech_segments = vad.detect(audio);
2571
2572 if speech_segments.is_empty() {
2573 return Ok(VadTranscriptionResult {
2574 text: String::new(),
2575 language: options.language.unwrap_or_else(|| "en".to_string()),
2576 segments: Vec::new(),
2577 speech_segments: Vec::new(),
2578 total_duration_secs: start_time.elapsed().as_secs_f32(),
2579 speech_duration_secs: 0.0,
2580 });
2581 }
2582
2583 let sample_rate = audio::SAMPLE_RATE as f32;
2585 let mut speech_audios = Vec::with_capacity(speech_segments.len());
2586 let mut speech_duration = 0.0f32;
2587
2588 for segment in &speech_segments {
2589 let start_sample = (segment.start * sample_rate) as usize;
2590 let end_sample = ((segment.end * sample_rate) as usize).min(audio.len());
2591
2592 if end_sample > start_sample {
2593 speech_audios.push((
2594 segment.start,
2595 segment.end,
2596 audio[start_sample..end_sample].to_vec(),
2597 ));
2598 speech_duration += segment.duration();
2599 }
2600 }
2601
2602 let mut all_segments = Vec::new();
2604 let mut full_text = String::new();
2605 let language = options.language.clone().unwrap_or_else(|| "en".to_string());
2606
2607 for (seg_start, seg_end, speech_audio) in &speech_audios {
2608 let result = self.transcribe(speech_audio, options.clone())?;
2610
2611 let segment = VadSpeechSegment {
2613 start: *seg_start,
2614 end: *seg_end,
2615 text: result.text.clone(),
2616 tokens: result
2617 .segments
2618 .first()
2619 .map(|s| s.tokens.clone())
2620 .unwrap_or_default(),
2621 };
2622
2623 if !full_text.is_empty() {
2624 full_text.push(' ');
2625 }
2626 full_text.push_str(&result.text);
2627
2628 all_segments.push(segment);
2629 }
2630
2631 let total_duration_secs = start_time.elapsed().as_secs_f32();
2632
2633 Ok(VadTranscriptionResult {
2634 text: full_text,
2635 language,
2636 segments: all_segments,
2637 speech_segments: speech_segments
2638 .into_iter()
2639 .map(|s| (s.start, s.end))
2640 .collect(),
2641 total_duration_secs,
2642 speech_duration_secs: speech_duration,
2643 })
2644 }
2645
2646 pub fn transcribe_with_silence_detection(
2662 &self,
2663 audio: &[f32],
2664 options: TranscribeOptions,
2665 silence_config: Option<vad::SilenceConfig>,
2666 ) -> WhisperResult<VadTranscriptionResult> {
2667 let start_time = std::time::Instant::now();
2668
2669 let config = silence_config.unwrap_or_default();
2671 let mut detector = vad::SilenceDetector::new(config, audio::SAMPLE_RATE);
2672
2673 let frame_size = 480; let silence_segments = detector.detect(audio, frame_size);
2676
2677 let speech_segments = self.invert_silence_segments(&silence_segments, audio.len());
2679
2680 if speech_segments.is_empty() {
2681 return Ok(VadTranscriptionResult {
2682 text: String::new(),
2683 language: options.language.unwrap_or_else(|| "en".to_string()),
2684 segments: Vec::new(),
2685 speech_segments: Vec::new(),
2686 total_duration_secs: start_time.elapsed().as_secs_f32(),
2687 speech_duration_secs: 0.0,
2688 });
2689 }
2690
2691 let sample_rate = audio::SAMPLE_RATE as f32;
2693 let mut all_segments = Vec::new();
2694 let mut full_text = String::new();
2695 let mut speech_duration = 0.0f32;
2696 let language = options.language.clone().unwrap_or_else(|| "en".to_string());
2697
2698 for (start, end) in &speech_segments {
2699 let start_sample = (start * sample_rate) as usize;
2700 let end_sample = ((end * sample_rate) as usize).min(audio.len());
2701
2702 if end_sample > start_sample {
2703 let speech_audio = &audio[start_sample..end_sample];
2704 let result = self.transcribe(speech_audio, options.clone())?;
2705
2706 let segment = VadSpeechSegment {
2707 start: *start,
2708 end: *end,
2709 text: result.text.clone(),
2710 tokens: result
2711 .segments
2712 .first()
2713 .map(|s| s.tokens.clone())
2714 .unwrap_or_default(),
2715 };
2716
2717 if !full_text.is_empty() {
2718 full_text.push(' ');
2719 }
2720 full_text.push_str(&result.text);
2721
2722 speech_duration += end - start;
2723 all_segments.push(segment);
2724 }
2725 }
2726
2727 let total_duration_secs = start_time.elapsed().as_secs_f32();
2728
2729 Ok(VadTranscriptionResult {
2730 text: full_text,
2731 language,
2732 segments: all_segments,
2733 speech_segments,
2734 total_duration_secs,
2735 speech_duration_secs: speech_duration,
2736 })
2737 }
2738
2739 fn invert_silence_segments(
2741 &self,
2742 silence_segments: &[vad::SilenceSegment],
2743 audio_len: usize,
2744 ) -> Vec<(f32, f32)> {
2745 let _ = self; let sample_rate = audio::SAMPLE_RATE as f32;
2747 let total_duration = audio_len as f32 / sample_rate;
2748 let mut speech_segments = Vec::new();
2749 let mut current_pos = 0.0f32;
2750
2751 for silence in silence_segments {
2752 if silence.start > current_pos {
2753 speech_segments.push((current_pos, silence.start));
2754 }
2755 current_pos = silence.end;
2756 }
2757
2758 if current_pos < total_duration {
2760 speech_segments.push((current_pos, total_duration));
2761 }
2762
2763 speech_segments
2764 }
2765
2766 pub fn transcribe_partial(
2803 &self,
2804 partial_audio: &[f32],
2805 options: TranscribeOptions,
2806 is_final: bool,
2807 ) -> WhisperResult<PartialTranscriptionResult> {
2808 let start_time = std::time::Instant::now();
2809
2810 let min_samples = (audio::SAMPLE_RATE as f32 * 0.5) as usize;
2812 if partial_audio.len() < min_samples {
2813 return Ok(PartialTranscriptionResult {
2814 text: String::new(),
2815 language: options.language.unwrap_or_else(|| "en".to_string()),
2816 is_final,
2817 confidence: 0.0,
2818 duration_secs: partial_audio.len() as f32 / audio::SAMPLE_RATE as f32,
2819 processing_time_secs: start_time.elapsed().as_secs_f32(),
2820 });
2821 }
2822
2823 let result = self.transcribe(partial_audio, options)?;
2825
2826 let processing_time = start_time.elapsed().as_secs_f32();
2827
2828 Ok(PartialTranscriptionResult {
2829 text: result.text,
2830 language: result.language,
2831 is_final,
2832 confidence: 1.0, duration_secs: partial_audio.len() as f32 / audio::SAMPLE_RATE as f32,
2834 processing_time_secs: processing_time,
2835 })
2836 }
2837
2838 #[must_use]
2875 pub fn create_streaming_session(
2876 &self,
2877 options: TranscribeOptions,
2878 input_sample_rate: u32,
2879 ) -> StreamingSession<'_> {
2880 let streaming_config = audio::StreamingConfig::with_sample_rate(input_sample_rate);
2881 let processor = audio::StreamingProcessor::new(streaming_config);
2882
2883 StreamingSession {
2884 whisper: self,
2885 processor,
2886 options,
2887 last_partial_text: String::new(),
2888 }
2889 }
2890
2891 pub fn transcribe_and_summarize(
2926 &self,
2927 audio: &[f32],
2928 transcribe_options: TranscribeOptions,
2929 summarize_options: SummarizeOptions<'_>,
2930 ) -> WhisperResult<TranscribeSummaryResult> {
2931 let transcription = self.transcribe(audio, transcribe_options)?;
2933
2934 if transcription.text.trim().is_empty() {
2936 return Ok(TranscribeSummaryResult {
2937 transcription,
2938 summary: String::new(),
2939 generation_stats: None,
2940 });
2941 }
2942
2943 let input_tokens = summarize_options.tokenizer.encode(&transcription.text);
2945
2946 let (output_tokens, stats) = summarize_options.model.generate_with_stats(
2948 &input_tokens,
2949 summarize_options.max_tokens,
2950 summarize_options.temperature,
2951 Some(|_token: u32, _idx: usize| true), )?;
2953
2954 let summary = summarize_options.tokenizer.decode(&output_tokens);
2956
2957 Ok(TranscribeSummaryResult {
2958 transcription,
2959 summary,
2960 generation_stats: Some(stats),
2961 })
2962 }
2963}
2964
2965#[derive(Debug, Clone)]
2967pub struct PartialTranscriptionResult {
2968 pub text: String,
2970 pub language: String,
2972 pub is_final: bool,
2974 pub confidence: f32,
2976 pub duration_secs: f32,
2978 pub processing_time_secs: f32,
2980}
2981
2982impl PartialTranscriptionResult {
2983 #[must_use]
2985 pub fn has_text(&self) -> bool {
2986 !self.text.is_empty()
2987 }
2988
2989 #[must_use]
2991 pub fn is_empty_interim(&self) -> bool {
2992 self.text.is_empty() && !self.is_final
2993 }
2994
2995 #[must_use]
2997 pub fn real_time_factor(&self) -> f32 {
2998 if self.duration_secs <= 0.0 {
2999 0.0
3000 } else {
3001 self.processing_time_secs / self.duration_secs
3002 }
3003 }
3004}
3005
3006#[derive(Debug)]
3011pub struct StreamingSession<'a> {
3012 whisper: &'a WhisperApr,
3014 processor: audio::StreamingProcessor,
3016 options: TranscribeOptions,
3018 last_partial_text: String,
3020}
3021
3022impl StreamingSession<'_> {
3023 pub fn push(&mut self, audio: &[f32]) -> WhisperResult<Option<PartialTranscriptionResult>> {
3034 self.processor.push_audio(audio);
3035 self.processor.process();
3036
3037 if self.processor.has_partial() {
3039 if let Some(partial_audio) = self.processor.get_partial() {
3040 let result =
3041 self.whisper
3042 .transcribe_partial(&partial_audio, self.options.clone(), false)?;
3043
3044 if result.text != self.last_partial_text {
3046 result.text.clone_into(&mut self.last_partial_text);
3047 return Ok(Some(result));
3048 }
3049 }
3050 }
3051
3052 Ok(None)
3053 }
3054
3055 #[must_use]
3057 pub fn has_chunk(&self) -> bool {
3058 self.processor.has_chunk()
3059 }
3060
3061 #[must_use]
3063 pub fn has_events(&self) -> bool {
3064 self.processor.has_events()
3065 }
3066
3067 pub fn drain_events(&mut self) -> Vec<audio::StreamingEvent> {
3069 self.processor.drain_events()
3070 }
3071
3072 pub fn finalize(&mut self) -> WhisperResult<PartialTranscriptionResult> {
3080 let chunk = self
3081 .processor
3082 .get_chunk()
3083 .ok_or_else(|| WhisperError::Audio("no chunk ready for finalization".into()))?;
3084
3085 let result = self
3086 .whisper
3087 .transcribe_partial(&chunk, self.options.clone(), true)?;
3088 self.last_partial_text.clear();
3089
3090 Ok(result)
3091 }
3092
3093 pub fn flush(&mut self) -> WhisperResult<Option<PartialTranscriptionResult>> {
3101 if let Some(chunk) = self.processor.flush() {
3102 let result = self
3103 .whisper
3104 .transcribe_partial(&chunk, self.options.clone(), true)?;
3105 self.last_partial_text.clear();
3106 Ok(Some(result))
3107 } else {
3108 Ok(None)
3109 }
3110 }
3111
3112 pub fn reset(&mut self) {
3114 self.processor.reset();
3115 self.last_partial_text.clear();
3116 }
3117
3118 #[must_use]
3120 pub fn state(&self) -> audio::ProcessorState {
3121 self.processor.state()
3122 }
3123
3124 #[must_use]
3126 pub fn chunk_progress(&self) -> f32 {
3127 self.processor.chunk_progress()
3128 }
3129
3130 #[must_use]
3132 pub fn partial_duration(&self) -> f32 {
3133 self.processor.partial_duration()
3134 }
3135
3136 pub fn set_partial_threshold(&mut self, seconds: f32) {
3138 self.processor.set_partial_threshold(seconds);
3139 }
3140}
3141
3142#[derive(Debug, Clone)]
3144pub struct VadTranscriptionResult {
3145 pub text: String,
3147 pub language: String,
3149 pub segments: Vec<VadSpeechSegment>,
3151 pub speech_segments: Vec<(f32, f32)>,
3153 pub total_duration_secs: f32,
3155 pub speech_duration_secs: f32,
3157}
3158
3159impl VadTranscriptionResult {
3160 #[must_use]
3162 pub fn num_segments(&self) -> usize {
3163 self.segments.len()
3164 }
3165
3166 #[must_use]
3168 pub fn has_speech(&self) -> bool {
3169 !self.segments.is_empty()
3170 }
3171
3172 #[must_use]
3174 pub fn silence_ratio(&self, audio_duration: f32) -> f32 {
3175 if audio_duration <= 0.0 {
3176 return 1.0;
3177 }
3178 1.0 - (self.speech_duration_secs / audio_duration)
3179 }
3180
3181 #[must_use]
3183 pub fn first_segment(&self) -> Option<&VadSpeechSegment> {
3184 self.segments.first()
3185 }
3186
3187 #[must_use]
3189 pub fn last_segment(&self) -> Option<&VadSpeechSegment> {
3190 self.segments.last()
3191 }
3192
3193 pub fn iter(&self) -> impl Iterator<Item = &VadSpeechSegment> {
3195 self.segments.iter()
3196 }
3197}
3198
3199#[derive(Debug, Clone)]
3201pub struct VadSpeechSegment {
3202 pub start: f32,
3204 pub end: f32,
3206 pub text: String,
3208 pub tokens: Vec<u32>,
3210}
3211
3212impl VadSpeechSegment {
3213 #[must_use]
3215 pub fn duration(&self) -> f32 {
3216 self.end - self.start
3217 }
3218
3219 #[must_use]
3221 pub fn has_text(&self) -> bool {
3222 !self.text.is_empty()
3223 }
3224}
3225
3226#[cfg(test)]
3227mod tests {
3228 use super::*;
3229
3230 #[test]
3231 fn test_default_options() {
3232 let options = TranscribeOptions::default();
3233 assert!(options.language.is_none());
3234 assert_eq!(options.task, Task::Transcribe);
3235 assert!(!options.word_timestamps);
3236 assert!(options.prompt.is_none());
3237 assert!(options.hotwords.is_empty());
3238 }
3239
3240 #[test]
3241 fn test_options_with_prompt() {
3242 let options = TranscribeOptions {
3243 prompt: Some("This lecture covers AWS, YAML, and Rust programming.".into()),
3244 ..TranscribeOptions::default()
3245 };
3246 assert_eq!(
3247 options.prompt.as_deref(),
3248 Some("This lecture covers AWS, YAML, and Rust programming.")
3249 );
3250 }
3251
3252 #[test]
3253 fn test_options_with_hotwords() {
3254 let options = TranscribeOptions {
3255 hotwords: vec!["AWS".into(), "YAML".into(), "SIMD".into(), "Rust".into()],
3256 ..TranscribeOptions::default()
3257 };
3258 assert_eq!(options.hotwords.len(), 4);
3259 assert_eq!(options.hotwords[0], "AWS");
3260 }
3261
3262 #[test]
3263 fn test_options_with_prompt_and_hotwords() {
3264 let options = TranscribeOptions {
3265 language: Some("en".into()),
3266 prompt: Some("Technical lecture on cloud computing".into()),
3267 hotwords: vec!["Kubernetes".into(), "Docker".into()],
3268 ..TranscribeOptions::default()
3269 };
3270 assert!(options.prompt.is_some());
3271 assert_eq!(options.hotwords.len(), 2);
3272 assert_eq!(options.language, Some("en".into()));
3273 }
3274
3275 #[test]
3276 fn test_decoding_strategy_default() {
3277 let strategy = DecodingStrategy::default();
3278 assert!(matches!(strategy, DecodingStrategy::Greedy));
3279 }
3280
3281 #[test]
3286 #[ignore = "Allocates large model - run with --ignored"]
3287 fn test_whisper_tiny() {
3288 let whisper = WhisperApr::tiny();
3289 assert_eq!(whisper.model_type(), ModelType::Tiny);
3290 assert_eq!(whisper.config().n_audio_layer, 4);
3291 }
3292
3293 #[test]
3294 #[ignore = "Allocates large model - run with --ignored"]
3295 fn test_whisper_base() {
3296 let whisper = WhisperApr::base();
3297 assert_eq!(whisper.model_type(), ModelType::Base);
3298 assert_eq!(whisper.config().n_audio_layer, 6);
3299 }
3300
3301 #[test]
3302 #[ignore = "Allocates large model - run with --ignored"]
3303 fn test_whisper_memory_size() {
3304 let tiny = WhisperApr::tiny();
3305 let base = WhisperApr::base();
3306
3307 assert!(tiny.memory_size() < base.memory_size());
3308 assert!(tiny.memory_size() > 100_000_000); }
3310
3311 #[test]
3312 #[ignore = "Allocates large model - run with --ignored"]
3313 fn test_whisper_initial_tokens() {
3314 let whisper = WhisperApr::tiny();
3315
3316 let tokens = whisper.get_initial_tokens("en", Task::Transcribe);
3317 assert_eq!(tokens[0], tokenizer::special_tokens::SOT);
3318 assert!(tokens.len() >= 3); let translate_tokens = whisper.get_initial_tokens("es", Task::Translate);
3321 assert!(translate_tokens.contains(&tokenizer::special_tokens::TRANSLATE));
3322 }
3323
3324 #[test]
3325 #[ignore = "Allocates large model - run with --ignored"]
3326 fn test_whisper_set_resampler() {
3327 let mut whisper = WhisperApr::tiny();
3328
3329 assert!(whisper.resampler.is_none());
3331
3332 whisper.set_resampler(44100).expect("should succeed");
3334 assert!(whisper.resampler.is_some());
3335
3336 whisper.set_resampler(16000).expect("should succeed");
3338 assert!(whisper.resampler.is_none());
3339 }
3340
3341 #[test]
3342 #[ignore = "Allocates large model - run with --ignored"]
3343 fn test_whisper_resample_passthrough() {
3344 let whisper = WhisperApr::tiny();
3345 let audio = vec![0.1, 0.2, 0.3, 0.4];
3346
3347 let resampled = whisper.resample(&audio).expect("should succeed");
3349 assert_eq!(resampled, audio);
3350 }
3351
3352 #[test]
3353 #[ignore = "Allocates large model - run with --ignored"]
3354 fn test_whisper_resample_with_resampler() {
3355 let mut whisper = WhisperApr::tiny();
3356 whisper.set_resampler(32000).expect("should succeed");
3357
3358 let n_samples = 16000;
3360 let audio: Vec<f32> = (0..n_samples)
3361 .map(|i| (2.0 * std::f32::consts::PI * 440.0 * i as f32 / 32000.0).sin())
3362 .collect();
3363
3364 let resampled = whisper.resample(&audio).expect("should succeed");
3365
3366 assert!(resampled.len() > n_samples / 3);
3368 assert!(resampled.len() < n_samples);
3369 }
3370
3371 #[test]
3372 #[ignore = "Allocates large model - run with --ignored"]
3373 fn test_whisper_tokenizer() {
3374 let whisper = WhisperApr::tiny();
3375 let tokenizer = whisper.tokenizer();
3376
3377 let vocab_size = tokenizer.vocab_size();
3379 assert!(
3380 vocab_size >= 256,
3381 "vocab_size should include base byte tokens"
3382 );
3383 }
3384
3385 #[test]
3386 fn test_transcribe_options_with_language() {
3387 let options = TranscribeOptions {
3388 language: Some("es".to_string()),
3389 task: Task::Transcribe,
3390 strategy: DecodingStrategy::Greedy,
3391 word_timestamps: false,
3392 profile: false,
3393 ..Default::default()
3394 };
3395
3396 assert_eq!(options.language, Some("es".to_string()));
3397 }
3398
3399 #[test]
3400 fn test_transcribe_options_beam_search() {
3401 let options = TranscribeOptions {
3402 language: None,
3403 task: Task::Transcribe,
3404 strategy: DecodingStrategy::BeamSearch {
3405 beam_size: 5,
3406 temperature: 0.0,
3407 patience: 1.0,
3408 },
3409 word_timestamps: false,
3410 profile: false,
3411 ..Default::default()
3412 };
3413
3414 assert!(matches!(
3415 options.strategy,
3416 DecodingStrategy::BeamSearch { .. }
3417 ));
3418 }
3419
3420 #[test]
3421 fn test_segment_struct() {
3422 let segment = Segment {
3423 start: 0.0,
3424 end: 2.5,
3425 text: "Hello world".to_string(),
3426 tokens: vec![1, 2, 3],
3427 };
3428
3429 assert!((segment.start - 0.0).abs() < f32::EPSILON);
3430 assert!((segment.end - 2.5).abs() < f32::EPSILON);
3431 assert_eq!(segment.text, "Hello world");
3432 assert_eq!(segment.tokens.len(), 3);
3433 }
3434
3435 #[test]
3436 fn test_transcription_result_struct() {
3437 let result = TranscriptionResult {
3438 text: "Test transcription".to_string(),
3439 language: "en".to_string(),
3440 segments: vec![],
3441 profiling: None,
3442 };
3443
3444 assert_eq!(result.text, "Test transcription");
3445 assert_eq!(result.language, "en");
3446 assert!(result.segments.is_empty());
3447 }
3448
3449 #[test]
3454 fn test_load_from_apr_basic() {
3455 let data = format::create_test_apr();
3457
3458 let result = WhisperApr::load_from_apr(&data);
3459 assert!(result.is_ok());
3460
3461 let whisper = result.expect("should load");
3462 assert_eq!(whisper.model_type(), ModelType::Tiny);
3463 }
3464
3465 #[test]
3466 fn test_load_from_apr_with_progress_callback() {
3467 let data = format::create_test_apr();
3468
3469 let mut progress_updates = Vec::new();
3470 let mut callback = |p: &progress::Progress| {
3471 progress_updates.push(p.percent());
3472 };
3473
3474 let result = WhisperApr::load_from_apr_with_progress(&data, &mut callback);
3475 assert!(result.is_ok());
3476
3477 assert!(!progress_updates.is_empty());
3479 }
3480
3481 #[test]
3482 fn test_load_from_apr_invalid_magic() {
3483 let mut data = format::create_test_apr();
3484 data[0] = b'X'; let result = WhisperApr::load_from_apr(&data);
3487 assert!(result.is_err());
3488 }
3489
3490 #[test]
3491 fn test_load_from_apr_too_short() {
3492 let data = vec![b'A', b'P', b'R', b'1']; let result = WhisperApr::load_from_apr(&data);
3495 assert!(result.is_err());
3496 }
3497
3498 #[test]
3499 #[ignore = "Allocates large model - run with --ignored"]
3500 fn test_encoder_mut_accessor() {
3501 let mut whisper = WhisperApr::tiny();
3502 let encoder = whisper.encoder_mut();
3503 assert_eq!(encoder.n_layers(), 4);
3504 }
3505
3506 #[test]
3507 #[ignore = "Allocates large model - run with --ignored"]
3508 fn test_decoder_mut_accessor() {
3509 let mut whisper = WhisperApr::tiny();
3510 let decoder = whisper.decoder_mut();
3511 assert_eq!(decoder.n_layers(), 4);
3512 }
3513
3514 #[test]
3519 #[ignore = "Allocates large model - run with --ignored"]
3520 fn test_full_pipeline_tiny_model() {
3521 let whisper = WhisperApr::tiny();
3523 assert_eq!(whisper.model_type(), ModelType::Tiny);
3524
3525 assert_eq!(whisper.config().n_audio_layer, 4);
3527 assert_eq!(whisper.config().n_text_layer, 4);
3528 assert_eq!(whisper.config().n_audio_state, 384);
3529 assert_eq!(whisper.config().n_text_state, 384);
3530 }
3531
3532 #[test]
3533 #[ignore = "Allocates large model - run with --ignored"]
3534 fn test_full_pipeline_base_model() {
3535 let whisper = WhisperApr::base();
3537 assert_eq!(whisper.model_type(), ModelType::Base);
3538
3539 assert_eq!(whisper.config().n_audio_layer, 6);
3541 assert_eq!(whisper.config().n_text_layer, 6);
3542 assert_eq!(whisper.config().n_audio_state, 512);
3543 assert_eq!(whisper.config().n_text_state, 512);
3544 }
3545
3546 #[test]
3547 fn test_transcribe_options_all_strategies() {
3548 let opts_greedy = TranscribeOptions::default();
3550 assert!(matches!(opts_greedy.strategy, DecodingStrategy::Greedy));
3551
3552 let opts_beam = TranscribeOptions {
3554 language: Some("en".to_string()),
3555 task: Task::Transcribe,
3556 strategy: DecodingStrategy::BeamSearch {
3557 beam_size: 5,
3558 temperature: 0.0,
3559 patience: 1.0,
3560 },
3561 word_timestamps: false,
3562 profile: false,
3563 ..Default::default()
3564 };
3565 assert!(matches!(
3566 opts_beam.strategy,
3567 DecodingStrategy::BeamSearch { .. }
3568 ));
3569
3570 let opts_sampling = TranscribeOptions {
3572 language: None,
3573 task: Task::Translate,
3574 strategy: DecodingStrategy::Sampling {
3575 temperature: 0.7,
3576 top_k: Some(50),
3577 top_p: Some(0.9),
3578 },
3579 word_timestamps: true,
3580 profile: false,
3581 ..Default::default()
3582 };
3583 assert!(matches!(
3584 opts_sampling.strategy,
3585 DecodingStrategy::Sampling { .. }
3586 ));
3587 }
3588
3589 #[test]
3590 #[ignore = "Allocates large model - run with --ignored"]
3591 fn test_memory_estimation_consistency() {
3592 let tiny = WhisperApr::tiny();
3593 let base = WhisperApr::base();
3594
3595 assert!(base.config().weights_memory_mb() > tiny.config().weights_memory_mb());
3597 assert!(base.config().peak_memory_mb() > tiny.config().peak_memory_mb());
3598 assert!(base.config().parameter_count() > tiny.config().parameter_count());
3599 }
3600
3601 #[test]
3602 fn test_simd_operations_integration() {
3603 use crate::simd;
3604
3605 let a = vec![1.0, 2.0, 3.0, 4.0];
3607 let b = vec![5.0, 6.0, 7.0, 8.0];
3608
3609 let sum = simd::add(&a, &b);
3611 assert_eq!(sum.len(), 4);
3612 assert!((sum[0] - 6.0).abs() < 1e-5);
3613 assert!((sum[3] - 12.0).abs() < 1e-5);
3614
3615 let softmax = simd::softmax(&a);
3617 let sum_softmax: f32 = softmax.iter().sum();
3618 assert!((sum_softmax - 1.0).abs() < 1e-5);
3619
3620 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);
3624 assert_eq!(result.len(), 4);
3625 }
3626
3627 #[test]
3628 fn test_memory_pool_integration() {
3629 use crate::memory::{get_buffer, pool_stats, return_buffer, MemoryPool};
3630
3631 let pool = MemoryPool::new();
3633
3634 let buf1 = pool.get(1024);
3636 assert_eq!(buf1.len(), 1024);
3637 pool.return_buffer(buf1);
3638
3639 let buf2 = pool.get(1024);
3640 assert_eq!(buf2.len(), 1024);
3641
3642 let stats = pool.stats();
3644 assert_eq!(stats.hits, 1);
3645
3646 let tlbuf = get_buffer(512);
3648 assert_eq!(tlbuf.len(), 512);
3649 return_buffer(tlbuf);
3650
3651 let tl_stats = pool_stats();
3652 assert!(tl_stats.allocations > 0);
3653 }
3654
3655 #[test]
3656 fn test_audio_resampling_integration() {
3657 use audio::SincResampler;
3658
3659 let resampler = SincResampler::new(44100, 16000).expect("resampler should work");
3661
3662 let duration_ms = 100;
3664 let samples_at_44100 = (44100 * duration_ms) / 1000;
3665 let input: Vec<f32> = (0..samples_at_44100)
3666 .map(|i| {
3667 let t = i as f32 / 44100.0;
3668 (2.0 * std::f32::consts::PI * 440.0 * t).sin()
3669 })
3670 .collect();
3671
3672 let output = resampler.resample(&input).expect("resample should work");
3673
3674 let expected_len = (input.len() as f32 * 16000.0 / 44100.0) as usize;
3676 assert!(
3677 (output.len() as i32 - expected_len as i32).abs() < 10,
3678 "output len {} vs expected ~{}",
3679 output.len(),
3680 expected_len
3681 );
3682 }
3683
3684 #[test]
3685 fn test_vad_integration() {
3686 use vad::VadConfig;
3687
3688 let config = VadConfig::default();
3690 assert!(config.energy_threshold > 0.0);
3691 assert!(config.zcr_threshold > 0.0);
3692 assert!(config.min_speech_frames > 0);
3693 }
3694
3695 #[test]
3700 fn test_vad_transcription_result_new() {
3701 let result = VadTranscriptionResult {
3702 text: "hello world".to_string(),
3703 language: "en".to_string(),
3704 segments: vec![],
3705 speech_segments: vec![],
3706 total_duration_secs: 1.0,
3707 speech_duration_secs: 0.5,
3708 };
3709
3710 assert_eq!(result.text, "hello world");
3711 assert_eq!(result.language, "en");
3712 assert!(!result.has_speech());
3713 assert_eq!(result.num_segments(), 0);
3714 }
3715
3716 #[test]
3717 fn test_vad_transcription_result_with_segments() {
3718 let result = VadTranscriptionResult {
3719 text: "hello world".to_string(),
3720 language: "en".to_string(),
3721 segments: vec![
3722 VadSpeechSegment {
3723 start: 0.0,
3724 end: 1.0,
3725 text: "hello".to_string(),
3726 tokens: vec![1, 2],
3727 },
3728 VadSpeechSegment {
3729 start: 1.5,
3730 end: 2.5,
3731 text: "world".to_string(),
3732 tokens: vec![3, 4],
3733 },
3734 ],
3735 speech_segments: vec![(0.0, 1.0), (1.5, 2.5)],
3736 total_duration_secs: 3.0,
3737 speech_duration_secs: 2.0,
3738 };
3739
3740 assert!(result.has_speech());
3741 assert_eq!(result.num_segments(), 2);
3742 assert!(result.first_segment().is_some());
3743 assert!(result.last_segment().is_some());
3744 assert_eq!(
3745 result.first_segment().map(|s| &s.text),
3746 Some(&"hello".to_string())
3747 );
3748 assert_eq!(
3749 result.last_segment().map(|s| &s.text),
3750 Some(&"world".to_string())
3751 );
3752 }
3753
3754 #[test]
3755 fn test_vad_transcription_result_silence_ratio() {
3756 let result = VadTranscriptionResult {
3757 text: String::new(),
3758 language: "en".to_string(),
3759 segments: vec![],
3760 speech_segments: vec![],
3761 total_duration_secs: 1.0,
3762 speech_duration_secs: 0.5,
3763 };
3764
3765 let ratio = result.silence_ratio(2.0);
3766 assert!((ratio - 0.75).abs() < 0.01); }
3768
3769 #[test]
3770 fn test_vad_transcription_result_silence_ratio_zero_duration() {
3771 let result = VadTranscriptionResult {
3772 text: String::new(),
3773 language: "en".to_string(),
3774 segments: vec![],
3775 speech_segments: vec![],
3776 total_duration_secs: 0.0,
3777 speech_duration_secs: 0.0,
3778 };
3779
3780 let ratio = result.silence_ratio(0.0);
3781 assert!((ratio - 1.0).abs() < 0.01);
3782 }
3783
3784 #[test]
3785 fn test_vad_transcription_result_iter() {
3786 let result = VadTranscriptionResult {
3787 text: "a b".to_string(),
3788 language: "en".to_string(),
3789 segments: vec![
3790 VadSpeechSegment {
3791 start: 0.0,
3792 end: 1.0,
3793 text: "a".to_string(),
3794 tokens: vec![1],
3795 },
3796 VadSpeechSegment {
3797 start: 1.0,
3798 end: 2.0,
3799 text: "b".to_string(),
3800 tokens: vec![2],
3801 },
3802 ],
3803 speech_segments: vec![(0.0, 1.0), (1.0, 2.0)],
3804 total_duration_secs: 2.0,
3805 speech_duration_secs: 2.0,
3806 };
3807
3808 let texts: Vec<_> = result.iter().map(|s| s.text.as_str()).collect();
3809 assert_eq!(texts, vec!["a", "b"]);
3810 }
3811
3812 #[test]
3813 fn test_vad_speech_segment_duration() {
3814 let segment = VadSpeechSegment {
3815 start: 1.5,
3816 end: 3.0,
3817 text: "test".to_string(),
3818 tokens: vec![1, 2, 3],
3819 };
3820
3821 assert!((segment.duration() - 1.5).abs() < 0.01);
3822 }
3823
3824 #[test]
3825 fn test_vad_speech_segment_has_text() {
3826 let with_text = VadSpeechSegment {
3827 start: 0.0,
3828 end: 1.0,
3829 text: "hello".to_string(),
3830 tokens: vec![1],
3831 };
3832 let empty = VadSpeechSegment {
3833 start: 0.0,
3834 end: 1.0,
3835 text: String::new(),
3836 tokens: vec![],
3837 };
3838
3839 assert!(with_text.has_text());
3840 assert!(!empty.has_text());
3841 }
3842
3843 #[test]
3844 #[ignore = "Allocates large model - run with --ignored"]
3845 fn test_transcribe_with_vad_silence_only() {
3846 let whisper = WhisperApr::tiny();
3847 let silence = vec![0.0; 16000]; let result = whisper
3850 .transcribe_with_vad(&silence, TranscribeOptions::default(), None)
3851 .expect("should succeed");
3852
3853 assert!(!result.has_speech());
3854 assert_eq!(result.num_segments(), 0);
3855 assert!(result.text.is_empty());
3856 }
3857
3858 #[test]
3859 fn test_transcribe_with_silence_detection_config() {
3860 let config = vad::SilenceConfig::new()
3862 .with_min_silence_duration(0.3)
3863 .with_max_silence_duration(2.0)
3864 .with_silence_threshold(0.001);
3865
3866 assert!((config.min_silence_duration - 0.3).abs() < 0.01);
3867 assert!((config.max_silence_duration - 2.0).abs() < 0.01);
3868 assert!((config.silence_threshold - 0.001).abs() < 0.001);
3869 }
3870
3871 #[test]
3872 #[ignore = "Allocates large model - run with --ignored"]
3873 fn test_invert_silence_segments_empty() {
3874 let whisper = WhisperApr::tiny();
3875 let audio_len = 16000; let silence_segments: Vec<vad::SilenceSegment> = vec![];
3878 let speech = whisper.invert_silence_segments(&silence_segments, audio_len);
3879
3880 assert_eq!(speech.len(), 1);
3882 assert!((speech[0].0 - 0.0).abs() < 0.01);
3883 assert!((speech[0].1 - 1.0).abs() < 0.01);
3884 }
3885
3886 #[test]
3887 #[ignore = "Allocates large model - run with --ignored"]
3888 fn test_invert_silence_segments_single() {
3889 let whisper = WhisperApr::tiny();
3890 let audio_len = 32000; let silence_segments = vec![vad::SilenceSegment {
3893 start: 0.5,
3894 end: 1.5,
3895 noise_floor: 0.001,
3896 }];
3897 let speech = whisper.invert_silence_segments(&silence_segments, audio_len);
3898
3899 assert_eq!(speech.len(), 2);
3901 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); }
3906
3907 #[test]
3908 #[ignore = "Allocates large model - run with --ignored"]
3909 fn test_invert_silence_segments_multiple() {
3910 let whisper = WhisperApr::tiny();
3911 let audio_len = 48000; let silence_segments = vec![
3914 vad::SilenceSegment {
3915 start: 0.5,
3916 end: 1.0,
3917 noise_floor: 0.001,
3918 },
3919 vad::SilenceSegment {
3920 start: 2.0,
3921 end: 2.5,
3922 noise_floor: 0.001,
3923 },
3924 ];
3925 let speech = whisper.invert_silence_segments(&silence_segments, audio_len);
3926
3927 assert_eq!(speech.len(), 3);
3929 }
3930
3931 #[test]
3932 fn test_tokenizer_integration() {
3933 use tokenizer::special_tokens;
3934
3935 assert!(special_tokens::SOT > 0);
3937 assert!(special_tokens::EOT > 0);
3938 assert!(special_tokens::TRANSCRIBE > 0);
3939 assert!(special_tokens::TRANSLATE > 0);
3940 assert!(special_tokens::NO_TIMESTAMPS > 0);
3941 }
3942
3943 #[test]
3944 fn test_inference_beam_search_integration() {
3945 use inference::BeamSearchDecoder;
3946
3947 let decoder = BeamSearchDecoder::new(5, 448); assert_eq!(decoder.beam_size(), 5);
3951 }
3952
3953 #[test]
3954 fn test_format_decompression_integration() {
3955 use format::Decompressor;
3956
3957 let mut decompressor = Decompressor::new();
3959
3960 assert!(decompressor.is_empty());
3962
3963 decompressor.reset();
3965 assert!(decompressor.is_empty());
3966 }
3967
3968 #[test]
3969 fn test_timestamps_generation() {
3970 let segment = Segment {
3971 start: 1.5,
3972 end: 3.25,
3973 text: "This is a test".to_string(),
3974 tokens: vec![1, 2, 3, 4],
3975 };
3976
3977 let duration = segment.end - segment.start;
3979 assert!((duration - 1.75).abs() < 1e-5);
3980 }
3981
3982 #[test]
3983 fn test_language_detection_integration() {
3984 use detection::{is_supported, language_name};
3985
3986 assert!(is_supported("en"));
3988 assert!(is_supported("es"));
3989 assert!(is_supported("ja"));
3990 assert!(!is_supported("invalid_lang"));
3991
3992 assert_eq!(language_name("en"), Some("English"));
3994 assert_eq!(language_name("es"), Some("Spanish"));
3995 assert_eq!(language_name("zh"), Some("Chinese"));
3996
3997 for lang in &["en", "es", "fr", "de", "it", "ja", "zh", "ko", "pt", "ru"] {
4000 assert!(
4001 language_name(lang).is_some(),
4002 "language {} should have a name",
4003 lang
4004 );
4005 }
4006 }
4007
4008 #[test]
4009 fn test_model_kv_cache_integration() {
4010 use model::{Decoder, ModelConfig};
4011
4012 let config = ModelConfig::tiny();
4013 let decoder = Decoder::new(&config);
4014
4015 let cache = decoder.create_kv_cache();
4017 assert!(cache.self_attn_cache.iter().all(|c| c.is_empty()));
4018 assert!(cache.cross_attn_cache.iter().all(|c| c.is_empty()));
4019 }
4020
4021 #[test]
4022 fn test_progress_tracking_integration() {
4023 use progress::{format_bytes, Progress};
4024
4025 let progress = Progress::new(50, 100);
4027 assert_eq!(progress.percent(), 50.0);
4028 assert_eq!(progress.current, 50);
4029 assert_eq!(progress.total, 100);
4030
4031 let kb_str = format_bytes(1024);
4033 assert!(kb_str.contains("KB"), "should contain KB: {}", kb_str);
4034 let mb_str = format_bytes(1024 * 1024);
4035 assert!(mb_str.contains("MB"), "should contain MB: {}", mb_str);
4036 }
4037
4038 #[test]
4043 fn test_model_type_variants() {
4044 let tiny = ModelType::Tiny;
4045 let tiny_en = ModelType::TinyEn;
4046 let base = ModelType::Base;
4047 let base_en = ModelType::BaseEn;
4048 let small = ModelType::Small;
4049
4050 assert!(format!("{tiny:?}").contains("Tiny"));
4052 assert!(format!("{tiny_en:?}").contains("TinyEn"));
4053 assert!(format!("{base:?}").contains("Base"));
4054 assert!(format!("{base_en:?}").contains("BaseEn"));
4055 assert!(format!("{small:?}").contains("Small"));
4056
4057 let tiny_clone = tiny;
4059 assert_eq!(tiny_clone, ModelType::Tiny);
4060
4061 assert_eq!(tiny, ModelType::Tiny);
4063 assert_ne!(tiny, base);
4064 }
4065
4066 #[test]
4067 fn test_task_variants() {
4068 let transcribe = Task::Transcribe;
4069 let translate = Task::Translate;
4070
4071 assert!(format!("{transcribe:?}").contains("Transcribe"));
4073 assert!(format!("{translate:?}").contains("Translate"));
4074
4075 let transcribe_clone = transcribe;
4077 assert_eq!(transcribe_clone, Task::Transcribe);
4078
4079 assert_eq!(transcribe, Task::Transcribe);
4081 assert_ne!(transcribe, translate);
4082
4083 assert_eq!(Task::default(), Task::Transcribe);
4085 }
4086
4087 #[test]
4088 fn test_decoding_strategy_sampling() {
4089 let sampling = DecodingStrategy::Sampling {
4090 temperature: 0.8,
4091 top_k: Some(50),
4092 top_p: Some(0.9),
4093 };
4094
4095 let debug_str = format!("{sampling:?}");
4097 assert!(debug_str.contains("Sampling"));
4098
4099 let cloned = sampling.clone();
4101 assert!(matches!(cloned, DecodingStrategy::Sampling { .. }));
4102 }
4103
4104 #[test]
4105 fn test_decoding_strategy_beam_search() {
4106 let beam = DecodingStrategy::BeamSearch {
4107 beam_size: 5,
4108 temperature: 0.0,
4109 patience: 1.0,
4110 };
4111
4112 let debug_str = format!("{beam:?}");
4113 assert!(debug_str.contains("BeamSearch"));
4114 }
4115
4116 #[test]
4117 fn test_transcribe_options_clone() {
4118 let options = TranscribeOptions {
4119 language: Some("fr".to_string()),
4120 task: Task::Translate,
4121 strategy: DecodingStrategy::Greedy,
4122 word_timestamps: true,
4123 profile: false,
4124 ..Default::default()
4125 };
4126
4127 let cloned = options.clone();
4128 assert_eq!(cloned.language, Some("fr".to_string()));
4129 assert_eq!(cloned.task, Task::Translate);
4130 assert!(cloned.word_timestamps);
4131 }
4132
4133 #[test]
4134 fn test_transcribe_options_debug() {
4135 let options = TranscribeOptions::default();
4136 let debug_str = format!("{options:?}");
4137 assert!(debug_str.contains("TranscribeOptions"));
4138 }
4139
4140 #[test]
4141 fn test_segment_clone() {
4142 let segment = Segment {
4143 start: 1.0,
4144 end: 2.0,
4145 text: "test".to_string(),
4146 tokens: vec![1, 2],
4147 };
4148
4149 let cloned = segment.clone();
4150 assert_eq!(cloned.text, "test");
4151 assert_eq!(cloned.tokens, vec![1, 2]);
4152 }
4153
4154 #[test]
4155 fn test_segment_debug() {
4156 let segment = Segment {
4157 start: 0.0,
4158 end: 1.0,
4159 text: "hello".to_string(),
4160 tokens: vec![],
4161 };
4162
4163 let debug_str = format!("{segment:?}");
4164 assert!(debug_str.contains("hello"));
4165 }
4166
4167 #[test]
4168 fn test_transcription_result_clone() {
4169 let result = TranscriptionResult {
4170 text: "hello world".to_string(),
4171 language: "en".to_string(),
4172 segments: vec![Segment {
4173 start: 0.0,
4174 end: 1.0,
4175 text: "hello".to_string(),
4176 tokens: vec![1],
4177 }],
4178 profiling: None,
4179 };
4180
4181 let cloned = result.clone();
4182 assert_eq!(cloned.text, "hello world");
4183 assert_eq!(cloned.segments.len(), 1);
4184 }
4185
4186 #[test]
4187 fn test_transcription_result_debug() {
4188 let result = TranscriptionResult {
4189 text: "test".to_string(),
4190 language: "en".to_string(),
4191 segments: vec![],
4192 profiling: None,
4193 };
4194
4195 let debug_str = format!("{result:?}");
4196 assert!(debug_str.contains("TranscriptionResult"));
4197 }
4198
4199 #[test]
4200 #[ignore = "Allocates large model - run with --ignored"]
4201 fn test_whisper_debug() {
4202 let whisper = WhisperApr::tiny();
4203 let debug_str = format!("{whisper:?}");
4204 assert!(debug_str.contains("WhisperApr"));
4205 }
4206
4207 #[test]
4208 #[ignore = "Allocates large model - run with --ignored"]
4209 fn test_whisper_memory_size_all_models() {
4210 let tiny = WhisperApr::tiny();
4211 let base = WhisperApr::base();
4212
4213 assert!(tiny.memory_size() < base.memory_size());
4215
4216 assert!(tiny.memory_size() > 0);
4218 assert!(base.memory_size() > 0);
4219 }
4220
4221 #[test]
4222 fn test_transcribe_options_sampling_strategy() {
4223 let options = TranscribeOptions {
4224 language: None,
4225 task: Task::Transcribe,
4226 strategy: DecodingStrategy::Sampling {
4227 temperature: 1.0,
4228 top_k: None,
4229 top_p: None,
4230 },
4231 word_timestamps: false,
4232 profile: false,
4233 ..Default::default()
4234 };
4235
4236 assert!(matches!(
4237 options.strategy,
4238 DecodingStrategy::Sampling { .. }
4239 ));
4240 }
4241
4242 #[test]
4243 fn test_transcribe_options_word_timestamps_enabled() {
4244 let options = TranscribeOptions {
4245 language: Some("en".to_string()),
4246 task: Task::Transcribe,
4247 strategy: DecodingStrategy::default(),
4248 word_timestamps: true,
4249 profile: false,
4250 ..Default::default()
4251 };
4252
4253 assert!(options.word_timestamps);
4254 }
4255
4256 #[test]
4261 fn test_model_type_small() {
4262 let small = ModelType::Small;
4263 assert!(format!("{small:?}").contains("Small"));
4264 }
4265
4266 #[test]
4267 fn test_model_type_medium() {
4268 let medium = ModelType::Medium;
4269 assert!(format!("{medium:?}").contains("Medium"));
4270 }
4271
4272 #[test]
4273 fn test_model_type_medium_en() {
4274 let medium_en = ModelType::MediumEn;
4275 assert!(format!("{medium_en:?}").contains("MediumEn"));
4276 }
4277
4278 #[test]
4279 fn test_model_type_large() {
4280 let large = ModelType::Large;
4281 assert!(format!("{large:?}").contains("Large"));
4282 }
4283
4284 #[test]
4285 fn test_model_type_large_v1() {
4286 let large_v1 = ModelType::LargeV1;
4287 assert!(format!("{large_v1:?}").contains("LargeV1"));
4288 }
4289
4290 #[test]
4291 fn test_model_type_large_v2() {
4292 let large_v2 = ModelType::LargeV2;
4293 assert!(format!("{large_v2:?}").contains("LargeV2"));
4294 }
4295
4296 #[test]
4297 fn test_model_type_large_v3() {
4298 let large_v3 = ModelType::LargeV3;
4299 assert!(format!("{large_v3:?}").contains("LargeV3"));
4300 }
4301
4302 #[test]
4303 #[ignore = "Allocates large model - run with --ignored"]
4304 fn test_whisper_small() {
4305 let whisper = WhisperApr::small();
4306 assert_eq!(whisper.model_type(), ModelType::Small);
4307 assert_eq!(whisper.config().n_audio_layer, 12);
4308 assert_eq!(whisper.config().n_audio_state, 768);
4309 }
4310
4311 #[test]
4312 #[ignore = "Allocates large model - run with --ignored"]
4313 fn test_whisper_medium() {
4314 let whisper = WhisperApr::medium();
4315 assert_eq!(whisper.model_type(), ModelType::Medium);
4316 assert_eq!(whisper.config().n_audio_layer, 24);
4317 assert_eq!(whisper.config().n_audio_state, 1024);
4318 }
4319
4320 #[test]
4321 #[ignore = "Allocates large model - run with --ignored"]
4322 fn test_whisper_large() {
4323 let whisper = WhisperApr::large();
4324 assert_eq!(whisper.model_type(), ModelType::Large);
4325 assert_eq!(whisper.config().n_audio_layer, 32);
4326 assert_eq!(whisper.config().n_audio_state, 1280);
4327 }
4328
4329 #[test]
4330 #[ignore = "Allocates large model - run with --ignored"]
4331 fn test_whisper_memory_size_all_extended_models() {
4332 let small = WhisperApr::small();
4333 let medium = WhisperApr::medium();
4334 let large = WhisperApr::large();
4335
4336 assert!(small.memory_size() < medium.memory_size());
4338 assert!(medium.memory_size() < large.memory_size());
4339
4340 assert!(small.memory_size() > 0);
4342 assert!(medium.memory_size() > 0);
4343 assert!(large.memory_size() > 0);
4344 }
4345
4346 #[test]
4347 #[ignore = "Allocates large model - run with --ignored"]
4348 fn test_extended_model_memory_size_estimates() {
4349 let small = WhisperApr::small();
4350 let medium = WhisperApr::medium();
4351 let large = WhisperApr::large();
4352
4353 assert!(small.memory_size() > 900_000_000);
4355 assert!(small.memory_size() < 1_200_000_000);
4356
4357 assert!(medium.memory_size() > 2_500_000_000);
4359 assert!(medium.memory_size() < 3_500_000_000);
4360
4361 assert!(large.memory_size() > 5_000_000_000);
4363 assert!(large.memory_size() < 7_000_000_000);
4364 }
4365
4366 #[test]
4371 fn test_batch_transcription_result_len() {
4372 let result = BatchTranscriptionResult {
4373 results: vec![
4374 TranscriptionResult {
4375 text: "Hello".to_string(),
4376 language: "en".to_string(),
4377 segments: vec![],
4378 profiling: None,
4379 },
4380 TranscriptionResult {
4381 text: "World".to_string(),
4382 language: "en".to_string(),
4383 segments: vec![],
4384 profiling: None,
4385 },
4386 ],
4387 total_duration_secs: 1.5,
4388 };
4389
4390 assert_eq!(result.len(), 2);
4391 assert!(!result.is_empty());
4392 }
4393
4394 #[test]
4395 fn test_batch_transcription_result_empty() {
4396 let result = BatchTranscriptionResult {
4397 results: vec![],
4398 total_duration_secs: 0.0,
4399 };
4400
4401 assert!(result.is_empty());
4402 assert_eq!(result.len(), 0);
4403 }
4404
4405 #[test]
4406 fn test_batch_transcription_result_get() {
4407 let result = BatchTranscriptionResult {
4408 results: vec![
4409 TranscriptionResult {
4410 text: "First".to_string(),
4411 language: "en".to_string(),
4412 segments: vec![],
4413 profiling: None,
4414 },
4415 TranscriptionResult {
4416 text: "Second".to_string(),
4417 language: "es".to_string(),
4418 segments: vec![],
4419 profiling: None,
4420 },
4421 ],
4422 total_duration_secs: 2.0,
4423 };
4424
4425 assert!(result.get(0).is_some());
4426 assert_eq!(result.get(0).map(|r| r.text.as_str()), Some("First"));
4427 assert_eq!(result.get(1).map(|r| r.language.as_str()), Some("es"));
4428 assert!(result.get(2).is_none());
4429 }
4430
4431 #[test]
4432 fn test_batch_transcription_result_texts() {
4433 let result = BatchTranscriptionResult {
4434 results: vec![
4435 TranscriptionResult {
4436 text: "One".to_string(),
4437 language: "en".to_string(),
4438 segments: vec![],
4439 profiling: None,
4440 },
4441 TranscriptionResult {
4442 text: "Two".to_string(),
4443 language: "en".to_string(),
4444 segments: vec![],
4445 profiling: None,
4446 },
4447 TranscriptionResult {
4448 text: "Three".to_string(),
4449 language: "en".to_string(),
4450 segments: vec![],
4451 profiling: None,
4452 },
4453 ],
4454 total_duration_secs: 3.0,
4455 };
4456
4457 let texts = result.texts();
4458 assert_eq!(texts, vec!["One", "Two", "Three"]);
4459 }
4460
4461 #[test]
4462 fn test_batch_transcription_result_iter() {
4463 let result = BatchTranscriptionResult {
4464 results: vec![
4465 TranscriptionResult {
4466 text: "A".to_string(),
4467 language: "en".to_string(),
4468 segments: vec![],
4469 profiling: None,
4470 },
4471 TranscriptionResult {
4472 text: "B".to_string(),
4473 language: "en".to_string(),
4474 segments: vec![],
4475 profiling: None,
4476 },
4477 ],
4478 total_duration_secs: 1.0,
4479 };
4480
4481 let collected: Vec<&str> = result.iter().map(|r| r.text.as_str()).collect();
4482 assert_eq!(collected, vec!["A", "B"]);
4483 }
4484
4485 #[test]
4486 #[ignore = "Allocates large model - run with --ignored"]
4487 fn test_transcribe_batch_empty() {
4488 let whisper = WhisperApr::tiny();
4489 let result = whisper.transcribe_batch(&[], TranscribeOptions::default());
4490 assert!(result.is_err());
4491 }
4492
4493 #[test]
4494 #[ignore = "Allocates large model - run with --ignored"]
4495 fn test_transcribe_audio_batch_empty() {
4496 let whisper = WhisperApr::tiny();
4497 let batch = audio::AudioBatch::with_default_config();
4498 let result = whisper.transcribe_audio_batch(&batch, TranscribeOptions::default());
4499 assert!(result.is_err());
4500 }
4501
4502 #[test]
4503 #[ignore = "Allocates large model - run with --ignored"]
4504 fn test_transcribe_batch_optimized_empty() {
4505 let whisper = WhisperApr::tiny();
4506 let result = whisper.transcribe_batch_optimized(&[], TranscribeOptions::default());
4507 assert!(result.is_err());
4508 }
4509
4510 #[test]
4511 fn test_create_audio_batch() {
4512 let segments = vec![vec![0.1_f32, 0.2, 0.3], vec![0.4_f32, 0.5]];
4513
4514 let batch = WhisperApr::create_audio_batch(&segments);
4515 assert_eq!(batch.len(), 2);
4516 assert!(!batch.is_empty());
4517 }
4518
4519 #[test]
4520 fn test_create_audio_batch_empty() {
4521 let segments: Vec<Vec<f32>> = vec![];
4522 let batch = WhisperApr::create_audio_batch(&segments);
4523 assert!(batch.is_empty());
4524 }
4525
4526 #[test]
4527 fn test_batch_transcription_result_duration() {
4528 let result = BatchTranscriptionResult {
4529 results: vec![],
4530 total_duration_secs: 5.25,
4531 };
4532
4533 assert!((result.total_duration_secs - 5.25).abs() < f32::EPSILON);
4534 }
4535
4536 #[test]
4541 fn test_partial_transcription_result_new() {
4542 let result = PartialTranscriptionResult {
4543 text: "hello".to_string(),
4544 language: "en".to_string(),
4545 is_final: false,
4546 confidence: 0.95,
4547 duration_secs: 1.5,
4548 processing_time_secs: 0.3,
4549 };
4550
4551 assert_eq!(result.text, "hello");
4552 assert_eq!(result.language, "en");
4553 assert!(!result.is_final);
4554 assert!((result.confidence - 0.95).abs() < 0.01);
4555 }
4556
4557 #[test]
4558 fn test_partial_transcription_result_has_text() {
4559 let with_text = PartialTranscriptionResult {
4560 text: "hello".to_string(),
4561 language: "en".to_string(),
4562 is_final: false,
4563 confidence: 1.0,
4564 duration_secs: 1.0,
4565 processing_time_secs: 0.1,
4566 };
4567 let empty = PartialTranscriptionResult {
4568 text: String::new(),
4569 language: "en".to_string(),
4570 is_final: false,
4571 confidence: 0.0,
4572 duration_secs: 0.5,
4573 processing_time_secs: 0.05,
4574 };
4575
4576 assert!(with_text.has_text());
4577 assert!(!empty.has_text());
4578 }
4579
4580 #[test]
4581 fn test_partial_transcription_result_is_empty_interim() {
4582 let empty_interim = PartialTranscriptionResult {
4583 text: String::new(),
4584 language: "en".to_string(),
4585 is_final: false,
4586 confidence: 0.0,
4587 duration_secs: 0.5,
4588 processing_time_secs: 0.05,
4589 };
4590 let empty_final = PartialTranscriptionResult {
4591 text: String::new(),
4592 language: "en".to_string(),
4593 is_final: true,
4594 confidence: 0.0,
4595 duration_secs: 0.5,
4596 processing_time_secs: 0.05,
4597 };
4598 let with_text = PartialTranscriptionResult {
4599 text: "hello".to_string(),
4600 language: "en".to_string(),
4601 is_final: false,
4602 confidence: 1.0,
4603 duration_secs: 1.0,
4604 processing_time_secs: 0.1,
4605 };
4606
4607 assert!(empty_interim.is_empty_interim());
4608 assert!(!empty_final.is_empty_interim()); assert!(!with_text.is_empty_interim()); }
4611
4612 #[test]
4613 fn test_partial_transcription_result_real_time_factor() {
4614 let result = PartialTranscriptionResult {
4615 text: "hello".to_string(),
4616 language: "en".to_string(),
4617 is_final: false,
4618 confidence: 1.0,
4619 duration_secs: 2.0,
4620 processing_time_secs: 0.5,
4621 };
4622
4623 assert!((result.real_time_factor() - 0.25).abs() < 0.01);
4625 }
4626
4627 #[test]
4628 fn test_partial_transcription_result_real_time_factor_zero_duration() {
4629 let result = PartialTranscriptionResult {
4630 text: String::new(),
4631 language: "en".to_string(),
4632 is_final: false,
4633 confidence: 0.0,
4634 duration_secs: 0.0,
4635 processing_time_secs: 0.0,
4636 };
4637
4638 assert!((result.real_time_factor() - 0.0).abs() < 0.01);
4639 }
4640
4641 #[test]
4642 fn test_partial_transcription_result_debug_clone() {
4643 let result = PartialTranscriptionResult {
4644 text: "test".to_string(),
4645 language: "en".to_string(),
4646 is_final: true,
4647 confidence: 0.9,
4648 duration_secs: 1.0,
4649 processing_time_secs: 0.1,
4650 };
4651
4652 let debug_str = format!("{result:?}");
4653 assert!(debug_str.contains("PartialTranscriptionResult"));
4654
4655 let cloned = result.clone();
4656 assert_eq!(cloned.text, "test");
4657 assert!(cloned.is_final);
4658 }
4659
4660 #[test]
4661 #[ignore = "Allocates large model - run with --ignored"]
4662 fn test_transcribe_partial_too_short() {
4663 let whisper = WhisperApr::tiny();
4664 let short_audio = vec![0.0; 4000]; let result = whisper
4667 .transcribe_partial(&short_audio, TranscribeOptions::default(), false)
4668 .expect("should succeed with empty result");
4669
4670 assert!(result.text.is_empty());
4671 assert!(!result.is_final);
4672 assert!((result.confidence - 0.0).abs() < 0.01);
4673 }
4674
4675 #[test]
4676 #[ignore = "Allocates large model - run with --ignored"]
4677 fn test_encode_3_second_chunk() {
4678 let whisper = WhisperApr::tiny();
4681 let audio = vec![0.0; 48000]; let mel = whisper
4685 .compute_mel(&audio)
4686 .expect("mel computation should succeed");
4687
4688 let result = whisper.encode(&mel);
4690 assert!(
4691 result.is_ok(),
4692 "encode should succeed for 3s audio mel: {:?}",
4693 result.err()
4694 );
4695
4696 let encoded = result.expect("encode should succeed");
4698 let d_model = whisper.config().n_text_state as usize; assert_eq!(
4700 encoded.len() % d_model,
4701 0,
4702 "encoded output should be multiple of d_model"
4703 );
4704 }
4705
4706 #[test]
4707 #[ignore = "Allocates large model - run with --ignored"]
4708 fn test_create_streaming_session() {
4709 let whisper = WhisperApr::tiny();
4710 let session = whisper.create_streaming_session(TranscribeOptions::default(), 44100);
4711
4712 assert_eq!(session.state(), audio::ProcessorState::WaitingForSpeech);
4713 assert!((session.chunk_progress() - 0.0).abs() < 0.01);
4714 assert!(!session.has_chunk());
4715 assert!(!session.has_events());
4716 }
4717
4718 #[test]
4719 #[ignore = "Allocates large model - run with --ignored"]
4720 fn test_streaming_session_reset() {
4721 let whisper = WhisperApr::tiny();
4722 let mut session = whisper.create_streaming_session(TranscribeOptions::default(), 16000);
4723
4724 session.push(&vec![0.1; 1000]).expect("push should work");
4726
4727 session.reset();
4729
4730 assert_eq!(session.state(), audio::ProcessorState::WaitingForSpeech);
4731 assert!((session.partial_duration() - 0.0).abs() < 0.01);
4732 }
4733
4734 #[test]
4735 #[ignore = "Allocates large model - run with --ignored"]
4736 fn test_streaming_session_set_partial_threshold() {
4737 let whisper = WhisperApr::tiny();
4738 let mut session = whisper.create_streaming_session(TranscribeOptions::default(), 16000);
4739
4740 session.set_partial_threshold(5.0);
4741 }
4744
4745 #[test]
4746 #[ignore = "Allocates large model - run with --ignored"]
4747 fn test_streaming_session_finalize_no_chunk() {
4748 let whisper = WhisperApr::tiny();
4749 let mut session = whisper.create_streaming_session(TranscribeOptions::default(), 16000);
4750
4751 let result = session.finalize();
4752 assert!(result.is_err());
4753 }
4754
4755 #[test]
4756 #[ignore = "Allocates large model - run with --ignored"]
4757 fn test_streaming_session_flush_empty() {
4758 let whisper = WhisperApr::tiny();
4759 let mut session = whisper.create_streaming_session(TranscribeOptions::default(), 16000);
4760
4761 let result = session.flush().expect("flush should work");
4762 assert!(result.is_none());
4763 }
4764
4765 #[test]
4766 #[ignore = "Allocates large model - run with --ignored"]
4767 fn test_streaming_session_drain_events() {
4768 let whisper = WhisperApr::tiny();
4769 let mut session = whisper.create_streaming_session(TranscribeOptions::default(), 16000);
4770
4771 session.reset();
4773
4774 let events = session.drain_events();
4775 assert!(!events.is_empty());
4776 assert!(events
4777 .iter()
4778 .any(|e| matches!(e, audio::StreamingEvent::Reset)));
4779 }
4780
4781 #[test]
4782 #[ignore = "Allocates large model - run with --ignored"]
4783 fn test_streaming_session_push_silence() {
4784 let whisper = WhisperApr::tiny();
4785 let mut session = whisper.create_streaming_session(TranscribeOptions::default(), 16000);
4786
4787 let result = session.push(&vec![0.0; 16000]).expect("push should work");
4789 assert!(result.is_none()); }
4791
4792 #[test]
4793 #[ignore = "Allocates large model - run with --ignored"]
4794 fn test_streaming_session_debug() {
4795 let whisper = WhisperApr::tiny();
4796 let session = whisper.create_streaming_session(TranscribeOptions::default(), 16000);
4797
4798 let debug_str = format!("{session:?}");
4799 assert!(debug_str.contains("StreamingSession"));
4800 }
4801
4802 #[test]
4807 #[ignore = "Allocates large model - run with --ignored"]
4808 fn test_streaming_session_state() {
4809 let whisper = WhisperApr::tiny();
4810 let session = whisper.create_streaming_session(TranscribeOptions::default(), 16000);
4811
4812 let state = session.state();
4814 assert_eq!(state, audio::ProcessorState::WaitingForSpeech);
4815 }
4816
4817 #[test]
4818 #[ignore = "Allocates large model - run with --ignored"]
4819 fn test_streaming_session_chunk_progress() {
4820 let whisper = WhisperApr::tiny();
4821 let session = whisper.create_streaming_session(TranscribeOptions::default(), 16000);
4822
4823 let progress = session.chunk_progress();
4824 assert!(progress >= 0.0 && progress <= 1.0);
4825 }
4826
4827 #[test]
4828 #[ignore = "Allocates large model - run with --ignored"]
4829 fn test_streaming_session_partial_duration() {
4830 let whisper = WhisperApr::tiny();
4831 let session = whisper.create_streaming_session(TranscribeOptions::default(), 16000);
4832
4833 let duration = session.partial_duration();
4834 assert!(duration >= 0.0);
4835 }
4836
4837 #[test]
4838 fn test_partial_transcription_result_rtf_with_zero_processing() {
4839 let result = PartialTranscriptionResult {
4840 text: "test".to_string(),
4841 language: "en".to_string(),
4842 is_final: true,
4843 confidence: 0.9,
4844 duration_secs: 5.0,
4845 processing_time_secs: 0.0,
4846 };
4847
4848 let rtf = result.real_time_factor();
4849 assert!(rtf >= 0.0);
4850 }
4851
4852 #[test]
4853 fn test_vad_transcription_result_methods_empty() {
4854 let result = VadTranscriptionResult {
4855 text: "Hello world".to_string(),
4856 language: "en".to_string(),
4857 segments: vec![],
4858 speech_segments: vec![],
4859 total_duration_secs: 5.0,
4860 speech_duration_secs: 0.0,
4861 };
4862
4863 assert_eq!(result.num_segments(), 0);
4864 assert!(!result.has_speech());
4865 }
4866
4867 #[test]
4868 fn test_batch_transcription_result_defaults() {
4869 let result = BatchTranscriptionResult {
4870 results: vec![],
4871 total_duration_secs: 0.0,
4872 };
4873
4874 assert_eq!(result.len(), 0);
4875 assert!(result.is_empty());
4876 assert!(result.get(0).is_none());
4877 assert!(result.texts().is_empty());
4878 }
4879
4880 #[test]
4881 #[ignore = "Allocates large model - run with --ignored"]
4882 fn test_whisper_config_accessors() {
4883 let whisper = WhisperApr::tiny();
4884 let config = whisper.config();
4885
4886 assert!(config.n_vocab > 0);
4887 assert!(config.n_audio_ctx > 0);
4888 assert!(config.n_text_ctx > 0);
4889 }
4890
4891 #[test]
4892 fn test_transcribe_options_all_fields() {
4893 let options = TranscribeOptions {
4894 language: Some("fr".to_string()),
4895 task: Task::Translate,
4896 strategy: DecodingStrategy::BeamSearch {
4897 beam_size: 3,
4898 temperature: 0.2,
4899 patience: 1.5,
4900 },
4901 word_timestamps: true,
4902 profile: false,
4903 prompt: Some("domain-specific prompt".into()),
4904 hotwords: vec!["hotword1".into(), "hotword2".into()],
4905 };
4906
4907 assert_eq!(options.language, Some("fr".to_string()));
4908 assert_eq!(options.task, Task::Translate);
4909 assert!(options.word_timestamps);
4910 }
4911
4912 #[test]
4913 fn test_segment_with_tokens() {
4914 let segment = Segment {
4915 text: "Hello".to_string(),
4916 start: 0.0,
4917 end: 1.0,
4918 tokens: vec![1, 2, 3, 4, 5],
4919 };
4920
4921 assert_eq!(segment.tokens.len(), 5);
4922 assert!((segment.end - segment.start - 1.0).abs() < f32::EPSILON);
4923 }
4924
4925 #[test]
4926 fn test_model_from_config() {
4927 let config = model::ModelConfig::tiny();
4928 let whisper = WhisperApr::from_config(config);
4929
4930 assert_eq!(whisper.model_type(), ModelType::Tiny);
4931 }
4932
4933 #[test]
4934 fn test_decoding_strategy_variants() {
4935 let greedy = DecodingStrategy::Greedy;
4936 assert!(matches!(greedy, DecodingStrategy::Greedy));
4937
4938 let sampling = DecodingStrategy::Sampling {
4939 temperature: 0.5,
4940 top_k: Some(40),
4941 top_p: Some(0.9),
4942 };
4943 if let DecodingStrategy::Sampling {
4944 temperature,
4945 top_k,
4946 top_p,
4947 } = sampling
4948 {
4949 assert!((temperature - 0.5).abs() < f32::EPSILON);
4950 assert_eq!(top_k, Some(40));
4951 assert_eq!(top_p, Some(0.9));
4952 }
4953 }
4954
4955 #[test]
4956 fn test_task_variants_eq() {
4957 assert_eq!(Task::Transcribe, Task::Transcribe);
4958 assert_ne!(Task::Transcribe, Task::Translate);
4959 }
4960
4961 #[test]
4962 fn test_model_type_all_variants() {
4963 let variants = vec![
4964 ModelType::Tiny,
4965 ModelType::TinyEn,
4966 ModelType::Base,
4967 ModelType::BaseEn,
4968 ModelType::Small,
4969 ModelType::SmallEn,
4970 ModelType::Medium,
4971 ModelType::MediumEn,
4972 ModelType::Large,
4973 ModelType::LargeV1,
4974 ModelType::LargeV2,
4975 ModelType::LargeV3,
4976 ];
4977
4978 for variant in variants {
4979 let debug_str = format!("{variant:?}");
4980 assert!(!debug_str.is_empty());
4981 }
4982 }
4983
4984 #[test]
4985 fn test_vad_speech_segment_empty_text() {
4986 let segment = VadSpeechSegment {
4987 start: 0.0,
4988 end: 1.0,
4989 text: String::new(),
4990 tokens: vec![],
4991 };
4992
4993 assert!(!segment.has_text());
4994 assert!((segment.duration() - 1.0).abs() < f32::EPSILON);
4995 }
4996
4997 #[test]
4998 fn test_transcription_result_empty() {
4999 let result = TranscriptionResult {
5000 text: String::new(),
5001 language: "en".to_string(),
5002 segments: vec![],
5003 profiling: None,
5004 };
5005
5006 assert!(result.text.is_empty());
5007 assert!(result.segments.is_empty());
5008 }
5009
5010 #[test]
5011 fn test_batch_transcription_result_iter_coverage() {
5012 let results = vec![
5013 TranscriptionResult {
5014 text: "First".to_string(),
5015 language: "en".to_string(),
5016 segments: vec![],
5017 profiling: None,
5018 },
5019 TranscriptionResult {
5020 text: "Second".to_string(),
5021 language: "en".to_string(),
5022 segments: vec![],
5023 profiling: None,
5024 },
5025 ];
5026
5027 let batch = BatchTranscriptionResult {
5028 results,
5029 total_duration_secs: 1.0,
5030 };
5031
5032 let mut count = 0;
5033 for result in batch.iter() {
5034 count += 1;
5035 assert!(!result.text.is_empty());
5036 }
5037 assert_eq!(count, 2);
5038 }
5039
5040 #[test]
5041 #[ignore = "Allocates large model - run with --ignored"]
5042 fn test_whisper_clone() {
5043 let whisper = WhisperApr::tiny();
5044 let cloned = whisper.clone();
5045
5046 assert_eq!(whisper.model_type(), cloned.model_type());
5047 assert_eq!(whisper.memory_size(), cloned.memory_size());
5048 }
5049
5050 #[test]
5051 fn test_vad_transcription_result_first_last_segment() {
5052 let result = VadTranscriptionResult {
5053 text: "hello world".to_string(),
5054 language: "en".to_string(),
5055 segments: vec![
5056 VadSpeechSegment {
5057 start: 0.0,
5058 end: 1.0,
5059 text: "hello".to_string(),
5060 tokens: vec![1],
5061 },
5062 VadSpeechSegment {
5063 start: 1.5,
5064 end: 2.5,
5065 text: "world".to_string(),
5066 tokens: vec![2],
5067 },
5068 ],
5069 speech_segments: vec![(0.0, 1.0), (1.5, 2.5)],
5070 total_duration_secs: 3.0,
5071 speech_duration_secs: 2.0,
5072 };
5073
5074 let first = result.first_segment().expect("first segment");
5075 assert_eq!(first.text, "hello");
5076
5077 let last = result.last_segment().expect("last segment");
5078 assert_eq!(last.text, "world");
5079 }
5080
5081 #[test]
5082 fn test_vad_transcription_result_iter_segments() {
5083 let result = VadTranscriptionResult {
5084 text: "test".to_string(),
5085 language: "en".to_string(),
5086 segments: vec![
5087 VadSpeechSegment {
5088 start: 0.0,
5089 end: 1.0,
5090 text: "a".to_string(),
5091 tokens: vec![1],
5092 },
5093 VadSpeechSegment {
5094 start: 1.0,
5095 end: 2.0,
5096 text: "b".to_string(),
5097 tokens: vec![2],
5098 },
5099 ],
5100 speech_segments: vec![(0.0, 1.0), (1.0, 2.0)],
5101 total_duration_secs: 2.0,
5102 speech_duration_secs: 2.0,
5103 };
5104
5105 let mut count = 0;
5106 for segment in result.iter() {
5107 count += 1;
5108 assert!(segment.has_text());
5109 }
5110 assert_eq!(count, 2);
5111 }
5112
5113 #[test]
5114 fn test_partial_transcription_result_methods() {
5115 let result = PartialTranscriptionResult {
5116 text: "hello".to_string(),
5117 language: "en".to_string(),
5118 is_final: false,
5119 confidence: 0.85,
5120 duration_secs: 2.0,
5121 processing_time_secs: 0.5,
5122 };
5123
5124 assert!(result.has_text());
5125 assert!((result.real_time_factor() - 0.25).abs() < f32::EPSILON);
5126 }
5127
5128 #[test]
5129 fn test_partial_transcription_result_zero_duration() {
5130 let result = PartialTranscriptionResult {
5131 text: "".to_string(),
5132 language: "en".to_string(),
5133 is_final: true,
5134 confidence: 0.0,
5135 duration_secs: 0.0,
5136 processing_time_secs: 0.1,
5137 };
5138
5139 assert!(!result.has_text());
5140 assert!((result.real_time_factor() - 0.0).abs() < f32::EPSILON);
5141 }
5142
5143 #[test]
5148 #[ignore = "Requires model file - run with --ignored"]
5149 fn test_e2e_transcribe_with_int8_model() {
5150 let model_path = std::path::Path::new("models/whisper-tiny-int8.apr");
5152 if !model_path.exists() {
5153 eprintln!(
5154 "Skipping E2E test: model file not found at {:?}",
5155 model_path
5156 );
5157 return;
5158 }
5159
5160 let model_data = std::fs::read(model_path).expect("Failed to read model file");
5161 eprintln!("Loaded model: {} bytes", model_data.len());
5162
5163 let whisper = WhisperApr::load_from_apr(&model_data).expect("Failed to load model");
5164 eprintln!("Model loaded: {:?}", whisper.model_type());
5165
5166 let sample_rate = 16000;
5169 let duration_secs = 3.0;
5170 let num_samples = (sample_rate as f32 * duration_secs) as usize;
5171
5172 let audio: Vec<f32> = (0..num_samples)
5174 .map(|i| {
5175 let noise = ((i as f32 * 0.1).sin() * 0.001) + ((i as f32 * 0.37).cos() * 0.001);
5177 noise
5178 })
5179 .collect();
5180
5181 eprintln!("Generated {} samples of test audio", audio.len());
5182
5183 let start = std::time::Instant::now();
5189 let result = whisper.transcribe(&audio, TranscribeOptions::default());
5190 let elapsed = start.elapsed();
5191
5192 eprintln!("Transcription completed in {:?}", elapsed);
5193
5194 match result {
5195 Ok(transcription) => {
5196 eprintln!("Result: '{}'", transcription.text);
5197 eprintln!("Language: {}", transcription.language);
5198 eprintln!("Segments: {}", transcription.segments.len());
5199
5200 assert!(
5203 transcription.text.len() < 100,
5204 "Unexpected long transcription for silence"
5205 );
5206 }
5207 Err(e) => {
5208 panic!("Transcription failed: {e:?}");
5209 }
5210 }
5211
5212 let rtf = elapsed.as_secs_f32() / duration_secs;
5214 eprintln!("Real-time factor: {rtf:.2}x");
5215
5216 assert!(rtf < 50.0, "RTF {rtf} is too slow, SIMD may not be working");
5218 }
5219
5220 #[test]
5225 fn test_summarize_options_default_params() {
5226 use model::lfm2::{Lfm2, Lfm2Tokenizer};
5227
5228 let config = format::apr2::Lfm2Config {
5230 hidden_size: 64,
5231 num_layers: 2,
5232 num_q_heads: 4,
5233 num_kv_heads: 2,
5234 intermediate_size: 128,
5235 vocab_size: 1000,
5236 max_seq_len: 512,
5237 rope_theta: 10000.0,
5238 conv_dimension: 32,
5239 layer_types: vec![
5240 format::apr2::LayerType::Convolution {
5241 kernel_size: 4,
5242 cache_len: 3,
5243 },
5244 format::apr2::LayerType::Attention { use_gqa: true },
5245 ],
5246 };
5247 let model = Lfm2::new(config).expect("Model creation should succeed");
5248 let tokenizer = Lfm2Tokenizer::new();
5249
5250 let options = SummarizeOptions::new(&model, &tokenizer);
5251
5252 assert_eq!(options.max_tokens, 256);
5253 assert!((options.temperature - 0.3).abs() < f32::EPSILON);
5254 }
5255
5256 #[test]
5257 fn test_summarize_options_builder() {
5258 use model::lfm2::{Lfm2, Lfm2Tokenizer};
5259
5260 let config = format::apr2::Lfm2Config {
5261 hidden_size: 64,
5262 num_layers: 2,
5263 num_q_heads: 4,
5264 num_kv_heads: 2,
5265 intermediate_size: 128,
5266 vocab_size: 1000,
5267 max_seq_len: 512,
5268 rope_theta: 10000.0,
5269 conv_dimension: 32,
5270 layer_types: vec![
5271 format::apr2::LayerType::Convolution {
5272 kernel_size: 4,
5273 cache_len: 3,
5274 },
5275 format::apr2::LayerType::Attention { use_gqa: true },
5276 ],
5277 };
5278 let model = Lfm2::new(config).expect("Model creation should succeed");
5279 let tokenizer = Lfm2Tokenizer::new();
5280
5281 let options = SummarizeOptions::new(&model, &tokenizer)
5282 .with_max_tokens(512)
5283 .with_temperature(0.7);
5284
5285 assert_eq!(options.max_tokens, 512);
5286 assert!((options.temperature - 0.7).abs() < f32::EPSILON);
5287 }
5288
5289 #[test]
5290 fn test_transcribe_summary_result_accessors() {
5291 let transcription = TranscriptionResult {
5292 text: "Hello world".to_string(),
5293 language: "en".to_string(),
5294 segments: vec![],
5295 profiling: None,
5296 };
5297
5298 let result = TranscribeSummaryResult {
5299 transcription,
5300 summary: "Summary text".to_string(),
5301 generation_stats: None,
5302 };
5303
5304 assert_eq!(result.transcript(), "Hello world");
5305 assert_eq!(result.summary(), "Summary text");
5306 assert!(result.has_summary());
5307 }
5308
5309 #[test]
5310 fn test_transcribe_summary_result_empty_summary() {
5311 let transcription = TranscriptionResult {
5312 text: "Hello world".to_string(),
5313 language: "en".to_string(),
5314 segments: vec![],
5315 profiling: None,
5316 };
5317
5318 let result = TranscribeSummaryResult {
5319 transcription,
5320 summary: String::new(),
5321 generation_stats: None,
5322 };
5323
5324 assert!(!result.has_summary());
5325 }
5326
5327 #[test]
5328 #[ignore = "Slow test - run with --ignored"]
5329 fn test_transcribe_and_summarize_empty_audio() {
5330 use model::lfm2::{Lfm2, Lfm2Tokenizer};
5331
5332 let whisper = WhisperApr::tiny();
5333 let config = format::apr2::Lfm2Config {
5335 hidden_size: 64,
5336 num_layers: 2,
5337 num_q_heads: 4,
5338 num_kv_heads: 2,
5339 intermediate_size: 128,
5340 vocab_size: 1000,
5341 max_seq_len: 2048, rope_theta: 10000.0,
5343 conv_dimension: 32,
5344 layer_types: vec![
5345 format::apr2::LayerType::Convolution {
5346 kernel_size: 4,
5347 cache_len: 3,
5348 },
5349 format::apr2::LayerType::Attention { use_gqa: true },
5350 ],
5351 };
5352 let lfm2 = Lfm2::new(config).expect("Model creation should succeed");
5353 let tokenizer = Lfm2Tokenizer::new();
5354
5355 let audio = vec![0.0f32; 16000]; let transcribe_options = TranscribeOptions::default();
5358 let summarize_options = SummarizeOptions::new(&lfm2, &tokenizer).with_max_tokens(8); let result = whisper
5361 .transcribe_and_summarize(&audio, transcribe_options, summarize_options)
5362 .expect("Should not fail");
5363
5364 if result.transcription.text.trim().is_empty() {
5367 assert!(!result.has_summary());
5368 }
5369 }
5370
5371 #[test]
5372 #[ignore = "Slow test - run with --ignored"]
5373 fn test_transcribe_and_summarize_integration() {
5374 use model::lfm2::{Lfm2, Lfm2Tokenizer};
5375
5376 let whisper = WhisperApr::tiny();
5377 let config = format::apr2::Lfm2Config {
5379 hidden_size: 64,
5380 num_layers: 2,
5381 num_q_heads: 4,
5382 num_kv_heads: 2,
5383 intermediate_size: 128,
5384 vocab_size: 1000,
5385 max_seq_len: 2048, rope_theta: 10000.0,
5387 conv_dimension: 32,
5388 layer_types: vec![
5389 format::apr2::LayerType::Convolution {
5390 kernel_size: 4,
5391 cache_len: 3,
5392 },
5393 format::apr2::LayerType::Attention { use_gqa: true },
5394 ],
5395 };
5396 let lfm2 = Lfm2::new(config).expect("Model creation should succeed");
5397 let tokenizer = Lfm2Tokenizer::new();
5398
5399 let sample_rate = 16000;
5401 let duration_secs = 1.0;
5402 let num_samples = (sample_rate as f32 * duration_secs) as usize;
5403 let audio: Vec<f32> = (0..num_samples)
5404 .map(|i| (2.0 * std::f32::consts::PI * 440.0 * i as f32 / sample_rate as f32).sin())
5405 .collect();
5406
5407 let transcribe_options = TranscribeOptions::default();
5408 let summarize_options = SummarizeOptions::new(&lfm2, &tokenizer).with_max_tokens(8); let result =
5412 whisper.transcribe_and_summarize(&audio, transcribe_options, summarize_options);
5413
5414 assert!(result.is_ok());
5417
5418 let result = result.expect("Should succeed");
5419 assert_eq!(result.transcription.language, "en");
5421 }
5422
5423 #[test]
5428 fn test_chunk_constants() {
5429 assert_eq!(WhisperApr::CHUNK_SAMPLES, 30 * 16000);
5431
5432 assert_eq!(WhisperApr::OVERLAP_SAMPLES, 5 * 16000);
5434 }
5435
5436 #[test]
5437 #[ignore = "Allocates large model - run with --ignored"]
5438 fn test_short_audio_uses_single_chunk() {
5439 let _whisper = WhisperApr::tiny();
5440
5441 let short_audio = vec![0.0_f32; 10 * 16000]; assert!(short_audio.len() <= WhisperApr::CHUNK_SAMPLES);
5444
5445 }
5448
5449 #[test]
5450 #[ignore = "Allocates large model - run with --ignored"]
5451 fn test_long_audio_uses_chunking() {
5452 let _whisper = WhisperApr::tiny();
5453
5454 let long_audio = vec![0.0_f32; 60 * 16000]; assert!(long_audio.len() > WhisperApr::CHUNK_SAMPLES);
5457
5458 }
5461
5462 #[test]
5463 #[ignore = "Allocates large model - run with --ignored"]
5464 fn test_merge_overlapping_segments_empty() {
5465 let whisper = WhisperApr::tiny();
5466 let segments: Vec<Segment> = vec![];
5467 let merged = whisper.merge_overlapping_segments(segments);
5468 assert!(merged.is_empty());
5469 }
5470
5471 #[test]
5472 #[ignore = "Allocates large model - run with --ignored"]
5473 fn test_merge_overlapping_segments_single() {
5474 let whisper = WhisperApr::tiny();
5475 let segments = vec![Segment {
5476 start: 0.0,
5477 end: 5.0,
5478 text: "Hello".to_string(),
5479 tokens: vec![1, 2, 3],
5480 }];
5481 let merged = whisper.merge_overlapping_segments(segments);
5482 assert_eq!(merged.len(), 1);
5483 assert_eq!(merged[0].text, "Hello");
5484 }
5485
5486 #[test]
5487 #[ignore = "Allocates large model - run with --ignored"]
5488 fn test_merge_overlapping_segments_no_overlap() {
5489 let whisper = WhisperApr::tiny();
5490 let segments = vec![
5491 Segment {
5492 start: 0.0,
5493 end: 5.0,
5494 text: "Hello".to_string(),
5495 tokens: vec![1],
5496 },
5497 Segment {
5498 start: 10.0,
5499 end: 15.0,
5500 text: "World".to_string(),
5501 tokens: vec![2],
5502 },
5503 ];
5504 let merged = whisper.merge_overlapping_segments(segments);
5505 assert_eq!(merged.len(), 2);
5506 assert_eq!(merged[0].text, "Hello");
5507 assert_eq!(merged[1].text, "World");
5508 }
5509
5510 #[test]
5511 #[ignore = "Allocates large model - run with --ignored"]
5512 fn test_merge_overlapping_segments_with_overlap() {
5513 let whisper = WhisperApr::tiny();
5514 let segments = vec![
5515 Segment {
5516 start: 0.0,
5517 end: 5.0,
5518 text: "Hello".to_string(),
5519 tokens: vec![1],
5520 },
5521 Segment {
5522 start: 4.9, end: 10.0,
5524 text: "World".to_string(),
5525 tokens: vec![2],
5526 },
5527 ];
5528 let merged = whisper.merge_overlapping_segments(segments);
5529 assert_eq!(merged.len(), 1);
5530 assert_eq!(merged[0].text, "Hello World");
5531 assert_eq!(merged[0].start, 0.0);
5532 assert_eq!(merged[0].end, 10.0);
5533 assert_eq!(merged[0].tokens, vec![1, 2]);
5534 }
5535
5536 #[test]
5537 #[ignore = "Allocates large model - run with --ignored"]
5538 fn test_merge_overlapping_segments_multiple_overlaps() {
5539 let whisper = WhisperApr::tiny();
5540 let segments = vec![
5541 Segment {
5542 start: 0.0,
5543 end: 5.0,
5544 text: "A".to_string(),
5545 tokens: vec![1],
5546 },
5547 Segment {
5548 start: 4.95,
5549 end: 10.0,
5550 text: "B".to_string(),
5551 tokens: vec![2],
5552 },
5553 Segment {
5554 start: 9.95,
5555 end: 15.0,
5556 text: "C".to_string(),
5557 tokens: vec![3],
5558 },
5559 Segment {
5560 start: 20.0, end: 25.0,
5562 text: "D".to_string(),
5563 tokens: vec![4],
5564 },
5565 ];
5566 let merged = whisper.merge_overlapping_segments(segments);
5567 assert_eq!(merged.len(), 2);
5568 assert_eq!(merged[0].text, "A B C");
5569 assert_eq!(merged[0].end, 15.0);
5570 assert_eq!(merged[1].text, "D");
5571 assert_eq!(merged[1].start, 20.0);
5572 }
5573
5574 #[test]
5575 fn test_chunk_boundary_calculation() {
5576 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);
5583 assert_eq!(chunk1_end, 35 * 16000);
5584
5585 let chunk2_start = chunk_size; let chunk2_end = (chunk2_start + chunk_size + overlap).min(total_samples);
5588 assert_eq!(chunk2_end, 65 * 16000);
5589
5590 let chunk3_start = 2 * chunk_size; let chunk3_end = (chunk3_start + chunk_size + overlap).min(total_samples);
5593 assert_eq!(chunk3_end, total_samples);
5594 }
5595
5596 #[test]
5597 #[ignore = "Allocates large model - run with --ignored"]
5598 fn test_falsification_point_25_long_audio_full_transcription() {
5599 let _whisper = WhisperApr::tiny();
5604
5605 let ten_minutes_samples = 10 * 60 * 16000;
5607 assert!(ten_minutes_samples > WhisperApr::CHUNK_SAMPLES);
5608
5609 let expected_chunks =
5611 (ten_minutes_samples + WhisperApr::CHUNK_SAMPLES - 1) / WhisperApr::CHUNK_SAMPLES;
5612 assert_eq!(expected_chunks, 20); let mut offset = 0;
5616 let mut chunk_count = 0;
5617 while offset < ten_minutes_samples {
5618 let chunk_end = (offset + WhisperApr::CHUNK_SAMPLES + WhisperApr::OVERLAP_SAMPLES)
5619 .min(ten_minutes_samples);
5620 let _chunk_len = chunk_end - offset;
5621 offset += WhisperApr::CHUNK_SAMPLES;
5622 chunk_count += 1;
5623 }
5624 assert_eq!(chunk_count, 20);
5625 }
5626
5627 #[test]
5628 fn test_falsification_point_30_streaming_consistency() {
5629 let overlap_seconds = WhisperApr::OVERLAP_SAMPLES as f32 / 16000.0;
5635 assert!((overlap_seconds - 5.0).abs() < 0.01);
5636
5637 let expected_overlap_words = overlap_seconds * 2.5;
5640 assert!(expected_overlap_words >= 10.0);
5641 }
5642
5643 #[test]
5648 #[ignore = "Allocates large model - run with --ignored"]
5649 fn test_moonshine_tiny_constructs() {
5650 let model = WhisperApr::moonshine_tiny();
5651 assert_eq!(model.config.n_audio_state, 288);
5652 assert_eq!(model.config.n_text_state, 288);
5653 assert_eq!(model.config.n_audio_layer, 6);
5654 assert_eq!(model.config.n_text_layer, 6);
5655 assert_eq!(model.config.n_audio_head, 8);
5656 assert_eq!(model.config.n_text_head, 8);
5657 }
5658
5659 #[test]
5660 #[ignore = "Allocates large model - run with --ignored"]
5661 fn test_moonshine_tiny_encoder_has_moonshine_blocks() {
5662 let model = WhisperApr::moonshine_tiny();
5663 assert!(model.encoder.moonshine_blocks().len() > 0);
5665 assert!(model.encoder.blocks().is_empty());
5666 assert!(model.encoder.rope().is_some());
5667 }
5668
5669 #[test]
5670 #[ignore = "Allocates large model - run with --ignored"]
5671 fn test_moonshine_tiny_decoder_has_moonshine_blocks() {
5672 let model = WhisperApr::moonshine_tiny();
5673 assert!(model.decoder.moonshine_blocks().len() > 0);
5675 assert!(model.decoder.blocks().is_empty());
5676 assert!(model.decoder.rope().is_some());
5677 }
5678
5679 #[test]
5680 #[ignore = "Allocates large model - run with --ignored"]
5681 fn test_moonshine_tiny_has_sentencepiece_tokenizer() {
5682 let model = WhisperApr::moonshine_tiny();
5683 match model.tokenizer() {
5684 tokenizer::Tokenizer::SentencePiece(_) => {} tokenizer::Tokenizer::Bpe(_) => panic!("Moonshine should use SentencePiece tokenizer"),
5686 }
5687 }
5688
5689 #[test]
5690 #[ignore = "Allocates large model - run with --ignored"]
5691 fn test_moonshine_tiny_has_conv_stem() {
5692 let model = WhisperApr::moonshine_tiny();
5693 assert!(model.conv_stem.is_some());
5694 assert!(model.mel_filters.is_none());
5695 }
5696
5697 #[test]
5698 #[ignore = "Allocates large model - run with --ignored"]
5699 fn test_moonshine_tiny_initial_tokens() {
5700 let model = WhisperApr::moonshine_tiny();
5701 let tokens = model.get_initial_tokens("en", crate::Task::Transcribe);
5702 assert_eq!(tokens, vec![1]);
5704 }
5705
5706 #[test]
5707 #[ignore = "Allocates large model - run with --ignored"]
5708 fn test_moonshine_tiny_eot_token() {
5709 let model = WhisperApr::moonshine_tiny();
5710 assert_eq!(model.eot_token(), 2); }
5712
5713 #[test]
5714 #[ignore = "Allocates large model - run with --ignored"]
5715 fn test_moonshine_encoder_forward_shape() {
5716 let model = WhisperApr::moonshine_tiny();
5717 let d_model = 288;
5718 let seq_len = 7; let features = vec![0.1_f32; seq_len * d_model];
5722 let output = model.encoder.forward(&features).expect("encoder forward");
5723
5724 assert_eq!(output.len(), seq_len * d_model);
5726 assert!(output.iter().all(|v| v.is_finite()));
5727 }
5728
5729 #[test]
5730 #[ignore = "Allocates large model - run with --ignored"]
5731 fn test_moonshine_decoder_forward_shape() {
5732 let model = WhisperApr::moonshine_tiny();
5733 let d_model = 288;
5734 let n_vocab = model.config.n_vocab as usize;
5735
5736 let enc_seq_len = 7;
5738 let encoder_output = vec![0.1_f32; enc_seq_len * d_model];
5739
5740 let tokens = vec![1_u32, 100];
5742 let logits = model
5743 .decoder
5744 .forward(&tokens, &encoder_output)
5745 .expect("decoder forward");
5746
5747 assert_eq!(logits.len(), tokens.len() * n_vocab);
5749 assert!(logits.iter().all(|v| v.is_finite()));
5750 }
5751
5752 #[test]
5753 #[ignore = "Allocates large model - run with --ignored"]
5754 fn test_moonshine_variable_length_input() {
5755 let model = WhisperApr::moonshine_tiny();
5756 let d_model = 288;
5757
5758 for seq_len in [3, 7, 15, 31] {
5760 let features = vec![0.1_f32; seq_len * d_model];
5761 let output = model.encoder.forward(&features).expect("encoder forward");
5762 assert_eq!(
5763 output.len(),
5764 seq_len * d_model,
5765 "seq_len={seq_len} output mismatch"
5766 );
5767 }
5768 }
5769
5770 #[test]
5771 #[ignore = "Allocates large model - run with --ignored"]
5772 fn test_whisper_tiny_unaffected() {
5773 let model = WhisperApr::tiny();
5775 assert!(model.encoder.blocks().len() > 0);
5776 assert!(model.encoder.moonshine_blocks().is_empty());
5777 assert!(model.encoder.rope().is_none());
5778 assert!(model.decoder.blocks().len() > 0);
5779 assert!(model.decoder.moonshine_blocks().is_empty());
5780 assert!(model.decoder.rope().is_none());
5781 assert!(model.mel_filters.is_some());
5782 assert!(model.conv_stem.is_none());
5783 match model.tokenizer() {
5784 tokenizer::Tokenizer::Bpe(_) => {} tokenizer::Tokenizer::SentencePiece(_) => {
5786 panic!("Whisper should use BPE tokenizer")
5787 }
5788 }
5789 }
5790
5791 #[test]
5796 fn test_moonshine_encoder_uses_rmsnorm_final() {
5797 let config = model::ModelConfig::moonshine_tiny();
5799 let encoder = model::Encoder::new(&config);
5800 assert!(
5801 encoder.ln_post_rms().is_some(),
5802 "Moonshine encoder must use RmsNorm for final layer norm"
5803 );
5804 }
5805
5806 #[test]
5807 fn test_whisper_encoder_uses_layernorm_final() {
5808 let config = model::ModelConfig::tiny();
5810 let encoder = model::Encoder::new(&config);
5811 assert!(
5812 encoder.ln_post_rms().is_none(),
5813 "Whisper encoder must use LayerNorm, not RmsNorm"
5814 );
5815 }
5816
5817 #[test]
5818 #[ignore = "Allocates large model - run with --ignored"]
5819 fn test_moonshine_decoder_forward_one_sets_cross_attn_cached() {
5820 let model = WhisperApr::moonshine_tiny();
5822 let d_model = 288;
5823 let n_layers = model.decoder.n_layers();
5824 let max_len = model.decoder.max_len();
5825 let enc_seq_len = 7;
5826 let encoder_output = vec![0.1_f32; enc_seq_len * d_model];
5827
5828 let mut cache = model::DecoderKVCache::new(n_layers, d_model, max_len);
5829 assert!(!cache.cross_attn_cached);
5830
5831 let _logits = model
5833 .decoder
5834 .forward_one(1, &encoder_output, &mut cache)
5835 .expect("forward_one");
5836
5837 assert!(
5838 cache.cross_attn_cached,
5839 "cross_attn_cached must be true after forward_one for Moonshine"
5840 );
5841 }
5842
5843 #[test]
5844 #[ignore = "Allocates large model - run with --ignored"]
5845 fn test_moonshine_decoder_forward_one_cache_has_layers() {
5846 let config = model::ModelConfig::moonshine_tiny();
5848 let decoder = model::Decoder::new(&config);
5849
5850 let cache =
5852 model::DecoderKVCache::new(decoder.n_layers(), decoder.d_model(), decoder.max_len());
5853 assert!(
5854 !cache.self_attn_cache.is_empty(),
5855 "Moonshine decoder cache must have at least 1 layer"
5856 );
5857 }
5858
5859 #[test]
5860 fn test_moonshine_decoder_is_finalized() {
5861 let config = model::ModelConfig::moonshine_tiny();
5863 let decoder = model::Decoder::new(&config);
5864 assert!(
5866 decoder.is_finalized(),
5867 "Moonshine decoder should be considered finalized"
5868 );
5869 }
5870
5871 #[test]
5872 fn test_whisper_decoder_is_finalized_initially_false() {
5873 let config = model::ModelConfig::tiny();
5875 let decoder = model::Decoder::new(&config);
5876 assert!(
5878 !decoder.is_finalized(),
5879 "Whisper decoder should NOT be finalized before finalize_weights()"
5880 );
5881 }
5882
5883 #[test]
5884 #[ignore = "Allocates large model - run with --ignored"]
5885 fn test_moonshine_decoder_forward_traced() {
5886 let model = WhisperApr::moonshine_tiny();
5888 let d_model = 288;
5889 let enc_seq_len = 7;
5890 let encoder_output = vec![0.1_f32; enc_seq_len * d_model];
5891
5892 let (logits, trace) = model
5893 .decoder
5894 .forward_traced(&[1, 2], &encoder_output)
5895 .expect("forward_traced");
5896
5897 let layer_traces: Vec<_> = trace
5899 .iter()
5900 .filter(|(name, _)| name.starts_with("layer_"))
5901 .collect();
5902 assert!(
5903 !layer_traces.is_empty(),
5904 "forward_traced must produce layer traces for Moonshine"
5905 );
5906
5907 assert!(!logits.is_empty());
5908 assert!(logits.iter().all(|v| v.is_finite()));
5909 }
5910
5911 #[test]
5912 fn test_moonshine_encoder_forward_batch_uses_forward() {
5913 let config = model::ModelConfig::moonshine_tiny();
5915 let encoder = model::Encoder::new(&config);
5916 let d_model = 288;
5917
5918 let batch = vec![vec![0.1_f32; 7 * d_model], vec![0.1_f32; 5 * d_model]];
5920 let results = encoder.forward_batch(&batch).expect("forward_batch");
5921 assert_eq!(results.len(), 2);
5922 assert_eq!(results[0].len(), 7 * d_model);
5923 assert_eq!(results[1].len(), 5 * d_model);
5924 }
5925}