Build masks and position id for left to right model.
(
data,
eod_token,
reset_position_ids,
reset_attention_mask,
)
| 10 | |
| 11 | |
| 12 | def get_ltor_masks_and_position_ids( |
| 13 | data, |
| 14 | eod_token, |
| 15 | reset_position_ids, |
| 16 | reset_attention_mask, |
| 17 | ): |
| 18 | """Build masks and position id for left to right model.""" |
| 19 | |
| 20 | # Extract batch size and sequence length. |
| 21 | micro_batch_size, seq_length = data.size() |
| 22 | |
| 23 | # Attention mask (lower triangular). |
| 24 | if reset_attention_mask: |
| 25 | att_mask_batch = micro_batch_size |
| 26 | else: |
| 27 | att_mask_batch = 1 |
| 28 | attention_mask = torch.tril( |
| 29 | torch.ones((att_mask_batch, seq_length, seq_length), device=data.device) |
| 30 | ).view(att_mask_batch, 1, seq_length, seq_length) |
| 31 | |
| 32 | # Position ids. |
| 33 | position_ids = torch.arange(seq_length, dtype=torch.long, device=data.device) |
| 34 | position_ids = position_ids.unsqueeze(0).expand_as(data) |
| 35 | # We need to clone as the ids will be modifed based on batch index. |
| 36 | if reset_position_ids: |
| 37 | position_ids = position_ids.clone() |
| 38 | |
| 39 | if reset_position_ids or reset_attention_mask: |
| 40 | # Loop through the batches: |
| 41 | for b in range(micro_batch_size): |
| 42 | |
| 43 | # Find indecies where EOD token is. |
| 44 | eod_index = position_ids[b, data[b] == eod_token] |
| 45 | # Detach indecies from positions if going to modify positions. |
| 46 | if reset_position_ids: |
| 47 | eod_index = eod_index.clone() |
| 48 | |
| 49 | # Loop through EOD indecies: |
| 50 | prev_index = 0 |
| 51 | for j in range(eod_index.size()[0]): |
| 52 | i = eod_index[j] |
| 53 | # Mask attention loss. |
| 54 | if reset_attention_mask: |
| 55 | attention_mask[b, 0, (i + 1) :, : (i + 1)] = 0 |
| 56 | # Reset positions. |
| 57 | if reset_position_ids: |
| 58 | position_ids[b, (i + 1) :] -= i + 1 - prev_index |
| 59 | prev_index = i + 1 |
| 60 | |
| 61 | # Convert attention mask to binary: |
| 62 | attention_mask = attention_mask < 0.5 |
| 63 | |
| 64 | return attention_mask, position_ids |
| 65 | |
| 66 | |
| 67 | def get_batch( |