(self,
*,
batch_size: int,
scfg: SamplingConfig,
sequence_lengths: torch.Tensor,
context_lengths: torch.Tensor,
host_context_lengths,
max_context_length: int,
beam_width: int,
cache_indirections: list,
input_ids: torch.Tensor,
hidden_states: torch.Tensor,
prompt_embedding_table: torch.Tensor,
tasks: torch.Tensor,
prompt_vocab_size: torch.Tensor,
ite: int,
sequence_limit_lengths: torch.Tensor,
stop_words_data,
bad_words_data,
output_sequence_lengths: bool = False,
output_generation_logits: bool = False,
return_dict: bool = False,
encoder_output: torch.Tensor = None,
encoder_input_lengths: torch.Tensor = None,
stopping_criteria: StoppingCriteria = None,
logits_processor: LogitsProcessor = None,
cross_attention_mask: List[torch.Tensor] = None,
**kwargs)
| 4153 | return None |
| 4154 | |
| 4155 | def decode_stream(self, |
| 4156 | *, |
| 4157 | batch_size: int, |
| 4158 | scfg: SamplingConfig, |
| 4159 | sequence_lengths: torch.Tensor, |
| 4160 | context_lengths: torch.Tensor, |
| 4161 | host_context_lengths, |
| 4162 | max_context_length: int, |
| 4163 | beam_width: int, |
| 4164 | cache_indirections: list, |
| 4165 | input_ids: torch.Tensor, |
| 4166 | hidden_states: torch.Tensor, |
| 4167 | prompt_embedding_table: torch.Tensor, |
| 4168 | tasks: torch.Tensor, |
| 4169 | prompt_vocab_size: torch.Tensor, |
| 4170 | ite: int, |
| 4171 | sequence_limit_lengths: torch.Tensor, |
| 4172 | stop_words_data, |
| 4173 | bad_words_data, |
| 4174 | output_sequence_lengths: bool = False, |
| 4175 | output_generation_logits: bool = False, |
| 4176 | return_dict: bool = False, |
| 4177 | encoder_output: torch.Tensor = None, |
| 4178 | encoder_input_lengths: torch.Tensor = None, |
| 4179 | stopping_criteria: StoppingCriteria = None, |
| 4180 | logits_processor: LogitsProcessor = None, |
| 4181 | cross_attention_mask: List[torch.Tensor] = None, |
| 4182 | **kwargs): |
| 4183 | kv_cache_block_offsets = None |
| 4184 | host_kv_cache_block_offsets = None |
| 4185 | cross_kv_cache_block_offsets = None |
| 4186 | host_cross_kv_cache_block_offsets = None |
| 4187 | attention_mask = None |
| 4188 | outputs_context_logits = None |
| 4189 | |
| 4190 | def get_outputs_dict(output_ids): |
| 4191 | outputs = {} |
| 4192 | outputs['output_ids'] = output_ids |
| 4193 | if output_sequence_lengths: |
| 4194 | outputs[ |
| 4195 | 'sequence_lengths'] = self.sequence_length_buffer.reshape( |
| 4196 | [batch_size, beam_width]) |
| 4197 | if self.gather_context_logits: |
| 4198 | outputs['context_logits'] = outputs_context_logits |
| 4199 | return outputs |
| 4200 | |
| 4201 | # prepare cross attention mask. |
| 4202 | cross_attention_mask_for_context = None |
| 4203 | cross_attention_mask_for_gen = None |
| 4204 | if cross_attention_mask is not None: |
| 4205 | cross_attention_mask_for_context, cross_attention_mask_for_gen = self._prepare_cross_attention_mask( |
| 4206 | batch_size, context_lengths, cross_attention_mask) |
| 4207 | |
| 4208 | next_step_tensors = None |
| 4209 | for step in range(0, self.max_new_tokens): |
| 4210 | |
| 4211 | should_stop, next_step_tensors, tasks, context_lengths, host_context_lengths, attention_mask, context_logits, generation_logits, encoder_input_lengths = self.handle_per_step( |
| 4212 | cache_indirections=cache_indirections, |
no test coverage detected