Replace token function with vocab function

Author: JamePeng

llama_cpp/_internals.py

@@ -112,48 +112,69 @@ def get_tensor(self, name: str) -> ctypes.c_void_p:
     # Vocab
 
     def token_get_text(self, token: int) -> str:
-        return llama_cpp.llama_token_get_text(self.vocab, token).decode("utf-8")
+        return llama_cpp.llama_vocab_get_text(self.vocab, token).decode("utf-8")
 
     def token_get_score(self, token: int) -> float:
-        return llama_cpp.llama_token_get_score(self.vocab, token)
+        return llama_cpp.llama_vocab_get_score(self.vocab, token)
 
     def token_get_attr(self, token: int) -> int:
-        return llama_cpp.llama_token_get_attr(self.vocab, token)
+        return llama_cpp.llama_vocab_get_attr(self.vocab, token)
+
+    def token_is_eog(self, token: int) -> bool:
+        return llama_cpp.llama_vocab_is_eog(self.vocab, token)
+
+    def token_is_control(self, token: int) -> bool:
+        return llama_cpp.llama_vocab_is_control(self.vocab, token)
 
     # Special tokens
 
     def token_bos(self) -> int:
-        return llama_cpp.llama_token_bos(self.vocab)
+        return llama_cpp.llama_vocab_bos(self.vocab)
 
     def token_eos(self) -> int:
-        return llama_cpp.llama_token_eos(self.vocab)
+        return llama_cpp.llama_vocab_eos(self.vocab)
 
-    def token_cls(self) -> int:
-        return llama_cpp.llama_token_cls(self.vocab)
+    def token_eot(self) -> int:
+        return llama_cpp.llama_vocab_eot(self.vocab)
 
     def token_sep(self) -> int:
-        return llama_cpp.llama_token_sep(self.vocab)
+        return llama_cpp.llama_vocab_sep(self.vocab)
 
     def token_nl(self) -> int:
-        return llama_cpp.llama_token_nl(self.vocab)
+        return llama_cpp.llama_vocab_nl(self.vocab)
 
-    def token_prefix(self) -> int:
-        raise NotImplementedError("token_prefix is not implemented in llama.cpp")
+    def token_pad(self) -> int:
+        return llama_cpp.llama_vocab_pad(self.vocab)
 
-    def token_middle(self) -> int:
-        raise NotImplementedError("token_middle is not implemented in llama.cpp")
+    def token_cls(self) -> int:
+        return llama_cpp.llama_vocab_cls(self.vocab)
 
-    def token_suffix(self) -> int:
-        raise NotImplementedError("token_suffix is not implemented in llama.cpp")
+    def token_fim_pre(self) -> int:
+        return llama_cpp.llama_vocab_fim_pre(self.vocab)
 
-    def token_eot(self) -> int:
-        return llama_cpp.llama_token_eot(self.vocab)
+    def token_fim_suf(self) -> int:
+        return llama_cpp.llama_vocab_fim_suf(self.vocab)
+
+    def token_fim_mid(self) -> int:
+        return llama_cpp.llama_vocab_fim_mid(self.vocab)
+
+    def token_fim_pad(self) -> int:
+        return llama_cpp.llama_vocab_fim_pad(self.vocab)
+
+    def token_fim_rep(self) -> int:
+        return llama_cpp.llama_vocab_fim_rep(self.vocab)
+
+    def token_fim_sep(self) -> int:
+        return llama_cpp.llama_vocab_fim_sep(self.vocab)
+
+    def get_add_bos(self) -> bool:
+        return llama_cpp.llama_vocab_get_add_bos(self.vocab)
 
-    def add_bos_token(self) -> bool:
-        return llama_cpp.llama_add_bos_token(self.vocab)
+    def get_add_eos(self) -> bool:
+        return llama_cpp.llama_vocab_get_add_eos(self.vocab)
 
-    def add_eos_token(self) -> bool:
-        return llama_cpp.llama_add_eos_token(self.vocab)
+    def get_add_sep(self) -> bool:
+        return llama_cpp.llama_vocab_get_add_sep(self.vocab)
 
     # Tokenization
 

llama_cpp/llama.py

@@ -1189,29 +1189,29 @@ def _create_completion(
 
         completion_id: str = f"cmpl-{str(uuid.uuid4())}"
         created: int = int(time.time())
-        bos_token_id: int = self.token_bos()
-        cls_token_id: int = self._model.token_cls()
+        bos_token_id: int = self._model.token_bos()
+        eos_token_id: int = self._model.token_eos()
         sep_token_id: int = self._model.token_sep()
-        prefix_token_id: int = 0 # self._model.token_prefix() # TODO: Fix
-        middle_token_id: int = 0 # self._model.token_middle() # TODO: Fix
-        suffix_token_id: int = 0 # self._model.token_suffix() # TODO: Fix
+        prefix_token_id: int = self._model.token_fim_pre()
+        middle_token_id: int = self._model.token_fim_mid()
+        suffix_token_id: int = self._model.token_fim_suf()
         add_space_prefix: bool = (
             self.metadata.get("tokenizer.ggml.add_space_prefix", "true") == "true"
         )
-        bos_tokens: List[int] = [cls_token_id if cls_token_id != -1 else bos_token_id]
+        bos_tokens: List[int] = [bos_token_id]
         eos_tokens: List[int] = [
-            sep_token_id if sep_token_id != -1 else self.token_eos()
+            sep_token_id if self._model.get_add_sep() else eos_token_id
         ]
 
         if (
             (isinstance(prompt, list) and suffix is None)
-            or not self._model.add_bos_token()
+            or not self._model.get_add_bos()
             or bos_tokens[:1] == [-1]
         ):
             bos_tokens = []
 
         if (isinstance(prompt, list) and suffix is None) or (
-            not self._model.add_eos_token() and sep_token_id == -1
+            not self._model.get_add_eos() and not self._model.get_add_sep()
         ):
             eos_tokens = []
 
@@ -2294,18 +2294,30 @@ def tokenizer(self) -> LlamaTokenizer:
         """Return the llama tokenizer for this model."""
         return LlamaTokenizer(self)
 
+    def token_bos(self) -> int:
+        """Return the beginning-of-sequence token."""
+        return self._model.token_bos()
+
     def token_eos(self) -> int:
         """Return the end-of-sequence token."""
         return self._model.token_eos()
 
-    def token_bos(self) -> int:
-        """Return the beginning-of-sequence token."""
-        return self._model.token_bos()
+    def token_eot(self) -> int:
+        """Return the end-of-turn token."""
+        return self._model.token_eot()
+
+    def token_sep(self) -> int:
+        """Return the sentence-separator token."""
+        return self._model.token_sep()
 
     def token_nl(self) -> int:
-        """Return the newline token."""
+        """Return the next-line token."""
         return self._model.token_nl()
 
+    def token_pad(self) -> int:
+        """Return the padding token."""
+        return self._model.token_pad()
+
     def pooling_type(self) -> str:
         """Return the pooling type."""
         return self._ctx.pooling_type()