Spaces:
Sleeping
Sleeping
| """ | |
| Shared NativityClassifier for training and inference. | |
| Variants: MLP (baseline), MLP + channel attention (Squeeze-and-Excitation style). | |
| """ | |
| import torch | |
| import torch.nn as nn | |
| class ImprovedAttentionBlock(nn.Module): | |
| """ | |
| Squeeze-and-Excitation style attention: squeeze 1024 dims to a bottleneck, | |
| then excite back to 1024 channel weights. Learns which channels matter for nativity. | |
| """ | |
| def __init__(self, input_dim=1024, reduction_ratio=8): | |
| super().__init__() | |
| reduced = max(input_dim // reduction_ratio, 32) | |
| self.fc1 = nn.Linear(input_dim, reduced) | |
| self.relu = nn.ReLU() | |
| self.fc2 = nn.Linear(reduced, input_dim) | |
| self.sigmoid = nn.Sigmoid() | |
| def forward(self, x): | |
| # x shape: (batch_size, 1024) | |
| attention_weights = self.sigmoid(self.fc2(self.relu(self.fc1(x)))) | |
| return x * attention_weights | |
| class NativityClassifier(nn.Module): | |
| """ | |
| MLP classifier on wav2vec embeddings -> Native vs Non-Native. | |
| Uses BatchNorm and dropout for better generalization. | |
| """ | |
| def __init__( | |
| self, | |
| input_dim=1024, | |
| hidden_dims=(384, 256, 128), | |
| output_dim=2, | |
| dropout=0.35, | |
| ): | |
| super(NativityClassifier, self).__init__() | |
| self.input_dim = input_dim | |
| self.output_dim = output_dim | |
| layers = [] | |
| prev = input_dim | |
| for h in hidden_dims: | |
| layers.extend([ | |
| nn.Linear(prev, h), | |
| nn.BatchNorm1d(h), | |
| nn.ReLU(inplace=True), | |
| nn.Dropout(dropout), | |
| ]) | |
| prev = h | |
| layers.append(nn.Linear(prev, output_dim)) | |
| self.net = nn.Sequential(*layers) | |
| def forward(self, x): | |
| return self.net(x) | |
| class NativityClassifierWithAttention(nn.Module): | |
| """ | |
| Squeeze-and-Excitation attention over 1024-d embedding then MLP. | |
| SE block learns which channels matter for non-native prosody; higher final_dropout | |
| reduces overfitting to majority class. | |
| """ | |
| def __init__( | |
| self, | |
| input_dim=1024, | |
| hidden_dims=(384, 256, 128), | |
| output_dim=2, | |
| dropout=0.35, | |
| final_dropout=0.6, | |
| reduction_ratio=8, | |
| ): | |
| super(NativityClassifierWithAttention, self).__init__() | |
| self.input_dim = input_dim | |
| self.output_dim = output_dim | |
| self.attention = ImprovedAttentionBlock(input_dim=input_dim, reduction_ratio=reduction_ratio) | |
| layers = [] | |
| prev = input_dim | |
| for i, h in enumerate(hidden_dims): | |
| layers.extend([ | |
| nn.Linear(prev, h), | |
| nn.BatchNorm1d(h), | |
| nn.ReLU(inplace=True), | |
| nn.Dropout(final_dropout if i == len(hidden_dims) - 1 else dropout), | |
| ]) | |
| prev = h | |
| layers.append(nn.Linear(prev, output_dim)) | |
| self.net = nn.Sequential(*layers) | |
| def forward(self, x): | |
| x = self.attention(x) | |
| return self.net(x) | |