Update custom_interface.py

custom_interface.py

@@ -8,22 +8,21 @@ class ASR(Pretrained):
 
     def encode_batch(self, wavs, wav_lens=None, normalize=False):
         wavs = wavs.to(self.device)
-        wav_lens = wav_lens.to(self.device)
+        self.wav_lens = wav_lens.to(self.device)
 
         # Forward pass
-        encoded_outputs = self.modules.encoder_w2v2(wavs.detach())
+        encoded_outputs = self.mods.encoder_w2v2(wavs.detach())
         # append
         tokens_bos = torch.zeros((wavs.size(0), 1), dtype=torch.long).to(self.device)
-        print(tokens_bos.size())
-        embedded_tokens = self.modules.embedding(tokens_bos)
-        decoder_outputs, _ = self.modules.decoder(embedded_tokens, encoded_outputs, self.sig_lens)
+        embedded_tokens = self.mods.embedding(tokens_bos)
+        decoder_outputs, _ = self.mods.decoder(embedded_tokens, encoded_outputs, self.wav_lens)
 
         # Output layer for seq2seq log-probabilities
-        logits = self.modules.seq_lin(decoder_outputs)
-        predictions = {"seq_logprobs": self.hparams.log_softmax(logits)}
-        predictions["tokens"], _, _, _ = self.hparams.test_search(encoded_outputs, self.sig_lens)
+        predictions = self.hparams.test_search(encoded_outputs, self.wav_lens)[0]
+        predicted_words = [self.hparams.tokenizer.decode_ids(prediction).split(" ") for prediction in predictions]
+        print(predicted_words)
 
-        return predictions
+        return predicted_words
 
 
     def classify_file(self, path):
@@ -31,7 +30,7 @@ class ASR(Pretrained):
         # Fake a batch:
         batch = waveform.unsqueeze(0)
         rel_length = torch.tensor([1.0])
-        outputs = self.encode_batch(batch, rel_length)["tokens"]
+        outputs = self.encode_batch(batch, rel_length)
        
         return outputs