| 2227 | ) |
| 2228 | |
| 2229 | def load_state(self, state: LlamaState) -> None: |
| 2230 | # Only filling in up to `n_tokens` and then zero-ing out the rest |
| 2231 | self.scores[: state.n_tokens, :] = state.scores.copy() |
| 2232 | rest = self.scores[state.n_tokens :, :] |
| 2233 | rest[rest > 0] = 0.0 |
| 2234 | self.input_ids = state.input_ids.copy() |
| 2235 | self.n_tokens = state.n_tokens |
| 2236 | self._requires_eval = True |
| 2237 | self._seed = state.seed |
| 2238 | state_size = state.llama_state_size |
| 2239 | LLamaStateArrayType = ctypes.c_uint8 * state_size |
| 2240 | llama_state = LLamaStateArrayType.from_buffer_copy(state.llama_state) |
| 2241 | |
| 2242 | if ( |
| 2243 | llama_cpp.llama_state_set_data(self._ctx.ctx, llama_state, state_size) |
| 2244 | != state_size |
| 2245 | ): |
| 2246 | raise RuntimeError("Failed to set llama state data") |
| 2247 | |
| 2248 | def n_ctx(self) -> int: |
| 2249 | """Return the context window size.""" |