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 Safetensors 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()