|
import torch
|
|
from transformers import AutoTokenizer, AutoModel, Wav2Vec2FeatureExtractor
|
|
|
|
class AstralQuantizer(torch.nn.Module):
|
|
def __init__(
|
|
self,
|
|
tokenizer_name: str,
|
|
ssl_model_name: str,
|
|
ssl_output_layer: int,
|
|
encoder: torch.nn.Module,
|
|
quantizer: torch.nn.Module,
|
|
skip_ssl: bool = False,
|
|
):
|
|
super().__init__()
|
|
self.encoder = encoder
|
|
self.quantizer = quantizer
|
|
self.tokenizer_name = tokenizer_name
|
|
self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
|
|
|
|
|
|
self.ssl_model_name = ssl_model_name
|
|
self.ssl_output_layer = ssl_output_layer
|
|
self.ssl_feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(ssl_model_name)
|
|
|
|
if skip_ssl:
|
|
self.ssl_model = None
|
|
else:
|
|
self.ssl_model = AutoModel.from_pretrained(ssl_model_name).eval()
|
|
self.ssl_model.encoder.layers = self.ssl_model.encoder.layers[:ssl_output_layer]
|
|
self.ssl_model.encoder.layer_norm = torch.nn.Identity()
|
|
|
|
def load_separate_checkpoint(self, checkpoint_path):
|
|
params = torch.load(checkpoint_path, map_location='cpu')['net']
|
|
for key in params.keys():
|
|
for k in list(params[key].keys()):
|
|
if k.startswith("module."):
|
|
params[key][k[len("module."):]] = params[key][k]
|
|
del params[key][k]
|
|
self.encoder.load_state_dict(params['encoder'])
|
|
self.quantizer.load_state_dict(params['vq'])
|
|
if self.decoder is not None:
|
|
self.decoder.load_state_dict(params['decoder'])
|
|
if self.asr_decoder is not None:
|
|
self.asr_decoder.load_state_dict(params['predictor'], strict=False)
|
|
|
|
def forward(self, waves_16k, wave_16k_lens, ssl_model=None):
|
|
ssl_fn = self.ssl_model if self.ssl_model else ssl_model
|
|
assert ssl_fn is not None, "In case in-class SSL model loading is skipped, external ssl_model must be provided"
|
|
waves_16k_input_list = [
|
|
waves_16k[bib, :wave_16k_lens[bib]].cpu().numpy()
|
|
for bib in range(len(waves_16k))
|
|
]
|
|
alt_inputs = self.ssl_feature_extractor(
|
|
waves_16k_input_list,
|
|
return_tensors='pt',
|
|
return_attention_mask=True,
|
|
padding=True,
|
|
sampling_rate=16000
|
|
).to(waves_16k.device)
|
|
feature_lens = alt_inputs.data['attention_mask'].sum(-1) // 320
|
|
|
|
outputs = ssl_fn(
|
|
alt_inputs.input_values,
|
|
attention_mask=alt_inputs.attention_mask,
|
|
)
|
|
last_hidden_states = outputs.last_hidden_state
|
|
last_hidden_states = last_hidden_states[:, :feature_lens.max(), :]
|
|
feature_lens = feature_lens.clamp(max=last_hidden_states.size(1))
|
|
last_hidden_states = last_hidden_states.transpose(1, 2)
|
|
x_hidden = self.encoder(last_hidden_states, feature_lens)
|
|
x_hidden = x_hidden.transpose(1, 2)
|
|
x_quantized, indices = self.quantizer(x_hidden)[:2]
|
|
return x_quantized, indices, feature_lens |
