Update custom_interface.py

custom_interface.py

@@ -140,6 +140,27 @@ class CustomEncoderWav2vec2Classifier(Pretrained):
         rel_length = torch.tensor([1.0])
         outputs = self.encode_batch(batch, rel_length)
         return outputs
+
+    def embed_sample(self, sample, sr):
+        """Returns embedding (last layer output) for the given audiofile.
+
+        Arguments
+        ---------
+        ample : torch tensor
+            wav tensor. ([T, 1])
+        sr: int 
+            sampling rate.
+
+        Returns
+        -------
+        embed
+            The log posterior probabilities of each class ([batch, embed_dim])
+        """
+        waveform = self.audio_normalizer(sample, sr)
+        batch = waveform.unsqueeze(0) 
+        rel_length = torch.tensor([1.0])
+        outputs = self.encode_batch(batch, rel_length)
+        return outputs
         
 
     def classify_file(self, path):
@@ -192,10 +213,7 @@ class CustomEncoderWav2vec2Classifier(Pretrained):
             (label encoder should be provided).
         """
         # Fake a batch:
-        waveform = self.audio_normalizer(sample, sr)
-        batch = waveform.unsqueeze(0) 
-        rel_length = torch.tensor([1.0])
-        outputs = self.encode_batch(batch, rel_length)
+        outputs = self.embed_sample(sample, sr)
         outputs = self.mods.output_mlp(outputs).squeeze(1)
         out_prob = self.hparams.softmax(outputs)
         score, index = torch.max(out_prob, dim=-1)