v1

app.py

@@ -0,0 +1,66 @@
+import gradio as gr
+
+import torch
+import safetensors
+# hack to load safetensors.torch
+from safetensors.torch import save_file
+from huggingface_hub import hf_hub_download
+
+def run(pr_number, model_id):
+    try:
+        st_weights_path = hf_hub_download(repo_id=model_id, filename="model.safetensors", revision=f"refs/pr/{pr_number}")
+        torch_weights_path = hf_hub_download(repo_id=model_id, filename="pytorch_model.bin")
+    except Exception as e:
+        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"
+
+
+    st_weights = safetensors.torch.load_file(st_weights_path)
+    torch_weights = torch.load(torch_weights_path)
+
+    # check if keys are the same
+    if st_weights.keys() != torch_weights.keys():
+        # retrieve different keys
+        unexpected_keys = st_weights.keys() - torch_weights.keys()
+        return f"keys are not the same ! Conversion failed - unexpected keys are: {unexpected_keys}"
+
+    total_errors = []
+
+    # check all weights are same
+    for key, value in st_weights.items():
+        # this automatically asserts that the weights are same and raises error if not
+        try:
+            torch.testing.assert_close(torch_weights[key], value, rtol=1e-5, atol=1e-5)
+        except Exception as e:
+            total_errors.append(e)
+
+
+    if len(total_errors) > 0:
+        return f"weights are not the same ! Conversion failed - {len(total_errors)} errors : {total_errors}"
+    
+    return "Safetensors and torch weights are the same! Conversion sucessfull - you can safely merge the PR"
+
+DESCRIPTION = """
+The steps are the following:
+- You got tagged in a safetensor PR? Check if it works!
+- Identify the PR number that you want to check.
+- Paste the model id and the PR number below
+- Click "Submit"
+- That's it! You'll get feedback if the user successfully converted a model in `safetensors` format or not!
+
+For now this app supports only `pytorch_model.bin` files, and we'll extend it in the future to support sharded formats.
+"""
+
+demo = gr.Interface(
+    title="SafeTensors Checker",
+    description=DESCRIPTION,
+    allow_flagging="never",
+    article="Check out the [Safetensors repo on GitHub](https://github.com/huggingface/safetensors)",
+    inputs=[
+        gr.Text(max_lines=1, label="PR number"),
+        gr.Text(max_lines=1, label="model_id"),
+    ],
+    outputs=[gr.Markdown(label="output")],
+    fn=run,
+).queue()
+
+demo.launch()

requirements.txt

@@ -0,0 +1,3 @@
+torch
+safetensors
+gradio