Update custom_interface.py

custom_interface.py

@@ -104,7 +104,7 @@ class CustomEncoderWav2vec2Classifier(Pretrained):
             (label encoder should be provided).
         """
         outputs = self.encode_batch(wavs, wav_lens)
-        outputs = self.mods.output_mlp(outputs)
+        outputs = self.mods.label_lin(outputs)
         out_prob = self.hparams.softmax(outputs)
         score, index = torch.max(out_prob, dim=-1)
         text_lab = self.hparams.label_encoder.decode_torch(index)
@@ -137,6 +137,17 @@ class CustomEncoderWav2vec2Classifier(Pretrained):
         out_prob = self.hparams.softmax(outputs)
         score, index = torch.max(out_prob, dim=-1)
         text_lab = self.hparams.label_encoder.decode_torch(index)
+        if text_lab[0] == "1":
+            text_lab = "neutral"
+        elif text_lab[0] == "2":
+            text_lab = "sadness"
+        elif text_lab[0] == "3":
+            text_lab = "joy"
+        elif text_lab[0] == "4":
+            text_lab = "anger"
+        elif text_lab[0] == "5":
+            text_lab = "affection"
+        
         return out_prob, score, index, text_lab
 
     def forward(self, wavs, wav_lens=None, normalize=False):