File size: 2,771 Bytes
378eaeb | 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 | from dataclasses import dataclass
from typing import Optional
import torch
from torch import nn
from transformers import AutoConfig, AutoModel, PreTrainedModel
from transformers.utils import ModelOutput
from .configuration_genaid import GenAIDConfig
@dataclass
class GenAIDOutput(ModelOutput):
embedding: torch.FloatTensor = None
accent_logits: Optional[torch.FloatTensor] = None
speaker_logits: Optional[torch.FloatTensor] = None
class GenAIDModel(PreTrainedModel):
config_class = GenAIDConfig
base_model_prefix = "genaid"
main_input_name = "input_values"
# GenAID has no tied parameters. Transformers 5.x expects custom models to
# expose this mapping while finalizing low-memory checkpoint loading.
all_tied_weights_keys = {}
def __init__(self, config):
super().__init__(config)
encoder_dict = dict(config.encoder_config)
model_type = encoder_dict.pop("model_type")
encoder_config = AutoConfig.for_model(model_type, **encoder_dict)
self.encoder = AutoModel.from_config(encoder_config)
hidden = encoder_config.hidden_size
dim = config.bottleneck_dim
self.bottleneck = nn.Sequential(
nn.Linear(hidden, dim), nn.GELU(), nn.Linear(dim, dim), nn.GELU()
)
self.accent_classifier = nn.Linear(dim, config.num_accents, bias=False)
self.speaker_classifier = nn.Linear(dim, config.num_speakers, bias=False)
@staticmethod
def masked_mean(hidden_states, attention_mask):
if attention_mask is None:
return hidden_states.mean(1)
lengths = attention_mask.sum(-1)
frame_lengths = (lengths * hidden_states.shape[1] / attention_mask.shape[1]).ceil().long()
frame_lengths = frame_lengths.clamp(1, hidden_states.shape[1])
frame_mask = torch.arange(hidden_states.shape[1], device=hidden_states.device)[None]
frame_mask = frame_mask < frame_lengths[:, None]
return (hidden_states * frame_mask.unsqueeze(-1)).sum(1) / frame_lengths.unsqueeze(-1)
def forward(self, input_values, attention_mask=None, return_dict=True, **kwargs):
encoded = self.encoder(
input_values=input_values,
attention_mask=attention_mask,
return_dict=True,
**kwargs,
).last_hidden_state
embedding = self.bottleneck(self.masked_mean(encoded, attention_mask))
accent_logits = self.accent_classifier(embedding)
speaker_logits = self.speaker_classifier(embedding)
if not return_dict:
return embedding, accent_logits, speaker_logits
return GenAIDOutput(
embedding=embedding,
accent_logits=accent_logits,
speaker_logits=speaker_logits,
)
|