File size: 3,520 Bytes
66662af |
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 109 110 111 112 113 |
import argparse
import json
import os
import torch
from huggingface_hub import CommitOperationAdd, HfApi, hf_hub_download
from safetensors.torch import save_file
def rename(pt_filename) -> str:
local = pt_filename.replace(".bin", ".safetensors")
local = local.replace("pytorch_model", "model")
return local
def convert_multi(model_id) -> str:
local_filenames = []
try:
filename = hf_hub_download(
repo_id=model_id, filename="pytorch_model.bin.index.json"
)
with open(filename, "r") as f:
data = json.load(f)
filenames = set(data["weight_map"].values())
for filename in filenames:
cached_filename = hf_hub_download(repo_id=model_id, filename=filename)
loaded = torch.load(cached_filename)
local = rename(filename)
save_file(loaded, local, metadata={"format": "pt"})
local_filenames.append(local)
index = "model.safetensors.index.json"
with open(index, "w") as f:
newdata = {k: v for k, v in data.items()}
newmap = {k: rename(v) for k, v in data["weight_map"].items()}
newdata["weight_map"] = newmap
json.dump(newdata, f)
local_filenames.append(index)
api = HfApi()
operations = [
CommitOperationAdd(path_in_repo=local, path_or_fileobj=local)
for local in local_filenames
]
return api.create_commit(
repo_id=model_id,
operations=operations,
commit_message="Adding `safetensors` variant of this model",
create_pr=True,
)
finally:
for local in local_filenames:
os.remove(local)
def convert_single(model_id) -> str:
local = "model.safetensors"
try:
filename = hf_hub_download(repo_id=model_id, filename="pytorch_model.bin")
loaded = torch.load(filename)
save_file(loaded, local, metadata={"format": "pt"})
api = HfApi()
return api.upload_file(
path_or_fileobj=local,
create_pr=True,
path_in_repo=local,
repo_id=model_id,
)
finally:
os.remove(local)
def convert(token: str, model_id: str) -> str:
"""
returns url to the PR
"""
api = HfApi(token=token)
info = api.model_info(model_id)
filenames = set(s.rfilename for s in info.siblings)
if "pytorch_model.bin" in filenames:
return convert_single(model_id)
elif "pytorch_model.bin.index.json" in filenames:
return convert_multi(model_id)
raise ValueError("repo does not seem to have a pytorch_model in it")
if __name__ == "__main__":
DESCRIPTION = """
Simple utility tool to convert automatically some weights on the hub to `safetensors` format.
It is PyTorch exclusive for now.
It works by downloading the weights (PT), converting them locally, and uploading them back
as a PR on the hub.
"""
parser = argparse.ArgumentParser(description=DESCRIPTION)
parser.add_argument(
"model_id",
type=str,
help="The name of the model on the hub to convert. E.g. `gpt2` or `facebook/wav2vec2-base-960h`",
)
args = parser.parse_args()
model_id = args.model_id
api = HfApi()
info = api.model_info(model_id)
filenames = set(s.rfilename for s in info.siblings)
if "pytorch_model.bin" in filenames:
convert_single(model_id)
else:
convert_multi(model_id)
|