<commit_before>import torch
import torch.nn as nn
import librosa
import numpy as np
from .clustering import GMM
from torch.utils.checkpoint import checkpoint


class MelProjection(nn.Module):
    def __init__(self, sample_rate, num_frequencies, num_mels, direction='forward', clamp=False, trainable=False):
        """
        MelProjection takes as input a time-frequency representation (e.g. a spectrogram, or a mask) and outputs a mel
        project that can be learned or fixed. The initialization uses librosa to get a mel filterbank. Direction
        controls whether it is a forward transform or the inverse transform (e.g. back to spectrogram).

        Args:
            sample_rate: (int) Sample rate of audio for computing the mel filters.
            num_frequencies: (int) Number of frequency bins in input spectrogram.
            num_mels: (int) Number of mel bins in output mel spectrogram. if num_mels < 0, this does nothing
                other than clamping the output of clamp is True
            direction: (str) Which direction to go in (either 'forward' - to mel, or 'backward' - to frequencies).
                Defaults to 'forward'.
            clamp: (bool) Whether to clamp the output values of the transform between 0.0 and 1.0. Used for transforming
                a mask in and out of the mel-domain. Defaults to False.
            trainable: (bool) Whether the mel transform can be adjusted by the optimizer. Defaults to False.
        """
        super(MelProjection, self).__init__()
        self.num_mels = num_mels
        if direction not in ['backward', 'forward']:
            raise ValueError("direction must be one of ['backward', 'forward']!")
        self.direction = direction
        self.clamp = clamp

        if self.num_mels > 0:
            shape = (num_frequencies, num_mels) if self.direction == 'forward' else (num_mels, num_frequencies)
            self.add_module('transform', nn.Linear(*shape))

            mel_filters = librosa.filters.mel(sample_rate, 2 * (num_frequencies - 1), num_mels)
            mel_filters = (mel_filters.T / mel_filters.sum(axis=1)).T
            filter_bank = mel_filters if self.direction == 'forward' else np.linalg.pinv(mel_filters)

            for name, param in self.transform.named_parameters():
                if 'bias' in name:
                    nn.init.constant_(param, 0.0)
                if 'weight' in name:
                    param.data = torch.from_numpy(filter_bank).float()
                param.requires_grad_(trainable)

    def forward(self, data):
        """
        Args:
            data: Representation - shape: (batch_size, sequence_length, num_frequencies or num_mels, num_sources)

        Returns:
            Mel-spectrogram or time-frequency representation of shape:
                (batch_size, sequence_length, num_mels or num_frequencies, num_sources).
        """

        if self.num_mels > 0:
            data = data.transpose(2, -1)
            data = self.transform(data)
            data = data.transpose(2, -1)

        if self.clamp:
            data = data.clamp(0.0, 1.0)
        return data


class Embedding(nn.Module):
    def __init__(self, num_features, hidden_size, embedding_size, activation):
        """
        Maps output from an audio representation module (e.g. RecurrentStack, DilatedConvolutionalStack) to an embedding
        space. The output shape is (batch_size, sequence_length, num_features, embedding_size). The embeddings can
        be passed through an activation function. If activation is 'softmax' or 'sigmoid', and embedding_size is equal
        to the number of sources, this module can be used to implement a mask inference network (or a mask inference
        head in a Chimera network setup).

        Args:
            num_features: (int) Number of features being mapped for each frame. Either num_frequencies, or if used with
                MelProjection, num_mels if using RecurrentStack. Should be 1 if using DilatedConvolutionalStack. 
            hidden_size: (int) Size of output from RecurrentStack (hidden_size) or DilatedConvolutionalStack (num_filters). If
                RecurrentStack is bidirectional, this should be set to 2 * hidden_size.
            embedding_size: (int) Dimensionality of embedding.
            activation: (list of str) Activation functions to be applied. Options are 'sigmoid', 'tanh', 'softmax'.
                Unit normalization can be applied by adding 'unit_norm' in list (e.g. ['sigmoid', unit_norm']).
        """
        super(Embedding, self).__init__()
        self.add_module('linear', nn.Linear(hidden_size, num_features * embedding_size))
        self.num_features = num_features
        self.activation = activation
        self.embedding_size = embedding_size

        for name, param in self.linear.named_parameters():
            if 'bias' in name:
                nn.init.constant_(param, 0.0)
            elif 'weight' in name:
                nn.init.xavier_normal_(param)


    def forward(self, data):
        """
        Args:
            data: output from RecurrentStack or ConvolutionalStack. Shape is:
                (num_batch, sequence_length, hidden_size or num_filters)
        Returns:
            An embedding (with an optional activation) for each point in the representation of shape
            (num_batch, sequence_length, num_features, embedding_size).
        """
        data = self.linear(data)

        if len(data.shape) < 4:
            # Then this is the output of RecurrentStack and needs to be reshaped a bit.
            # Came in as [num_batch, sequence_length, num_features * embedding_size]
            # Goes out as [num_batch, sequence_length, num_features, embedding_size]
            data = data.view(
                data.shape[0], data.shape[1], self.num_features,  self.embedding_size
            )

        if 'sigmoid' in self.activation:
            data = torch.sigmoid(data)
        elif 'tanh' in self.activation:
            data = torch.tanh(data)
        elif 'relu' in self.activation:
            data = torch.relu(data)

        if 'unit_norm' in self.activation:
            data = nn.functional.normalize(data, dim=-1, p=2)

        return data

class Mask(nn.Module):
    def __init__(self):
        super(Mask, self).__init__()

    def forward(self, mask, magnitude_spectrogram):
        return mask * magnitude_spectrogram.unsqueeze(-1)

class RecurrentStack(nn.Module):
    def __init__(self, num_features, hidden_size, num_layers, bidirectional, dropout, rnn_type='lstm'):
        """
        Creates a stack of RNNs used to process an audio sequence represented as (sequence_length, num_features). With
        bidirectional = True, hidden_size = 600, num_layers = 4, rnn_type='lstm', and dropout = .3, this becomes
        the state-of-the-art audio processor used in deep clustering networks, deep attractor networks, etc. Note that
        batch_first is set to True here.

        Args:
            num_features: (int) Number of features being mapped for each frame. Either num_frequencies, or if used with
                MelProjection, num_mels.
            hidden_size: (int) Hidden size of recurrent stack for each layer.
            num_layers: (int) Number of layers in stack.
            bidirectional: (int) True makes this a BiLSTM or a BiGRU. Note that this doubles the hidden size.
            dropout: (float) Dropout between layers.
            rnn_type: (str) LSTM ('lstm') or GRU ('gru').
        """
        super(RecurrentStack, self).__init__()
        if rnn_type not in ['lstm', 'gru']:
            raise ValueError("rnn_type must be one of ['lstm', 'gru']!")

        if rnn_type == 'lstm':
            self.add_module('rnn', nn.LSTM(num_features,
                                           hidden_size,
                                           num_layers,
                                           batch_first=True,<commit_msg>fixing a bug in mask inference<commit_after>import torch
import torch.nn as nn
import librosa
import numpy as np
from .clustering import GMM
from torch.utils.checkpoint import checkpoint


class MelProjection(nn.Module):
    def __init__(self, sample_rate, num_frequencies, num_mels, direction='forward', clamp=False, trainable=False):
        """
        MelProjection takes as input a time-frequency representation (e.g. a spectrogram, or a mask) and outputs a mel
        project that can be learned or fixed. The initialization uses librosa to get a mel filterbank. Direction
        controls whether it is a forward transform or the inverse transform (e.g. back to spectrogram).

        Args:
            sample_rate: (int) Sample rate of audio for computing the mel filters.
            num_frequencies: (int) Number of frequency bins in input spectrogram.
            num_mels: (int) Number of mel bins in output mel spectrogram. if num_mels < 0, this does nothing
                other than clamping the output of clamp is True
            direction: (str) Which direction to go in (either 'forward' - to mel, or 'backward' - to frequencies).
                Defaults to 'forward'.
            clamp: (bool) Whether to clamp the output values of the transform between 0.0 and 1.0. Used for transforming
                a mask in and out of the mel-domain. Defaults to False.
            trainable: (bool) Whether the mel transform can be adjusted by the optimizer. Defaults to False.
        """
        super(MelProjection, self).__init__()
        self.num_mels = num_mels
        if direction not in ['backward', 'forward']:
            raise ValueError("direction must be one of ['backward', 'forward']!")
        self.direction = direction
        self.clamp = clamp

        if self.num_mels > 0:
            shape = (num_frequencies, num_mels) if self.direction == 'forward' else (num_mels, num_frequencies)
            self.add_module('transform', nn.Linear(*shape))

            mel_filters = librosa.filters.mel(sample_rate, 2 * (num_frequencies - 1), num_mels)
            mel_filters = (mel_filters.T / mel_filters.sum(axis=1)).T
            filter_bank = mel_filters if self.direction == 'forward' else np.linalg.pinv(mel_filters)

            for name, param in self.transform.named_parameters():
                if 'bias' in name:
                    nn.init.constant_(param, 0.0)
                if 'weight' in name:
                    param.data = torch.from_numpy(filter_bank).float()
                param.requires_grad_(trainable)

    def forward(self, data):
        """
        Args:
            data: Representation - shape: (batch_size, sequence_length, num_frequencies or num_mels, num_sources)

        Returns:
            Mel-spectrogram or time-frequency representation of shape:
                (batch_size, sequence_length, num_mels or num_frequencies, num_sources).
        """

        if self.num_mels > 0:
            data = data.transpose(2, -1)
            data = self.transform(data)
            data = data.transpose(2, -1)

        if self.clamp:
            data = data.clamp(0.0, 1.0)
        return data


class Embedding(nn.Module):
    def __init__(self, num_features, hidden_size, embedding_size, activation):
        """
        Maps output from an audio representation module (e.g. RecurrentStack, DilatedConvolutionalStack) to an embedding
        space. The output shape is (batch_size, sequence_length, num_features, embedding_size). The embeddings can
        be passed through an activation function. If activation is 'softmax' or 'sigmoid', and embedding_size is equal
        to the number of sources, this module can be used to implement a mask inference network (or a mask inference
        head in a Chimera network setup).

        Args:
            num_features: (int) Number of features being mapped for each frame. Either num_frequencies, or if used with
                MelProjection, num_mels if using RecurrentStack. Should be 1 if using DilatedConvolutionalStack. 
            hidden_size: (int) Size of output from RecurrentStack (hidden_size) or DilatedConvolutionalStack (num_filters). If
                RecurrentStack is bidirectional, this should be set to 2 * hidden_size.
            embedding_size: (int) Dimensionality of embedding.
            activation: (list of str) Activation functions to be applied. Options are 'sigmoid', 'tanh', 'softmax'.
                Unit normalization can be applied by adding 'unit_norm' in list (e.g. ['sigmoid', unit_norm']).
        """
        super(Embedding, self).__init__()
        self.add_module('linear', nn.Linear(hidden_size, num_features * embedding_size))
        self.num_features = num_features
        self.activation = activation
        self.embedding_size = embedding_size

        for name, param in self.linear.named_parameters():
            if 'bias' in name:
                nn.init.constant_(param, 0.0)
            elif 'weight' in name:
                nn.init.xavier_normal_(param)


    def forward(self, data):
        """
        Args:
            data: output from RecurrentStack or ConvolutionalStack. Shape is:
                (num_batch, sequence_length, hidden_size or num_filters)
        Returns:
            An embedding (with an optional activation) for each point in the representation of shape
            (num_batch, sequence_length, num_features, embedding_size).
        """
        data = self.linear(data)

        if len(data.shape) < 4:
            # Then this is the output of RecurrentStack and needs to be reshaped a bit.
            # Came in as [num_batch, sequence_length, num_features * embedding_size]
            # Goes out as [num_batch, sequence_length, num_features, embedding_size]
            data = data.view(
                data.shape[0], data.shape[1], self.num_features,  self.embedding_size
            )

        if 'sigmoid' in self.activation:
            data = torch.sigmoid(data)
        elif 'tanh' in self.activation:
            data = torch.tanh(data)
        elif 'relu' in self.activation:
            data = torch.relu(data)

        if 'unit_norm' in self.activation:
            data = nn.functional.normalize(data, dim=-1, p=2)

        return data

class Mask(nn.Module):
    def __init__(self):
        super(Mask, self).__init__()

    def forward(self, mask, magnitude_spectrogram):
        return mask * magnitude_spectrogram

class RecurrentStack(nn.Module):
    def __init__(self, num_features, hidden_size, num_layers, bidirectional, dropout, rnn_type='lstm'):
        """
        Creates a stack of RNNs used to process an audio sequence represented as (sequence_length, num_features). With
        bidirectional = True, hidden_size = 600, num_layers = 4, rnn_type='lstm', and dropout = .3, this becomes
        the state-of-the-art audio processor used in deep clustering networks, deep attractor networks, etc. Note that
        batch_first is set to True here.

        Args:
            num_features: (int) Number of features being mapped for each frame. Either num_frequencies, or if used with
                MelProjection, num_mels.
            hidden_size: (int) Hidden size of recurrent stack for each layer.
            num_layers: (int) Number of layers in stack.
            bidirectional: (int) True makes this a BiLSTM or a BiGRU. Note that this doubles the hidden size.
            dropout: (float) Dropout between layers.
            rnn_type: (str) LSTM ('lstm') or GRU ('gru').
        """
        super(RecurrentStack, self).__init__()
        if rnn_type not in ['lstm', 'gru']:
            raise ValueError("rnn_type must be one of ['lstm', 'gru']!")

        if rnn_type == 'lstm':
            self.add_module('rnn', nn.LSTM(num_features,
                                           hidden_size,
                                           num_layers,
                                           batch_first=True,
{'fraction_non_alphanumeric': 0.08141949285849116, 'fraction_numerical': 0.004341534008683068, 'mean_word_length': 2.3084929225645294, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 12, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2205527', 'n_tokens_mistral': 4465, 'n_tokens_neox': 4259, 'n_words': 1504}