import torch
from torch import nn
from transformers.modeling_outputs import BaseModelOutput, CausalLMOutput
from transformers.models.whisper.modeling_whisper import WhisperEncoder, WhisperEncoderLayer, WhisperAttention
from .FDDT import FDDT
from .config import DiCoWConfig
from .layers import CustomLinear, CustomDiagonalLinear, Gate, SpeakerCommunicationBlock
class DiCoWEncoder(WhisperEncoder):
config_class = DiCoWConfig
def __init__(self, config: DiCoWConfig):
super().__init__(config)
self.ctc_weight = config.ctc_weight
if config.additional_layer and self.ctc_weight > 0.0:
self.additional_layer = WhisperEncoderLayer(config)
if config.additional_self_attention_layer and self.ctc_weight > 0.0:
self.additional_self_attention_layer = WhisperAttention(
embed_dim=config.d_model,
num_heads=config.encoder_attention_heads,
dropout=config.attention_dropout,
config=config,
)
if config.pre_ctc_sub_sample and self.ctc_weight > 0.0:
self.subsample_conv1 = nn.Conv1d(
in_channels=config.d_model,
out_channels=config.d_model,
kernel_size=3,
stride=2,
padding=1,
bias=False,
)
self.subsample_conv2 = nn.Conv1d(
in_channels=config.d_model,
out_channels=config.d_model,
kernel_size=3,
stride=2,
padding=1,
bias=False,
)
if self.ctc_weight > 0.0:
self.lm_head = nn.Linear(config.d_model, config.vocab_size + 1, bias=False)
self.final_dropout = nn.Dropout(config.final_dropout)
if config.use_fddt:
num_fddts = self.config.apply_fddt_to_n_layers if self.config.apply_fddt_to_n_layers != -1 else len(
self.layers)
self.fddts = nn.ModuleList([
FDDT(
d_model=config.d_model,
non_target_rate=1.0,
fddt_init=config.fddt_init,
is_diagonal=config.fddt_is_diagonal,
bias_only=config.fddt_bias_only,
use_silence=config.fddt_use_silence,
use_target=config.fddt_use_target,
use_overlap=config.fddt_use_overlap,
use_non_target=config.fddt_use_non_target,
)
for _ in range(num_fddts)
])
if config.use_pre_pos_fddt:
self.initial_fddt = FDDT(
d_model=config.d_model,
non_target_rate=config.non_target_fddt_value,
fddt_init=config.fddt_init,
is_diagonal=config.fddt_is_diagonal,
bias_only=config.fddt_bias_only,
use_silence=config.fddt_use_silence,
use_target=config.fddt_use_target,
use_overlap=config.fddt_use_overlap,
use_non_target=config.fddt_use_non_target,
)
if config.use_enrollments and config.scb_layers is not None:
self.ca_enrolls = nn.ModuleList([SpeakerCommunicationBlock(config) for _ in range(config.scb_layers)])
self.first_task_token = self.config.vocab_size - 30 * 50 - 1 - 6 # 30 seconds of 50 Hz timestamps -1 to get to 0.0 and -6 number of tasks
self.post_init()
def _init_weights(self, module):
super()._init_weights(module)
if isinstance(module, CustomLinear) or isinstance(module, CustomDiagonalLinear) or isinstance(module, Gate):
module.reset_parameters()
def get_output_embeddings(self):
return None
def possibly_update_last_hidden_states(self, hidden_states):
if hasattr(self, "additional_layer"):
hidden_states, = self.additional_layer(
hidden_states,
attention_mask=None,
output_attentions=False,
layer_head_mask=None,
)
elif hasattr(self, "additional_self_attention_layer"):
hidden_states, _ = self.additional_self_attention_layer(
hidden_states,
attention_mask=None,
output_attentions=False,
layer_head_mask=None,
)
hidden_states = self.final_dropout(hidden_states)
if hasattr(self, "subsample_conv2"):
hidden_states = self.subsample_conv2(self.subsample_conv1(hidden_states.transpose(1, 2))).transpose(1, 2)
return hidden_states
def get_loss(self, logits, labels):
if labels.max() >= self.config.vocab_size:
raise ValueError(f"Label values must be <= vocab_size: {self.config.vocab_size}")
if self.config.remove_timestamps_from_ctc:
labels = torch.nn.utils.rnn.pad_sequence([label[label < self.first_task_token] for label in labels],
padding_value=-100).T
input_lengths = torch.full((logits.shape[0],), fill_value=logits.shape[1],
device=logits.device)
# assuming that padded tokens are filled with -100
# when not being attended to
labels_mask = labels >= 0
target_lengths = labels_mask.sum(-1)
# ctc_loss doesn't support fp16
log_probs = nn.functional.log_softmax(logits, dim=-1, dtype=torch.float32).transpose(0, 1)
with torch.backends.cudnn.flags(enabled=True):
ctc_loss = nn.functional.ctc_loss(
log_probs,
labels,
input_lengths,
target_lengths,
blank=logits.shape[-1] - 1,
reduction=self.config.ctc_loss_reduction,
zero_infinity=True,
)
return ctc_loss
def get_max_len(self):
return self.config.max_source_positions * self.conv1.stride[0] * self.conv2.stride[0]
def forward(
self,
input_features,
attention_mask=None,
head_mask=None,
output_attentions=None,
output_hidden_states=None,
return_dict=None,
stno_mask=None,
return_logits=False,
enrollments=None
):
if enrollments is not None:
input_features = torch.stack((input_features, enrollments['input_features']), dim=1).flatten(0,1)
stno_mask = torch.stack((stno_mask, enrollments['stno_mask']),dim=1).flatten(0,1)
expected_seq_length = self.config.max_source_positions * self.conv1.stride[0] * self.conv2.stride[0]
if input_features.shape[-1] != expected_seq_length:
raise ValueError(
f"Whisper expects the mel input features to be of length {expected_seq_length}, but found {input_features.shape[-1]}. Make sure to pad the input mel features to {expected_seq_length}."
)
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
output_hidden_states = (
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
)
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
inputs_embeds = nn.functional.gelu(self.conv1(input_features))
inputs_embeds = nn.functional.gelu(self.conv2(inputs_embeds))
inputs_embeds = inputs_embeds.permute(0, 2, 1)
""""""
if self.config.use_fddt and self.config.use_pre_pos_fddt:
inputs_embeds = self.initial_fddt(inputs_embeds, stno_mask)
""""""
all_positions = torch.arange(self.embed_positions.num_embeddings, device=inputs_embeds.device)
hidden_states = inputs_embeds + self.embed_positions(all_positions)
hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
encoder_states = () if output_hidden_states else None
all_attentions = () if output_attentions else None
# check if head_mask has a correct number of layers specified if desired
if head_mask is not None:
assert head_mask.size()[0] == (len(self.layers)), (
f"The head_mask should be specified for {len(self.layers)} layers, but it is for {head_mask.size()[0]}."
)
for idx, encoder_layer in enumerate(self.layers):
if output_hidden_states:
encoder_states = encoder_states + (hidden_states,)
# add LayerDrop (see https://huggingface.co/papers/1909.11556 for description)
to_drop = False
if self.training:
dropout_probability = torch.rand([])
if dropout_probability < self.layerdrop: # skip the layer
to_drop = True
if to_drop:
layer_outputs = (None, None)
else:
""""""
if self.config.use_fddt and idx < len(self.fddts):
hidden_states = self.fddts[idx](hidden_states, stno_mask)
if self.config.use_enrollments and idx < self.config.scb_layers:
hidden_states = self.ca_enrolls[idx](hidden_states)
if idx == self.config.scb_layers -1:
# enrollment representations are not longer needed
hidden_states = hidden_states[::2]
stno_mask = stno_mask[::2]
""""""
layer_outputs = encoder_layer(
hidden_states,
None,
layer_head_mask=(head_mask[idx] if head_mask is not None else None),
output_attentions=output_attentions,
)
hidden_states = layer_outputs[0]
if output_attentions:
all_attentions = all_attentions + (layer_outputs[1],)
hidden_states = self.layer_norm(hidden_states)
if output_hidden_states:
encoder_states = encoder_states + (hidden_states,)
if return_logits:
hidden_states = hidden_states
hidden_states = self.possibly_update_last_hidden_states(hidden_states)
logits = self.lm_head(hidden_states)
return CausalLMOutput(
loss=None, logits=logits, hidden_states=hidden_states,
)
if not return_dict:
return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None)
return BaseModelOutput(
last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions
)