diff --git a/src/download_model.py b/src/download_model.py index 086a2a0..f0f2d79 100644 --- a/src/download_model.py +++ b/src/download_model.py @@ -1,5 +1,6 @@ import os import logging +from transformers import AutoTokenizer from vllm.model_executor.weight_utils import prepare_hf_model_weights if __name__ == "__main__": @@ -17,6 +18,12 @@ if __name__ == "__main__": model_name_or_path=model, cache_dir=download_dir, ) + + tokenizer = os.getenv("TOKENIZER_NAME", model) + AutoTokenizer.from_pretrained( + pretrained_model_name_or_path=tokenizer, + cache_dir=download_dir, + ) logging.info(f"Finished downloading model {model} to {download_dir}")