(input_ids, model, tree_attn_mask, past_key_values, logits_processor)
| 225 | |
| 226 | |
| 227 | def initialize_tree(input_ids, model, tree_attn_mask, past_key_values, logits_processor): |
| 228 | position_ids = torch.arange(input_ids.shape[1], dtype=torch.long, device=input_ids.device) |
| 229 | tree_logits, outputs, logits, hidden_state, sample_token = model( |
| 230 | input_ids, past_key_values=past_key_values, output_orig=True, logits_processor=logits_processor, |
| 231 | position_ids=position_ids |
| 232 | ) |
| 233 | return tree_logits, logits, hidden_state, sample_token |
| 234 | |
| 235 | |
| 236 | def reset_tree_mode( |
no outgoing calls
no test coverage detected