Update modeling_dicow.py

modeling_dicow.py

@@ -25,8 +25,7 @@ from .encoder import DiCoWEncoder
 from .FDDT import FDDT
 from .layers import CustomLinear, CustomDiagonalLinear, Gate
 from .generation import DiCoWGenerationMixin
-from .contrastive_loss import ContrastiveLoss
-import wandb
+
 logging.set_verbosity_debug()
 logger = logging.get_logger("transformers")
 
@@ -334,21 +333,6 @@ class DiCoWForConditionalGeneration(DiCoWGenerationMixin, WhisperForConditionalG
                 wandb.log({"ctc_loss": ctc_loss})
             loss = (1 - self.config.ctc_weight) * dec_loss + self.config.ctc_weight * ctc_loss
 
-            if self.config.contrastive_loss_weight > 0.0:
-                loss_fct = ContrastiveLoss(distance_metric="cosine")
-                stno_per_spk_pair = stno_mask.view(-1, self.config.mt_num_speakers, stno_mask.shape[1], stno_mask.shape[2])
-                positive_mask = ((stno_per_spk_pair[:, :, 1, :] + stno_per_spk_pair[:, :, 3, :]) > 0.5).flatten(1)
-                intermediate_states = outputs.encoder_hidden_states[8].view(-1, self.config.mt_num_speakers, stno_mask.shape[2],
-                                                          outputs.encoder_hidden_states[8].shape[-1]).flatten(1, 2)
-                valid_pairs = is_valid.view((-1, self.config.mt_num_speakers)).all(dim=-1)
-                contrastive_loss = loss_fct(
-                    intermediate_states[valid_pairs],
-                    positive_mask[valid_pairs])
-                # print(contrastive_loss)
-                if wandb.run is not None:
-                    wandb.log({"contrastive_loss": contrastive_loss})
-                if contrastive_loss != 0.0 and loss < 0.5:
-                    loss += self.config.contrastive_loss_weight * contrastive_loss
         if not return_dict:
             output = (dec_lm_logits,) + outputs[1:]
             return ((loss,) + output) if loss is not None else output