Spaces:
Sleeping
Sleeping
| import torch | |
| import torch.nn as nn | |
| from transformers import Wav2Vec2PreTrainedModel, Wav2Vec2Model, Wav2Vec2Config | |
| class Wav2Vec2PhonemeEmbedder(Wav2Vec2PreTrainedModel): | |
| def __init__(self, config): | |
| super().__init__(config) | |
| # 1. The Base Wav2Vec2 Audio Encoder | |
| self.wav2vec2 = Wav2Vec2Model(config) | |
| # 2. Audio projection to the embedding space | |
| self.proj_size = getattr(config, "classifier_proj_size", 256) | |
| self.audio_proj = nn.Linear(config.hidden_size, self.proj_size) | |
| # 2b. Dropout for anti-collapse (expert recommendation) | |
| self.proj_dropout = nn.Dropout(p=0.1) | |
| # 3. The Learnable Phoneme Dictionary | |
| self.vocab_size = config.vocab_size | |
| self.phoneme_embeddings = nn.Parameter(torch.randn(self.vocab_size, self.proj_size)) | |
| # 4. Temperature parameter for cosine similarity (learnable) | |
| self.logit_scale = nn.Parameter(torch.ones([]) * torch.log(torch.tensor(1 / 0.07))) | |
| # 5. Optional CTC class weights (set externally by training script) | |
| self.ctc_class_weights = None | |
| # Initialize weights | |
| self.post_init() | |
| def forward(self, input_values, attention_mask=None, labels=None, output_attentions=None, output_hidden_states=None, return_dict=None): | |
| # Force plain Python bool — newer transformers configs or internal states can return Tensors here | |
| return_dict = True if return_dict is None else bool(return_dict) | |
| # Extract audio features (Shape: [Batch, Time, Hidden_Size]) | |
| outputs = self.wav2vec2( | |
| input_values, | |
| attention_mask=attention_mask, | |
| output_attentions=output_attentions, | |
| output_hidden_states=output_hidden_states, | |
| return_dict=return_dict, | |
| ) | |
| hidden_states = outputs[0] | |
| # Project audio features to the embedding space + dropout | |
| audio_features = self.proj_dropout(self.audio_proj(hidden_states)) | |
| # L2 Normalize audio and phoneme embeddings for Cosine Similarity | |
| audio_features = audio_features / (audio_features.norm(dim=-1, keepdim=True) + 1e-8) | |
| phoneme_features = self.phoneme_embeddings / (self.phoneme_embeddings.norm(dim=-1, keepdim=True) + 1e-8) | |
| # Calculate Cosine Similarity Logits (Shape: [Batch, Time, Vocab_Size]) | |
| logit_scale = self.logit_scale.exp() | |
| logits = logit_scale * torch.matmul(audio_features, phoneme_features.t()) | |
| loss = None | |
| if labels is not None: | |
| # 1. Prepare targets for CTC | |
| labels_mask = labels >= 0 | |
| target_lengths = labels_mask.sum(dim=-1) | |
| flattened_targets = labels[labels_mask] | |
| # 2. Calculate input lengths for CTC (downsampled by CNN layers) | |
| if attention_mask is not None: | |
| input_lengths = self._get_feat_extract_output_lengths(attention_mask.sum(-1)).to(torch.long) | |
| else: | |
| input_lengths = torch.full((logits.shape[0],), logits.shape[1], device=logits.device, dtype=torch.long) | |
| # 3. Apply class weighting (anti-collapse: penalizes schwa dominance) | |
| if self.ctc_class_weights is not None: | |
| weights = self.ctc_class_weights.to(logits.device) | |
| log_probs = nn.functional.log_softmax(logits + weights.log(), dim=-1).transpose(0, 1) | |
| else: | |
| log_probs = nn.functional.log_softmax(logits, dim=-1).transpose(0, 1) | |
| # 4. Calculate CTC Loss | |
| loss_fn = nn.CTCLoss(blank=self.config.pad_token_id or 0, zero_infinity=True) | |
| loss = loss_fn(log_probs, flattened_targets, input_lengths, target_lengths) | |
| if not return_dict: | |
| output = (logits,) + outputs[2:] | |
| return ((loss,) + output) if loss is not None else output | |
| return { | |
| "loss": loss, | |
| "logits": logits, | |
| "hidden_states": outputs.hidden_states, | |
| "attentions": outputs.attentions, | |
| } | |
| def _get_feat_extract_output_lengths(self, input_lengths): | |
| """Helper to compute downsampled lengths through the CNN layers.""" | |
| # This implementation matches the one in Wav2Vec2ForCTC | |
| def _conv_out_length(input_length, kernel_size, stride): | |
| return torch.div(input_length - kernel_size, stride, rounding_mode="floor") + 1 | |
| for kernel_size, stride in zip(self.config.conv_kernel, self.config.conv_stride): | |
| input_lengths = _conv_out_length(input_lengths, kernel_size, stride) | |
| return input_lengths | |