| 10257 | |
| 10258 | |
| 10259 | class Sampler: |
| 10260 | TOKEN_DATA_DTYPE = np.dtype( |
| 10261 | [("id", np.intc), ("logit", np.single), ("p", np.single)], |
| 10262 | align=True, |
| 10263 | ) |
| 10264 | |
| 10265 | def __init__( |
| 10266 | self, |
| 10267 | *, |
| 10268 | seed: int, |
| 10269 | vocab: llama_cpp.llama_vocab_p, |
| 10270 | n_vocab: int, |
| 10271 | top_p: float, |
| 10272 | temperature: float, |
| 10273 | frequency_penalty: float, |
| 10274 | presence_penalty: float, |
| 10275 | logit_bias: Optional[Dict[int, float]], |
| 10276 | grammar_text: Optional[str] = None, |
| 10277 | grammar_root: str = "root", |
| 10278 | ) -> None: |
| 10279 | params = llama_cpp.llama_sampler_chain_default_params() |
| 10280 | self._sampler = llama_cpp.llama_sampler_chain_init(params) |
| 10281 | self._closed = False |
| 10282 | self._sample_logits_size = 0 |
| 10283 | self._sample_logits_token_data: Optional[Any] = None |
| 10284 | self._sample_logits_token_array: Optional[Any] = None |
| 10285 | self._sample_logits_recarray: Optional[np.recarray] = None |
| 10286 | if logit_bias: |
| 10287 | bias_array = (llama_cpp.llama_logit_bias * len(logit_bias))() |
| 10288 | for index, (token, bias) in enumerate(logit_bias.items()): |
| 10289 | bias_array[index].token = ctypes.c_int32(token) |
| 10290 | bias_array[index].bias = float(bias) |
| 10291 | llama_cpp.llama_sampler_chain_add( |
| 10292 | self._sampler, |
| 10293 | llama_cpp.llama_sampler_init_logit_bias( |
| 10294 | n_vocab, len(logit_bias), bias_array |
| 10295 | ), |
| 10296 | ) |
| 10297 | self.bias_array = bias_array |
| 10298 | if frequency_penalty != 0.0 or presence_penalty != 0.0: |
| 10299 | llama_cpp.llama_sampler_chain_add( |
| 10300 | self._sampler, |
| 10301 | llama_cpp.llama_sampler_init_penalties( |
| 10302 | 64, |
| 10303 | 1.0, |
| 10304 | frequency_penalty, |
| 10305 | presence_penalty, |
| 10306 | ), |
| 10307 | ) |
| 10308 | if grammar_text is not None: |
| 10309 | grammar_sampler = llama_cpp.llama_sampler_init_grammar( |
| 10310 | vocab, |
| 10311 | grammar_text.encode("utf-8"), |
| 10312 | grammar_root.encode("utf-8"), |
| 10313 | ) |
| 10314 | if grammar_sampler is None: |
| 10315 | raise RuntimeError("failed to initialize grammar sampler") |
| 10316 | llama_cpp.llama_sampler_chain_add(self._sampler, grammar_sampler) |
no outgoing calls
no test coverage detected
searching dependent graphs…