Upload 4 files

hyperparams.yaml

@@ -4,7 +4,7 @@
 # ############################################################################
 
 # Hparams NEEDED
-HPARAMS_NEEDED: ["encoder_dim", "out_n_neurons", "accent_encoder", "softmax"]
+HPARAMS_NEEDED: ["encoder_dim", "out_n_neurons", "label_encoder", "softmax"]
 # Modules Needed
 MODULES_NEEDED: ["wav2vec2", "avg_pool", "output_mlp"]
 
@@ -27,7 +27,7 @@ wav2vec2: !new:speechbrain.lobes.models.huggingface_wav2vec.HuggingFaceWav2Vec2
     save_path: wav2vec2_checkpoints
 
 # Mean and std normalization of the input features
-mean_var_norm_input: !new:speechbrain.processing.features.InputNormalization
+mean_var_norm_emb: !new:speechbrain.processing.features.InputNormalization
     norm_type: sentence
     std_norm: False
 
@@ -43,25 +43,25 @@ model: !new:torch.nn.ModuleList
     - [!ref <output_mlp>]
 
 modules:
-    mean_var_norm_input: !ref <mean_var_norm_input>
+    mean_var_norm_emb: !ref <mean_var_norm_emb>
     wav2vec2: !ref <wav2vec2>
     output_mlp: !ref <output_mlp>
     avg_pool: !ref <avg_pool>
 
 softmax: !new:speechbrain.nnet.activations.Softmax
         
-accent_encoder: !new:speechbrain.dataio.encoder.CategoricalEncoder
+label_encoder: !new:speechbrain.dataio.encoder.CategoricalEncoder
 
 
 pretrainer: !new:speechbrain.utils.parameter_transfer.Pretrainer
     loadables:
-        mean_var_norm_input: !ref <mean_var_norm_input>
+        mean_var_norm_emb: !ref <mean_var_norm_emb>
         wav2vec2: !ref <wav2vec2>
         model: !ref <model>
-        accent_encoder: !ref <accent_encoder>
+        label_encoder: !ref <label_encoder>
     paths:
-        mean_var_norm_input: !ref <pretrained_path>/normalizer_input.ckpt
+        mean_var_norm_emb: !ref <pretrained_path>/normalizer_input.ckpt
         wav2vec2: !ref <pretrained_path>/wav2vec2.ckpt
         model: !ref <pretrained_path>/model.ckpt
-        accent_encoder: !ref <pretrained_path>/accent_encoder.txt
+        label_encoder: !ref <pretrained_path>/label_encoder.txt
 

config.json

@@ -1,3 +1,4 @@
 {
-    "speechbrain_interface": "EncoderClassifier"
+    "speechbrain_interface": "CustomEncoderWav2vec2Classifier",
+    "model_type": "wav2vec2"
 }

custom_interface.py

@@ -0,0 +1,158 @@
+import torch
+from speechbrain.pretrained import Pretrained
+
+
+class CustomEncoderWav2vec2Classifier(Pretrained):
+    """A ready-to-use class for utterance-level classification (e.g, speaker-id,
+    language-id, emotion recognition, keyword spotting, etc).
+
+    The class assumes that an self-supervised encoder like wav2vec2/hubert and a classifier model
+    are defined in the yaml file. If you want to
+    convert the predicted index into a corresponding text label, please
+    provide the path of the label_encoder in a variable called 'lab_encoder_file'
+    within the yaml.
+
+    The class can be used either to run only the encoder (encode_batch()) to
+    extract embeddings or to run a classification step (classify_batch()).
+    ```
+
+    Example
+    -------
+    >>> import torchaudio
+    >>> from speechbrain.pretrained import EncoderClassifier
+    >>> # Model is downloaded from the speechbrain HuggingFace repo
+    >>> tmpdir = getfixture("tmpdir")
+    >>> classifier = EncoderClassifier.from_hparams(
+    ...     source="speechbrain/spkrec-ecapa-voxceleb",
+    ...     savedir=tmpdir,
+    ... )
+
+    >>> # Compute embeddings
+    >>> signal, fs = torchaudio.load("samples/audio_samples/example1.wav")
+    >>> embeddings =  classifier.encode_batch(signal)
+
+    >>> # Classification
+    >>> prediction =  classifier .classify_batch(signal)
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def encode_batch(self, wavs, wav_lens=None, normalize=False):
+        """Encodes the input audio into a single vector embedding.
+
+        The waveforms should already be in the model's desired format.
+        You can call:
+        ``normalized = <this>.normalizer(signal, sample_rate)``
+        to get a correctly converted signal in most cases.
+
+        Arguments
+        ---------
+        wavs : torch.tensor
+            Batch of waveforms [batch, time, channels] or [batch, time]
+            depending on the model. Make sure the sample rate is fs=16000 Hz.
+        wav_lens : torch.tensor
+            Lengths of the waveforms relative to the longest one in the
+            batch, tensor of shape [batch]. The longest one should have
+            relative length 1.0 and others len(waveform) / max_length.
+            Used for ignoring padding.
+        normalize : bool
+            If True, it normalizes the embeddings with the statistics
+            contained in mean_var_norm_emb.
+
+        Returns
+        -------
+        torch.tensor
+            The encoded batch
+        """
+        # Manage single waveforms in input
+        if len(wavs.shape) == 1:
+            wavs = wavs.unsqueeze(0)
+
+        # Assign full length if wav_lens is not assigned
+        if wav_lens is None:
+            wav_lens = torch.ones(wavs.shape[0], device=self.device)
+
+        # Storing waveform in the specified device
+        wavs, wav_lens = wavs.to(self.device), wav_lens.to(self.device)
+        wavs = wavs.float()
+
+        # Computing features and embeddings
+        outputs = self.mods.wav2vec2(wavs)
+
+        # last dim will be used for AdaptativeAVG pool
+        outputs = self.mods.avg_pool(outputs, wav_lens)
+        outputs = outputs.view(outputs.shape[0], -1)
+        return outputs
+
+    def classify_batch(self, wavs, wav_lens=None):
+        """Performs classification on the top of the encoded features.
+
+        It returns the posterior probabilities, the index and, if the label
+        encoder is specified it also the text label.
+
+        Arguments
+        ---------
+        wavs : torch.tensor
+            Batch of waveforms [batch, time, channels] or [batch, time]
+            depending on the model. Make sure the sample rate is fs=16000 Hz.
+        wav_lens : torch.tensor
+            Lengths of the waveforms relative to the longest one in the
+            batch, tensor of shape [batch]. The longest one should have
+            relative length 1.0 and others len(waveform) / max_length.
+            Used for ignoring padding.
+
+        Returns
+        -------
+        out_prob
+            The log posterior probabilities of each class ([batch, N_class])
+        score:
+            It is the value of the log-posterior for the best class ([batch,])
+        index
+            The indexes of the best class ([batch,])
+        text_lab:
+            List with the text labels corresponding to the indexes.
+            (label encoder should be provided).
+        """
+        outputs = self.encode_batch(wavs, wav_lens)
+        outputs = self.mods.output_mlp(outputs)
+        out_prob = self.hparams.softmax(outputs)
+        score, index = torch.max(out_prob, dim=-1)
+        text_lab = self.hparams.label_encoder.decode_torch(index)
+        return out_prob, score, index, text_lab
+
+    def classify_file(self, path):
+        """Classifies the given audiofile into the given set of labels.
+
+        Arguments
+        ---------
+        path : str
+            Path to audio file to classify.
+
+        Returns
+        -------
+        out_prob
+            The log posterior probabilities of each class ([batch, N_class])
+        score:
+            It is the value of the log-posterior for the best class ([batch,])
+        index
+            The indexes of the best class ([batch,])
+        text_lab:
+            List with the text labels corresponding to the indexes.
+            (label encoder should be provided).
+        """
+        waveform = self.load_audio(path)
+        # Fake a batch:
+        batch = waveform.unsqueeze(0)
+        rel_length = torch.tensor([1.0])
+        outputs = self.encode_batch(batch, rel_length)
+        outputs = self.mods.output_mlp(outputs).squeeze(1)
+        out_prob = self.hparams.softmax(outputs)
+        score, index = torch.max(out_prob, dim=-1)
+        text_lab = self.hparams.label_encoder.decode_torch(index)
+        return out_prob, score, index, text_lab
+
+    def forward(self, wavs, wav_lens=None, normalize=False):
+        return self.encode_batch(
+            wavs=wavs, wav_lens=wav_lens, normalize=normalize
+        )

label_encoder.txt

@@ -0,0 +1,8 @@
+'ESPANA SUR PENINSULAR (ANDALUCIA EXTREMADURA MURCIA)' => 0
+'MEXICO' => 1
+'ANDINOPACIFICO COLOMBIA PERU ECUADOR OESTE DE BOLIVIA Y VENEZUELA ANDINA' => 2
+'CARIBE CUBA VENEZUELA PUERTO RICO REPUBLICA DOMINICANA PANAMA COLOMBIA CARIBENA MEXICO CARIBENO COSTA DEL GOLFO DE MEXICO' => 3
+'RIOPLATENSE ARGENTINA URUGUAY ESTE DE BOLIVIA PARAGUAY' => 4
+'CHILENO CHILE CUYO' => 5
+================
+'starting_index' => 0