| import math
|
| import random
|
| from typing import Optional, Tuple
|
| from fairseq.checkpoint_utils import load_model_ensemble_and_task
|
| import numpy as np
|
| import torch
|
| import torch.nn.functional as F
|
|
|
|
|
| from fairseq.utils import index_put
|
|
|
|
|
|
|
| def pad_to_multiple(x, multiple, dim=-1, value=0):
|
|
|
| if x is None:
|
| return None, 0
|
| tsz = x.size(dim)
|
| m = tsz / multiple
|
| remainder = math.ceil(m) * multiple - tsz
|
| if int(tsz % multiple) == 0:
|
| return x, 0
|
| pad_offset = (0,) * (-1 - dim) * 2
|
|
|
| return F.pad(x, (*pad_offset, 0, remainder), value=value), remainder
|
|
|
|
|
| def extract_features(
|
| self,
|
| x,
|
| padding_mask=None,
|
| tgt_layer=None,
|
| min_layer=0,
|
| ):
|
| if padding_mask is not None:
|
| x = index_put(x, padding_mask, 0)
|
|
|
| x_conv = self.pos_conv(x.transpose(1, 2))
|
| x_conv = x_conv.transpose(1, 2)
|
| x = x + x_conv
|
|
|
| if not self.layer_norm_first:
|
| x = self.layer_norm(x)
|
|
|
|
|
| x, pad_length = pad_to_multiple(x, self.required_seq_len_multiple, dim=-2, value=0)
|
| if pad_length > 0 and padding_mask is None:
|
| padding_mask = x.new_zeros((x.size(0), x.size(1)), dtype=torch.bool)
|
| padding_mask[:, -pad_length:] = True
|
| else:
|
| padding_mask, _ = pad_to_multiple(
|
| padding_mask, self.required_seq_len_multiple, dim=-1, value=True
|
| )
|
| x = F.dropout(x, p=self.dropout, training=self.training)
|
|
|
|
|
| x = x.transpose(0, 1)
|
|
|
| layer_results = []
|
| r = None
|
| for i, layer in enumerate(self.layers):
|
| dropout_probability = np.random.random() if self.layerdrop > 0 else 1
|
| if not self.training or (dropout_probability > self.layerdrop):
|
| x, (z, lr) = layer(
|
| x, self_attn_padding_mask=padding_mask, need_weights=False
|
| )
|
| if i >= min_layer:
|
| layer_results.append((x, z, lr))
|
| if i == tgt_layer:
|
| r = x
|
| break
|
|
|
| if r is not None:
|
| x = r
|
|
|
|
|
| x = x.transpose(0, 1)
|
|
|
|
|
| if pad_length > 0:
|
| x = x[:, :-pad_length]
|
|
|
| def undo_pad(a, b, c):
|
| return (
|
| a[:-pad_length],
|
| b[:-pad_length] if b is not None else b,
|
| c[:-pad_length],
|
| )
|
|
|
| layer_results = [undo_pad(*u) for u in layer_results]
|
|
|
| return x, layer_results
|
|
|
|
|
| def compute_mask_indices(
|
| shape: Tuple[int, int],
|
| padding_mask: Optional[torch.Tensor],
|
| mask_prob: float,
|
| mask_length: int,
|
| mask_type: str = "static",
|
| mask_other: float = 0.0,
|
| min_masks: int = 0,
|
| no_overlap: bool = False,
|
| min_space: int = 0,
|
| require_same_masks: bool = True,
|
| mask_dropout: float = 0.0,
|
| ) -> torch.Tensor:
|
| """
|
| Computes random mask spans for a given shape
|
|
|
| Args:
|
| shape: the the shape for which to compute masks.
|
| should be of size 2 where first element is batch size and 2nd is timesteps
|
| padding_mask: optional padding mask of the same size as shape, which will prevent masking padded elements
|
| mask_prob: probability for each token to be chosen as start of the span to be masked. this will be multiplied by
|
| number of timesteps divided by length of mask span to mask approximately this percentage of all elements.
|
| however due to overlaps, the actual number will be smaller (unless no_overlap is True)
|
| mask_type: how to compute mask lengths
|
| static = fixed size
|
| uniform = sample from uniform distribution [mask_other, mask_length*2]
|
| normal = sample from normal distribution with mean mask_length and stdev mask_other. mask is min 1 element
|
| poisson = sample from possion distribution with lambda = mask length
|
| min_masks: minimum number of masked spans
|
| no_overlap: if false, will switch to an alternative recursive algorithm that prevents spans from overlapping
|
| min_space: only used if no_overlap is True, this is how many elements to keep unmasked between spans
|
| require_same_masks: if true, will randomly drop out masks until same amount of masks remains in each sample
|
| mask_dropout: randomly dropout this percentage of masks in each example
|
| """
|
|
|
| bsz, all_sz = shape
|
| mask = torch.full((bsz, all_sz), False)
|
|
|
| all_num_mask = int(
|
|
|
| mask_prob * all_sz / float(mask_length)
|
| + torch.rand([1]).item()
|
| )
|
|
|
| all_num_mask = max(min_masks, all_num_mask)
|
|
|
| mask_idcs = []
|
| for i in range(bsz):
|
| if padding_mask is not None:
|
| sz = all_sz - padding_mask[i].long().sum().item()
|
| num_mask = int(mask_prob * sz / float(mask_length) + np.random.rand())
|
| num_mask = max(min_masks, num_mask)
|
| else:
|
| sz = all_sz
|
| num_mask = all_num_mask
|
|
|
| if mask_type == "static":
|
| lengths = torch.full([num_mask], mask_length)
|
| elif mask_type == "uniform":
|
| lengths = torch.randint(mask_other, mask_length * 2 + 1, size=[num_mask])
|
| elif mask_type == "normal":
|
| lengths = torch.normal(mask_length, mask_other, size=[num_mask])
|
| lengths = [max(1, int(round(x))) for x in lengths]
|
| else:
|
| raise Exception("unknown mask selection " + mask_type)
|
|
|
| if sum(lengths) == 0:
|
| lengths[0] = min(mask_length, sz - 1)
|
|
|
| if no_overlap:
|
| mask_idc = []
|
|
|
| def arrange(s, e, length, keep_length):
|
| span_start = torch.randint(low=s, high=e - length, size=[1]).item()
|
| mask_idc.extend(span_start + i for i in range(length))
|
|
|
| new_parts = []
|
| if span_start - s - min_space >= keep_length:
|
| new_parts.append((s, span_start - min_space + 1))
|
| if e - span_start - length - min_space > keep_length:
|
| new_parts.append((span_start + length + min_space, e))
|
| return new_parts
|
|
|
| parts = [(0, sz)]
|
| min_length = min(lengths)
|
| for length in sorted(lengths, reverse=True):
|
| t = [e - s if e - s >= length + min_space else 0 for s, e in parts]
|
| lens = torch.asarray(t, dtype=torch.int)
|
| l_sum = torch.sum(lens)
|
| if l_sum == 0:
|
| break
|
| probs = lens / torch.sum(lens)
|
| c = torch.multinomial(probs.float(), len(parts)).item()
|
| s, e = parts.pop(c)
|
| parts.extend(arrange(s, e, length, min_length))
|
| mask_idc = torch.asarray(mask_idc)
|
| else:
|
| min_len = min(lengths)
|
| if sz - min_len <= num_mask:
|
| min_len = sz - num_mask - 1
|
| mask_idc = torch.asarray(
|
| random.sample([i for i in range(sz - min_len)], num_mask)
|
| )
|
| mask_idc = torch.asarray(
|
| [
|
| mask_idc[j] + offset
|
| for j in range(len(mask_idc))
|
| for offset in range(lengths[j])
|
| ]
|
| )
|
|
|
| mask_idcs.append(torch.unique(mask_idc[mask_idc < sz]))
|
|
|
| min_len = min([len(m) for m in mask_idcs])
|
| for i, mask_idc in enumerate(mask_idcs):
|
| if isinstance(mask_idc, torch.Tensor):
|
| mask_idc = torch.asarray(mask_idc, dtype=torch.float)
|
| if len(mask_idc) > min_len and require_same_masks:
|
| mask_idc = torch.asarray(
|
| random.sample([i for i in range(mask_idc)], min_len)
|
| )
|
| if mask_dropout > 0:
|
| num_holes = int(round(len(mask_idc) * mask_dropout))
|
| mask_idc = torch.asarray(
|
| random.sample([i for i in range(mask_idc)], len(mask_idc) - num_holes)
|
| )
|
|
|
| mask[i, mask_idc.int()] = True
|
|
|
| return mask
|
|
|
|
|
| def apply_mask(self, x, padding_mask, target_list):
|
| B, T, C = x.shape
|
| torch.zeros_like(x)
|
| if self.mask_prob > 0:
|
| mask_indices = compute_mask_indices(
|
| (B, T),
|
| padding_mask,
|
| self.mask_prob,
|
| self.mask_length,
|
| self.mask_selection,
|
| self.mask_other,
|
| min_masks=2,
|
| no_overlap=self.no_mask_overlap,
|
| min_space=self.mask_min_space,
|
| )
|
| mask_indices = mask_indices.to(x.device)
|
| x[mask_indices] = self.mask_emb
|
| else:
|
| mask_indices = None
|
|
|
| if self.mask_channel_prob > 0:
|
| mask_channel_indices = compute_mask_indices(
|
| (B, C),
|
| None,
|
| self.mask_channel_prob,
|
| self.mask_channel_length,
|
| self.mask_channel_selection,
|
| self.mask_channel_other,
|
| no_overlap=self.no_mask_channel_overlap,
|
| min_space=self.mask_channel_min_space,
|
| )
|
| mask_channel_indices = (
|
| mask_channel_indices.to(x.device).unsqueeze(1).expand(-1, T, -1)
|
| )
|
| x[mask_channel_indices] = 0
|
|
|
| return x, mask_indices
|
|
|
|
|
| def get_hubert_model(
|
| model_path="assets/hubert/hubert_base.pt", device=torch.device("cpu")
|
| ):
|
| models, _, _ = load_model_ensemble_and_task(
|
| [model_path],
|
| suffix="",
|
| )
|
| hubert_model = models[0]
|
| hubert_model = hubert_model.to(device)
|
|
|
| def _apply_mask(x, padding_mask, target_list):
|
| return apply_mask(hubert_model, x, padding_mask, target_list)
|
|
|
| hubert_model.apply_mask = _apply_mask
|
|
|
| def _extract_features(
|
| x,
|
| padding_mask=None,
|
| tgt_layer=None,
|
| min_layer=0,
|
| ):
|
| return extract_features(
|
| hubert_model.encoder,
|
| x,
|
| padding_mask=padding_mask,
|
| tgt_layer=tgt_layer,
|
| min_layer=min_layer,
|
| )
|
|
|
| hubert_model.encoder.extract_features = _extract_features
|
|
|
| hubert_model._forward = hubert_model.forward
|
|
|
| def hubert_extract_features(
|
| self,
|
| source: torch.Tensor,
|
| padding_mask: Optional[torch.Tensor] = None,
|
| mask: bool = False,
|
| ret_conv: bool = False,
|
| output_layer: Optional[int] = None,
|
| ) -> Tuple[torch.Tensor, torch.Tensor]:
|
| res = self._forward(
|
| source,
|
| padding_mask=padding_mask,
|
| mask=mask,
|
| features_only=True,
|
| output_layer=output_layer,
|
| )
|
| feature = res["features"] if ret_conv else res["x"]
|
| return feature, res["padding_mask"]
|
|
|
| def _hubert_extract_features(
|
| source: torch.Tensor,
|
| padding_mask: Optional[torch.Tensor] = None,
|
| mask: bool = False,
|
| ret_conv: bool = False,
|
| output_layer: Optional[int] = None,
|
| ) -> Tuple[torch.Tensor, torch.Tensor]:
|
| return hubert_extract_features(
|
| hubert_model, source, padding_mask, mask, ret_conv, output_layer
|
| )
|
|
|
| hubert_model.extract_features = _hubert_extract_features
|
|
|
| def infer(source, padding_mask, output_layer: torch.Tensor):
|
| output_layer = output_layer.item()
|
| logits = hubert_model.extract_features(
|
| source=source, padding_mask=padding_mask, output_layer=output_layer
|
| )
|
| feats = hubert_model.final_proj(logits[0]) if output_layer == 9 else logits[0]
|
| return feats
|
|
|
| hubert_model.infer = infer
|
|
|
|
|
|
|
| return hubert_model
|
|
|
