text: import sys
  import librosa
  import numpy as np
  import soundfile as sf
  import functools
  import torch
  from torch.nn.functional import cosine_similarity
  
  from correloss import corre_loss
  
  
  def logme(f):
      @functools.wraps(f)
      def wrapped(*args, **kwargs):
          print('\n-----------------\n')
          print('   MODEL: {}'.format(f.__name__.upper()))
          print('\n-----------------\n')
          return f(*args, **kwargs)
      return wrapped
  
  
  class ProgressBar:
      """Progress bar
      
      """
      def __init__ (self, valmax, maxbar, title):
          if valmax == 0:  valmax = 1
          if maxbar > 200: maxbar = 200
          self.valmax = valmax
          self.maxbar = maxbar
          self.title  = title
          print ('')
  
      def update(self, val, avg_loss=0):
          # format
          if val > self.valmax: val = self.valmax
  
          # process
          perc  = round((float(val) / float(self.valmax)) * 100)
          scale = 100.0 / float(self.maxbar)
          bar   = int(perc / scale)
  
          # render
          if avg_loss:
              # out = '\r %20s [%s%s] %3d / %3d  cost: %.2f  r_loss: %.0f  l_loss: %.4f  clf_loss: %.4f' % (
              out = '\r %20s [%s%s] %3d / %3d  loss: %.5f' % (
                  self.title, 
                  '=' * bar, ' ' * (self.maxbar - bar), 
                  val, 
                  self.valmax, 
                  avg_loss, 
                  )
          else:
              out = '\r %20s [%s%s] %3d / %3d ' % (self.title, '=' * bar, ' ' * (self.maxbar - bar), val, self.valmax)
  
          sys.stdout.write(out)
          sys.stdout.flush()
  
  
  def pad(l, sr):
      # 0-Pad 10 sec at fs hz and add little noise
      z = np.zeros(10*sr, dtype='float32')
      z[:l.size] = l
      z = z + 5*1e-4*np.random.rand(z.size).astype('float32')
      return z
  
  
  def compute_spectrogram(filename, sr=22000, n_mels=96):
      # zero pad and compute log mel spec
      try:
          audio, sr = librosa.load(filename, sr=sr, res_type='kaiser_fast')
      except:
          audio, o_sr = sf.read(filename)
          audio = librosa.core.resample(audio, o_sr, sr)
      try:
          x = pad(audio, sr)
      except ValueError:
          x = audio
      audio_rep = librosa.feature.melspectrogram(y=x, sr=sr, hop_length=512, n_fft=1024, n_mels=n_mels, power=1.)
      audio_rep = np.log(audio_rep + np.finfo(np.float32).eps)
      return audio_rep
  
  
  def return_spectrogram_max_nrg_frame(spectrogram):
      frames = librosa.util.frame(np.asfortranarray(spectrogram), frame_length=96, hop_length=12)
      idx_max_nrg = np.argmax(np.sum(np.sum(frames, axis=0), axis=0))
      return frames[:,:,idx_max_nrg]
  
  
  def return_spectrogram_3_max_nrg_frames(spectrogram):
      frames = librosa.util.frame(np.asfortranarray(spectrogram), frame_length=96, hop_length=12)
      idxes_max_nrg = (-np.sum(np.sum(frames, axis=0), axis=0)).argsort()[:3]
      return frames[:,:,idxes_max_nrg]
  
  
  def spectrogram_to_audio(filename, y, sr=22000):
      y = np.exp(y)
      x = librosa.feature.inverse.mel_to_audio(y, sr=sr, n_fft=1024, hop_length=512, power=1.)
      librosa.output.write_wav(filename, x, sr)
  
  
  def kullback_leibler(y_hat, y):
      """Generalized Kullback Leibler divergence.
      :param y_hat: The predicted distribution.
      :type y_hat: torch.Tensor
      :param y: The true distribution.
      :type y: torch.Tensor
      :return: The generalized Kullback Leibler divergence\
               between predicted and true distributions.
      :rtype: torch.Tensor
      """
      return (y * (y.add(1e-5).log() - y_hat.add(1e-5).log()) + (y_hat - y)).sum(dim=-1).mean()
  
  
  def embeddings_to_cosine_similarity_matrix(z):
      """Converts a a tensor of n embeddings to an (n, n) tensor of similarities.
      """
      cosine_similarity = torch.matmul(z, z.t())
      embedding_norms = torch.norm(z, p=2, dim=1)
      embedding_norms_mat = embedding_norms.unsqueeze(0)*embedding_norms.unsqueeze(1)
      cosine_similarity = cosine_similarity / (embedding_norms_mat)
      return cosine_similarity
  
  
  def contrastive_loss(z_audio, z_tag, t=1):
      """Computes contrastive loss following the paper:
          A Simple Framework for Contrastive Learning of Visual Representations
          https://arxiv.org/pdf/2002.05709v1.pdf
          TODO: make it robust to NaN (with low values of t it happens). 
          e.g Cast to double float for exp calculation.
      """
      z = torch.cat((z_audio, z_tag), dim=0)
      s = embeddings_to_cosine_similarity_matrix(z)
      N = int(s.shape[0]/2)
      s = torch.exp(s/t)
      try:
          s = s * (1 - torch.eye(len(s), len(s)).cuda())
          # s[range(len(s)), range(len(s))] = torch.zeros((len(s),)).cuda()
      except AssertionError:
          s = s * (1 - torch.eye(len(s), len(s)))
      denom = s.sum(dim=-1)
      num = torch.cat((s[:N,N:].diag(), s[N:,:N].diag()), dim=0)
      return torch.log((num / denom) + 1e-5).neg().mean()
  
  
  def contrastive_corr_loss(z_audio, z_tag, t=1):
      s = corre_loss(z_audio, z_tag)
      N = int(s.shape[0]/2)
      try:
          s = s * (1 - torch.eye(len(s), len(s)).cuda())
          # s[range(len(s)), range(len(s))] = torch.zeros((len(s),)).cuda()
      except AssertionError:
          s = s * (1 - torch.eye(len(s), len(s)))
      denom = s.sum(dim=-1)
      num = torch.cat((s[:N,N:].diag(), s[N:,:N].diag()), dim=0)
      return torch.log((num / denom) + 1e-5).neg().mean()
  