test_online_whisper_fbank.py (1396B)
1 #!/usr/bin/env python3 2 # 3 # Copyright (c) 2023 Xiaomi Corporation (authors: Fangjun Kuang) 4 5 import pickle 6 7 import kaldi_native_fbank as knf 8 import torch 9 10 11 def test(): 12 opts = knf.WhisperFeatureOptions() 13 14 # Use 128 for whisper large v3 15 opts.dim = 128 16 online_whisper_fbank = knf.OnlineWhisperFbank(opts) 17 18 audio = torch.rand(100000) 19 # audio should be normalized into the range [-1, 1] 20 print(audio.shape, audio.max(), audio.min()) 21 online_whisper_fbank.accept_waveform(sampling_rate=16000, waveform=audio.numpy()) 22 online_whisper_fbank.input_finished() 23 print(online_whisper_fbank.num_frames_ready) 24 25 features = [] 26 for i in range(online_whisper_fbank.num_frames_ready): 27 f = online_whisper_fbank.get_frame(i) 28 f = torch.from_numpy(f) 29 features.append(f) 30 31 features = torch.stack(features) 32 print(features.shape) 33 34 log_spec = torch.clamp(features, min=1e-10).log10() 35 log_spec = torch.maximum(log_spec, log_spec.max() - 8.0) 36 mel = (log_spec + 4.0) / 4.0 37 target = 3000 38 mel = torch.nn.functional.pad(mel, (0, 0, 0, target - mel.shape[0]), "constant", 0) 39 mel = mel.t().unsqueeze(0) 40 print(mel.shape) 41 42 assert mel.shape == (1, opts.dim, 3000), mel.shape 43 # Now you can input 'mel' to whisper.encoder model 44 45 46 def main(): 47 test() 48 49 50 if __name__ == "__main__": 51 torch.manual_seed(20240112) 52 main()