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.")