Spaces:
Runtime error
Runtime error
File size: 1,521 Bytes
ec0c335 |
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 |
"""
Upload weights to huggingface.
Usage:
python3 -m fastchat.model.upload_hub --model-path ~/model_weights/vicuna-13b --hub-repo-id lmsys/vicuna-13b-v1.3
"""
import argparse
import tempfile
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
def upload_hub(model_path, hub_repo_id, component, private):
if component == "all":
components = ["model", "tokenizer"]
else:
components = [component]
kwargs = {"push_to_hub": True, "repo_id": hub_repo_id, "private": args.private}
if "model" in components:
model = AutoModelForCausalLM.from_pretrained(
model_path, torch_dtype=torch.float16, low_cpu_mem_usage=True
)
with tempfile.TemporaryDirectory() as tmp_path:
model.save_pretrained(tmp_path, **kwargs)
if "tokenizer" in components:
tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
with tempfile.TemporaryDirectory() as tmp_path:
tokenizer.save_pretrained(tmp_path, **kwargs)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--model-path", type=str, required=True)
parser.add_argument("--hub-repo-id", type=str, required=True)
parser.add_argument(
"--component", type=str, choices=["all", "model", "tokenizer"], default="all"
)
parser.add_argument("--private", action="store_true")
args = parser.parse_args()
upload_hub(args.model_path, args.hub_repo_id, args.component, args.private)
|