File size: 4,731 Bytes
65a0eff ee8c043 65a0eff 1057217 65a0eff 1057217 ee8c043 1057217 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 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 |
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)
@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:
# move labels to correct device to enable model parallelism
labels = labels.to(logits.device)
if self.config.problem_type is None:
if self.num_labels == 1:
self.config.problem_type = "regression"
elif self.num_labels > 1 and (
labels.dtype == torch.long or labels.dtype == torch.int
):
self.config.problem_type = "single_label_classification"
else:
self.config.problem_type = "multi_label_classification"
if self.config.problem_type == "regression":
loss_fct = F.mse_loss
if self.num_labels == 1:
loss = loss_fct(logits.squeeze(), labels.squeeze())
else:
loss = loss_fct(logits, labels)
elif self.config.problem_type == "single_label_classification":
loss_fct = F.cross_entropy
loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
elif self.config.problem_type == "multi_label_classification":
loss_fct = F.binary_cross_entropy_with_logits
loss = loss_fct(logits, labels)
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
)
|