| 2197 | self.__init__(**state) |
| 2198 | |
| 2199 | def save_state(self) -> LlamaState: |
| 2200 | if self.verbose: |
| 2201 | print("Llama.save_state: saving llama state", file=sys.stderr) |
| 2202 | state_size = llama_cpp.llama_state_get_size(self._ctx.ctx) |
| 2203 | if self.verbose: |
| 2204 | print(f"Llama.save_state: got state size: {state_size}", file=sys.stderr) |
| 2205 | llama_state = (ctypes.c_uint8 * int(state_size))() |
| 2206 | if self.verbose: |
| 2207 | print("Llama.save_state: allocated state", file=sys.stderr) |
| 2208 | n_bytes = llama_cpp.llama_state_get_data(self._ctx.ctx, llama_state, state_size) |
| 2209 | if self.verbose: |
| 2210 | print(f"Llama.save_state: copied llama state: {n_bytes}", file=sys.stderr) |
| 2211 | if int(n_bytes) > int(state_size): |
| 2212 | raise RuntimeError("Failed to copy llama state data") |
| 2213 | llama_state_compact = (ctypes.c_uint8 * int(n_bytes))() |
| 2214 | llama_cpp.ctypes.memmove(llama_state_compact, llama_state, int(n_bytes)) |
| 2215 | if self.verbose: |
| 2216 | print( |
| 2217 | f"Llama.save_state: saving {n_bytes} bytes of llama state", |
| 2218 | file=sys.stderr, |
| 2219 | ) |
| 2220 | return LlamaState( |
| 2221 | scores=self._scores.copy(), |
| 2222 | input_ids=self.input_ids.copy(), |
| 2223 | n_tokens=self.n_tokens, |
| 2224 | llama_state=bytes(llama_state_compact), |
| 2225 | llama_state_size=n_bytes, |
| 2226 | seed=self._seed, |
| 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 |