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)
| 21 | |
| 22 | |
| 23 | def pad_batch_data(insts, |
| 24 | pretraining_task='seq2seq', |
| 25 | pad_idx=1, |
| 26 | sent_b_starts=None, |
| 27 | return_pos=False, |
| 28 | return_input_mask=False, |
| 29 | return_max_len=False, |
| 30 | return_num_token=False, |
| 31 | return_seq_lens=False): |
| 32 | """ |
| 33 | Pad the instances to the max sequence length in batch, and generate the |
| 34 | corresponding position data and attention bias. |
| 35 | """ |
| 36 | return_list = [] |
| 37 | max_len = max(len(inst) for inst in insts) |
| 38 | # Any token included in dict can be used to pad, since the paddings' loss |
| 39 | # will be masked out by weights and make no effect on parameter gradients. |
| 40 | |
| 41 | inst_data = np.array( |
| 42 | [inst + list([pad_idx] * (max_len - len(inst))) for inst in insts]) |
| 43 | return_list += [inst_data.astype('int64').reshape([-1, max_len, 1])] |
| 44 | |
| 45 | # position data |
| 46 | if return_pos: |
| 47 | inst_pos = np.array([ |
| 48 | list(range(0, len(inst))) + [pad_idx] * (max_len - len(inst)) |
| 49 | for inst in insts |
| 50 | ]) |
| 51 | |
| 52 | return_list += [inst_pos.astype('int64').reshape([-1, max_len, 1])] |
| 53 | |
| 54 | if return_input_mask: |
| 55 | if pretraining_task is 'seq2seq': |
| 56 | assert sent_b_starts is not None, \ |
| 57 | "[FATAL] For seq2seq lanugae model loss," \ |
| 58 | " sent_b_starts should not be None" |
| 59 | # This is used to avoid attention on paddings and subsequent words. |
| 60 | input_mask_data = np.zeros((inst_data.shape[0], max_len, max_len)) |
| 61 | for index, mask_data in enumerate(input_mask_data): |
| 62 | start = sent_b_starts[index] |
| 63 | end = len(insts[index]) |
| 64 | mask_data[:end, :start] = 1.0 |
| 65 | # Generate the lower triangular matrix using the slice of matrix |
| 66 | b = np.tril(np.ones([end - start, end - start]), 0) |
| 67 | mask_data[start:end, start:end] = b |
| 68 | input_mask_data = np.array(input_mask_data).reshape([-1, max_len, max_len]) |
| 69 | else: |
| 70 | # This is used to avoid attention on paddings. |
| 71 | input_mask_data = np.array([[1] * len(inst) + [0] * |
| 72 | (max_len - len(inst)) for inst in insts]) |
| 73 | input_mask_data = np.expand_dims(input_mask_data, axis=-1) |
| 74 | # input_mask_data = np.matmul(input_mask_data, np.transpose(input_mask_data, (0, 2, 1))) |
| 75 | return_list += [input_mask_data.astype("float32")] |
| 76 | |
| 77 | if return_max_len: |
| 78 | return_list += [max_len] |
| 79 | |
| 80 | if return_num_token: |
no outgoing calls
no test coverage detected