import argparse import json import os import shutil from tempfile import TemporaryDirectory from collections import defaultdict from inspect import signature from typing import Optional, List import torch from huggingface_hub import CommitOperationAdd, HfApi, hf_hub_download, get_repo_discussions from huggingface_hub.file_download import repo_folder_name from transformers import AutoConfig from transformers.pipelines.base import infer_framework_load_model from safetensors.torch import save_file class AlreadyExists(Exception): pass def shared_pointers(tensors): ptrs = defaultdict(list) for k, v in tensors.items(): ptrs[v.data_ptr()].append(k) failing = [] for ptr, names in ptrs.items(): if len(names) > 1: failing.append(names) return failing def check_file_size(sf_filename: str, pt_filename: str): sf_size = os.stat(sf_filename).st_size pt_size = os.stat(pt_filename).st_size if (sf_size - pt_size) / pt_size > 0.01: raise RuntimeError(f"""The file size different is more than 1%: - {sf_filename}: {sf_size} - {pt_filename}: {pt_size} """) def rename(pt_filename: str) -> str: local = pt_filename.replace(".bin", ".safetensors") local = local.replace("pytorch_model", "model") return local def convert_multi(model_id: str) -> List["CommitOperationAdd"]: 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) sf_filename = rename(filename) local = os.path.join(folder, sf_filename) save_file(loaded, local, metadata={"format": "pt"}) check_file_size(local, cached_filename) local_filenames.append(local) 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) local_filenames.append(index) operations = [CommitOperationAdd(path_in_repo=local.split("/")[-1], path_or_fileobj=local) for local in local_filenames] return operations def convert_single(model_id: str, folder: str) -> List["CommitOperationAdd"]: sf_filename = "model.safetensors" filename = hf_hub_download(repo_id=model_id, filename="pytorch_model.bin") loaded = torch.load(filename) local = os.path.join(folder, sf_filename) shared = shared_pointers(loaded) for shared_weights in shared: for name in shared_weights[1:]: loaded.pop(name) # For tensors to be contiguous loaded = {k: v.contiguous() for k, v in loaded.items()} save_file(loaded, local, metadata={"format": "pt"}) check_file_size(local, filename) operations = [CommitOperationAdd(path_in_repo=sf_filename, path_or_fileobj=local)] return operations def check_final_model(model_id: str, folder: str): config = hf_hub_download(repo_id=model_id, filename="config.json") shutil.copy(config, os.path.join(folder, "config.json")) config = AutoConfig.from_pretrained(folder) _, pt_model = infer_framework_load_model(model_id, config) _, sf_model = infer_framework_load_model(folder, config) pt_model = pt_model sf_model = sf_model pt_params = pt_model.state_dict() sf_params = sf_model.state_dict() pt_shared = shared_pointers(pt_params) sf_shared = shared_pointers(sf_params) if pt_shared != sf_shared: raise RuntimeError("The reconstructed model is wrong, shared tensors are different {shared_pt} != {shared_tf}") sig = signature(pt_model.forward) input_ids = torch.arange(10).unsqueeze(0) pixel_values = torch.randn(1, 3, 224, 224) input_values = torch.arange(1000).float().unsqueeze(0) kwargs = {} if "input_ids" in sig.parameters: kwargs["input_ids"] = input_ids if "decoder_input_ids" in sig.parameters: kwargs["decoder_input_ids"] = input_ids if "pixel_values" in sig.parameters: kwargs["pixel_values"] = pixel_values if "input_values" in sig.parameters: kwargs["input_values"] = input_values if "bbox" in sig.parameters: kwargs["bbox"] = torch.zeros((1, 10, 4)).long() if "image" in sig.parameters: kwargs["image"] = pixel_values if torch.cuda.is_available(): pt_model = pt_model.cuda() sf_model = sf_model.cuda() kwargs = {k: v.cuda() for k, v in kwargs.items()} pt_logits = pt_model(**kwargs)[0] sf_logits = sf_model(**kwargs)[0] torch.testing.assert_close(sf_logits, pt_logits) print(f"Model {model_id} is ok !") def previous_pr(api: "HfApi", model_id: str, pr_title: str) -> Optional["Discussion"]: try: 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: return discussion def convert(api: "HfApi", model_id: str, force: bool=False) -> Optional["CommitInfo"]: pr_title = "Adding `safetensors` variant of this model" 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) if ("model.safetensors" in filenames or "model_index.safetensors.index.json" 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 "pytorch_model.bin" in filenames: operations = convert_single(model_id, folder) elif "pytorch_model.bin.index.json" in filenames: operations = convert_multi(model_id, folder) else: raise RuntimeError(f"Model {model_id} doesn't seem to be a valid pytorch model. Cannot convert") if operations: check_final_model(model_id, folder) # new_pr = api.create_commit( # repo_id=model_id, # operations=operations, # commit_message=pr_title, # create_pr=True, # ) finally: shutil.rmtree(folder) return new_pr 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.", ) args = parser.parse_args() model_id = args.model_id api = HfApi() convert(api, model_id, force=args.force)