(data,
loss_mask=None,
attention_mask=None, args=None)
| 14 | |
| 15 | |
| 16 | def get_masks_and_position_ids(data, |
| 17 | loss_mask=None, |
| 18 | attention_mask=None, args=None): |
| 19 | # Extract batch size and sequence length. |
| 20 | batch_size, seq_length = data.size() |
| 21 | |
| 22 | # Attention mask (lower triangular). |
| 23 | if attention_mask is None: |
| 24 | attention_mask = torch.ones((batch_size, seq_length, seq_length), device=data.device) |
| 25 | attention_mask.tril_() |
| 26 | attention_mask.unsqueeze_(1) |
| 27 | |
| 28 | # Loss mask. |
| 29 | if loss_mask is None: |
| 30 | loss_mask = torch.ones(data.size(), dtype=data.dtype, device=data.device) |
| 31 | |
| 32 | # Position ids. |
| 33 | position_ids = torch.arange(seq_length, dtype=torch.long, |
| 34 | device=data.device) |
| 35 | position_ids = position_ids.unsqueeze(0).expand_as(data) |
| 36 | |
| 37 | return attention_mask, loss_mask, position_ids |
| 38 | |
| 39 | |
| 40 | def get_batch(data_iterator, args, timers): |
no outgoing calls
no test coverage detected