| 10339 | return int(llama_cpp.llama_sampler_sample(self._sampler, ctx, output_index)) |
| 10340 | |
| 10341 | def _ensure_sample_logits_buffer(self, size: int) -> None: |
| 10342 | if size == self._sample_logits_size and self._sample_logits_recarray is not None: |
| 10343 | return |
| 10344 | token_data = (llama_cpp.llama_token_data * size)() |
| 10345 | token_data_address = ctypes.addressof(token_data) |
| 10346 | recarray = np.recarray( |
| 10347 | shape=(size,), |
| 10348 | dtype=self.TOKEN_DATA_DTYPE, |
| 10349 | buf=cast( |
| 10350 | Any, |
| 10351 | (llama_cpp.llama_token_data * size).from_address(token_data_address), |
| 10352 | ), |
| 10353 | ) |
| 10354 | recarray.id[:] = np.arange(size, dtype=np.intc) |
| 10355 | token_array = llama_cpp.llama_token_data_array( |
| 10356 | data=token_data, |
| 10357 | size=size, |
| 10358 | selected=-1, |
| 10359 | sorted=False, |
| 10360 | ) |
| 10361 | self._sample_logits_size = size |
| 10362 | self._sample_logits_token_data = token_data |
| 10363 | self._sample_logits_token_array = token_array |
| 10364 | self._sample_logits_recarray = recarray |
| 10365 | |
| 10366 | def sample_logits(self, logits: np.ndarray) -> int: |
| 10367 | self._ensure_sample_logits_buffer(len(logits)) |