(beam_size, output_ids, parent_ids, out_seq_lens, end_id, max_seq_len=None, args=None)
| 182 | |
| 183 | |
| 184 | def finalize(beam_size, output_ids, parent_ids, out_seq_lens, end_id, max_seq_len=None, args=None): |
| 185 | out_seq_lens = torch.reshape(out_seq_lens, (-1, beam_size)) |
| 186 | max_lens = torch.max(out_seq_lens, 1)[0] |
| 187 | if max_seq_len: |
| 188 | shape = (max_seq_len, -1, beam_size) |
| 189 | else: |
| 190 | shape = (torch.max(max_lens), -1, beam_size) |
| 191 | output_ids = torch.reshape(output_ids, shape) |
| 192 | parent_ids = torch.reshape(parent_ids, shape) |
| 193 | if output_ids.is_cuda: |
| 194 | torch.classes.load_library("./lib/libth_transformer.so") |
| 195 | batch_size = output_ids.shape[1] |
| 196 | end_ids = end_id * torch.ones(batch_size, dtype=torch.int32, device=output_ids.device) |
| 197 | ids = torch.ops.fastertransformer.gather_tree(output_ids.to(torch.int32), parent_ids.to(torch.int32), out_seq_lens.to(torch.int32), end_ids) |
| 198 | else: |
| 199 | ids = gather_tree(output_ids, parent_ids, max_lens, end_id) |
| 200 | ids = torch.einsum('ijk->jki', ids) # batch_size, beam_size, max_seq_len |
| 201 | lengths = torch.eq(ids, end_id) |
| 202 | lengths = 1 - lengths.to(output_ids.dtype) |
| 203 | lengths = torch.sum(lengths, -1) |
| 204 | return ids, lengths |
| 205 | |
| 206 | |
| 207 |
no test coverage detected