def greedy_search(self, h: torch.Tensor) -> List[Hypothesis]: """Greedy search implementation for transformer-transducer. Args: h: Encoded speech features (T_max, D_enc) Returns: hyp: 1-best decoding results """ dec_state = self.decoder.init_state(1) hyp = Hypothesis(score=0.0, yseq=[self.blank], dec_state=dec_state) cache = {} y, state, _ = self.decoder.score(hyp, cache) for i, hi in enumerate(h): ytu = torch.log_softmax(self.joint_network(hi, y), dim=-1) logp, pred = torch.max(ytu, dim=-1) if pred != self.blank: hyp.yseq.append(int(pred)) hyp.score += float(logp) hyp.dec_state = state y, state, _ = self.decoder.score(hyp, cache) return [hyp]
def greedy_search(self, enc_out: torch.Tensor) -> List[Hypothesis]: """Greedy search implementation. Args: enc_out: Encoder output sequence. (T, D_enc) Returns: hyp: 1-best hypotheses. """ dec_state = self.decoder.init_state(1) hyp = Hypothesis(score=0.0, yseq=[self.blank_id], dec_state=dec_state) cache = {} dec_out, state, _ = self.decoder.score(hyp, cache) for enc_out_t in enc_out: logp = torch.log_softmax( self.joint_network(enc_out_t, dec_out, quantization=self.quantization) / self.softmax_temperature, dim=-1, ) top_logp, pred = torch.max(logp, dim=-1) if pred != self.blank_id: hyp.yseq.append(int(pred)) hyp.score += float(top_logp) hyp.dec_state = state dec_out, state, _ = self.decoder.score(hyp, cache) return [hyp]
def align_length_sync_decoding(self, h: torch.Tensor) -> List[Hypothesis]: """Alignment-length synchronous beam search implementation. Based on https://ieeexplore.ieee.org/document/9053040 Args: h: Encoded speech features (T_max, D_enc) Returns: nbest_hyps: N-best decoding results """ beam = min(self.beam_size, self.vocab_size) h_length = int(h.size(0)) u_max = min(self.u_max, (h_length - 1)) beam_state = self.decoder.init_state(beam) B = [ Hypothesis( yseq=[self.blank], score=0.0, dec_state=self.decoder.select_state(beam_state, 0), ) ] final = [] cache = {} if self.use_lm and not self.is_wordlm: B[0].lm_state = init_lm_state(self.lm_predictor) for i in range(h_length + u_max): A = [] B_ = [] h_states = [] for hyp in B: u = len(hyp.yseq) - 1 t = i - u + 1 if t > (h_length - 1): continue B_.append(hyp) h_states.append((t, h[t])) if B_: beam_y, beam_state, beam_lm_tokens = self.decoder.batch_score( B_, beam_state, cache, self.use_lm, ) h_enc = torch.stack([h[1] for h in h_states]) beam_logp = torch.log_softmax(self.joint_network( h_enc, beam_y), dim=-1) beam_topk = beam_logp[:, 1:].topk(beam, dim=-1) if self.use_lm: beam_lm_states = create_lm_batch_state( [b.lm_state for b in B_], self.lm_layers, self.is_wordlm) beam_lm_states, beam_lm_scores = self.lm.buff_predict( beam_lm_states, beam_lm_tokens, len(B_)) for i, hyp in enumerate(B_): new_hyp = Hypothesis( score=(hyp.score + float(beam_logp[i, 0])), yseq=hyp.yseq[:], dec_state=hyp.dec_state, lm_state=hyp.lm_state, ) A.append(new_hyp) if h_states[i][0] == (h_length - 1): final.append(new_hyp) for logp, k in zip(beam_topk[0][i], beam_topk[1][i] + 1): new_hyp = Hypothesis( score=(hyp.score + float(logp)), yseq=(hyp.yseq[:] + [int(k)]), dec_state=self.decoder.select_state(beam_state, i), lm_state=hyp.lm_state, ) if self.use_lm: new_hyp.score += self.lm_weight * beam_lm_scores[i, k] new_hyp.lm_state = select_lm_state( beam_lm_states, i, self.lm_layers, self.is_wordlm) A.append(new_hyp) B = sorted(A, key=lambda x: x.score, reverse=True)[:beam] B = recombine_hyps(B) if final: return self.sort_nbest(final) else: return B
def time_sync_decoding(self, h: torch.Tensor) -> List[Hypothesis]: """Time synchronous beam search implementation. Based on https://ieeexplore.ieee.org/document/9053040 Args: h: Encoded speech features (T_max, D_enc) Returns: nbest_hyps: N-best decoding results """ beam = min(self.beam_size, self.vocab_size) beam_state = self.decoder.init_state(beam) B = [ Hypothesis( yseq=[self.blank], score=0.0, dec_state=self.decoder.select_state(beam_state, 0), ) ] cache = {} if self.use_lm and not self.is_wordlm: B[0].lm_state = init_lm_state(self.lm_predictor) for hi in h: A = [] C = B h_enc = hi.unsqueeze(0) for v in range(self.max_sym_exp): D = [] beam_y, beam_state, beam_lm_tokens = self.decoder.batch_score( C, beam_state, cache, self.use_lm, ) beam_logp = torch.log_softmax(self.joint_network( h_enc, beam_y), dim=-1) beam_topk = beam_logp[:, 1:].topk(beam, dim=-1) seq_A = [h.yseq for h in A] for i, hyp in enumerate(C): if hyp.yseq not in seq_A: A.append( Hypothesis( score=(hyp.score + float(beam_logp[i, 0])), yseq=hyp.yseq[:], dec_state=hyp.dec_state, lm_state=hyp.lm_state, )) else: dict_pos = seq_A.index(hyp.yseq) A[dict_pos].score = np.logaddexp( A[dict_pos].score, (hyp.score + float(beam_logp[i, 0]))) if v < (self.max_sym_exp - 1): if self.use_lm: beam_lm_states = create_lm_batch_state( [c.lm_state for c in C], self.lm_layers, self.is_wordlm) beam_lm_states, beam_lm_scores = self.lm.buff_predict( beam_lm_states, beam_lm_tokens, len(C)) for i, hyp in enumerate(C): for logp, k in zip(beam_topk[0][i], beam_topk[1][i] + 1): new_hyp = Hypothesis( score=(hyp.score + float(logp)), yseq=(hyp.yseq + [int(k)]), dec_state=self.decoder.select_state( beam_state, i), lm_state=hyp.lm_state, ) if self.use_lm: new_hyp.score += self.lm_weight * beam_lm_scores[ i, k] new_hyp.lm_state = select_lm_state( beam_lm_states, i, self.lm_layers, self.is_wordlm) D.append(new_hyp) C = sorted(D, key=lambda x: x.score, reverse=True)[:beam] B = sorted(A, key=lambda x: x.score, reverse=True)[:beam] return self.sort_nbest(B)
def default_beam_search(self, h: torch.Tensor) -> List[Hypothesis]: """Beam search implementation. Args: x: Encoded speech features (T_max, D_enc) Returns: nbest_hyps: N-best decoding results """ beam = min(self.beam_size, self.vocab_size) beam_k = min(beam, (self.vocab_size - 1)) dec_state = self.decoder.init_state(1) kept_hyps = [ Hypothesis(score=0.0, yseq=[self.blank], dec_state=dec_state) ] cache = {} for hi in h: hyps = kept_hyps kept_hyps = [] while True: max_hyp = max(hyps, key=lambda x: x.score) hyps.remove(max_hyp) y, state, lm_tokens = self.decoder.score(max_hyp, cache) ytu = torch.log_softmax(self.joint_network(hi, y), dim=-1) top_k = ytu[1:].topk(beam_k, dim=-1) kept_hyps.append( Hypothesis( score=(max_hyp.score + float(ytu[0:1])), yseq=max_hyp.yseq[:], dec_state=max_hyp.dec_state, lm_state=max_hyp.lm_state, )) if self.use_lm: lm_state, lm_scores = self.lm.predict( max_hyp.lm_state, lm_tokens) else: lm_state = max_hyp.lm_state for logp, k in zip(*top_k): score = max_hyp.score + float(logp) if self.use_lm: score += self.lm_weight * lm_scores[0][k + 1] hyps.append( Hypothesis( score=score, yseq=max_hyp.yseq[:] + [int(k + 1)], dec_state=state, lm_state=lm_state, )) hyps_max = float(max(hyps, key=lambda x: x.score).score) kept_most_prob = sorted( [hyp for hyp in kept_hyps if hyp.score > hyps_max], key=lambda x: x.score, ) if len(kept_most_prob) >= beam: kept_hyps = kept_most_prob break return self.sort_nbest(kept_hyps)
def default_beam_search(self, enc_out: torch.Tensor) -> List[Hypothesis]: """Beam search implementation. Modified from https://arxiv.org/pdf/1211.3711.pdf Args: enc_out: Encoder output sequence. (T, D) Returns: nbest_hyps: N-best hypothesis. """ beam = min(self.beam_size, self.vocab_size) beam_k = min(beam, (self.vocab_size - 1)) dec_state = self.decoder.init_state(1) kept_hyps = [Hypothesis(score=0.0, yseq=[self.blank_id], dec_state=dec_state)] cache = {} for enc_out_t in enc_out: hyps = kept_hyps kept_hyps = [] while True: max_hyp = max(hyps, key=lambda x: x.score) hyps.remove(max_hyp) dec_out, state, lm_tokens = self.decoder.score(max_hyp, cache) logp = torch.log_softmax( self.joint_network( enc_out_t, dec_out, quantization=self.quantization ) / self.softmax_temperature, dim=-1, ) top_k = logp[1:].topk(beam_k, dim=-1) kept_hyps.append( Hypothesis( score=(max_hyp.score + float(logp[0:1])), yseq=max_hyp.yseq[:], dec_state=max_hyp.dec_state, lm_state=max_hyp.lm_state, ) ) if self.use_lm: lm_state, lm_scores = self.lm.predict(max_hyp.lm_state, lm_tokens) else: lm_state = max_hyp.lm_state for logp, k in zip(*top_k): score = max_hyp.score + float(logp) if self.use_lm: score += self.lm_weight * lm_scores[0][k + 1] hyps.append( Hypothesis( score=score, yseq=max_hyp.yseq[:] + [int(k + 1)], dec_state=state, lm_state=lm_state, ) ) hyps_max = float(max(hyps, key=lambda x: x.score).score) kept_most_prob = sorted( [hyp for hyp in kept_hyps if hyp.score > hyps_max], key=lambda x: x.score, ) if len(kept_most_prob) >= beam: kept_hyps = kept_most_prob break return self.sort_nbest(kept_hyps)