Pad the instances to the max sequence length in batch, and generate the corresponding position data and attention bias.
(insts,
pretraining_task='seq2seq',
pad_idx=1,
sent_b_starts=None,
return_pos=False,
return_input_mask=False,
return_max_len=False,
return_num_token=False,
return_seq_lens=False)
| 457 | |
| 458 | |
| 459 | def pad_batch_data(insts, |
| 460 | pretraining_task='seq2seq', |
| 461 | pad_idx=1, |
| 462 | sent_b_starts=None, |
| 463 | return_pos=False, |
| 464 | return_input_mask=False, |
| 465 | return_max_len=False, |
| 466 | return_num_token=False, |
| 467 | return_seq_lens=False): |
| 468 | """ |
| 469 | Pad the instances to the max sequence length in batch, and generate the |
| 470 | corresponding position data and attention bias. |
| 471 | """ |
| 472 | return_list = [] |
| 473 | max_len = max(len(inst) for inst in insts) |
| 474 | # Any token included in dict can be used to pad, since the paddings' loss |
| 475 | # will be masked out by weights and make no effect on parameter gradients. |
| 476 | |
| 477 | inst_data = np.array( |
| 478 | [inst + list([pad_idx] * (max_len - len(inst))) for inst in insts]) |
| 479 | return_list += [inst_data.astype('int64').reshape([-1, max_len, 1])] |
| 480 | |
| 481 | # position data |
| 482 | if return_pos: |
| 483 | inst_pos = np.array([ |
| 484 | list(range(0, len(inst))) + [pad_idx] * (max_len - len(inst)) |
| 485 | for inst in insts |
| 486 | ]) |
| 487 | |
| 488 | return_list += [inst_pos.astype('int64').reshape([-1, max_len, 1])] |
| 489 | |
| 490 | if return_input_mask: |
| 491 | if pretraining_task is 'seq2seq': |
| 492 | assert sent_b_starts is not None, \ |
| 493 | "[FATAL] For seq2seq lanugae model loss," \ |
| 494 | " sent_b_starts should not be None" |
| 495 | # This is used to avoid attention on paddings and subsequent words. |
| 496 | input_mask_data = np.zeros((inst_data.shape[0], max_len, max_len)) |
| 497 | for index, mask_data in enumerate(input_mask_data): |
| 498 | start = sent_b_starts[index] |
| 499 | end = len(insts[index]) |
| 500 | mask_data[:end, :start] = 1.0 |
| 501 | # Generate the lower triangular matrix using the slice of matrix |
| 502 | b = np.tril(np.ones([end - start, end - start]), 0) |
| 503 | mask_data[start:end, start:end] = b |
| 504 | input_mask_data = np.array(input_mask_data).reshape([-1, max_len, max_len]) |
| 505 | else: |
| 506 | # This is used to avoid attention on paddings. |
| 507 | input_mask_data = np.array([[1] * len(inst) + [0] * |
| 508 | (max_len - len(inst)) for inst in insts]) |
| 509 | input_mask_data = np.expand_dims(input_mask_data, axis=1) |
| 510 | # input_mask_data = np.matmul(input_mask_data, np.transpose(input_mask_data, (0, 2, 1))) |
| 511 | return_list += [input_mask_data.astype("float32")] |
| 512 | |
| 513 | if return_max_len: |
| 514 | return_list += [max_len] |
| 515 | |
| 516 | if return_num_token: |
no outgoing calls
no test coverage detected