From cedb4870c660875d7d0b6fd98015960a27dd302e Mon Sep 17 00:00:00 2001
From: klosax <131523366+klosax@users.noreply.github.com>
Date: Mon, 14 Aug 2023 22:08:40 +0200
Subject: [PATCH] gguf.py : add token types

---
 gguf.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/gguf.py b/gguf.py
index 93591d5c7..d3b95963a 100644
--- a/gguf.py
+++ b/gguf.py
@@ -297,6 +297,9 @@ class GGUFWriter:
     def add_token_merges(self, merges: List):
         self.add_array(constants.KEY_TOKENIZER_MERGES, merges)
 
+    def add_token_types(self, types: List[int]):
+        self.add_array(constants.KEY_TOKENIZER_TOKEN_TYPE, types)
+
     def add_token_scores(self, scores: List[float]):
         self.add_array(constants.KEY_TOKENIZER_SCORES, scores)