get_batch subdivides the source data into chunks of length args.seq_length. If source is equal to the example output of the data loading example, with a seq_length limit of 2, we'd get the following two Variables for i = 0: ┌ a g m s ┐ ┌ b h n t ┐ └ b h n t ┘ └ c i o u ┘ Not
(data, args)
| 142 | |
| 143 | |
| 144 | def get_batch(data, args): |
| 145 | ''' get_batch subdivides the source data into chunks of |
| 146 | length args.seq_length. If source is equal to the example |
| 147 | output of the data loading example, with a seq_length limit |
| 148 | of 2, we'd get the following two Variables for i = 0: |
| 149 | ┌ a g m s ┐ ┌ b h n t ┐ |
| 150 | └ b h n t ┘ └ c i o u ┘ |
| 151 | Note that despite the name of the function, the subdivison of data is not |
| 152 | done along the batch dimension (i.e. dimension 1), since that was handled |
| 153 | by the data loader. The chunks are along dimension 0, corresponding |
| 154 | to the seq_len dimension in the LSTM. A Variable representing an appropriate |
| 155 | shard reset mask of the same dimensions is also returned. |
| 156 | ''' |
| 157 | # Items and their type. |
| 158 | keys = ['text', 'loss_mask'] |
| 159 | if args.transformer_xl or args.block_lm: |
| 160 | keys += ['target', 'attention_mask'] |
| 161 | if args.block_lm: |
| 162 | keys += ['position_id'] |
| 163 | datatype = torch.int64 |
| 164 | |
| 165 | # Broadcast data. |
| 166 | data_b = mpu.broadcast_data(keys, data, datatype) |
| 167 | # Unpack. |
| 168 | if args.transformer_xl: |
| 169 | tokens = data_b['text'].long() |
| 170 | labels = data_b['target'].long() |
| 171 | attention_mask = data_b['attention_mask'].float() |
| 172 | loss_mask = data_b['loss_mask'].float() |
| 173 | elif args.block_lm: |
| 174 | tokens = data_b['text'].long() |
| 175 | labels = data_b['target'].long() |
| 176 | attention_mask = data_b['attention_mask'].long() |
| 177 | loss_mask = data_b['loss_mask'].float() |
| 178 | position_ids = data_b['position_id'].long() |
| 179 | else: |
| 180 | tokens_ = data_b['text'].long() |
| 181 | loss_mask = data_b['loss_mask'].float() |
| 182 | labels = tokens_[:, 1:].contiguous() |
| 183 | loss_mask = loss_mask[:, 1:].contiguous() |
| 184 | tokens = tokens_[:, :-1].contiguous() |
| 185 | attention_mask = None |
| 186 | |
| 187 | # Get the masks and postition ids. |
| 188 | if not args.block_lm: |
| 189 | attention_mask, loss_mask, position_ids = get_masks_and_position_ids( |
| 190 | tokens, |
| 191 | args.eod_token, |
| 192 | args.reset_position_ids, |
| 193 | args.reset_attention_mask, |
| 194 | loss_mask=loss_mask, |
| 195 | attention_mask=attention_mask, |
| 196 | mem_length=args.mem_length, |
| 197 | set_loss_mask=not args.transformer_xl) |
| 198 | # Convert |
| 199 | if args.fp16: |
| 200 | attention_mask = attention_mask.half() |
| 201 | return tokens, labels, loss_mask, attention_mask, position_ids |
no test coverage detected