Add HF tokenizer helper, refactor example

eval/longctx.py

@@ -2,7 +2,6 @@ import sys, os
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
 import argparse
-from transformers import AutoTokenizer
 from exllamav3.util.progress import ProgressBar
 from exllamav3 import Config, Model, Cache, Tokenizer, model_init, Generator, Job, GreedySampler
 import torch
@@ -27,9 +26,6 @@ def main(args):
     print(f" -- Model: {args.model_dir}")
     print(f" -- Bitrate: {bpw_layer:.2f} bpw / {bpw_head:.2f} bpw (head)")
 
-    # Load Transformers tokenizers
-    t_tokenizer = AutoTokenizer.from_pretrained(args.model_dir)
-
     # Get
     texts_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "eval_texts")
     with open(os.path.join(texts_dir, "illustrious_client.txt"), "r") as file:
@@ -55,8 +51,7 @@ def main(args):
             "role": "user",
             "content": instruction
         }]
-        input_ids = t_tokenizer.apply_chat_template(chat, add_generation_prompt = True)
-        input_ids = torch.tensor(input_ids, dtype = torch.long).unsqueeze(0)
+        input_ids = tokenizer.hf_chat_template(chat, add_generation_prompt = True)
         job = Job(
             input_ids = input_ids,
             max_new_tokens = 768,

exllamav3/tokenizer/tokenizer.py

@@ -37,6 +37,7 @@ class Tokenizer:
         self.unspecial_piece_to_id = {}
         self.unspecial_id_to_piece = {}
         self.vocab = None
+        self.hf_tokenizer = None
 
         # Regex
         self.ord_exp = re.compile(r"^<0x([0-9A-Fa-f]+)>$")
@@ -621,3 +622,35 @@ class Tokenizer:
         }
     def get_vocab_dict(self):
         return self._get_vocab_dict
+
+    def hf_chat_template(
+        self,
+        messages: list,
+        add_generation_prompt: bool = True
+    ):
+        """
+        Tokenize with HF tokenizer. Requires `transformers`
+
+        :param messages:
+            HF-formatted list of messages, e.g.
+                [{
+                    "role": "user",
+                    "content": "Hello."
+                }]
+
+        :param add_generation_prompt:
+            bool, add generation prompt
+
+        :return:
+            Topken IDs tensor, shape (1, num_tokens)
+        """
+
+        from transformers import AutoTokenizer
+        if self.hf_tokenizer is None:
+            self.hf_tokenizer = AutoTokenizer.from_pretrained(self.config.directory)
+        ids = self.hf_tokenizer.apply_chat_template(
+            messages,
+            add_generation_prompt = add_generation_prompt
+        )["input_ids"]
+        ids = torch.tensor(ids, dtype = torch.long).unsqueeze(0)
+        return ids