Update model.py

model.py

@@ -4,6 +4,7 @@ from typing import Optional, Tuple, Union
 from torch.nn import MSELoss
 import torch
 import torch.nn as nn
+import math
 
 class Wav2Vec2BertForSequenceClassification(Wav2Vec2BertPreTrainedModel):
     # Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2ForSequenceClassification.__init__ with Wav2Vec2->Wav2Vec2Bert,wav2vec2->wav2vec2_bert
@@ -19,7 +20,19 @@ class Wav2Vec2BertForSequenceClassification(Wav2Vec2BertPreTrainedModel):
         if config.use_weighted_layer_sum:
             self.layer_weights = nn.Parameter(torch.ones(num_layers) / num_layers)
         self.projector = nn.Linear(config.hidden_size, config.classifier_proj_size)
-        self.classifier = nn.Linear(config.classifier_proj_size, config.num_labels)
+        self.pooled_conv = nn.Sequential(nn.Conv1d(config.hidden_size, config.hidden_size // 2, kernel_size=15, stride=3, padding=30),
+                                         nn.AvgPool1d(2, 2),
+                                         nn.BatchNorm1d(config.hidden_size // 2),
+                                         nn.Conv1d(config.hidden_size // 2, config.classifier_proj_size, kernel_size=7, stride=2, padding=0),
+                                         nn.ReLU()
+                                        )
+        
+        self.classifier = nn.Sequential(nn.Dropout(p=0.091,),
+                                        nn.Linear(config.classifier_proj_size, config.classifier_proj_size // 2), 
+                                        nn.ReLU(),
+                                        nn.Linear(config.classifier_proj_size // 2, config.num_labels),
+                                        nn.ReLU(),
+                                       )
 
         # Initialize weights and apply final processing
         self.post_init()
@@ -69,20 +82,14 @@ class Wav2Vec2BertForSequenceClassification(Wav2Vec2BertPreTrainedModel):
         else:
             hidden_states = outputs[0]
 
-        hidden_states = self.projector(hidden_states)
-        if attention_mask is None:
-            pooled_output = hidden_states.mean(dim=1)
-        else:
-            padding_mask = self._get_feature_vector_attention_mask(hidden_states.shape[1], attention_mask)
-            hidden_states[~padding_mask] = 0.0
-            pooled_output = hidden_states.sum(dim=1) / padding_mask.sum(dim=1).view(-1, 1)
-
-        logits = self.classifier(pooled_output)
-        logits = nn.functional.relu(logits)
+        hidden_states = hidden_states.permute(0, 2, 1)
+        hidden_states = self.pooled_conv(hidden_states)
+        hidden_states = torch.mean(hidden_states, dim=2)
+        logits        = self.classifier(hidden_states)
 
         loss = None
         if labels is not None:
-            loss_fct = MSELoss()
+            loss_fct = nn.L1Loss()
             loss = loss_fct(logits.view(-1, self.config.num_labels), labels.view(-1, self.config.num_labels))
 
         if not return_dict: