| | 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 |
