Update custom_interface.py

custom_interface.py

@@ -11,25 +11,14 @@ class ASR(Pretrained):
         wavs = wavs.to(device)
         wav_lens = wav_lens.to(device)
 
-        # Forward pass
-        encoded_outputs = self.mods.encoder_w2v2(wavs.detach())
-        # append
-        tokens_bos = torch.zeros((wavs.size(0), 1), dtype=torch.long).to(device)
-        embedded_tokens = self.mods.embedding(tokens_bos)
-        decoder_outputs, _ = self.mods.decoder(embedded_tokens, encoded_outputs, wav_lens)
+        # Forward encoder + decoder
+        tokens = torch.tensor([[1, 1]]) * self.mods.whisper.config.decoder_start_token_id
+        tokens = tokens.to(device)
+        enc_out, logits, _ = self.mods.whisper(wavs, tokens)
+        log_probs = self.hparams.log_softmax(logits)
 
-        # Output layer for seq2seq log-probabilities
-        predictions = self.hparams.test_search(encoded_outputs, wav_lens)[0]
-        # predicted_words = [self.hparams.tokenizer.decode_ids(prediction).split(" ") for prediction in predictions]
-        predicted_words = []
-        for prediction in predictions:
-            prediction = [token for token in prediction if token != 0]
-            predicted_words.append(self.hparams.tokenizer.decode_ids(prediction).split(" "))
-        prediction = []
-        for sent in predicted_words:
-            sent = self.filter_repetitions(sent, 3)
-            prediction.append(sent)
-        predicted_words = prediction
+        hyps, _, _, _ = self.hparams.test_search(enc_out.detach(), wav_lens)
+        predicted_words = [self.mods.whisper.tokenizer.decode(token, skip_special_tokens=True).strip() for token in hyps]
         return predicted_words
 
     def filter_repetitions(self, seq, max_repetition_length):