| 253 | return z_all, ys_in_lens |
| 254 | |
| 255 | def init_state(self, x): |
| 256 | # to support mutiple encoder asr mode, in single encoder mode, |
| 257 | # convert torch.Tensor to List of torch.Tensor |
| 258 | if self.num_encs == 1: |
| 259 | x = [x] |
| 260 | |
| 261 | c_list = [self.zero_state(x[0].unsqueeze(0))] |
| 262 | z_list = [self.zero_state(x[0].unsqueeze(0))] |
| 263 | for _ in range(1, self.dlayers): |
| 264 | c_list.append(self.zero_state(x[0].unsqueeze(0))) |
| 265 | z_list.append(self.zero_state(x[0].unsqueeze(0))) |
| 266 | # TODO(karita): support strm_index for `asr_mix` |
| 267 | strm_index = 0 |
| 268 | att_idx = min(strm_index, len(self.att_list) - 1) |
| 269 | if self.num_encs == 1: |
| 270 | a = None |
| 271 | self.att_list[att_idx].reset() # reset pre-computation of h |
| 272 | else: |
| 273 | a = [None] * (self.num_encs + 1) # atts + han |
| 274 | for idx in range(self.num_encs + 1): |
| 275 | # reset pre-computation of h in atts and han |
| 276 | self.att_list[idx].reset() |
| 277 | return dict( |
| 278 | c_prev=c_list[:], |
| 279 | z_prev=z_list[:], |
| 280 | a_prev=a, |
| 281 | workspace=(att_idx, z_list, c_list), |
| 282 | ) |
| 283 | |
| 284 | def score(self, yseq, state, x): |
| 285 | # to support mutiple encoder asr mode, in single encoder mode, |