| 156 | |
| 157 | |
| 158 | def get_token_stream( |
| 159 | model, |
| 160 | tokenizer, |
| 161 | seq_length, |
| 162 | out_seq_length, |
| 163 | context_tokens, |
| 164 | return_scores: bool = False, |
| 165 | prompt_length: int = None, |
| 166 | micro_batch_size: int = None, |
| 167 | bad_ids: List = None, |
| 168 | temperature: float = 1.0, |
| 169 | topp: float = 1.0, |
| 170 | topk: int = 0.0, |
| 171 | greedy: bool = False, |
| 172 | recompute: bool = False, |
| 173 | ): |
| 174 | context_tokens, context_lengths = pad_batch(context_tokens, tokenizer.eos_token_id, seq_length) |
| 175 | |
| 176 | context_tokens_tensor = torch.cuda.LongTensor(context_tokens) |
| 177 | context_length_tensor = torch.cuda.LongTensor(context_lengths) |
| 178 | context_length = context_length_tensor.min().item() |
| 179 | tokens, attention_mask, position_ids = get_batch( |
| 180 | context_tokens_tensor, |
| 181 | micro_batch_size, |
| 182 | tokenizer.eos_token_id, |
| 183 | ) |
| 184 | |
| 185 | batch_token_iterator = sample_sequence_batch( |
| 186 | model, |
| 187 | tokenizer, |
| 188 | context_tokens_tensor, |
| 189 | context_length_tensor, |
| 190 | attention_mask, |
| 191 | position_ids, |
| 192 | seq_length=seq_length, |
| 193 | out_seq_length=out_seq_length, |
| 194 | return_scores=return_scores, |
| 195 | prompt_length=prompt_length, |
| 196 | bad_ids=bad_ids, |
| 197 | temperature=temperature, |
| 198 | topp=topp, |
| 199 | topk=topk, |
| 200 | greedy=greedy, |
| 201 | recompute=recompute, |
| 202 | ) |
| 203 | |
| 204 | for tokens, lengths in batch_token_iterator: |
| 205 | context_length += 1 |
| 206 | if tokens is not None: |
| 207 | yield tokens[:, :context_length], lengths |
| 208 | else: |
| 209 | yield None, None |
| 210 | |
| 211 | |
| 212 | def switch(val1, val2, boolean): |