Narsil HF staff commited on
Commit
dede4f0
1 Parent(s): 3e03394

Make the modifications to harden everything

Browse files
Files changed (1) hide show
  1. convert.py +93 -71
convert.py CHANGED
@@ -1,91 +1,118 @@
1
  import argparse
2
  import json
3
  import os
 
4
 
5
  import torch
6
 
7
- from huggingface_hub import CommitOperationAdd, HfApi, hf_hub_download, CommitInfo
 
 
 
8
  from safetensors.torch import save_file
9
 
10
 
 
 
 
 
 
 
 
 
 
 
 
11
  def rename(pt_filename) -> str:
12
  local = pt_filename.replace(".bin", ".safetensors")
13
  local = local.replace("pytorch_model", "model")
14
  return local
15
 
16
 
17
- def convert_multi(api: HfApi, model_id) -> CommitInfo:
18
- local_filenames = []
19
- try:
20
- filename = hf_hub_download(
21
- repo_id=model_id, filename="pytorch_model.bin.index.json"
22
- )
23
- with open(filename, "r") as f:
24
- data = json.load(f)
25
-
26
- filenames = set(data["weight_map"].values())
27
- for filename in filenames:
28
- cached_filename = hf_hub_download(repo_id=model_id, filename=filename)
29
- loaded = torch.load(cached_filename)
30
- local = rename(filename)
31
- save_file(loaded, local, metadata={"format": "pt"})
32
- local_filenames.append(local)
33
-
34
- index = "model.safetensors.index.json"
35
- with open(index, "w") as f:
36
- newdata = {k: v for k, v in data.items()}
37
- newmap = {k: rename(v) for k, v in data["weight_map"].items()}
38
- newdata["weight_map"] = newmap
39
- json.dump(newdata, f)
40
- local_filenames.append(index)
41
-
42
- operations = [
43
- CommitOperationAdd(path_in_repo=local, path_or_fileobj=local)
44
- for local in local_filenames
45
- ]
46
- return api.create_commit(
47
- repo_id=model_id,
48
- operations=operations,
49
- commit_message="Adding `safetensors` variant of this model",
50
- commit_description="Converted from this Space: https://huggingface.co/spaces/safetensors/convert",
51
- create_pr=True,
52
- )
53
- finally:
54
- for local in local_filenames:
55
- os.remove(local)
56
 
 
 
 
 
 
57
 
58
- def convert_single(api: HfApi, model_id) -> CommitInfo:
59
- local = "model.safetensors"
60
- try:
61
- filename = hf_hub_download(repo_id=model_id, filename="pytorch_model.bin")
62
- loaded = torch.load(filename)
63
  save_file(loaded, local, metadata={"format": "pt"})
 
 
64
 
65
- operations = [CommitOperationAdd(path_in_repo=local, path_or_fileobj=local)]
66
- return api.create_commit(
67
- repo_id=model_id,
68
- operations=operations,
69
- commit_message="Adding `safetensors` variant of this model",
70
- commit_description="Converted from this Space: https://huggingface.co/spaces/safetensors/convert",
71
- create_pr=True,
72
- )
73
- finally:
74
- os.remove(local)
75
 
 
76
 
77
- def convert(token: str, model_id: str) -> CommitInfo:
78
- """
79
- returns url to the PR
80
- """
81
- api = HfApi(token=token)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
82
  info = api.model_info(model_id)
83
  filenames = set(s.rfilename for s in info.siblings)
84
- if "pytorch_model.bin" in filenames:
85
- return convert_single(api, model_id)
86
- elif "pytorch_model.bin.index.json" in filenames:
87
- return convert_multi(api, model_id)
88
- raise ValueError("repo does not seem to have a pytorch_model in it")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
89
 
90
 
91
  if __name__ == "__main__":
@@ -104,9 +131,4 @@ if __name__ == "__main__":
104
  args = parser.parse_args()
105
  model_id = args.model_id
106
  api = HfApi()
107
- info = api.model_info(model_id)
108
- filenames = set(s.rfilename for s in info.siblings)
109
- if "pytorch_model.bin" in filenames:
110
- convert_single(model_id)
111
- else:
112
- convert_multi(model_id)
 
1
  import argparse
2
  import json
3
  import os
4
+ import shutil
5
 
6
  import torch
7
 
8
+ from huggingface_hub import CommitOperationAdd, HfApi, hf_hub_download
9
+ from huggingface_hub.file_download import repo_folder_name
10
+ from transformers import AutoConfig
11
+ from transformers.pipelines.base import infer_framework_load_model
12
  from safetensors.torch import save_file
13
 
14
 
15
+ def check_file_size(sf_filename, pt_filename):
16
+ sf_size = os.stat(sf_filename).st_size
17
+ pt_size = os.stat(pt_filename).st_size
18
+
19
+ if (sf_size - pt_size) / pt_size > 0.01:
20
+ raise RuntimeError(f"""The file size different is more than 1%:
21
+ - {sf_filename}: {sf_size}
22
+ - {pt_filename}: {pt_size}
23
+ """)
24
+
25
+
26
  def rename(pt_filename) -> str:
27
  local = pt_filename.replace(".bin", ".safetensors")
28
  local = local.replace("pytorch_model", "model")
29
  return local
30
 
31
 
32
+ def convert_multi(model_id):
33
+ filename = hf_hub_download(repo_id=model_id, filename="pytorch_model.bin.index.json")
34
+ with open(filename, "r") as f:
35
+ data = json.load(f)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
36
 
37
+ filenames = set(data["weight_map"].values())
38
+ for filename in filenames:
39
+ cached_filename = hf_hub_download(repo_id=model_id, filename=filename)
40
+ loaded = torch.load(cached_filename)
41
+ sf_filename = rename(filename)
42
 
43
+ local = os.path.join(folder, sf_filename)
 
 
 
 
44
  save_file(loaded, local, metadata={"format": "pt"})
45
+ check_file_size(local, cached_filename)
46
+ local_filenames.append(local)
47
 
48
+ index = os.path.join(folder, "model.safetensors.index.json")
49
+ with open(index, "w") as f:
50
+ newdata = {k: v for k, v in data.items()}
51
+ newmap = {k: rename(v) for k, v in data["weight_map"].items()}
52
+ newdata["weight_map"] = newmap
53
+ json.dump(newdata, f)
54
+ local_filenames.append(index)
 
 
 
55
 
56
+ operations = [CommitOperationAdd(path_in_repo=local.split("/")[-1], path_or_fileobj=local) for local in local_filenames]
57
 
58
+ return operations
59
+
60
+
61
+ def convert_single(model_id, folder):
62
+ sf_filename = "model.safetensors"
63
+ filename = hf_hub_download(repo_id=model_id, filename="pytorch_model.bin")
64
+ loaded = torch.load(filename)
65
+
66
+ local = os.path.join(folder, sf_filename)
67
+ save_file(loaded, local, metadata={"format": "pt"})
68
+
69
+ check_file_size(local, filename)
70
+
71
+ operations = [CommitOperationAdd(path_in_repo=sf_filename, path_or_fileobj=local)]
72
+ return operations
73
+
74
+ def check_final_model(model_id, folder):
75
+ config = hf_hub_download(repo_id=model_id, filename="config.json")
76
+ shutil.copy(config, os.path.join(folder, "config.json"))
77
+ config = AutoConfig.from_pretrained(folder)
78
+ _, sf_model = infer_framework_load_model(folder, config)
79
+ _, pt_model = infer_framework_load_model(model_id, config)
80
+
81
+ input_ids = torch.arange(10).long().unsqueeze(0)
82
+ sf_logits = sf_model(input_ids)
83
+ pt_logits = pt_model(input_ids)
84
+ torch.testing.assert_close(sf_logits, pt_logits)
85
+ print(f"Model {model_id} is ok !")
86
+
87
+
88
+ def convert(api, model_id):
89
  info = api.model_info(model_id)
90
  filenames = set(s.rfilename for s in info.siblings)
91
+
92
+ folder = repo_folder_name(repo_id=model_id, repo_type="models")
93
+ os.makedirs(folder)
94
+ try:
95
+ operations = None
96
+ if "model.safetensors" in filenames or "model_index.safetensors.index.json" in filenames:
97
+ print(f"Model {model_id} is already converted, skipping..")
98
+ elif "pytorch_model.bin" in filenames:
99
+ operations = convert_single(model_id, folder)
100
+ elif "pytorch_model.bin.index.json" in filenames:
101
+ operations = convert_multi(model_id, folder)
102
+ else:
103
+ raise RuntimeError(f"Model {model_id} doesn't seem to be a valid pytorch model. Cannot convert")
104
+
105
+ if operations:
106
+ check_final_model(model_id, folder)
107
+ api.create_commit(
108
+ repo_id=model_id,
109
+ operations=operations,
110
+ commit_message="Adding `safetensors` variant of this model",
111
+ create_pr=True,
112
+ )
113
+ finally:
114
+ shutil.rmtree(folder)
115
+ return 1
116
 
117
 
118
  if __name__ == "__main__":
 
131
  args = parser.parse_args()
132
  model_id = args.model_id
133
  api = HfApi()
134
+ convert(api, model_id)