from transformers import PreTrainedModel, PretrainedConfig, Wav2Vec2ForCTC import json import torch from torch import nn from torch.nn.utils.rnn import pad_sequence import math from typing import Optional # x: torch.FloatTensor [T, B, D] # mask: torch.BoolTensor [B, T], where True indicates padding # returns: torch.LongTensor [B] def get_lengths(x, mask=None): if mask is not None: return (~mask).long().sum(dim=1) else: return torch.LongTensor([x.size(0)] * x.size(1)).to(x.device) # lens: torch.LongTensor [B] # returns: torch.BoolTensor [B, max_lens], where True indicates padding def lengths_to_padding_mask(lens): bsz, max_lens = lens.size(0), torch.max(lens).item() mask = torch.arange(max_lens).to(lens.device).view(1, max_lens) mask = mask.expand(bsz, -1) >= lens.view(bsz, 1).expand(-1, max_lens) return mask # input_lengths: torch.LongTensor [B] def get_output_lengths(input_lengths): conv_feature_layers = "[(512, 10, 5)] + [(512, 3, 2)] * 4 + [(512,2,2)] + [(512,2,2)]" conv_cfg_list = eval(conv_feature_layers) def _conv_out_length(input_length, kernel_size, stride): return torch.floor((input_length - kernel_size) / stride + 1) for i in range(len(conv_cfg_list)): input_lengths = _conv_out_length( input_lengths, conv_cfg_list[i][1], conv_cfg_list[i][2] ) return input_lengths.to(torch.long) class ZeroSwotEncoderConfig(PretrainedConfig): model_type = "zero_swot_encoder" def __init__( self, wav2vec2_model_name_or_path="", compression_adapter=None, embed_dim=1024, **kwargs ): super().__init__(**kwargs) self.wav2vec2_model_name_or_path = wav2vec2_model_name_or_path self.compression_adapter = compression_adapter self.embed_dim = embed_dim @classmethod def from_json_file(cls, json_file): with open(json_file, "r") as reader: text = reader.read() config_dict = json.loads(text) return cls(**config_dict) class ZeroSwotEncoderModel(PreTrainedModel): config_class = ZeroSwotEncoderConfig model_type = "zero_swot_encoder" def __init__(self, config): super().__init__(config) self.wav2vec2 = Wav2Vec2ForCTC.from_pretrained(config.wav2vec2_model_name_or_path) self.compression_adapter = CompressionAdapter(config.compression_adapter) self.speech_embedder = SpeechEmbedder(config.embed_dim) def forward(self, input_values, attention_mask=None): input_lens = get_lengths(input_values, ~attention_mask) # Forward pass through wav2vec2 encoder x = self.wav2vec2.wav2vec2(input_values, attention_mask)[0] # [B, T, D] # CTC predictions preds = self.wav2vec2.lm_head(x).argmax(-1) # [B, T] # Get output lengths for x output_lens = get_output_lengths(input_lens) # Compression x, mask, _ = self.compression_adapter(x, preds, output_lens) # [B, N, D] with N << T # BOS and EOS embeddings x, mask = self.speech_embedder(x, mask) # [B, N+2, D] return x, ~mask class SpeechEmbedder(nn.Module): def __init__(self, embed_dim): super().__init__() self.embed_dim = embed_dim self.bos_emb = nn.Parameter(torch.empty(embed_dim)) self.eos_emb = nn.Parameter(torch.empty(embed_dim)) self.scale = self.embed_dim ** 0.5 def forward(self, x, padding_mask=None): """Add special embedding and positional embedding. Args: x (FloatTensor): (B, T, C) padding_mask (ByteTensor): (B, T) Outputs: x (FloatTensor): (B, T+2, C) padding_mask (ByteTensor): (B, T+2) """ B = x.size(0) lengths = get_lengths(x.transpose(0, 1), padding_mask) assert B == len(lengths) if padding_mask is not None: x = x * (1 - padding_mask.unsqueeze(-1).type_as(x)) # prepend bos x = torch.cat([self.bos_emb.view(1, 1, -1).expand(B, 1, -1), x], dim=1) lengths += 1 # append padding (zeros) and then convert first padding to eos x = torch.cat([x, torch.zeros(B, 1, x.size(-1), device=x.device, dtype=x.dtype)], dim=1) for i in range(B): x[i, lengths[i], :] = self.eos_emb lengths += 1 padding_mask = lengths_to_padding_mask(lengths) x = x * self.scale return x, padding_mask class PositionalEmbedding(nn.Module): def __init__(self, num_embeddings, embedding_dim, padding_idx): super().__init__() self.embedding_dim = embedding_dim self.padding_idx = padding_idx if padding_idx is not None else 0 num_embeddings += padding_idx + 1 self.weights = PositionalEmbedding.get_embedding( num_embeddings, embedding_dim, padding_idx ) self.register_buffer("_float_tensor", torch.FloatTensor(1)) self.max_positions = int(1e5) @staticmethod def get_embedding( num_embeddings: int, embedding_dim: int, padding_idx: Optional[int] = None ): half_dim = embedding_dim // 2 emb = math.log(10000) / (half_dim - 1) emb = torch.exp(torch.arange(half_dim, dtype=torch.float) * -emb) emb = torch.arange(num_embeddings, dtype=torch.float).unsqueeze(1) * emb.unsqueeze(0) emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1).view(num_embeddings, -1) if embedding_dim % 2 == 1: # zero pad emb = torch.cat([emb, torch.zeros(num_embeddings, 1)], dim=1) if padding_idx is not None: emb[padding_idx, :] = 0 return emb def make_positions(self, x, padding_idx: int): mask = x.ne(padding_idx).int() return (torch.cumsum(mask, dim=1).type_as(mask) * mask).long() + padding_idx def forward(self, input): """Input is expected to be of size [bsz x seqlen].""" bsz, seq_len = input.size() max_pos = self.padding_idx + 1 + seq_len if self.weights is None or max_pos > self.weights.size(0): # recompute/expand embeddings if needed self.weights = PositionalEmbedding.get_embedding( max_pos, self.embedding_dim, self.padding_idx ) self.weights = self.weights.to(self._float_tensor) positions = self.make_positions(input, self.padding_idx) return ( self.weights.index_select(0, positions.view(-1)) .view(bsz, seq_len, -1) .detach() ) class CLSPooling(nn.Module): def __init__(self, embed_dim, num_transformer_layers, dropout_rate): super().__init__() self.cls_token = nn.Parameter(torch.empty(1, 1, embed_dim)) nn.init.normal_(self.cls_token, mean=0.0, std=0.25) self.transformer = nn.TransformerEncoder( nn.TransformerEncoderLayer( embed_dim, nhead=16 if embed_dim == 1024 else 8, dim_feedforward=4*embed_dim, dropout=dropout_rate, activation="relu", batch_first=True, norm_first=True ), num_layers=num_transformer_layers, ) self.pos_emb = PositionalEmbedding(512, embed_dim, 1) self.scale = math.sqrt(embed_dim) def forward(self, x, lens): # x: [B, N, D] # lens: [B] # prepend cls token x = torch.cat( [ self.cls_token.to(dtype=x.dtype, device=x.device).repeat(x.size(0), 1, 1), # B x 1 x D x ], dim=1) # [B, N+1, D] mask = lengths_to_padding_mask(lens+1) x = x + self.pos_emb(mask.long()) / self.scale x = self.transformer(x, src_key_padding_mask=mask) # [B, N+1, D] x = x[:, 0] # [B, D] return x class CompressionAdapter(nn.Module): def __init__(self, cfg): super().__init__() self.embed_dim = cfg["embed_dim"] self.transformer_layers = cfg["transformer_layers"] self.dropout = cfg["dropout"] self.blank_idx = cfg["blank_idx"] self.sep_idx = cfg["sep_idx"] self.token_pooling_module = CLSPooling( self.embed_dim, self.transformer_layers, self.dropout ) def char_compression(self, x, preds, lens): # x: B x T x D # preds: B x T # lens: B B, T, D = x.size() device = x.device dtype = x.dtype # zero-out the padding mask = lengths_to_padding_mask(lens) # B x T x = x.masked_fill(mask.unsqueeze(-1), 0) preds = preds.masked_fill(mask, self.blank_idx) # add a vector of -1 to know where each example ends after flattening the batch preds = torch.cat([-torch.ones(B, 1, device=device, dtype=torch.long), preds], dim=1).view(-1) x = torch.cat([torch.zeros(B, 1, D, device=device, dtype=dtype), x], dim=1).view(-1, D) # get points of consecutive preds preds, counts = preds.unique_consecutive(return_counts=True) # split in representations of same chars x = torch.split(x, counts.tolist()) # remove blanks valid_mask = preds != self.blank_idx preds = preds[valid_mask] counts = counts[valid_mask] # [N] x = [x_i for x_i, v_i in zip(x, valid_mask) if v_i] # pack into tensor x = pad_sequence(x, batch_first=True, padding_value=0) # char pooling x = torch.sum(x, dim=1) / counts.to(dtype=x.dtype).unsqueeze(1) # [B, N, D] -> [B, D] # find split points for retrieving the examples split_points = (preds == -1).nonzero(as_tuple=True)[0] split_points = torch.cat([split_points, torch.tensor([len(preds)], device=device)]) split_points = (split_points[1:] - split_points[:-1]).tolist() # split into examples x = torch.split(x, split_points) preds = torch.split(preds, split_points) lens = torch.tensor([len(x_i) for x_i in x], device=device) # pack into tensors x = pad_sequence(x, batch_first=True, padding_value=0) preds = pad_sequence(preds, batch_first=True, padding_value=self.blank_idx) # remove the parts we add to identify the bounds for each example x = x[:, 1:] preds = preds[:, 1:] lens -= 1 mask = lengths_to_padding_mask(lens) # account for empty examples (just a sep token) empty_examples = lens == 0 num_empty_examples = empty_examples.sum() if num_empty_examples > 0: mask[empty_examples, 0] = True lens[empty_examples] = 1 preds[empty_examples, 0] = self.sep_idx return x, mask, lens, preds, num_empty_examples def token_compression(self, x, preds, lens): # x: B x T x D # preds: B x T # lens: B B, T, D = x.size() device = x.device dtype = x.dtype # new lengths after compression new_lens = preds.eq(self.sep_idx).sum(dim=1) # unpad and unpack to list of tensors preds = [preds[i, :lens[i]] for i in range(B)] x = [x[i, :lens[i]] for i in range(B)] # make sure every example ends with a separator num_examples_without_ending_sep = torch.tensor(0, device=device, dtype=torch.long) for i in range(B): if preds[i][-1] != self.sep_idx: preds[i] = torch.cat([preds[i], torch.tensor([self.sep_idx], device=device, dtype=torch.long)]) x[i] = torch.cat([x[i], torch.zeros(1, D, device=device, dtype=dtype)]) new_lens[i] += 1 num_examples_without_ending_sep += 1 # flatten preds = torch.cat(preds) x = torch.cat(x) # split points according to separators split_points = preds.eq(self.sep_idx).nonzero(as_tuple=True)[0] + 1 split_points = torch.cat([torch.tensor([0], device=device, dtype=torch.long), split_points]) split_points = (split_points[1:] - split_points[:-1]).tolist() # re-arrange in 3d [total_num_tokens x max(count) x D] x = torch.split(x, split_points) # Tuple[2d tensor] counts = torch.tensor([len(x_i) for x_i in x], device=device, dtype=torch.long) x = pad_sequence(x, batch_first=True, padding_value=0) # reduce dim 1 x = self.token_pooling_module(x, counts) # reconstruct the batch split_points = new_lens.cumsum(dim=0) split_points = torch.cat([torch.tensor([0], device=device, dtype=torch.long), split_points]) split_points = (split_points[1:] - split_points[:-1]).tolist() x = torch.split(x, split_points) x = pad_sequence(x, batch_first=True, padding_value=0) # B x ? x D mask = lengths_to_padding_mask(new_lens) return x, mask, new_lens, num_examples_without_ending_sep def forward(self, x, preds, lens): x, mask, lens, preds, _ = self.char_compression(x, preds, lens) x, mask, lens, _ = self.token_compression(x, preds, lens) return x, mask, lens