def score(self, hyp, cache):
"""Forward one step.
Args:
hyp (dataclass): hypothesis
cache (dict): states cache
Returns:
y (torch.Tensor): decoder outputs (1, dec_dim)
(list): decoder states
[L x (1, max_len, dec_dim)]
lm_tokens (torch.Tensor): token id for LM (1)
"""
device = next(self.parameters()).device
tgt = torch.tensor(hyp.yseq).unsqueeze(0).to(device=device)
lm_tokens = tgt[:, -1]
str_yseq = "".join([str(x) for x in hyp.yseq])
if str_yseq in cache:
y, new_state = cache[str_yseq]
else:
tgt_mask = subsequent_mask(len(
hyp.yseq)).unsqueeze(0).to(device=device)
state = check_state(hyp.dec_state, (tgt.size(1) - 1), self.blank)
tgt = self.embed(tgt)
new_state = []
for s, decoder in zip(state, self.decoders):
tgt, tgt_mask = decoder(tgt, tgt_mask, cache=s)
new_state.append(tgt)
y = self.after_norm(tgt[:, -1])
cache[str_yseq] = (y, new_state)
return y[0], new_state, lm_tokens
def score(
self, hyp: Hypothesis, cache: Dict[str, Any]
) -> Tuple[torch.Tensor, List[Optional[torch.Tensor]], torch.Tensor]:
"""One-step forward hypothesis.
Args:
hyp: Hypothesis.
cache: Pairs of (dec_out, dec_state) for each label sequence. (key)
Returns:
dec_out: Decoder output sequence. (1, D_dec)
dec_state: Decoder hidden states. [N x (1, U, D_dec)]
lm_label: Label ID for LM. (1,)
"""
labels = torch.tensor([hyp.yseq], device=self.device)
lm_label = labels[:, -1]
str_labels = "_".join(list(map(str, hyp.yseq)))
if str_labels in cache:
dec_out, dec_state = cache[str_labels]
else:
dec_out_mask = subsequent_mask(len(hyp.yseq)).unsqueeze_(0)
new_state = check_state(hyp.dec_state, (labels.size(1) - 1),
self.blank_id)
dec_out = self.embed(labels)
dec_state = []
for s, decoder in zip(new_state, self.decoders):
dec_out, dec_out_mask = decoder(dec_out, dec_out_mask, cache=s)
dec_state.append(dec_out)
dec_out = self.after_norm(dec_out[:, -1])
cache[str_labels] = (dec_out, dec_state)
return dec_out[0], dec_state, lm_label
