Generate batch from context tokens.
(context_tokens)
| 54 | |
| 55 | |
| 56 | def get_batch_(context_tokens): |
| 57 | """Generate batch from context tokens.""" |
| 58 | args = get_args() |
| 59 | tokenizer = get_tokenizer() |
| 60 | |
| 61 | # Move to GPU. |
| 62 | tokens = context_tokens.contiguous().cuda() |
| 63 | # Get the attention mask and postition ids. |
| 64 | attention_mask, _, position_ids = get_ltor_masks_and_position_ids( |
| 65 | tokens, |
| 66 | tokenizer.eod, |
| 67 | args.reset_position_ids, |
| 68 | args.reset_attention_mask, |
| 69 | args.eod_mask_loss, |
| 70 | ) |
| 71 | |
| 72 | return tokens, attention_mask, position_ids |
| 73 | |
| 74 | |
| 75 | def top_k_logits(logits, top_k=0, top_p=0.0, filter_value=-float("Inf")): |
no test coverage detected