From a7ba85834ffef7de128d5fda5a1d29f8f001b049 Mon Sep 17 00:00:00 2001 From: Andrei Betlen Date: Sat, 20 May 2023 08:13:41 -0400 Subject: [PATCH] Add n_ctx, n_vocab, and n_embd properties --- llama_cpp/llama.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/llama_cpp/llama.py b/llama_cpp/llama.py index 564c6c3..e399028 100644 --- a/llama_cpp/llama.py +++ b/llama_cpp/llama.py @@ -1291,6 +1291,24 @@ class Llama: if llama_cpp.llama_set_state_data(self.ctx, state.llama_state) != state_size: raise RuntimeError("Failed to set llama state data") + @property + def n_ctx(self) -> int: + """Return the context window size.""" + assert self.ctx is not None + return llama_cpp.llama_n_ctx(self.ctx) + + @property + def n_embd(self) -> int: + """Return the embedding size.""" + assert self.ctx is not None + return llama_cpp.llama_n_embd(self.ctx) + + @property + def n_vocab(self) -> int: + """Return the vocabulary size.""" + assert self.ctx is not None + return llama_cpp.llama_n_vocab(self.ctx) + @staticmethod def token_eos() -> int: """Return the end-of-sequence token."""