Update convert.py

convert.py

@@ -11,9 +11,7 @@ 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
-from transformers import AutoConfig
-from transformers.pipelines.base import infer_framework_load_model
+from safetensors.torch import load_file, save_file, _remove_duplicate_names
 
 
 COMMIT_DESCRIPTION = """
@@ -40,30 +38,6 @@ 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:
     filename, ext = os.path.splitext(pt_filename)
     local = f"{filename}.safetensors"
@@ -77,29 +51,40 @@ def convert_multi(model_id: str, folder: str) -> ConversionResult:
         data = json.load(f)
 
     filenames = set(data["weight_map"].values())
-    local_filenames = []
-    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)
-        local_filenames.append(sf_filename)
 
+    
     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)
-    local_filenames.append(index)
 
-    operations = [
-        CommitOperationAdd(path_in_repo=local.split("/")[-1], path_or_fileobj=local) for local in local_filenames
-    ]
-    errors: List[Tuple[str, "Exception"]] = []
+    
+    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=repo_id,
+                commit_message=f"Adds {sf_filename}",
+                revision=new_pr.git_reference,
+                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 operations, errors
 
 
 def convert_single(model_id: str, folder: str) -> ConversionResult:
@@ -108,9 +93,15 @@ def convert_single(model_id: str, folder: str) -> ConversionResult:
     sf_name = "model.safetensors"
     sf_filename = os.path.join(folder, sf_name)
     convert_file(pt_filename, sf_filename)
-    operations = [CommitOperationAdd(path_in_repo=sf_name, path_or_fileobj=sf_filename)]
-    errors: List[Tuple[str, "Exception"]] = []
-    return operations, errors
+
+    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(
@@ -120,18 +111,13 @@ def convert_file(
     loaded = torch.load(pt_filename, map_location="cpu")
     if "state_dict" in loaded:
         loaded = loaded["state_dict"]
-    shared = shared_pointers(loaded)
-    for shared_weights in shared:
-        for name in shared_weights[1:]:
-            loaded.pop(name)
-
+    loaded = _remove_duplicate_names(loaded)
     # 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={"format": "pt"})
-    check_file_size(sf_filename, pt_filename)
     reloaded = load_file(sf_filename)
     for k in loaded:
         pt_tensor = loaded[k]
@@ -156,56 +142,6 @@ def create_diff(pt_infos: Dict[str, List[str]], sf_infos: Dict[str, List[str]])
     return "\n".join(errors)
 
 
-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, pt_infos) = infer_framework_load_model(model_id, config, output_loading_info=True)
-    _, (sf_model, sf_infos) = infer_framework_load_model(folder, config, output_loading_info=True)
-
-    if pt_infos != sf_infos:
-        error_string = create_diff(pt_infos, sf_infos)
-        raise ValueError(f"Different infos when reloading the model: {error_string}")
-
-    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:
         main_commit = api.list_repo_commits(model_id)[0].commit_id
@@ -226,6 +162,8 @@ def convert_generic(model_id: str, folder: str, filenames: Set[str]) -> Conversi
     errors = []
 
     extensions = set([".bin", ".ckpt"])
+
+    new_pr = None
     for filename in filenames:
         prefix, ext = os.path.splitext(filename)
         if ext in extensions:
@@ -240,10 +178,28 @@ def convert_generic(model_id: str, folder: str, filenames: Set[str]) -> Conversi
             sf_filename = os.path.join(folder, sf_in_repo)
             try:
                 convert_file(pt_filename, sf_filename)
-                operations.append(CommitOperationAdd(path_in_repo=sf_in_repo, path_or_fileobj=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=repo_id,
+                        commit_message=f"Adds {sf_filename}",
+                        revision=new_pr.git_reference,
+                        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 operations, errors
+    return new_pr
 
 
 def convert(api: "HfApi", model_id: str, force: bool = False) -> Tuple["CommitInfo", List["Exception"]]:
@@ -268,26 +224,15 @@ def convert(api: "HfApi", model_id: str, force: bool = False) -> Tuple["CommitIn
                 raise AlreadyExists(f"Model {model_id} already has an open PR check out {url}")
             elif library_name == "transformers":
                 if "pytorch_model.bin" in filenames:
-                    operations, errors = convert_single(model_id, folder)
+                    new_pr = convert_single(model_id, folder)
                 elif "pytorch_model.bin.index.json" in filenames:
-                    operations, errors = convert_multi(model_id, folder)
+                    new_pr = convert_multi(model_id, folder)
                 else:
                     raise RuntimeError(f"Model {model_id} doesn't seem to be a valid pytorch model. Cannot convert")
-                check_final_model(model_id, folder)
-            else:
-                operations, errors = convert_generic(model_id, folder, filenames)
-
-            if operations:
-                new_pr = api.create_commit(
-                    repo_id=model_id,
-                    operations=operations,
-                    commit_message=pr_title,
-                    commit_description=COMMIT_DESCRIPTION,
-                    create_pr=True,
-                )
-                print(f"Pr created at {new_pr.pr_url}")
             else:
-                print("No files to convert")
+                new_pr = convert_generic(model_id, folder, filenames)
+
+            print(f"Pr created at {new_pr.pr_url}")
         finally:
             shutil.rmtree(folder)
         return new_pr, errors