mirror of
https://github.com/lilakk/BLEUBERI.git
synced 2026-04-19 12:58:12 +00:00
45 lines
1.5 KiB
Python
45 lines
1.5 KiB
Python
"""
|
|
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)
|