File size: 3,854 Bytes
65a0eff |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 |
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers.models.seamless_m4t.modeling_seamless_m4t import (
_compute_new_attention_mask,
)
from transformers.models.seamless_m4t_v2.modeling_seamless_m4t_v2 import (
SeamlessM4Tv2SpeechEncoder,
SeamlessM4Tv2PreTrainedModel,
)
from .configuration_seamless_m4t_v2_speech_encoder import (
MODEL_TYPE,
SeamlessM4Tv2EncoderConfig,
)
from transformers.modeling_outputs import SequenceClassifierOutput
from transformers.models.auto import AutoModel, AutoModelForAudioClassification, AutoModelForSequenceClassification
class SeamlessM4Tv2SpeechEncoder(SeamlessM4Tv2SpeechEncoder):
model_type = MODEL_TYPE
config_class = SeamlessM4Tv2EncoderConfig
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
def _compute_sub_sample_lengths_from_attention_mask(self, attention_mask):
pad = self.kernel_size // 2
seq_lens = attention_mask.size(1) - (1 - attention_mask.int()).sum(1)
seq_lens = ((seq_lens + 2 * pad - self.kernel_size) / self.stride) + 1
return seq_lens.floor()
@staticmethod
def mean_pooling(
hidden_states: torch.Tensor, attention_mask: torch.Tensor
) -> torch.Tensor:
# hidden_states shape: (batch_size, sequence_length, hidden_size)
# attention_mask shape: (batch_size, sequence_length)
# Apply attention mask and avoid division by zero
input_mask_expanded = (
attention_mask.unsqueeze(-1).expand(hidden_states.size()).float()
)
sum_hidden_states = torch.sum(hidden_states * input_mask_expanded, 1)
sum_mask = input_mask_expanded.sum(1)
return sum_hidden_states / torch.clamp(sum_mask, min=1e-9)
class SeamlessM4Tv2ForAudioClassification(SeamlessM4Tv2PreTrainedModel):
model_type = MODEL_TYPE
base_model_prefix = "model"
config_class = SeamlessM4Tv2EncoderConfig
def __init__(self, config, *args, **kwargs):
super().__init__(config)
self.num_labels = config.num_labels
self.model = SeamlessM4Tv2SpeechEncoder(config)
self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
def forward(
self,
input_features: torch.Tensor,
attention_mask: torch.Tensor,
labels: None | torch.Tensor,
*args,
**kwargs,
):
output_hidden_states = kwargs.pop("output_hidden_states", False)
outputs = self.model(
input_features,
attention_mask,
output_hidden_states=output_hidden_states,
*args,
**kwargs,
)
hidden_states = outputs.last_hidden_state
if attention_mask is not None:
sub_sampled_lengths = self._compute_sub_sample_lengths_from_attention_mask(
attention_mask
).to(outputs.last_hidden_state.device)
attention_mask = _compute_new_attention_mask(
hidden_states=hidden_states, seq_lens=sub_sampled_lengths
)
hidden_states = self.model.mean_pooling(
outputs.last_hidden_state, attention_mask
)
logits = self.score(hidden_states)
if labels is not None:
loss = F.cross_entropy(logits, labels)
else:
loss = None
return SequenceClassifierOutput(
loss=loss, # type: ignore
logits=logits,
hidden_states=outputs.hidden_states if output_hidden_states else None,
)
AutoModel.register(SeamlessM4Tv2EncoderConfig, SeamlessM4Tv2SpeechEncoder)
AutoModelForAudioClassification.register(
SeamlessM4Tv2EncoderConfig, SeamlessM4Tv2ForAudioClassification
)
AutoModelForSequenceClassification.register(
SeamlessM4Tv2EncoderConfig, SeamlessM4Tv2ForAudioClassification
)
|