diff --git a/include/speech_core/models/moss_whisper_features.h b/include/speech_core/models/moss_whisper_features.h index 1e5e923..fe49e05 100644 --- a/include/speech_core/models/moss_whisper_features.h +++ b/include/speech_core/models/moss_whisper_features.h @@ -13,7 +13,9 @@ struct MossLogMelFeatures { int time_frames = 0; }; -/// Exact 400-point Whisper frontend used by MOSS-Transcribe-Diarize. +/// Exact 400-point Whisper frontend used by MOSS-Transcribe-Diarize and other +/// Whisper-family models. MOSS uses the default 80 mel bins; Whisper large-v3 +/// uses 128. The FFT and normalization contract is otherwise identical. /// /// The ordinary speech-core FFT helper intentionally pads non-power-of-two /// transforms. MOSS was trained and exported with a true 400-point DFT, so @@ -28,16 +30,17 @@ class MossWhisperFeatureExtractor { static constexpr int kTimeFrames = 3000; static constexpr int kEncoderStrideSamples = 1280; - MossWhisperFeatureExtractor(); + explicit MossWhisperFeatureExtractor(int mel_bins = kMelBins); static std::size_t audio_token_count(std::size_t sample_count); - /// Extract exactly `[80, 3000]` features from one non-empty, at-most + /// Extract exactly `[mel_bins, 3000]` features from one non-empty, at-most /// 30-second 16 kHz mono chunk. Short chunks are right-padded with zeroes. MossLogMelFeatures extract_padded_chunk( const float* audio, std::size_t length) const; private: + int mel_bins_ = kMelBins; std::vector hann_window_; /// Row-major `[mel_bins, fft_bins]`. std::vector mel_filterbank_; diff --git a/src/models/moss/moss_whisper_features.cpp b/src/models/moss/moss_whisper_features.cpp index ef842c0..3c24205 100644 --- a/src/models/moss/moss_whisper_features.cpp +++ b/src/models/moss/moss_whisper_features.cpp @@ -42,7 +42,11 @@ double slaney_mel_to_hertz(double mel) { } // namespace -MossWhisperFeatureExtractor::MossWhisperFeatureExtractor() { +MossWhisperFeatureExtractor::MossWhisperFeatureExtractor(int mel_bins) + : mel_bins_(mel_bins) { + if (mel_bins_ <= 0) { + throw std::invalid_argument("Whisper mel-bin count must be positive"); + } hann_window_.resize(kFftSize); for (int index = 0; index < kFftSize; ++index) { hann_window_[static_cast(index)] = static_cast( @@ -52,21 +56,21 @@ MossWhisperFeatureExtractor::MossWhisperFeatureExtractor() { } constexpr int kFrequencyBins = kFftSize / 2 + 1; - std::vector points(kMelBins + 2); + std::vector points(static_cast(mel_bins_ + 2)); const double mel_minimum = hertz_to_slaney_mel(0.0); const double mel_maximum = hertz_to_slaney_mel(static_cast(kSampleRate) / 2.0); - for (int index = 0; index < kMelBins + 2; ++index) { + for (int index = 0; index < mel_bins_ + 2; ++index) { const double mel = mel_minimum + static_cast(index) * (mel_maximum - mel_minimum) - / static_cast(kMelBins + 1); + / static_cast(mel_bins_ + 1); points[static_cast(index)] = slaney_mel_to_hertz(mel); } mel_filterbank_.assign( - static_cast(kMelBins * kFrequencyBins), 0.0f); - for (int mel = 0; mel < kMelBins; ++mel) { + static_cast(mel_bins_ * kFrequencyBins), 0.0f); + for (int mel = 0; mel < mel_bins_; ++mel) { const double left = points[static_cast(mel)]; const double center = points[static_cast(mel + 1)]; const double right = points[static_cast(mel + 2)]; @@ -95,11 +99,11 @@ std::size_t MossWhisperFeatureExtractor::audio_token_count( MossLogMelFeatures MossWhisperFeatureExtractor::extract_padded_chunk( const float* audio, std::size_t length) const { if (!audio || length == 0) { - throw std::invalid_argument("MOSS audio chunk is empty"); + throw std::invalid_argument("Whisper audio chunk is empty"); } if (length > static_cast(kChunkSamples)) { throw std::invalid_argument( - "MOSS audio chunk exceeds the 30-second encoder input"); + "Whisper audio chunk exceeds the 30-second encoder input"); } std::vector fixed_audio( @@ -126,7 +130,7 @@ MossLogMelFeatures MossWhisperFeatureExtractor::extract_padded_chunk( KissFftrPlan plan( kiss_fftr_alloc(kFftSize, /*inverse_fft=*/0, nullptr, nullptr)); if (!plan) { - throw std::runtime_error("MOSS KISS FFT plan allocation failed"); + throw std::runtime_error("Whisper KISS FFT plan allocation failed"); } constexpr int kFrequencyBins = kFftSize / 2 + 1; @@ -135,7 +139,7 @@ MossLogMelFeatures MossWhisperFeatureExtractor::extract_padded_chunk( static_cast(kFrequencyBins)); std::vector power(static_cast(kFrequencyBins)); std::vector mel_by_frame( - static_cast(kTimeFrames * kMelBins)); + static_cast(kTimeFrames * mel_bins_)); for (int time = 0; time < kTimeFrames; ++time) { const std::size_t start = @@ -151,7 +155,7 @@ MossLogMelFeatures MossWhisperFeatureExtractor::extract_padded_chunk( power[static_cast(bin)] = value.r * value.r + value.i * value.i; } - for (int mel = 0; mel < kMelBins; ++mel) { + for (int mel = 0; mel < mel_bins_; ++mel) { const float* filter = mel_filterbank_.data() + static_cast(mel * kFrequencyBins); double sum = 0.0; @@ -161,7 +165,7 @@ MossLogMelFeatures MossWhisperFeatureExtractor::extract_padded_chunk( * static_cast(filter[bin]); } mel_by_frame[ - static_cast(time * kMelBins + mel)] = + static_cast(time * mel_bins_ + mel)] = static_cast(std::max(sum, 1e-10)); } } @@ -174,13 +178,13 @@ MossLogMelFeatures MossWhisperFeatureExtractor::extract_padded_chunk( const float floor = peak - 8.0f; MossLogMelFeatures result; - result.mel_bins = kMelBins; + result.mel_bins = mel_bins_; result.time_frames = kTimeFrames; - result.data.resize(static_cast(kMelBins * kTimeFrames)); + result.data.resize(static_cast(mel_bins_ * kTimeFrames)); for (int time = 0; time < kTimeFrames; ++time) { - for (int mel = 0; mel < kMelBins; ++mel) { + for (int mel = 0; mel < mel_bins_; ++mel) { const float value = mel_by_frame[ - static_cast(time * kMelBins + mel)]; + static_cast(time * mel_bins_ + mel)]; result.data[ static_cast(mel * kTimeFrames + time)] = (std::max(value, floor) + 4.0f) * 0.25f; diff --git a/tests/test_moss_whisper_features.cpp b/tests/test_moss_whisper_features.cpp index f916542..e562102 100644 --- a/tests/test_moss_whisper_features.cpp +++ b/tests/test_moss_whisper_features.cpp @@ -58,6 +58,36 @@ int main() { assert(std::fabs(actual - fixture.value) <= 5e-4f); } + // Whisper large-v3 uses the same exact 400-point frontend with 128 mel + // filters. These probes come from transformers 4.57.6 + // WhisperFeatureExtractor(feature_size=128) on the same deterministic + // waveform. They prevent a regression to speech_core::audio::fft_real, + // whose documented non-power-of-two behavior pads 400 to 512 and changes + // the model input materially. + const auto large_v3 = + MossWhisperFeatureExtractor(128).extract_padded_chunk( + audio.data(), audio.size()); + assert(large_v3.mel_bins == 128); + assert(large_v3.time_frames == 3000); + assert(large_v3.data.size() == 384000); + const std::vector large_v3_fixtures = { + {0, 0, 1.03680742f}, + {0, 10, 0.92322218f}, + {5, 50, 0.85368174f}, + {20, 10, 0.60412019f}, + {40, 99, 0.39996243f}, + {79, 0, -0.72825670f}, + {100, 100, 0.16020703f}, + {127, 0, -0.86562812f}, + {127, 2999, -0.86562812f}, + }; + for (const auto& fixture : large_v3_fixtures) { + const float actual = large_v3.data[ + static_cast( + fixture.mel * large_v3.time_frames + fixture.frame)]; + assert(std::fabs(actual - fixture.value) <= 5e-4f); + } + std::cout << "MOSS Whisper frontend tests passed\n"; return 0; }