Spaces:
Runtime error
Runtime error
import argparse | |
import json | |
import os | |
import shutil | |
from collections import defaultdict | |
from inspect import signature | |
from tempfile import TemporaryDirectory | |
from typing import Dict, List, Optional, Set, Tuple | |
import torch | |
from huggingface_hub import CommitInfo, CommitOperationAdd, Discussion, HfApi, hf_hub_download | |
from huggingface_hub.file_download import repo_folder_name | |
from safetensors.torch import load_file, save_file, _remove_duplicate_names | |
COMMIT_DESCRIPTION = """ | |
This is an automated PR created with https://huggingface.co/spaces/safetensors/convert | |
This new file is equivalent to `pytorch_model.bin` but safe in the sense that | |
no arbitrary code can be put into it. | |
These files also happen to load much faster than their pytorch counterpart: | |
https://colab.research.google.com/github/huggingface/notebooks/blob/main/safetensors_doc/en/speed.ipynb | |
The widgets on your model page will run using this model even if this is not merged | |
making sure the file actually works. | |
If you find any issues: please report here: https://huggingface.co/spaces/safetensors/convert/discussions | |
Feel free to ignore this PR. | |
""" | |
PR_TITLE = "Adding `safetensors` variant of this model" | |
ConversionResult = Tuple[List["CommitOperationAdd"], List[Tuple[str, "Exception"]]] | |
class AlreadyExists(Exception): | |
pass | |
def rename(pt_filename: str) -> str: | |
filename, ext = os.path.splitext(pt_filename) | |
local = f"{filename}.safetensors" | |
local = local.replace("pytorch_model", "model") | |
return local | |
def convert_multi(model_id: str, folder: str, api: "HfApi") -> ConversionResult: | |
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()) | |
index = os.path.join(folder, "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, indent=4) | |
new_pr = api.create_commit( | |
repo_id=model_id, | |
operations=[CommitOperationAdd(path_in_repo=index.split("/")[-1], path_or_fileobj=index)], | |
commit_message=PR_TITLE, | |
commit_description=COMMIT_DESCRIPTION, | |
create_pr=True, | |
) | |
for filename in filenames: | |
pt_filename = hf_hub_download(repo_id=model_id, filename=filename) | |
sf_filename = rename(pt_filename) | |
sf_filename = os.path.join(folder, sf_filename) | |
convert_file(pt_filename, sf_filename) | |
api.create_commit( | |
repo_id=model_id, | |
commit_message=f"Adds {sf_filename}", | |
revision=new_pr.pr_revision, | |
operations=[CommitOperationAdd(path_in_repo=sf_filename.split("/")[-1], path_or_fileobj=sf_filename)], | |
create_pr=False, | |
) | |
os.remove(pt_filename) | |
os.remove(sf_filename) | |
return new_pr, [] | |
def convert_single(model_id: str, folder: str, api: "HfApi") -> ConversionResult: | |
pt_filename = hf_hub_download(repo_id=model_id, filename="pytorch_model.bin") | |
sf_name = "model.safetensors" | |
sf_filename = os.path.join(folder, sf_name) | |
convert_file(pt_filename, sf_filename) | |
new_pr = api.create_commit( | |
repo_id=model_id, | |
operations=[CommitOperationAdd(path_in_repo=sf_name, path_or_fileobj=sf_filename)], | |
commit_message=PR_TITLE, | |
commit_description=COMMIT_DESCRIPTION, | |
create_pr=True, | |
) | |
return new_pr, [] | |
def convert_file( | |
pt_filename: str, | |
sf_filename: str, | |
): | |
loaded = torch.load(pt_filename, map_location="cpu") | |
if "state_dict" in loaded: | |
loaded = loaded["state_dict"] | |
to_removes = _remove_duplicate_names(loaded) | |
metadata = {"format": "pt"} | |
for kept_name, to_remove_group in to_removes.items(): | |
for to_remove in to_remove_group: | |
if to_remove not in metadata: | |
metadata[to_remove] = kept_name | |
del loaded[to_remove] | |
# For tensors to be contiguous | |
loaded = {k: v.contiguous() for k, v in loaded.items()} | |
dirname = os.path.dirname(sf_filename) | |
os.makedirs(dirname, exist_ok=True) | |
save_file(loaded, sf_filename, metadata=metadata) | |
reloaded = load_file(sf_filename) | |
for k in loaded: | |
pt_tensor = loaded[k] | |
sf_tensor = reloaded[k] | |
if not torch.equal(pt_tensor, sf_tensor): | |
raise RuntimeError(f"The output tensors do not match for key {k}") | |
def create_diff(pt_infos: Dict[str, List[str]], sf_infos: Dict[str, List[str]]) -> str: | |
errors = [] | |
for key in ["missing_keys", "mismatched_keys", "unexpected_keys"]: | |
pt_set = set(pt_infos[key]) | |
sf_set = set(sf_infos[key]) | |
pt_only = pt_set - sf_set | |
sf_only = sf_set - pt_set | |
if pt_only: | |
errors.append(f"{key} : PT warnings contain {pt_only} which are not present in SF warnings") | |
if sf_only: | |
errors.append(f"{key} : SF warnings contain {sf_only} which are not present in PT warnings") | |
return "\n".join(errors) | |
def previous_pr(api: "HfApi", model_id: str, pr_title: str) -> Optional["Discussion"]: | |
try: | |
main_commit = api.list_repo_commits(model_id)[0].commit_id | |
discussions = api.get_repo_discussions(repo_id=model_id) | |
except Exception: | |
return None | |
for discussion in discussions: | |
if discussion.status == "open" and discussion.is_pull_request and discussion.title == pr_title: | |
commits = api.list_repo_commits(model_id, revision=discussion.git_reference) | |
if main_commit == commits[1].commit_id: | |
return discussion | |
return None | |
def convert_generic(model_id: str, folder: str, filenames: Set[str], api: "HfApi") -> ConversionResult: | |
operations = [] | |
errors = [] | |
extensions = set([".bin", ".ckpt"]) | |
new_pr = None | |
for filename in filenames: | |
prefix, ext = os.path.splitext(filename) | |
if ext in extensions: | |
pt_filename = hf_hub_download(model_id, filename=filename) | |
dirname, raw_filename = os.path.split(filename) | |
if raw_filename == "pytorch_model.bin": | |
# XXX: This is a special case to handle `transformers` and the | |
# `transformers` part of the model which is actually loaded by `transformers`. | |
sf_in_repo = os.path.join(dirname, "model.safetensors") | |
else: | |
sf_in_repo = f"{prefix}.safetensors" | |
sf_filename = os.path.join(folder, sf_in_repo) | |
try: | |
convert_file(pt_filename, sf_filename) | |
if new_pr is None: | |
new_pr = api.create_commit( | |
repo_id=model_id, | |
operations=[CommitOperationAdd(path_in_repo=sf_in_repo, path_or_fileobj=sf_filename)], | |
commit_message=PR_TITLE, | |
commit_description=COMMIT_DESCRIPTION, | |
create_pr=True, | |
) | |
else: | |
api.create_commit( | |
repo_id=model_id, | |
commit_message=f"Adds {sf_filename}", | |
revision=new_pr.pr_revision, | |
operations=[CommitOperationAdd(path_in_repo=sf_in_repo, path_or_fileobj=sf_filename)], | |
create_pr=False, | |
) | |
os.remove(pt_filename) | |
os.remove(sf_filename) | |
except Exception as e: | |
errors.append((pt_filename, e)) | |
return new_pr, errors | |
def convert(api: "HfApi", model_id: str, force: bool = False) -> Tuple["CommitInfo", List["Exception"]]: | |
info = api.model_info(model_id) | |
filenames = set(s.rfilename for s in info.siblings) | |
with TemporaryDirectory() as d: | |
folder = os.path.join(d, repo_folder_name(repo_id=model_id, repo_type="models")) | |
os.makedirs(folder) | |
new_pr = None | |
try: | |
operations = None | |
pr = previous_pr(api, model_id, PR_TITLE) | |
library_name = getattr(info, "library_name", None) | |
if any(filename.endswith(".safetensors") for filename in filenames) and not force: | |
raise AlreadyExists(f"Model {model_id} is already converted, skipping..") | |
elif pr is not None and not force: | |
url = f"https://huggingface.co/{model_id}/discussions/{pr.num}" | |
new_pr = pr | |
raise AlreadyExists(f"Model {model_id} already has an open PR check out {url}") | |
elif library_name == "transformers": | |
if "pytorch_model.bin" in filenames: | |
new_pr, errors = convert_single(model_id, folder, api) | |
elif "pytorch_model.bin.index.json" in filenames: | |
new_pr, errors = convert_multi(model_id, folder, api) | |
else: | |
raise RuntimeError(f"Model {model_id} doesn't seem to be a valid pytorch model. Cannot convert") | |
else: | |
new_pr, errors = convert_generic(model_id, folder, filenames, api) | |
print(f"Pr created at {new_pr.pr_url}") | |
finally: | |
shutil.rmtree(folder) | |
return new_pr, errors | |
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`", | |
) | |
parser.add_argument( | |
"--force", | |
action="store_true", | |
help="Create the PR even if it already exists of if the model was already converted.", | |
) | |
parser.add_argument( | |
"-y", | |
action="store_true", | |
help="Ignore safety prompt", | |
) | |
args = parser.parse_args() | |
model_id = args.model_id | |
api = HfApi() | |
if args.y: | |
txt = "y" | |
else: | |
txt = input( | |
"This conversion script will unpickle a pickled file, which is inherently unsafe. If you do not trust this file, we invite you to use" | |
" https://huggingface.co/spaces/safetensors/convert or google colab or other hosted solution to avoid potential issues with this file." | |
" Continue [Y/n] ?" | |
) | |
if txt.lower() in {"", "y"}: | |
_commit_info, _errors = convert(api, model_id, force=args.force) | |
else: | |
print(f"Answer was `{txt}` aborting.") | |