Generate position_id and attention_mask according to input_ids considering eod reset Inputs: input_ids: the input token ids eod_id: the id for rank: the current rank dis: the slice value for each rank eod_reset: whether to open eod reset or not
(input_ids, eod_id, rank, dis, eod_reset)
| 33 | |
| 34 | |
| 35 | def get_input_data_batch_slice_map(input_ids, eod_id, rank, dis, eod_reset): |
| 36 | """ |
| 37 | Generate position_id and attention_mask according to input_ids considering eod reset |
| 38 | Inputs: |
| 39 | input_ids: the input token ids |
| 40 | eod_id: the id for <EOD> |
| 41 | rank: the current rank |
| 42 | dis: the slice value for each rank |
| 43 | eod_reset: whether to open eod reset or not |
| 44 | returns: |
| 45 | input_ids: the input token ids |
| 46 | position_id: the position ids cosidering eod reset |
| 47 | attention_mask: the attention mask considering eod reset |
| 48 | """ |
| 49 | # rank = int(rank) |
| 50 | # input_ids = input_ids[rank * dis : (rank + 1) * dis] |
| 51 | if np.any(input_ids > 60000): |
| 52 | raise ValueError("==exceed error") |
| 53 | # print("===input_ids tpye: ", input_ids.dtype, flush=True) |
| 54 | if not eod_reset: |
| 55 | return input_ids |
| 56 | seq_length = input_ids.shape[1] - 1 |
| 57 | # Initialize position_ids and attention_mask |
| 58 | batch_input_ids = deepcopy(input_ids) |
| 59 | batch_position_ids = np.ones((dis, seq_length)) |
| 60 | batch_attention_mask = np.ones((dis, seq_length, seq_length)) |
| 61 | |
| 62 | # Loop through batches |
| 63 | for bs_i in range(len(input_ids)): |
| 64 | # Get normal position_ids and attention_mask |
| 65 | local_ids = input_ids[bs_i] |
| 66 | batch_attention_mask[bs_i] = np.tril(np.ones(shape=(seq_length, seq_length))) |
| 67 | batch_position_ids[bs_i] = np.arange(seq_length) |
| 68 | # Find eod_of_document |
| 69 | eod_index = batch_position_ids[bs_i, local_ids[:-1] == eod_id].astype(np.int32) |
| 70 | prev_index = 0 |
| 71 | for i in range(eod_index.size): |
| 72 | # Reset position_ids and attention_mask considering EOD |
| 73 | index = eod_index[i] |
| 74 | batch_attention_mask[bs_i, (index + 1):, :(index + 1)] = 0 |
| 75 | batch_position_ids[bs_i, (index + 1):] -= (index + 1 - prev_index) |
| 76 | prev_index = index + 1 |
| 77 | return batch_input_ids, batch_position_ids, batch_attention_mask |
| 78 | |
| 79 | |
| 80 | def create_dataset(batch_size, data_path, args_opt, device_num=1, rank=0, drop=True, full_batch=False, |