File size: 3,758 Bytes
88a170b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 |
#!/usr/bin/env python
import os
import shutil
import argparse
import logging
import sys
import torch
from distutils.dir_util import copy_tree
from pathlib import Path
from tempfile import TemporaryDirectory
from huggingface_hub import snapshot_download, login
from tensorizer import TensorSerializer
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
from tensorize_model import tensorize_model
logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO, stream=sys.stdout)
def download_model_from_hf_hub(
model_name: str,
model_path: str,
rm_existing_model: bool = True,
) -> dict:
"""
This function downloads a model from the Hugging Face Hub and saves it locally.
It also saves the tokenizer in a separate location so that it can be easely included in a docker Image
without including the model weights.
Args:
model_name (str): Name of model on hugging face hub
path (str): Local path where model is saved
rm_existing_model (bool, optional): Whether to remove the existing model or not. Defaults to False.
Returns:
dict: Dictionary containing the model name and path
"""
# model_weights_path = os.path.join(os.getcwd(), "model_weights/torch_weights")
# model_path = os.path.join(model_weights_path, model_name)
if rm_existing_model:
logger.info(f"Removing existing model at {model_path}")
if os.path.exists(model_path):
shutil.rmtree(model_path)
# setup temporary directory
with TemporaryDirectory() as tmpdir:
logger.info(f"Downloading {model_name} weights to temp...")
snapshot_dir = snapshot_download(
repo_id=model_name,
cache_dir=tmpdir,
allow_patterns=["*.bin", "*.json", "*.md", "*.model", "*.py"],
)
# copy snapshot to model dir
logger.info(f"Copying weights to {model_path}...")
copy_tree(snapshot_dir, str(model_path))
return {"model_name": model_name, "model_path": model_path}
def download_hf_model_and_copy_tokenizer(
model_name: str,
model_path: str,
tokenizer_path: str,
rm_existing_model: bool = True,
):
model_info = download_model_from_hf_hub(model_name, model_path)
if tokenizer_path:
# Move tokenizer to separate location
logging.info(f"Copying tokenizer and model config to {tokenizer_path}...")
tokenizer = AutoTokenizer.from_pretrained(model_path, padding_side="left")
tokenizer.save_pretrained(tokenizer_path)
# Set the source and destination file paths
config_path = os.path.join(model_path, "config.json")
# Use the shutil.copy() function to copy the file to the destination directory
shutil.copy(config_path, tokenizer_path)
return model_info
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--model_name", type=str)
parser.add_argument("--model_path", type=str)
parser.add_argument("--tokenizer_path", type=str, default=None)
parser.add_argument("--hf_token", type=str, default=None)
parser.add_argument("--tensorize", action="store_true", default=False)
parser.add_argument("--dtype", type=str, default="fp32")
args = parser.parse_args()
if args.hf_token is not None:
login(token=args.hf_token)
# download_hf_model_and_copy_tokenizer(args.model_name, model_path=args.model_path, tokenizer_path=args.tokenizer_path)
tensorizer_path = os.path.join(args.model_path, "model.tensors")
if args.tensorize:
model = tensorize_model(args.model_name, model_path=args.model_path, dtype=args.dtype, tensorizer_path=tensorizer_path)
|