student wav2small 17K params

app.py

@@ -17,21 +17,6 @@ plt.style.use('seaborn-v0_8-whitegrid')
 
 
 
-def _prenorm(x, attention_mask=None):
-    '''mean/var'''
-    if attention_mask is not None:
-        N = attention_mask.sum(1, keepdim=True)  # 0=ignored 1=valid
-        x -= x.sum(1, keepdim=True) / N
-        var = (x * x).sum(1, keepdim=True) / N
-
-    else:
-        x -= x.mean(1, keepdim=True)  # mean is an onnx operator reducemean saves some ops compared to casting integer N to float and the div
-        var = (x * x).mean(1, keepdim=True)
-    return x / torch.sqrt(var + 1e-7)
-
-
-
-
 class ADV(nn.Module):
 
     def __init__(self, config):
@@ -96,16 +81,275 @@ dawn = Dawn.from_pretrained(
 ).to(device).eval()
 
 
-def wav2small(x):
-    return .5 * dawn(x) + .5 * base(x)
 
+# Wav2Small
+
+
+
+import torch
+import numpy as np
+import torch.nn.functional as F
+import librosa
+from transformers.models.wav2vec2.modeling_wav2vec2 import Wav2Vec2PreTrainedModel, Wav2Vec2Model
+from torch import nn
+from transformers import PretrainedConfig
+
+
+def _prenorm(x, attention_mask=None):
+    '''mean/var'''
+    if attention_mask is not None:
+        N = attention_mask.sum(1, keepdim=True)  # here attn msk is unprocessed just the original input
+        x -= x.sum(1, keepdim=True) / N
+        var = (x * x).sum(1, keepdim=True) / N
+
+    else:
+        x -= x.mean(1, keepdim=True)  # mean is an onnx operator reducemean saves some ops compared to casting integer N to float and the div
+        var = (x * x).mean(1, keepdim=True)
+    return x / torch.sqrt(var + 1e-7)
+
+
+
+
+class Spectrogram(nn.Module):
+    def __init__(self,
+        n_fft=64,   # num cols of DFT
+        n_time=64,  # num rows of DFT matrix
+        hop_length=32,
+        freeze_parameters=True):
+
+
+        super().__init__()
+
+        fft_window = librosa.filters.get_window('hann', n_time, fftbins=True)
+
+        fft_window = librosa.util.pad_center(fft_window, size=n_time)
+
+        
+        
+        
+
+        out_channels = n_fft // 2 + 1
+        
+        (x, y) = np.meshgrid(np.arange(n_time), np.arange(n_fft))
+        omega = np.exp(-2 * np.pi * 1j / n_time)
+        dft_matrix = np.power(omega, x * y)  # (n_fft, n_time)
+        dft_matrix = dft_matrix * fft_window[None, :]
+        dft_matrix = dft_matrix[0 : out_channels, :]
+        dft_matrix = dft_matrix[:, None, :]
+        
+        # ---- Assymetric DFT Non Square
+
+        self.conv_real = nn.Conv1d(1, out_channels, n_fft, stride=hop_length, padding=0, bias=False)
+        self.conv_imag = nn.Conv1d(1, out_channels, n_fft, stride=hop_length, padding=0, bias=False)
+        self.conv_real.weight.data = torch.tensor(np.real(dft_matrix), dtype=self.conv_real.weight.dtype).to(self.conv_real.weight.device)
+        self.conv_imag.weight.data = torch.tensor(np.imag(dft_matrix), dtype=self.conv_imag.weight.dtype).to(self.conv_imag.weight.device)
+        if freeze_parameters:
+            for param in self.parameters():
+                param.requires_grad = False
+
+    def forward(self, input):
+        x = input[:, None, :]
+
+        real = self.conv_real(x)
+        imag = self.conv_imag(x)
+        return real ** 2 + imag ** 2  # bs, mel, time-frames
+
+
+class LogmelFilterBank(nn.Module):
+    def __init__(self,
+                 sr=16000,
+                 n_fft=64,
+                 n_mels=26,  # maxpool
+                 fmin=0.0,
+                 freeze_parameters=True):
+
+        super().__init__()
+
+        fmax = sr//2
+
+        W2 = librosa.filters.mel(sr=sr,
+                                        n_fft=n_fft,
+                                        n_mels=n_mels,
+                                        fmin=fmin,
+                                        fmax=fmax).T
+
+        self.register_buffer('melW', torch.Tensor(W2))
+        self.register_buffer('amin', torch.Tensor([1e-10]))
+
+    def forward(self, x):
+
+        x = torch.matmul(x[:, None, :, :].transpose(2, 3), self.melW)   # changes melf not num frames
+
+        x = torch.where(x > self.amin, x, self.amin)  # not in place
+
+        x = 10 * torch.log10(x)
+        return x
+
+
+
+
+
+def length_after_conv_layer(_length, k=None, pad=None, stride=None):
+        return torch.floor( (_length + 2*pad - k) / stride + 1 )
+
+
+
+
+
+
+class Conv(nn.Module):
+
+    def __init__(self, c_in, c_out, k=3, stride=1, padding=1):
+
+        super().__init__()
+
+        self.conv = nn.Conv2d(c_in, c_out, k, stride=stride, padding=padding, bias=False)
+        self.norm = nn.BatchNorm2d(c_out)
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.norm(x)
+        return torch.relu_(x)
 
-fig_error, ax = plt.subplots(figsize=(8, 6))
 
-# Set the text to display
-error_message = "Error: No .wav or Mic. audio provided."
 
-# Add the text to the plot. We'll place it in the center of the plot
+
+
+class Vgg7(nn.Module):
+
+    def __init__(self):
+
+        super().__init__()
+
+        self.l1 = Conv( 1, 13)
+        self.l2 = Conv(13, 13)
+        self.l3 = Conv(13, 13)
+        self.maxpool_A = nn.MaxPool2d(3,
+                                      stride=2,
+                                      padding=1)
+        self.l4 = Conv(13, 13)
+        self.l5 = Conv(13, 13)
+        self.l6 = Conv(13, 13)
+        self.l7 = Conv(13, 13)
+        self.lin = nn.Conv2d(13, 13, 1, padding=0, stride=1)
+        self.sof = nn.Conv2d(13, 13, 1, padding=0, stride=1)  # pool time - reshape mel into channels after pooling
+        self.spectrogram_extractor = Spectrogram()
+        self.logmel_extractor = LogmelFilterBank()
+
+    def final_length(self, L):
+            conv_kernel = [64, 3]  # [nfft, maxpool]
+            conv_stride = [32, 2]  # [hop_len, maxpool_stride]    # consider only layers of stride > 1
+            conv_pad =    [0,  1]   # [pad_stft, pad_maxpool]
+            for k, stride, pad in zip(conv_kernel, conv_stride, conv_pad):
+                L = length_after_conv_layer(L, k=k, stride=stride, pad=pad)
+            return L
+
+    def final_attention_mask(self, feature_vector_length, attention_mask=None):
+            non_padded_lengths = attention_mask.sum(1)
+            out_lengths = self.final_length(non_padded_lengths)      # how can non_padded_lengths get exact 0 here DOES IT MEAN ATTNMASK WAS NOT FILLED?
+            out_lengths = out_lengths.to(torch.long)
+            bs, _ = attention_mask.shape
+            attention_mask = torch.ones((bs, feature_vector_length),
+                                        dtype=attention_mask.dtype,
+                                        device=attention_mask.device)
+            for b, _len in enumerate(out_lengths):
+                attention_mask[b, _len:] = 0
+            return attention_mask
+
+    def forward(self, x, attention_mask=None):
+        x = _prenorm(x,
+                     attention_mask=attention_mask)
+        x = self.spectrogram_extractor(x)
+        x = self.logmel_extractor(x)
+        x = self.l1(x)
+        x = self.l2(x)
+        x = self.l3(x)
+        x = self.maxpool_A(x)  # reshape here? so these conv will have large kernel
+        x = self.l4(x)
+        x = self.l5(x)
+        x = self.l6(x)
+        x = self.l7(x)
+        if attention_mask is not None:
+            bs, _, t, _ = x.shape
+            a = self.final_attention_mask(feature_vector_length=t,
+                                          attention_mask=attention_mask)[:, None, :, None]
+            #print(a.shape, x.shape, '\n\n\n\n')
+            x = torch.masked_fill(x, a < 1, 0)
+            # mask also affects lin !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
+            x = self.lin(x) * (  self.sof(x) -10000. * torch.logical_not(a)  ).softmax(2)
+        else:
+           x = self.lin(x) * self.sof(x).softmax(2)
+
+        x = x.sum(2)  # bs, ch, time-frames, HALF_MEL -> bs, ch, HALF_MEL
+        # --
+        xT = x.transpose(1,2)
+        x = torch.cat([x,
+                       torch.bmm(x, xT),   # corr  (chxmel) x (melxCH)
+                    #    torch.bmm(x, x),   # corr ch * ch
+                    #    torch.bmm(xT, xT)  # corr mel * mel
+                       ], 2)
+        # --
+        return x.reshape(-1, 338)
+
+
+class Wav2SmallConfig(PretrainedConfig):
+    model_type = "wav2vec2"
+
+    def __init__(self,
+                 **kwargs):
+        super().__init__(**kwargs)
+        self.half_mel = 13
+        self.n_fft = 64
+        self.n_time = 64
+        self.hidden = 2 * self.half_mel * self.half_mel
+        self.hop = self.n_time // 2
+
+
+class Wav2Small(Wav2Vec2PreTrainedModel):
+
+    def __init__(self,
+                 config):
+        super().__init__(config)
+        self.vgg7 = Vgg7()
+        self.adv     = nn.Linear(config.hidden, 3)   # 0=arousal, 1=dominance, 2=valence
+
+    def forward(self, x, attention_mask=None):
+        x = self.vgg7(x, attention_mask=attention_mask)
+        return self.adv(x)
+    
+    
+    
+
+
+
+
+def _ccc(x, y):
+    '''if len(x) = len(y) = 1 we have 0/0 as a&b can both be negative we should add 1e-7 to denominator protecting sign of denominator
+       to find sign of denominator and add 1e-7 if sgn>=0 or -1e-7 if sgn<0'''
+
+    mean_y = y.mean()
+    mean_x = x.mean()
+    a = x - mean_x
+    b = y - mean_y
+    L = (mean_x - mean_y).abs() * .1 * x.shape[0]
+    #print(L / ((mean_x - mean_y) **2  * x.shape[0]))
+    numerator = torch.dot(a, b)    # L term  if both a,b scalars dissallows 0 numerator [OFFICIAL CCC HAS L ONLY IN D]
+    denominator = torch.dot(a, a) + torch.dot(b, b) + L  # if both a,b are equalscalars then the dots are all zero and ccc=1
+    denominator = torch.where(denominator.sign() < 0,
+                              denominator - 1e-7,
+                              denominator + 1e-7)
+    ccc = numerator / denominator
+
+    return -ccc #+ F.l1_loss(a, b)
+
+
+
+wav2small = Wav2Small.from_pretrained('audeering/wav2small').to(device).eval()
+
+
+# Error figure for the first plot
+fig_error, ax = plt.subplots(figsize=(8, 6))
+error_message = "Error: No .wav or Mic. audio provided."
 ax.text(0.5, 0.5, error_message,
         ha='center',
         va='center',
@@ -113,125 +357,164 @@ ax.text(0.5, 0.5, error_message,
         color='gray',
         fontweight='bold',
         transform=ax.transAxes)
-
-# Hide the axis ticks and labels for a cleaner look
 ax.set_xticks([])
 ax.set_yticks([])
 ax.set_xticklabels([])
 ax.set_yticklabels([])
-
-# Optional: Add a border around the text to make it stand out more
 ax.set_frame_on(True)
 ax.spines['top'].set_visible(False)
 ax.spines['right'].set_visible(False)
 ax.spines['bottom'].set_visible(False)
 ax.spines['left'].set_visible(False)
 
-
-
-
-
 def process_audio(audio_filepath):
     if audio_filepath is None:
-        return fig_error
+        
+        return fig_error, fig_error
 
-    # Load the audio file
-    waveform, sample_rate = librosa.load(audio_filepath)
+    waveform, sample_rate = librosa.load(audio_filepath, sr=None)
 
-    # Ensure audio is mono: if stereo, take the mean across channels
-
-    # Resample audio to 16kHz if necessary
+    # Resample audio to 16kHz if the sample rate is different
     if sample_rate != 16000:
         resampled_waveform_np = audresample.resample(waveform, sample_rate, 16000)
-        x = torch.from_numpy(resampled_waveform_np)
-    x = x[:, :64000]  # 4s
+    else:
+        resampled_waveform_np = waveform[None, :]
+
+    x = torch.from_numpy(resampled_waveform_np).to(torch.float)
+
     with torch.no_grad():
+
         logits_dawn = dawn(x).cpu().numpy()[0, :]
-        logits_wavlm = base(x).cpu().numpy()[0, :]
 
-        logits_wav2small = .5 * logits_dawn + .5 * logits_wavlm
+        logits_wavlm = base(x).cpu().numpy()[0, :]        
+
+        # 17K params
+        logits_wav2small = wav2small(x).cpu().numpy()[0, :]
+
+
+    # --- Plot 1: Wav2Vec2 vs Wav2Small Teacher Outputs ---
+    
+    fig, ax = plt.subplots(figsize=(10, 6))
 
-    # left_bars_data = np.array([0.75, 0.5, 0.9])
-    # right_bars_data = np.array([0.3, 0.8, 0.65])
     left_bars_data = logits_dawn.clip(0, 1)
     right_bars_data = logits_wav2small.clip(0, 1)
 
-
     bar_labels = ['\nArousal', '\nDominance', '\nValence']
     y_pos = np.arange(len(bar_labels))
 
-    # Define the base colormaps for each category to ensure a different color per row
-    # Using Greys for Dominance as requested
+    # Define colormaps for each category to ensure distinct colors
     category_colormaps = [plt.cm.Blues, plt.cm.Greys, plt.cm.Oranges]
 
-    # Define color shades for left and right for each category
     left_filled_colors = []
     right_filled_colors = []
     background_colors = []
 
+    # Assign specific shades for filled bars and background bars
     for i, cmap in enumerate(category_colormaps):
-        # Pick a darker shade for the left filled bar
-        left_filled_colors.append(cmap(0.74))  # 0.7
-        # Pick a slightly lighter shade for the right filled bar
-        right_filled_colors.append(cmap(0.64))  # 0.5
-        # Pick a very light shade for the transparent background bar
+        left_filled_colors.append(cmap(0.74))
+        right_filled_colors.append(cmap(0.64))
         background_colors.append(cmap(0.1))
 
-    # Set up the figure and axes
-    fig, ax = plt.subplots(figsize=(10, 6))
-
-    # Plot the background bars with transparency
+    # Plot transparent background bars
     for i in range(len(bar_labels)):
-        # Left background bar (transparent, light shade of category color)
         ax.barh(y_pos[i], -1, color=background_colors[i], alpha=0.3, height=0.6)
-        # Right background bar (transparent, light shade of category color)
         ax.barh(y_pos[i], 1, color=background_colors[i], alpha=0.3, height=0.6)
 
-    # Plot the filled bars for the left and right side
+    # Plot the filled bars for actual data
     for i in range(len(bar_labels)):
-        # Left filled bar (opaque, darker shade of category color)
         ax.barh(y_pos[i], -left_bars_data[i], color=left_filled_colors[i], alpha=1, height=0.6)
-        # Right filled bar (opaque, lighter shade of category color)
         ax.barh(y_pos[i], right_bars_data[i], color=right_filled_colors[i], alpha=1, height=0.6)
 
-    # Add a central axis divider
+    # Add a central vertical axis divider
     ax.axvline(0, color='black', linewidth=0.8, linestyle='--')
 
-    # Set x-axis limits and y-axis ticks
+    # Set x-axis limits and y-axis ticks/labels
     ax.set_xlim(-1, 1)
     ax.set_yticks(y_pos)
     ax.set_yticklabels(bar_labels, fontsize=12)
 
-    
+    # Custom formatter for x-axis to show absolute percentage values
     def abs_tick_formatter(x, pos):
         return f'{int(abs(x) * 100)}%'
     ax.xaxis.set_major_formatter(plt.FuncFormatter(abs_tick_formatter))
 
-    # Add a clean title and labels
+    # Set plot title and x-axis label
     ax.set_title('', fontsize=16, pad=20)
-    ax.set_xlabel('Outputs of Wav2Vev2                                               Outputs of Wav2Small Teacher', fontsize=12)
+    ax.set_xlabel('Wav2Vev2 (Dawn)                                               Wav2Small (17K param.)', fontsize=12)
 
-    # Remove the top and right spines for a cleaner look
+    # Remove top, right, and left spines for a cleaner look
     ax.spines['top'].set_visible(False)
     ax.spines['right'].set_visible(False)
     ax.spines['left'].set_visible(False)
 
-    # Add annotations to the filled bars for clarity
+    # Add annotations (percentage values) to the filled bars
     for i in range(len(bar_labels)):
-        # Left annotation (uses left_filled_colors for text color)
         ax.text(-left_bars_data[i] - 0.05, y_pos[i], f'{int(left_bars_data[i] * 100)}%',
                 va='center', ha='right', color=left_filled_colors[i], fontweight='bold')
-        # Right annotation (uses right_filled_colors for text color)
         ax.text(right_bars_data[i] + 0.05, y_pos[i], f'{int(right_bars_data[i] * 100)}%',
                 va='center', ha='left', color=right_filled_colors[i], fontweight='bold')
 
 
-    return fig
 
+    # -- PLOT 2 : WavLM / Wav2Small Teacher
+
+    fig_2, ax_2 = plt.subplots(figsize=(10, 6))
+
+    
+    left_bars_data = logits_wavlm.clip(0, 1)
+    right_bars_data = (.5 * logits_dawn + .5 * logits_wavlm).clip(0, 1)
+
+    bar_labels = ['\nArousal', '\nDominance', '\nValence']
+    y_pos = np.arange(len(bar_labels))
+
+    # Define colormaps for each category to ensure distinct colors
+    category_colormaps = [plt.cm.Blues, plt.cm.Greys, plt.cm.Oranges]
+
+    left_filled_colors = []
+    right_filled_colors = []
+    background_colors = []
+
+    # Assign specific shades for filled bars and background bars
+    for i, cmap in enumerate(category_colormaps):
+        left_filled_colors.append(cmap(0.74))
+        right_filled_colors.append(cmap(0.64))
+        background_colors.append(cmap(0.1))
 
+    # Plot transparent background bars
+    for i in range(len(bar_labels)):
+        ax_2.barh(y_pos[i], -1, color=background_colors[i], alpha=0.3, height=0.6)
+        ax_2.barh(y_pos[i], 1, color=background_colors[i], alpha=0.3, height=0.6)
 
+    # Plot the filled bars for actual data
+    for i in range(len(bar_labels)):
+        ax_2.barh(y_pos[i], -left_bars_data[i], color=left_filled_colors[i], alpha=1, height=0.6)
+        ax_2.barh(y_pos[i], right_bars_data[i], color=right_filled_colors[i], alpha=1, height=0.6)
+
+    # Add a central vertical axis divider
+    ax_2.axvline(0, color='black', linewidth=0.8, linestyle='--')
 
+    # Set x-axis limits and y-axis ticks/labels
+    ax_2.set_xlim(-1, 1)
+    ax_2.set_yticks(y_pos)
+    ax_2.set_yticklabels(bar_labels, fontsize=12)
 
+    # Custom formatter for x-axis to show absolute percentage values
+    def abs_tick_formatter(x, pos):
+        return f'{int(abs(x) * 100)}%'
+    ax_2.xaxis.set_major_formatter(plt.FuncFormatter(abs_tick_formatter))
+    ax_2.set_title('', fontsize=16, pad=20)
+    ax_2.set_xlabel('WavLM (Baseline)                                              Wav2Small Teacher (0.4B param.)', fontsize=12)
+    ax_2.spines['top'].set_visible(False)
+    ax_2.spines['right'].set_visible(False)
+    ax_2.spines['left'].set_visible(False)
+
+    # Add annotations (percentage values) to the filled bars
+    for i in range(len(bar_labels)):
+        ax_2.text(-left_bars_data[i] - 0.05, y_pos[i], f'{int(left_bars_data[i] * 100)}%',
+                va='center', ha='right', color=left_filled_colors[i], fontweight='bold')
+        ax_2.text(right_bars_data[i] + 0.05, y_pos[i], f'{int(right_bars_data[i] * 100)}%',
+                va='center', ha='left', color=right_filled_colors[i], fontweight='bold')
+    return fig, fig_2
 
 
 iface = gr.Interface(
@@ -242,25 +525,27 @@ iface = gr.Interface(
         label=''
     ),
     outputs=[
-        gr.Plot(label="Arousal / Dominance / Valence Plots"),
+        gr.Plot(label="Wav2Vec2 vs Wav2Small (17K params) Plot"),   # First plot output
+        gr.Plot(label="WavLM vs Wav2Small Teacher Plot"),         # Second plot output
     ],
     title='',
     description='',
-    flagging_mode="never",  # save audio and .csv in the machine ?
+    flagging_mode="never",  # Disables flagging feature
     examples=[
                         "female-46-neutral.wav",
                         "female-20-happy.wav",
                         "male-60-angry.wav",
                         "male-27-sad.wav",
             ],
-    css="footer {visibility: hidden}"
+    css="footer {visibility: hidden}" # Hides the Gradio footer
 )
 
+# Gradio Blocks for tabbed interface
 with gr.Blocks() as demo:
-
-    # https://discuss.huggingface.co/t/how-to-get-the-microphone-streaming-input-file-when-using-blocks/37204/3
+    # First tab for the existing Arousal/Dominance/Valence plots
     with gr.Tab(label="Arousal / Dominance / Valence"):
         iface.render()
+    # Second tab for CCC (Concordance Correlation Coefficient) information
     with gr.Tab(label="CCC"):
        gr.Markdown('''<table style="width:500px"><tr><th colspan=5 >CCC MSP Podcast v1.7</th></tr>
   <tr> <td> </td><td>Arousal</td> <td>Dominance</td> <td>Valence</td> <td> Associated Paper </td> </tr>
@@ -269,5 +554,6 @@ with gr.Blocks() as demo:
 </table>
 ''')
 
+# Launch the Gradio application
 if __name__ == "__main__":
     demo.launch(share=False)