| 39 | |
| 40 | # For toolllama's predictions |
| 41 | def prepare_logits_processor( |
| 42 | temperature: float, repetition_penalty: float, top_p: float, top_k: int |
| 43 | ) -> LogitsProcessorList: |
| 44 | processor_list = LogitsProcessorList() |
| 45 | # TemperatureLogitsWarper doesn't accept 0.0, 1.0 makes it a no-op so we skip two cases. |
| 46 | if temperature >= 1e-5 and temperature != 1.0: |
| 47 | processor_list.append(TemperatureLogitsWarper(temperature)) |
| 48 | if repetition_penalty > 1.0: |
| 49 | processor_list.append(RepetitionPenaltyLogitsProcessor(repetition_penalty)) |
| 50 | if 1e-8 <= top_p < 1.0: |
| 51 | processor_list.append(TopPLogitsWarper(top_p)) |
| 52 | if top_k > 0: |
| 53 | processor_list.append(TopKLogitsWarper(top_k)) |
| 54 | return processor_list |
| 55 | |
| 56 | @torch.inference_mode() |
| 57 | def generate_stream( |