From e926c03b3d028ebad854e753517894fa6c637c3e Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Tue, 6 Aug 2024 19:41:18 -0700
Subject: [PATCH] Add a --tokenizer-dir command-line flag for llamacpp_HF

---
 modules/models.py | 33 ++++++++++++++++++++++-----------
 modules/shared.py |  1 +
 2 files changed, 23 insertions(+), 11 deletions(-)

diff --git a/modules/models.py b/modules/models.py
index ecef9060..b0e2346e 100644
--- a/modules/models.py
+++ b/modules/models.py
@@ -98,7 +98,7 @@ def load_model(model_name, loader=None):
         if model is None:
             return None, None
         else:
-            tokenizer = load_tokenizer(model_name, model)
+            tokenizer = load_tokenizer(model_name)
 
     shared.settings.update({k: v for k, v in metadata.items() if k in shared.settings})
     if loader.lower().startswith('exllama') or loader.lower().startswith('tensorrt'):
@@ -113,9 +113,13 @@ def load_model(model_name, loader=None):
     return model, tokenizer
 
 
-def load_tokenizer(model_name, model):
+def load_tokenizer(model_name, tokenizer_dir=None):
+    if tokenizer_dir:
+        path_to_model = Path(tokenizer_dir)
+    else:
+        path_to_model = Path(f"{shared.args.model_dir}/{model_name}/")
+
     tokenizer = None
-    path_to_model = Path(f"{shared.args.model_dir}/{model_name}/")
     if path_to_model.exists():
         if shared.args.no_use_fast:
             logger.info('Loading the tokenizer with use_fast=False.')
@@ -278,17 +282,24 @@ def llamacpp_loader(model_name):
 def llamacpp_HF_loader(model_name):
     from modules.llamacpp_hf import LlamacppHF
 
-    path = Path(f'{shared.args.model_dir}/{model_name}')
-
-    # Check if a HF tokenizer is available for the model
-    if all((path / file).exists() for file in ['tokenizer_config.json']):
-        logger.info(f'Using tokenizer from: \"{path}\"')
+    if shared.args.tokenizer_dir:
+        logger.info(f'Using tokenizer from: \"{shared.args.tokenizer_dir}\"')
     else:
-        logger.error("Could not load the model because a tokenizer in Transformers format was not found.")
-        return None, None
+        path = Path(f'{shared.args.model_dir}/{model_name}')
+        # Check if a HF tokenizer is available for the model
+        if all((path / file).exists() for file in ['tokenizer_config.json']):
+            logger.info(f'Using tokenizer from: \"{path}\"')
+        else:
+            logger.error("Could not load the model because a tokenizer in Transformers format was not found.")
+            return None, None
 
     model = LlamacppHF.from_pretrained(model_name)
-    return model
+
+    if shared.args.tokenizer_dir:
+        tokenizer = load_tokenizer(model_name, tokenizer_dir=shared.args.tokenizer_dir)
+        return model, tokenizer
+    else:
+        return model
 
 
 def AutoGPTQ_loader(model_name):
diff --git a/modules/shared.py b/modules/shared.py
index c27657ff..43533a14 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -132,6 +132,7 @@ group.add_argument('--cache-capacity', type=str, help='Maximum cache capacity (l
 group.add_argument('--row_split', action='store_true', help='Split the model by rows across GPUs. This may improve multi-gpu performance.')
 group.add_argument('--streaming-llm', action='store_true', help='Activate StreamingLLM to avoid re-evaluating the entire prompt when old messages are removed.')
 group.add_argument('--attention-sink-size', type=int, default=5, help='StreamingLLM: number of sink tokens. Only used if the trimmed prompt does not share a prefix with the old prompt.')
+group.add_argument('--tokenizer-dir', type=str, help='Load the tokenizer from this folder. Meant to be used with llamacpp_HF through the command-line.')
 
 # ExLlamaV2
 group = parser.add_argument_group('ExLlamaV2')