Fork of https://github.com/alokprasad/fastspeech_squeezewave to also fix denoising in squeezewave
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
100 lines
3.0 KiB
100 lines
3.0 KiB
import torch
|
|
import numpy as np
|
|
import transformer.Constants as Constants
|
|
|
|
|
|
class Beam():
|
|
''' Beam search '''
|
|
|
|
def __init__(self, size, device=False):
|
|
|
|
self.size = size
|
|
self._done = False
|
|
|
|
# The score for each translation on the beam.
|
|
self.scores = torch.zeros((size,), dtype=torch.float, device=device)
|
|
self.all_scores = []
|
|
|
|
# The backpointers at each time-step.
|
|
self.prev_ks = []
|
|
|
|
# The outputs at each time-step.
|
|
self.next_ys = [torch.full(
|
|
(size,), Constants.PAD, dtype=torch.long, device=device)]
|
|
self.next_ys[0][0] = Constants.BOS
|
|
|
|
def get_current_state(self):
|
|
"Get the outputs for the current timestep."
|
|
return self.get_tentative_hypothesis()
|
|
|
|
def get_current_origin(self):
|
|
"Get the backpointers for the current timestep."
|
|
return self.prev_ks[-1]
|
|
|
|
@property
|
|
def done(self):
|
|
return self._done
|
|
|
|
def advance(self, word_prob):
|
|
"Update beam status and check if finished or not."
|
|
num_words = word_prob.size(1)
|
|
|
|
# Sum the previous scores.
|
|
if len(self.prev_ks) > 0:
|
|
beam_lk = word_prob + self.scores.unsqueeze(1).expand_as(word_prob)
|
|
else:
|
|
beam_lk = word_prob[0]
|
|
|
|
flat_beam_lk = beam_lk.view(-1)
|
|
|
|
best_scores, best_scores_id = flat_beam_lk.topk(
|
|
self.size, 0, True, True) # 1st sort
|
|
best_scores, best_scores_id = flat_beam_lk.topk(
|
|
self.size, 0, True, True) # 2nd sort
|
|
|
|
self.all_scores.append(self.scores)
|
|
self.scores = best_scores
|
|
|
|
# bestScoresId is flattened as a (beam x word) array,
|
|
# so we need to calculate which word and beam each score came from
|
|
prev_k = best_scores_id / num_words
|
|
self.prev_ks.append(prev_k)
|
|
self.next_ys.append(best_scores_id - prev_k * num_words)
|
|
|
|
# End condition is when top-of-beam is EOS.
|
|
if self.next_ys[-1][0].item() == Constants.EOS:
|
|
self._done = True
|
|
self.all_scores.append(self.scores)
|
|
|
|
return self._done
|
|
|
|
def sort_scores(self):
|
|
"Sort the scores."
|
|
return torch.sort(self.scores, 0, True)
|
|
|
|
def get_the_best_score_and_idx(self):
|
|
"Get the score of the best in the beam."
|
|
scores, ids = self.sort_scores()
|
|
return scores[1], ids[1]
|
|
|
|
def get_tentative_hypothesis(self):
|
|
"Get the decoded sequence for the current timestep."
|
|
|
|
if len(self.next_ys) == 1:
|
|
dec_seq = self.next_ys[0].unsqueeze(1)
|
|
else:
|
|
_, keys = self.sort_scores()
|
|
hyps = [self.get_hypothesis(k) for k in keys]
|
|
hyps = [[Constants.BOS] + h for h in hyps]
|
|
dec_seq = torch.LongTensor(hyps)
|
|
|
|
return dec_seq
|
|
|
|
def get_hypothesis(self, k):
|
|
""" Walk back to construct the full hypothesis. """
|
|
hyp = []
|
|
for j in range(len(self.prev_ks) - 1, -1, -1):
|
|
hyp.append(self.next_ys[j+1][k])
|
|
k = self.prev_ks[j][k]
|
|
|
|
return list(map(lambda x: x.item(), hyp[::-1]))
|