|
import math |
|
from typing import Optional, Tuple, Union |
|
|
|
import numpy as np |
|
import torch |
|
import torch.nn.functional as F |
|
import torch.utils.checkpoint |
|
from torch import nn |
|
from torch.nn import CrossEntropyLoss |
|
|
|
from transformers.modeling_outputs import ( |
|
SequenceClassifierOutput |
|
) |
|
from transformers.modeling_utils import PreTrainedModel |
|
from transformers.utils import ( |
|
add_start_docstrings, |
|
add_start_docstrings_to_model_forward, |
|
logging, |
|
replace_return_docstrings, |
|
) |
|
from transformers.models.whisper.configuration_whisper import WhisperConfig |
|
from transformers.models.whisper.generation_whisper import WhisperGenerationMixin |
|
from transformers.models.whisper.modeling_whisper import WhisperPreTrainedModel, WHISPER_ENCODER_INPUTS_DOCSTRING, _CONFIG_FOR_DOC, WhisperEncoder |
|
|
|
class ViSpeechClassification(WhisperPreTrainedModel): |
|
def __init__(self, config): |
|
super().__init__(config) |
|
|
|
self.encoder = WhisperEncoder(config) |
|
num_layers = config.num_hidden_layers + 1 |
|
self.projector = nn.Sequential( |
|
nn.Linear(self.encoder.config.hidden_size, 1024), |
|
nn.ReLU(), |
|
nn.Dropout(config.dropout), |
|
|
|
nn.Linear(1024, 512), |
|
nn.ReLU(), |
|
nn.Dropout(config.dropout), |
|
|
|
nn.Linear(512, config.classifier_proj_size), |
|
nn.ReLU(), |
|
nn.Dropout(config.dropout) |
|
) |
|
|
|
self.classifier = nn.Linear(config.classifier_proj_size, config.num_labels) |
|
self.config.use_weighted_layer_sum = False |
|
|
|
|
|
self.post_init() |
|
|
|
def freeze_encoder(self): |
|
""" |
|
Calling this function will disable the gradient computation for the Whisper encoder so that its parameters will |
|
not be updated during training. Only the projection layers and classification head will be updated. |
|
""" |
|
self.encoder._freeze_parameters() |
|
|
|
def get_input_embeddings(self) -> nn.Module: |
|
return self.encoder.get_input_embeddings() |
|
|
|
def set_input_embeddings(self, value: nn.Module): |
|
self.encoder.set_input_embeddings(value) |
|
|
|
@add_start_docstrings_to_model_forward(WHISPER_ENCODER_INPUTS_DOCSTRING) |
|
@replace_return_docstrings(output_type=SequenceClassifierOutput, config_class=_CONFIG_FOR_DOC) |
|
def forward( |
|
self, |
|
input_features: Optional[torch.LongTensor] = None, |
|
head_mask: Optional[torch.Tensor] = None, |
|
encoder_outputs: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, |
|
labels: Optional[torch.LongTensor] = None, |
|
output_attentions: Optional[bool] = None, |
|
output_hidden_states: Optional[bool] = None, |
|
return_dict: Optional[bool] = None, |
|
) -> Union[Tuple[torch.Tensor], SequenceClassifierOutput]: |
|
r""" |
|
labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): |
|
Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., |
|
config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If |
|
`config.num_labels > 1` a classification loss is computed (Cross-Entropy). |
|
|
|
Returns: |
|
|
|
Example: |
|
|
|
```python |
|
>>> import torch |
|
>>> from transformers import AutoFeatureExtractor, WhisperForAudioClassification |
|
>>> from datasets import load_dataset |
|
|
|
>>> feature_extractor = AutoFeatureExtractor.from_pretrained("sanchit-gandhi/whisper-medium-fleurs-lang-id") |
|
>>> model = WhisperForAudioClassification.from_pretrained("sanchit-gandhi/whisper-medium-fleurs-lang-id") |
|
|
|
>>> ds = load_dataset("google/fleurs", "all", split="validation", streaming=True) |
|
>>> sample = next(iter(ds)) |
|
|
|
>>> inputs = feature_extractor( |
|
... sample["audio"]["array"], sampling_rate=sample["audio"]["sampling_rate"], return_tensors="pt" |
|
... ) |
|
>>> input_features = inputs.input_features |
|
|
|
>>> with torch.no_grad(): |
|
... logits = model(input_features).logits |
|
|
|
>>> predicted_class_ids = torch.argmax(logits).item() |
|
>>> predicted_label = model.config.id2label[predicted_class_ids] |
|
>>> predicted_label |
|
'Afrikaans' |
|
```""" |
|
|
|
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 |
|
) |
|
if self.config.use_weighted_layer_sum: |
|
output_hidden_states = True |
|
elif output_hidden_states is None: |
|
output_hidden_states = self.config.output_hidden_states |
|
|
|
return_dict = return_dict if return_dict is not None else self.config.use_return_dict |
|
|
|
if encoder_outputs is None: |
|
encoder_outputs = self.encoder( |
|
input_features, |
|
head_mask=head_mask, |
|
output_attentions=output_attentions, |
|
output_hidden_states=output_hidden_states, |
|
return_dict=return_dict, |
|
) |
|
|
|
if self.config.use_weighted_layer_sum: |
|
hidden_states = encoder_outputs[_HIDDEN_STATES_START_POSITION] |
|
hidden_states = torch.stack(hidden_states, dim=1) |
|
norm_weights = nn.functional.softmax(self.layer_weights, dim=-1) |
|
hidden_states = (hidden_states * norm_weights.view(-1, 1, 1)).sum(dim=1) |
|
else: |
|
hidden_states = encoder_outputs[0] |
|
|
|
hidden_states = self.projector(hidden_states) |
|
pooled_output = hidden_states.mean(dim=1) |
|
|
|
logits = self.classifier(pooled_output) |
|
|
|
loss = None |
|
|
|
if labels is not None: |
|
loss_fct = CrossEntropyLoss() |
|
|
|
labels = labels.to(logits.device) |
|
loss = loss_fct(logits.view(-1, self.config.num_labels), labels.view(-1)) |
|
|
|
if not return_dict: |
|
output = (logits,) + encoder_outputs[1:] |
|
return ((loss,) + output) if loss is not None else output |
|
|
|
return SequenceClassifierOutput( |
|
loss=loss, |
|
logits=logits, |
|
hidden_states=encoder_outputs.hidden_states, |
|
attentions=encoder_outputs.attentions, |
|
) |
