(
model,
tokens,
seq_length,
position_ids,
attention_mask,
layer_past=None,
get_key_value=None,
prompt_length=None,
context_length=None,
)
| 125 | |
| 126 | |
| 127 | def forward_step( |
| 128 | model, |
| 129 | tokens, |
| 130 | seq_length, |
| 131 | position_ids, |
| 132 | attention_mask, |
| 133 | layer_past=None, |
| 134 | get_key_value=None, |
| 135 | prompt_length=None, |
| 136 | context_length=None, |
| 137 | ): |
| 138 | # Forward pass through the model. |
| 139 | output_tensor = model( |
| 140 | tokens, |
| 141 | position_ids, |
| 142 | attention_mask, |
| 143 | layer_past=layer_past, |
| 144 | get_key_value=get_key_value, |
| 145 | prompt_length=prompt_length, |
| 146 | context_length=context_length, |
| 147 | ) |
| 148 | |
| 149 | if get_key_value: |
| 150 | output_tensor, layer_past = output_tensor |
| 151 | |
| 152 | if get_key_value: |
| 153 | return output_tensor, layer_past |
| 154 | |
| 155 | return output_tensor |
| 156 | |
| 157 | |
| 158 | def get_token_stream( |
nothing calls this directly
no outgoing calls
no test coverage detected