Move self._vocab to llama.py

Author: JamePeng

llama_cpp/_internals.py

@@ -43,7 +43,6 @@ def __init__(
         self._exit_stack = ExitStack()
 
         model = None
-        vocab = None
 
         if not os.path.exists(path_model):
             raise ValueError(f"Model path does not exist: {path_model}")
@@ -58,24 +57,12 @@ def __init__(
 
         self.model = model
 
-        vocab = llama_cpp.llama_model_get_vocab(self.model)
-
-        if vocab is None:
-            raise ValueError(f"Failed to load vocab from file: {path_model}")
-
-        self.vocab = vocab
-
         def free_model():
             if self.model is None:
                 return
             llama_cpp.llama_model_free(self.model)
             self.model = None
 
-            if self.vocab is None:
-                return
-            llama_cpp.llama_model_free(self.vocab)
-            self.vocab = None
-
         self._exit_stack.callback(free_model)
 
     def close(self):
@@ -84,11 +71,11 @@ def close(self):
     def __del__(self):
         self.close()
 
-    def vocab_type(self) -> int:
-        return llama_cpp.llama_vocab_type(self.vocab)
+    def vocab_type(self, _vocab:llama_cpp.llama_vocab_p) -> int:
+        return llama_cpp.llama_vocab_type(_vocab)
 
-    def n_vocab(self) -> int:
-        return llama_cpp.llama_vocab_n_tokens(self.vocab)
+    def n_vocab(self, _vocab:llama_cpp.llama_vocab_p) -> int:
+        return llama_cpp.llama_vocab_n_tokens(_vocab)
 
     def n_ctx_train(self) -> int:
         return llama_cpp.llama_model_n_ctx_train(self.model)
@@ -112,66 +99,66 @@ def n_params(self) -> int:
 
     # Vocab
 
-    def token_get_text(self, token: int) -> str:
-        return llama_cpp.llama_vocab_get_text(self.vocab, token).decode("utf-8")
+    def token_get_text(self, _vocab:llama_cpp.llama_vocab_p, token: int) -> str:
+        return llama_cpp.llama_vocab_get_text(_vocab, token).decode("utf-8")
 
-    def token_get_score(self, token: int) -> float:
-        return llama_cpp.llama_vocab_get_score(self.vocab, token)
+    def token_get_score(self, _vocab:llama_cpp.llama_vocab_p, token: int) -> float:
+        return llama_cpp.llama_vocab_get_score(_vocab, token)
 
-    def token_get_attr(self, token: int) -> int:
-        return llama_cpp.llama_vocab_get_attr(self.vocab, token)
+    def token_get_attr(self, _vocab:llama_cpp.llama_vocab_p, token: int) -> int:
+        return llama_cpp.llama_vocab_get_attr(_vocab, token)
 
     # Special tokens
 
-    def token_bos(self) -> int:
-        return llama_cpp.llama_vocab_bos(self.vocab)
+    def token_bos(self, _vocab:llama_cpp.llama_vocab_p) -> int:
+        return llama_cpp.llama_vocab_bos(_vocab)
 
-    def token_eos(self) -> int:
-        return llama_cpp.llama_vocab_eos(self.vocab)
+    def token_eos(self, _vocab:llama_cpp.llama_vocab_p) -> int:
+        return llama_cpp.llama_vocab_eos(_vocab)
 
-    def token_eot(self) -> int:
-        return llama_cpp.llama_vocab_eot(self.vocab)
+    def token_eot(self, _vocab:llama_cpp.llama_vocab_p) -> int:
+        return llama_cpp.llama_vocab_eot(_vocab)
 
-    def token_cls(self) -> int:
-        return llama_cpp.llama_vocab_cls(self.vocab)
+    def token_cls(self, _vocab:llama_cpp.llama_vocab_p) -> int:
+        return llama_cpp.llama_vocab_cls(_vocab)
 
-    def token_sep(self) -> int:
-        return llama_cpp.llama_vocab_sep(self.vocab)
+    def token_sep(self, _vocab:llama_cpp.llama_vocab_p) -> int:
+        return llama_cpp.llama_vocab_sep(_vocab)
 
-    def token_nl(self) -> int:
-        return llama_cpp.llama_vocab_nl(self.vocab)
+    def token_nl(self, _vocab:llama_cpp.llama_vocab_p) -> int:
+        return llama_cpp.llama_vocab_nl(_vocab)
 
-    def token_pad(self) -> int:
-        return llama_cpp.llama_vocab_pad(self.vocab)
+    def token_pad(self, _vocab:llama_cpp.llama_vocab_p) -> int:
+        return llama_cpp.llama_vocab_pad(_vocab)
 
-    def token_prefix(self) -> int:
-        return llama_cpp.llama_vocab_fim_pre(self.vocab)
+    def token_prefix(self, _vocab:llama_cpp.llama_vocab_p) -> int:
+        return llama_cpp.llama_vocab_fim_pre(_vocab)
 
-    def token_middle(self) -> int:
-        return llama_cpp.llama_vocab_fim_mid(self.vocab)
+    def token_middle(self, _vocab:llama_cpp.llama_vocab_p) -> int:
+        return llama_cpp.llama_vocab_fim_mid(_vocab)
 
-    def token_suffix(self) -> int:
-        return llama_cpp.llama_vocab_fim_suf(self.vocab)
+    def token_suffix(self, _vocab:llama_cpp.llama_vocab_p) -> int:
+        return llama_cpp.llama_vocab_fim_suf(_vocab)
 
-    def add_bos_token(self) -> bool:
-        return llama_cpp.llama_vocab_get_add_bos(self.vocab)
+    def add_bos_token(self, _vocab:llama_cpp.llama_vocab_p) -> bool:
+        return llama_cpp.llama_vocab_get_add_bos(_vocab)
 
-    def add_eos_token(self) -> bool:
-        return llama_cpp.llama_vocab_get_add_eos(self.vocab)
+    def add_eos_token(self, _vocab:llama_cpp.llama_vocab_p) -> bool:
+        return llama_cpp.llama_vocab_get_add_eos(_vocab)
 
     # Tokenization
 
-    def tokenize(self, text: bytes, add_bos: bool, special: bool):
+    def tokenize(self, _vocab:llama_cpp.llama_vocab_p, text: bytes, add_bos: bool, special: bool):
         n_ctx = self.n_ctx_train()
         tokens = (llama_cpp.llama_token * n_ctx)()
         n_tokens = llama_cpp.llama_tokenize(
-            self.vocab, text, len(text), tokens, n_ctx, add_bos, special
+            _vocab, text, len(text), tokens, n_ctx, add_bos, special
         )
         if n_tokens < 0:
             n_tokens = abs(n_tokens)
             tokens = (llama_cpp.llama_token * n_tokens)()
             n_tokens = llama_cpp.llama_tokenize(
-                self.vocab, text, len(text), tokens, n_tokens, add_bos, special
+                _vocab, text, len(text), tokens, n_tokens, add_bos, special
             )
             if n_tokens < 0:
                 raise RuntimeError(
@@ -618,10 +605,11 @@ def prev_str(self, ctx_main: LlamaContext, n: int) -> str:
     def sample(
         self,
         ctx_main: LlamaContext,
+        _vocab:llama_cpp.llama_vocab_p,
         idx: int = 0,
         logits_array: Optional[npt.NDArray[np.single]] = None,
     ):
-        n_vocab = ctx_main.model.n_vocab()
+        n_vocab = ctx_main.model.n_vocab(_vocab)
         id: int = 0
 
         if logits_array is None:

llama_cpp/llama.py

@@ -374,6 +374,8 @@ def __init__(
             )
         )
 
+        self._vocab = llama_cpp.llama_model_get_vocab(self._model.model)
+
         # Override tokenizer
         self.tokenizer_ = tokenizer or LlamaTokenizer(self)
 
@@ -2171,23 +2173,23 @@ def n_embd(self) -> int:
 
     def n_vocab(self) -> int:
         """Return the vocabulary size."""
-        return self._model.n_vocab()
+        return self._model.n_vocab(self._vocab)
 
     def tokenizer(self) -> LlamaTokenizer:
         """Return the llama tokenizer for this model."""
         return LlamaTokenizer(self)
 
     def token_eos(self) -> int:
         """Return the end-of-sequence token."""
-        return self._model.token_eos()
+        return self._model.token_eos(self._vocab)
 
     def token_bos(self) -> int:
         """Return the beginning-of-sequence token."""
-        return self._model.token_bos()
+        return self._model.token_bos(self._vocab)
 
     def token_nl(self) -> int:
         """Return the newline token."""
-        return self._model.token_nl()
+        return self._model.token_nl(self._vocab)
 
     def pooling_type(self) -> str:
         """Return the pooling type."""

llama_cpp/llama_cpp.py

@@ -1302,7 +1302,7 @@ def llama_pooling_type(ctx: llama_context_p, /) -> int:
     ...
 
 # LLAMA_API const struct llama_vocab * llama_model_get_vocab(const struct llama_model * model);
-@ctypes_function("llama_model_get_vocab", [llama_model_p_ctypes], llama_vocab_p)
+@ctypes_function("llama_model_get_vocab", [llama_model_p_ctypes], llama_vocab_p_ctypes)
 def llama_model_get_vocab(model: llama_model_p, /) -> Optional[llama_vocab_p]:
     ...