Generate batch from context tokens.
(context_tokens, micro_batch_size=None)
| 33 | |
| 34 | |
| 35 | def get_batch(context_tokens, micro_batch_size=None): |
| 36 | """Generate batch from context tokens.""" |
| 37 | args = get_args() |
| 38 | tokenizer = get_tokenizer() |
| 39 | |
| 40 | # Move to GPU. |
| 41 | if micro_batch_size is None: |
| 42 | micro_batch_size = args.micro_batch_size |
| 43 | tokens = context_tokens.view(micro_batch_size, -1).contiguous().cuda() |
| 44 | # Get the attention mask and postition ids. |
| 45 | attention_mask, _, position_ids = get_ltor_masks_and_position_ids( |
| 46 | tokens, |
| 47 | tokenizer.eod, |
| 48 | args.reset_position_ids, |
| 49 | args.reset_attention_mask, |
| 50 | args.eod_mask_loss, |
| 51 | ) |
| 52 | |
| 53 | return tokens, attention_mask, position_ids |
| 54 | |
| 55 | |
| 56 | def get_batch_(context_tokens): |
no test coverage detected