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