Sample a token from the model. Args: top_k: The top-k sampling parameter. top_p: The top-p sampling parameter. temp: The temperature parameter. repeat_penalty: The repeat penalty parameter. Returns: The sampled token.
(
self,
top_k: int = 40,
top_p: float = 0.95,
min_p: float = 0.05,
typical_p: float = 1.0,
temp: float = 0.80,
repeat_penalty: float = 1.0,
frequency_penalty: float = 0.0,
presence_penalty: float = 0.0,
tfs_z: float = 1.0,
mirostat_mode: int = 0,
mirostat_eta: float = 0.1,
mirostat_tau: float = 5.0,
penalize_nl: bool = True,
logits_processor: Optional[LogitsProcessorList] = None,
grammar: Optional[LlamaGrammar] = None,
idx: Optional[int] = None,
)
| 783 | return sampler |
| 784 | |
| 785 | def sample( |
| 786 | self, |
| 787 | top_k: int = 40, |
| 788 | top_p: float = 0.95, |
| 789 | min_p: float = 0.05, |
| 790 | typical_p: float = 1.0, |
| 791 | temp: float = 0.80, |
| 792 | repeat_penalty: float = 1.0, |
| 793 | frequency_penalty: float = 0.0, |
| 794 | presence_penalty: float = 0.0, |
| 795 | tfs_z: float = 1.0, |
| 796 | mirostat_mode: int = 0, |
| 797 | mirostat_eta: float = 0.1, |
| 798 | mirostat_tau: float = 5.0, |
| 799 | penalize_nl: bool = True, |
| 800 | logits_processor: Optional[LogitsProcessorList] = None, |
| 801 | grammar: Optional[LlamaGrammar] = None, |
| 802 | idx: Optional[int] = None, |
| 803 | ): |
| 804 | """Sample a token from the model. |
| 805 | |
| 806 | Args: |
| 807 | top_k: The top-k sampling parameter. |
| 808 | top_p: The top-p sampling parameter. |
| 809 | temp: The temperature parameter. |
| 810 | repeat_penalty: The repeat penalty parameter. |
| 811 | |
| 812 | Returns: |
| 813 | The sampled token. |
| 814 | """ |
| 815 | assert self.n_tokens > 0 |
| 816 | |
| 817 | tmp_sampler = False |
| 818 | |
| 819 | if self._sampler is None: |
| 820 | tmp_sampler = True |
| 821 | self._sampler = self._init_sampler( |
| 822 | top_k=top_k, |
| 823 | top_p=top_p, |
| 824 | min_p=min_p, |
| 825 | typical_p=typical_p, |
| 826 | temp=temp, |
| 827 | repeat_penalty=repeat_penalty, |
| 828 | frequency_penalty=frequency_penalty, |
| 829 | presence_penalty=presence_penalty, |
| 830 | tfs_z=tfs_z, |
| 831 | mirostat_mode=mirostat_mode, |
| 832 | mirostat_tau=mirostat_tau, |
| 833 | mirostat_eta=mirostat_eta, |
| 834 | penalize_nl=penalize_nl, |
| 835 | logits_processor=logits_processor, |
| 836 | grammar=grammar, |
| 837 | ) |
| 838 | |
| 839 | ridx = idx - self.n_tokens if idx is not None else -1 |
| 840 | |
| 841 | assert self.ctx is not None |
| 842 | token = self._sampler.sample(self._ctx, ridx) |