|
import os |
|
import io |
|
import glob |
|
import math |
|
import tarfile |
|
import torch |
|
import torchaudio |
|
import safetensors |
|
from .configuration_whisper import WhisperVQConfig |
|
from .modeling_whisper import WhisperVQEncoder, WhisperVQForConditionalGeneration |
|
from transformers import WhisperFeatureExtractor, WhisperTokenizerFast |
|
|
|
|
|
def load_quantize_encoder(model_path): |
|
config = WhisperVQConfig.from_pretrained(model_path) |
|
config.quantize_encoder_only = True |
|
model = WhisperVQEncoder(config) |
|
state_dict = {} |
|
for path in glob.glob(os.path.join(model_path, "model*.safetensors")): |
|
with safetensors.safe_open(path, framework="pt", device="cpu") as f: |
|
for key in f.keys(): |
|
if key.startswith("model.encoder."): |
|
new_key = key[len("model.encoder."):] |
|
if new_key.startswith("layer_norm"): |
|
continue |
|
if new_key.startswith("layers"): |
|
layer_id = int(new_key.split(".")[1]) |
|
if layer_id >= config.quantize_position: |
|
continue |
|
state_dict[new_key] = f.get_tensor(key) |
|
model.load_state_dict(state_dict) |
|
model.eval() |
|
model.cuda() |
|
return model |
|
|
|
|
|
_resample_buffer: dict[int, torchaudio.transforms.Resample] = {} |
|
|
|
|
|
def extract_speech_token(model: WhisperVQEncoder, feature_extractor: WhisperFeatureExtractor, utts): |
|
with torch.no_grad(): |
|
audios, indices = [], [] |
|
for idx, utt in enumerate(utts): |
|
if isinstance(utt, tuple): |
|
audio, sample_rate = utt |
|
else: |
|
audio, sample_rate = torchaudio.load(utt) |
|
audio = audio.cuda() |
|
if sample_rate != 16000: |
|
if sample_rate not in _resample_buffer: |
|
_resample_buffer[sample_rate] = torchaudio.transforms.Resample( |
|
orig_freq=sample_rate, |
|
new_freq=16000 |
|
).to('cuda') |
|
audio = _resample_buffer[sample_rate](audio) |
|
|
|
|
|
audio = audio[0] |
|
audio = audio.cpu().numpy() |
|
time_step = 0 |
|
while time_step * 16000 < audio.shape[0]: |
|
audio_segment = audio[time_step * 16000: (time_step + 30) * 16000] |
|
audios.append(audio_segment) |
|
indices.append(idx) |
|
time_step += 30 |
|
pooling_kernel_size = model.config.pooling_kernel_size or 1 |
|
stride = model.conv1.stride[0] * model.conv2.stride[0] * pooling_kernel_size * feature_extractor.hop_length |
|
all_speech_tokens = [[] for _ in range(len(utts))] |
|
batch_size = 128 |
|
for start in range(0, len(audios), batch_size): |
|
features = feature_extractor(audios[start: start + batch_size], sampling_rate=16000, |
|
return_attention_mask=True, return_tensors="pt", device='cuda', |
|
padding="longest", pad_to_multiple_of=stride) |
|
features = features.to(device="cuda") |
|
outputs = model(**features) |
|
speech_tokens = outputs.quantized_token_ids |
|
attention_mask = features.attention_mask[:, ::model.conv1.stride[0] * model.conv2.stride[0]] |
|
attention_mask = attention_mask[:, ::model.config.pooling_kernel_size] |
|
assert attention_mask.shape == speech_tokens.shape |
|
for i in range(len(speech_tokens)): |
|
idx = indices[start + i] |
|
speech_token = speech_tokens[i][attention_mask[i].bool()].tolist() |
|
all_speech_tokens[idx].extend(speech_token) |
|
return all_speech_tokens |
|
|
