ybelkada commited on
Commit
5adcd02
1 Parent(s): 5d5344b
Files changed (2) hide show
  1. app.py +66 -0
  2. requirements.txt +3 -0
app.py ADDED
@@ -0,0 +1,66 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+
3
+ import torch
4
+ import safetensors
5
+ # hack to load safetensors.torch
6
+ from safetensors.torch import save_file
7
+ from huggingface_hub import hf_hub_download
8
+
9
+ def run(pr_number, model_id):
10
+ try:
11
+ st_weights_path = hf_hub_download(repo_id=model_id, filename="model.safetensors", revision=f"refs/pr/{pr_number}")
12
+ torch_weights_path = hf_hub_download(repo_id=model_id, filename="pytorch_model.bin")
13
+ except Exception as e:
14
+ return f"Error: {e} | \n Maybe you specified model ids or PRs that does not exist or does not contain any `model.safetensors` or `pytorch_model.bin` files"
15
+
16
+
17
+ st_weights = safetensors.torch.load_file(st_weights_path)
18
+ torch_weights = torch.load(torch_weights_path)
19
+
20
+ # check if keys are the same
21
+ if st_weights.keys() != torch_weights.keys():
22
+ # retrieve different keys
23
+ unexpected_keys = st_weights.keys() - torch_weights.keys()
24
+ return f"keys are not the same ! Conversion failed - unexpected keys are: {unexpected_keys}"
25
+
26
+ total_errors = []
27
+
28
+ # check all weights are same
29
+ for key, value in st_weights.items():
30
+ # this automatically asserts that the weights are same and raises error if not
31
+ try:
32
+ torch.testing.assert_close(torch_weights[key], value, rtol=1e-5, atol=1e-5)
33
+ except Exception as e:
34
+ total_errors.append(e)
35
+
36
+
37
+ if len(total_errors) > 0:
38
+ return f"weights are not the same ! Conversion failed - {len(total_errors)} errors : {total_errors}"
39
+
40
+ return "Safetensors and torch weights are the same! Conversion sucessfull - you can safely merge the PR"
41
+
42
+ DESCRIPTION = """
43
+ The steps are the following:
44
+ - You got tagged in a safetensor PR? Check if it works!
45
+ - Identify the PR number that you want to check.
46
+ - Paste the model id and the PR number below
47
+ - Click "Submit"
48
+ - That's it! You'll get feedback if the user successfully converted a model in `safetensors` format or not!
49
+
50
+ For now this app supports only `pytorch_model.bin` files, and we'll extend it in the future to support sharded formats.
51
+ """
52
+
53
+ demo = gr.Interface(
54
+ title="SafeTensors Checker",
55
+ description=DESCRIPTION,
56
+ allow_flagging="never",
57
+ article="Check out the [Safetensors repo on GitHub](https://github.com/huggingface/safetensors)",
58
+ inputs=[
59
+ gr.Text(max_lines=1, label="PR number"),
60
+ gr.Text(max_lines=1, label="model_id"),
61
+ ],
62
+ outputs=[gr.Markdown(label="output")],
63
+ fn=run,
64
+ ).queue()
65
+
66
+ demo.launch()
requirements.txt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ torch
2
+ safetensors
3
+ gradio