add inference class

inference.py

@@ -1,120 +1,4 @@
 import torch
+from speechbrain.pretrained import Pretrained
 
-def forward(self, wavs, wav_lens=None):
-    """Runs the classification"""
-    return self.classify_batch(wavs, wav_lens)
-
-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
-    feats = self.mods.compute_features(wavs)
-    feats = self.mods.mean_var_norm(feats, wav_lens)
-    embeddings = self.mods.embedding_model(feats, wav_lens)
-    if normalize:
-        embeddings = self.hparams.mean_var_norm_emb(
-            embeddings, torch.ones(embeddings.shape[0], device=self.device)
-        )
-    return embeddings
-
-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).
-    """
-    emb = self.encode_batch(wavs, wav_lens)
-    out_prob = self.mods.classifier(emb).squeeze(1)
-    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])
-    emb = self.encode_batch(batch, rel_length)
-    out_prob = self.mods.classifier(emb).squeeze(1)
-    score, index = torch.max(out_prob, dim=-1)
-    text_lab = self.hparams.label_encoder.decode_torch(index)
-    return out_prob, score, index, text_lab
+class EncoderClassifier(Pretrained):