add py_module + hparams

custom_interface.py

@@ -0,0 +1,100 @@
+import torch
+from speechbrain.pretrained import Pretrained
+
+
+class CustomSLUDecoder(Pretrained):
+    """A end-to-end SLU model using hubert self-supervised encoder.
+
+    The class can be used either to run only the encoder (encode()) to extract
+    features or to run the entire model (decode()) to map the speech to its semantics.
+
+    Example
+    -------
+    >>> from speechbrain.pretrained.interfaces import foreign_class
+    >>> slu_model = foreign_class(source="speechbrain/slu-timers-and-such-direct-librispeech-asr", 
+                    pymodule_file="custom_interface.py", classname="CustomEncoderWav2vec2Classifier")
+    >>> slu_model.decode_file("samples/audio_samples/example6.wav")
+    "{'intent': 'SimpleMath', 'slots': {'number1': 37.67, 'number2': 75.7, 'op': ' minus '}}"
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.tokenizer = self.hparams.tokenizer
+
+    def decode_file(self, path):
+        """Maps the given audio file to a string representing the
+        semantic dictionary for the utterance.
+
+        Arguments
+        ---------
+        path : str
+            Path to audio file to decode.
+
+        Returns
+        -------
+        str
+            The predicted semantics.
+        """
+        waveform = self.load_audio(path)
+        waveform = waveform.to(self.device)
+        # Fake a batch:
+        batch = waveform.unsqueeze(0)
+        rel_length = torch.tensor([1.0])
+        predicted_words, predicted_tokens = self.decode_batch(batch, rel_length)
+        return predicted_words[0]
+    
+    def encode_batch(self, wavs):
+        """Encodes the input audio into a sequence of hidden states
+
+        Arguments
+        ---------
+        wavs : torch.tensor
+            Batch of waveforms [batch, time, channels] or [batch, time]
+            depending on the model.
+
+        Returns
+        -------
+        torch.tensor
+            The encoded batch
+        """
+        wavs = wavs.float()
+        wavs = wavs.to(self.device)
+        encoder_out = self.mods.hubert(wavs.detach())
+        return encoder_out
+
+    def decode_batch(self, wavs, wav_lens):
+        """Maps the input audio to its semantics
+
+        Arguments
+        ---------
+        wavs : torch.tensor
+            Batch of waveforms [batch, time, channels] or [batch, time]
+            depending on the model.
+        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
+        -------
+        list
+            Each waveform in the batch decoded.
+        tensor
+            Each predicted token id.
+        """
+        with torch.no_grad():
+            wavs, wav_lens = wavs.to(self.device), wav_lens.to(self.device)
+            encoder_out = self.encode_batch(wavs)
+            predicted_tokens, scores = self.mods.beam_searcher(
+                encoder_out, wav_lens
+            )
+            predicted_words = [
+                self.tokenizer.decode_ids(token_seq)
+                for token_seq in predicted_tokens
+            ]
+        return predicted_words, predicted_tokens
+
+    def forward(self, wavs, wav_lens):
+        """Runs full decoding - note: no gradients through decoding"""
+        return self.decode_batch(wavs, wav_lens)

hyperparams.yaml

@@ -0,0 +1,86 @@
+# ############################################################################
+# Model: WAV2VEC base for Emotion Recognition
+# ############################################################################
+
+
+# Hparams NEEDED
+HPARAMS_NEEDED: ["beam_searcher"]
+# Modules Needed
+MODULES_NEEDED: ["hubert", "decoder", "seq_lin"]
+
+# URL for the wav2vec2 model, you can change to benchmark diffrenet models
+wav2vec2_hub: facebook/hubert-base-ls960
+
+# Pretrain folder (HuggingFace)
+pretrained_path: speechbrain/SLU-direct-SLURP-hubert-enc
+
+# parameters
+encoder_dim: 768
+output_neurons: 58
+emb_size: 128
+dec_neurons: 512
+dec_attn_dim: 512
+dec_layer: 3
+
+hubert: !new:speechbrain.lobes.models.huggingface_wav2vec.HuggingFaceWav2Vec2
+    source: !ref <wav2vec2_hub>
+    output_norm: True
+    freeze: True
+    pretrain: False
+    save_path: wav2vec2_checkpoints
+
+output_emb: !new:speechbrain.nnet.embedding.Embedding
+  num_embeddings: !ref <output_neurons>
+  embedding_dim: !ref <emb_size>
+
+dec: !new:speechbrain.nnet.RNN.AttentionalRNNDecoder
+  enc_dim: !ref <encoder_dim>
+  input_size: !ref <emb_size>
+  rnn_type: lstm
+  attn_type: content
+  hidden_size: !ref <dec_neurons>
+  attn_dim: !ref <dec_attn_dim>
+  num_layers: !ref <dec_layer>
+  scaling: 1.0
+  dropout: 0.0
+
+seq_lin: !new:speechbrain.nnet.linear.Linear
+  input_size: !ref <dec_neurons>
+  n_neurons: !ref <output_neurons>
+
+beam_searcher: !new:speechbrain.decoders.S2SRNNBeamSearcher
+  embedding: !ref <output_emb>
+  decoder: !ref <dec>
+  linear: !ref <seq_lin>
+  bos_index: 0
+  eos_index: 0
+  min_decode_ratio: 0.0
+  max_decode_ratio: 10.0
+  beam_size: 80
+  eos_threshold: 1.5
+  temperature: 1.25
+  using_max_attn_shift: false
+  max_attn_shift: 30
+  coverage_penalty: 0.
+
+model: !new:torch.nn.ModuleList
+    - [!ref <output_emb>, !ref <dec>, !ref <seq_lin>]
+
+modules:
+    hubert: !ref <hubert>
+    beam_searcher: !ref <beam_searcher>
+
+tokenizer: !new:sentencepiece.SentencePieceProcessor
+
+
+pretrainer: !new:speechbrain.utils.parameter_transfer.Pretrainer
+    loadables:
+        hubert: !ref <hubert>
+        model: !ref <model>
+        tokenizer: !ref <tokenizer>
+    paths:
+        hubert: !ref <pretrained_path>/hubert.ckpt
+        model: !ref <pretrained_path>/model.ckpt
+        tokenizer: !ref <pretrained_path>/tokenizer_58_unigram.model
+
+