From 8f906c859ac71bd779e33468b4798fed0fd3be6a Mon Sep 17 00:00:00 2001
From: Anchen <li.anchen.au@gmail.com>
Date: Thu, 21 Mar 2024 15:39:39 +1100
Subject: [PATCH] chore(mlx-lm): enable to apply default chat template (#577)

* chore(mlx-lm): enable to apply default chat template

* Add option to use default chat template

* chore: rename the flag to use default chat template
---
 llms/mlx_lm/generate.py | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/llms/mlx_lm/generate.py b/llms/mlx_lm/generate.py
index 394c8e15..dfdffa1b 100644
--- a/llms/mlx_lm/generate.py
+++ b/llms/mlx_lm/generate.py
@@ -61,6 +61,11 @@ def setup_arg_parser():
         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(
         "--colorize",
         action="store_true",
@@ -108,6 +113,10 @@ def main(args):
         args.model, adapter_file=args.adapter_file, tokenizer_config=tokenizer_config
     )
 
+    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 (
         hasattr(tokenizer, "apply_chat_template")
         and tokenizer.chat_template is not None