Update convert.py

convert.py

@@ -5,7 +5,7 @@ import shutil
 from collections import defaultdict
 from inspect import signature
 from tempfile import TemporaryDirectory
-from typing import Dict, List, Optional, Set
+from typing import Dict, List, Optional, Set, Tuple
 
 import torch
 
@@ -33,6 +33,8 @@ If you find any issues: please report here: https://huggingface.co/spaces/safete
 Feel free to ignore this PR.
 """
 
+ConversionResult = Tuple[List["CommitOperationAdd"], List[Tuple[str, "Exception"]]]
+
 
 class AlreadyExists(Exception):
     pass
@@ -69,7 +71,7 @@ def rename(pt_filename: str) -> str:
     return local
 
 
-def convert_multi(model_id: str, folder: str) -> List["CommitOperationAdd"]:
+def convert_multi(model_id: str, folder: str) -> 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)
@@ -95,18 +97,20 @@ def convert_multi(model_id: str, folder: str) -> List["CommitOperationAdd"]:
     operations = [
         CommitOperationAdd(path_in_repo=local.split("/")[-1], path_or_fileobj=local) for local in local_filenames
     ]
+    errors: List[Tuple[str, "Exception"]] = []
 
-    return operations
+    return operations, errors
 
 
-def convert_single(model_id: str, folder: str) -> List["CommitOperationAdd"]:
+def convert_single(model_id: str, folder: str) -> 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)
     operations = [CommitOperationAdd(path_in_repo=sf_name, path_or_fileobj=sf_filename)]
-    return operations
+    errors: List[Tuple[str, "Exception"]] = []
+    return operations, errors
 
 
 def convert_file(
@@ -204,18 +208,22 @@ def check_final_model(model_id: str, folder: str):
 
 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:
-            details = api.get_discussion_details(repo_id=model_id, discussion_num=discussion.num)
-            if details.target_branch == "refs/heads/main":
+            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]) -> List["CommitOperationAdd"]:
+def convert_generic(model_id: str, folder: str, filenames: Set[str]) -> ConversionResult:
     operations = []
+    errors = []
 
     extensions = set([".bin", ".ckpt"])
     for filename in filenames:
@@ -230,12 +238,15 @@ def convert_generic(model_id: str, folder: str, filenames: Set[str]) -> List["Co
             else:
                 sf_in_repo = f"{prefix}.safetensors"
             sf_filename = os.path.join(folder, sf_in_repo)
-            convert_file(pt_filename, sf_filename)
-            operations.append(CommitOperationAdd(path_in_repo=sf_in_repo, path_or_fileobj=sf_filename))
-    return operations
+            try:
+                convert_file(pt_filename, sf_filename)
+                operations.append(CommitOperationAdd(path_in_repo=sf_in_repo, path_or_fileobj=sf_filename))
+            except Exception as e:
+                errors.append((pt_filename, e))
+    return operations, errors
 
 
-def convert(api: "HfApi", model_id: str, force: bool = False) -> Optional["CommitInfo"]:
+def convert(api: "HfApi", model_id: str, force: bool = False) -> Tuple["CommitInfo", List["Exception"]]:
     pr_title = "Adding `safetensors` variant of this model"
     info = api.model_info(model_id)
     filenames = set(s.rfilename for s in info.siblings)
@@ -257,14 +268,14 @@ def convert(api: "HfApi", model_id: str, force: bool = False) -> Optional["Commi
                 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 = convert_single(model_id, folder)
+                    operations, errors = convert_single(model_id, folder)
                 elif "pytorch_model.bin.index.json" in filenames:
-                    operations = convert_multi(model_id, folder)
+                    operations, errors = 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 = convert_generic(model_id, folder, filenames)
+                operations, errors = convert_generic(model_id, folder, filenames)
 
             if operations:
                 new_pr = api.create_commit(
@@ -279,7 +290,7 @@ def convert(api: "HfApi", model_id: str, force: bool = False) -> Optional["Commi
                 print("No files to convert")
         finally:
             shutil.rmtree(folder)
-        return new_pr
+        return new_pr, errors
 
 
 if __name__ == "__main__":
@@ -300,7 +311,23 @@ if __name__ == "__main__":
         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()
-    convert(api, model_id, force=args.force)
+    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.")