(
self, *, n_tokens: int, embd: int, n_seq_max: int, verbose: bool = True
)
| 466 | |
| 467 | class LlamaBatch: |
| 468 | def __init__( |
| 469 | self, *, n_tokens: int, embd: int, n_seq_max: int, verbose: bool = True |
| 470 | ): |
| 471 | self._n_tokens = n_tokens |
| 472 | self.embd = embd |
| 473 | self.n_seq_max = n_seq_max |
| 474 | self.verbose = verbose |
| 475 | self._exit_stack = ExitStack() |
| 476 | |
| 477 | batch = llama_cpp.llama_batch_init(self._n_tokens, self.embd, self.n_seq_max) |
| 478 | |
| 479 | if batch is None: |
| 480 | raise ValueError("Failed to create llama_batch") |
| 481 | |
| 482 | self.batch = batch |
| 483 | self.sampler = None # LlamaBatch doesn't use samplers, but some cleanup code expects this attribute |
| 484 | |
| 485 | def free_batch(): |
| 486 | if self.batch is None: |
| 487 | return |
| 488 | llama_cpp.llama_batch_free(self.batch) |
| 489 | self.batch = None |
| 490 | |
| 491 | self._exit_stack.callback(free_batch) |
| 492 | |
| 493 | def close(self): |
| 494 | self._exit_stack.close() |
nothing calls this directly
no outgoing calls
no test coverage detected