(ort_session, wav_path, pool_executor=None)
| 19 | import torch |
| 20 | |
| 21 | def extract_speech_token(ort_session, wav_path, pool_executor=None): |
| 22 | def tokenizer(audio_segment): |
| 23 | feat = whisper.log_mel_spectrogram(audio_segment, n_mels=128) |
| 24 | speech_token = ort_session.run(None, {ort_session.get_inputs()[0].name: feat.detach().cpu().numpy(), |
| 25 | ort_session.get_inputs()[1].name: np.array([feat.shape[2]], dtype=np.int32)})[ |
| 26 | 0].flatten().tolist() |
| 27 | return speech_token |
| 28 | |
| 29 | if isinstance(wav_path, str): |
| 30 | audio, sample_rate = torchaudio.load(wav_path, backend='soundfile') |
| 31 | else: |
| 32 | audio = wav_path |
| 33 | sample_rate = 16000 |
| 34 | if sample_rate != 16000: |
| 35 | audio = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=16000)(audio) |
| 36 | # print(f"audio.shape: {audio.shape}") |
| 37 | if audio.shape[0] > 1: |
| 38 | audio = audio.mean(dim=0, keepdim=True) |
| 39 | |
| 40 | time_step = 0 |
| 41 | audios =[] |
| 42 | |
| 43 | # Step 1: Split audio into 30-second segments |
| 44 | while time_step * 16000 < audio.shape[1]: |
| 45 | start = time_step * 16000 |
| 46 | end = min((time_step + 30) * 16000, audio.shape[1]) |
| 47 | audio_segment = audio[:, start:end] |
| 48 | audios.append(audio_segment) |
| 49 | time_step += 30 |
| 50 | |
| 51 | # Step 2: Handle last segment if too short |
| 52 | if len(audios) > 1 and audios[-1].shape[1] < 16000: # Less than 1 second |
| 53 | # Remove last two segments |
| 54 | last_segment = audios.pop() |
| 55 | second_last_segment = audios.pop() |
| 56 | |
| 57 | # Merge last two segments |
| 58 | merged_audio = torch.cat([second_last_segment, last_segment], dim=1) |
| 59 | total_length = merged_audio.shape[1] |
| 60 | |
| 61 | # Split merged audio into two equal parts |
| 62 | split_point = total_length // 2 |
| 63 | first_half = merged_audio[:, :split_point] |
| 64 | second_half = merged_audio[:, split_point:] |
| 65 | |
| 66 | # Add new segments back to list |
| 67 | audios.append(first_half) |
| 68 | audios.append(second_half) |
| 69 | |
| 70 | all_speech_tokens = [] |
| 71 | if pool_executor is not None: |
| 72 | # 提交所有任务 |
| 73 | futures = [ |
| 74 | pool_executor.submit( |
| 75 | tokenizer, |
| 76 | item, |
| 77 | ) |
| 78 | for item in audios |
no test coverage detected