Generate batch from context tokens.
(
context_tokens,
micro_batch_size,
eod_token,
reset_position_ids=False,
reset_attention_mask=False,
)
| 65 | |
| 66 | |
| 67 | def get_batch( |
| 68 | context_tokens, |
| 69 | micro_batch_size, |
| 70 | eod_token, |
| 71 | reset_position_ids=False, |
| 72 | reset_attention_mask=False, |
| 73 | ): |
| 74 | """Generate batch from context tokens.""" |
| 75 | tokens = context_tokens.view(micro_batch_size, -1).contiguous().cuda() |
| 76 | # Get the attention mask and postition ids. |
| 77 | attention_mask, position_ids = get_ltor_masks_and_position_ids( |
| 78 | tokens, |
| 79 | eod_token, |
| 80 | reset_position_ids, |
| 81 | reset_attention_mask, |
| 82 | ) |
| 83 | |
| 84 | return tokens, attention_mask, position_ids |
| 85 | |
| 86 | |
| 87 | def top_k_logits(logits, top_k=0, top_p=0.0, filter_value=-float("Inf")): |
no test coverage detected