custom dsv32 chat template (#693)

* custom dsv32 chat template

* use has_chat_template

Author: awni

mlx_lm/cache_prompt.py

@@ -41,16 +41,6 @@ def setup_arg_parser():
         default=None,
         help="End of sequence token for tokenizer",
     )
-    parser.add_argument(
-        "--ignore-chat-template",
-        action="store_true",
-        help="Use the raw prompt without the tokenizer's chat template.",
-    )
-    parser.add_argument(
-        "--use-default-chat-template",
-        action="store_true",
-        help="Use the default chat template",
-    )
     parser.add_argument(
         "--max-kv-size",
         type=int,
@@ -107,11 +97,7 @@ def main():
 
     args.prompt = sys.stdin.read() if args.prompt == "-" else args.prompt
 
-    if args.use_default_chat_template:
-        if tokenizer.chat_template is None:
-            tokenizer.chat_template = tokenizer.default_chat_template
-
-    if not args.ignore_chat_template and tokenizer.chat_template is not None:
+    if tokenizer.has_chat_template:
         messages = [{"role": "user", "content": args.prompt}]
         prompt = tokenizer.apply_chat_template(
             messages,
@@ -155,7 +141,6 @@ def callback(processed, total_tokens):
     print("Saving...")
     metadata = {}
     metadata["model"] = args.model
-    metadata["chat_template"] = json.dumps(tokenizer.chat_template)
     metadata["tokenizer_config"] = json.dumps(tokenizer_config)
     save_prompt_cache(args.prompt_cache_file, cache, metadata)
 

mlx_lm/generate.py

@@ -1302,15 +1302,9 @@ def main():
     if args.chat_template_config is not None:
         template_kwargs = json.loads(args.chat_template_config)
 
-    if args.use_default_chat_template:
-        if tokenizer.chat_template is None:
-            tokenizer.chat_template = tokenizer.default_chat_template
-    elif using_cache:
-        tokenizer.chat_template = json.loads(metadata["chat_template"])
-
     prompt = args.prompt.replace("\\n", "\n").replace("\\t", "\t")
     prompt = sys.stdin.read() if prompt == "-" else prompt
-    if not args.ignore_chat_template and tokenizer.chat_template is not None:
+    if not args.ignore_chat_template and tokenizer.has_chat_template:
         if args.system_prompt is not None:
             messages = [{"role": "system", "content": args.system_prompt}]
         else:

mlx_lm/tokenizer_utils.py

@@ -1,3 +1,4 @@
+import importlib
 import json
 from functools import partial
 from json import JSONDecodeError
@@ -248,7 +249,11 @@ class TokenizerWrapper:
     """
 
     def __init__(
-        self, tokenizer, detokenizer_class=NaiveStreamingDetokenizer, eos_token_ids=None
+        self,
+        tokenizer,
+        detokenizer_class=NaiveStreamingDetokenizer,
+        eos_token_ids=None,
+        chat_template=None,
     ):
         self._tokenizer = tokenizer
         self._detokenizer_class = detokenizer_class
@@ -261,6 +266,10 @@ def __init__(
         self._think_end = None
         self._tool_call_start = None
         self._tool_call_end = None
+        self._chat_template = chat_template
+        self.has_chat_template = (
+            tokenizer.chat_template is not None or chat_template is not None
+        )
 
         THINK_TOKENS = [("<think>", "</think>")]
         TOOL_CALL_TOKENS = [("<tool_call>", "</tool_call>")]
@@ -278,9 +287,15 @@ def __init__(
                     self._tool_call_end = tool_call_end
                     break
 
-    def apply_chat_template(self, *args, **kwargs):
+    def apply_chat_template(self, *args, tokenize=True, **kwargs):
+        if self._chat_template is not None:
+            out = self._chat_template(*args, **kwargs)
+            if tokenize:
+                out = self._tokenizer.encode(out, add_special_tokens=False)
+            return out
+
         kwargs["return_dict"] = False
-        return self._tokenizer.apply_chat_template(*args, **kwargs)
+        return self._tokenizer.apply_chat_template(*args, tokenize=tokenize, **kwargs)
 
     def add_eos_token(self, token: str):
         token_id = None
@@ -450,12 +465,28 @@ def load(
     if isinstance(eos_token_ids, int):
         eos_token_ids = [eos_token_ids]
 
+    tokenizer_config_file = model_path / "tokenizer_config.json"
+    custom_tokenizer = None
+    if tokenizer_config_file.exists():
+        with open(tokenizer_config_file, "r", encoding="utf-8") as fid:
+            try:
+                tokenizer_config = json.load(fid)
+            except JSONDecodeError as e:
+                raise JSONDecodeError(
+                    "Failed to parse tokenizer_config.json", e.doc, e.pos
+                )
+        if tokenizer_type := tokenizer_config.get("tokenizer_type", False):
+            custom_tokenizer = importlib.import_module(
+                f"mlx_lm.tokenizers.{tokenizer_type}"
+            )
+
     if return_tokenizer:
         kwargs = tokenizer_config_extra or {}
         return TokenizerWrapper(
             AutoTokenizer.from_pretrained(model_path, **kwargs),
             detokenizer_class,
             eos_token_ids=eos_token_ids,
+            chat_template=getattr(custom_tokenizer, "apply_chat_template", None),
         )
     else:
         return detokenizer_class