def reload_from_low_level_state_file(model: Llama, state_file: str):
    from llama_cpp import llama_load_session_file

    n_ctx = model.n_ctx()
    session_tokens = (llama_cpp.llama_token * n_ctx)()
    n_tok = llama_cpp.ctypes.c_size_t()

    # Load context from file
    status = llama_load_session_file(
        model.ctx,
        state_file.encode("utf-8"),
        session_tokens,
        n_ctx,
        llama_cpp.ctypes.byref(n_tok),
    )

    if status != 1:
        raise ValueError(f"Failed to load context from '{state_file}'")

    n_tok_int = int(n_tok.value)
    print(f"n_tok: {n_tok_int}")

    assert (
        session_tokens[:n_tok_int] == state.input_ids[:n_tok_int].tolist()
    ), "Tokens should match"

    # Set various objects like in `load_state`
    n_vocab = model.n_vocab()
    logits_from_low_level = np.ctypeslib.as_array(
        model._ctx.get_logits(), (n_tok_int, n_vocab)
    )

    model.scores[:n_tok_int, :] = logits_from_low_level.copy()
    model.scores[n_tok_int:, :] = 0.0

    for i in range(0, n_tok_int):
        model.input_ids[i] = session_tokens[i]

    model.n_tokens = n_tok_int

# Loads and resets to state created by doing `save_state` after `model.eval(prompt_tokens)`
model = get_nekomata7b()
model.load_state(state)

reload_from_low_level_state_file(model, "nekomata-7b-ctx.bin") 
out3 = model.create_chat_completion(
    [
        {"role": "system", "content": DEFAULT_SYSTEM_PROMPT},
        {"role": "user", "content": prompt},
    ],
    temperature=0.0,
    max_tokens=300,
    seed=1432
)
out3
pred_logits3 = model.scores[n_tok_prompt:model.n_tokens, :] 
assert (pred_logits == pred_logits3).all(), "All prediction logits should match"