gguf-py : add constants and method related to <|eom_id|> token

fairydreaming · Aug 4, 2024 · f10b0e2 · f10b0e2
1 parent cc50e78
commit f10b0e2
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 0 deletions.
diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py
@@ -161,6 +161,7 @@ class Tokenizer:
         SUFFIX_ID            = "tokenizer.ggml.suffix_token_id"
         MIDDLE_ID            = "tokenizer.ggml.middle_token_id"
         EOT_ID               = "tokenizer.ggml.eot_token_id"
+        EOM_ID               = "tokenizer.ggml.eom_token_id"
 
     class Adapter:
         TYPE       = "adapter.type"
@@ -1327,3 +1328,4 @@ def get_type(val: Any) -> GGUFValueType:
 KEY_TOKENIZER_SUFFIX_ID  = Keys.Tokenizer.SUFFIX_ID
 KEY_TOKENIZER_MIDDLE_ID  = Keys.Tokenizer.MIDDLE_ID
 KEY_TOKENIZER_EOT_ID     = Keys.Tokenizer.EOT_ID
+KEY_TOKENIZER_EOM_ID     = Keys.Tokenizer.EOM_ID
diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py
@@ -826,6 +826,9 @@ def add_middle_token_id(self, id: int) -> None:
     def add_eot_token_id(self, id: int) -> None:
         self.add_uint32(Keys.Tokenizer.EOT_ID, id)
 
+    def add_eom_token_id(self, id: int) -> None:
+        self.add_uint32(Keys.Tokenizer.EOM_ID, id)
+
     def _pack(self, fmt: str, value: Any, skip_pack_prefix: bool = False) -> bytes:
         pack_prefix = ''
         if not skip_pack_prefix: