Revert "Update app.py"

This reverts commit d49670bbab0c38913b141c315347c62c919a9d9f.

app.py

@@ -1,285 +1,94 @@
-import argparse
-import json
+import csv
+from datetime import datetime
 import os
-import shutil
-from collections import defaultdict
-from inspect import signature
-from tempfile import TemporaryDirectory
-from typing import Dict, List, Optional, Set
+from typing import Optional
+import gradio as gr
 
-import torch
+from convert import convert
+from huggingface_hub import HfApi, Repository
 
-from huggingface_hub import CommitInfo, CommitOperationAdd, Discussion, HfApi, hf_hub_download
-from huggingface_hub.file_download import repo_folder_name
-from safetensors.torch import load_file, save_file
-from transformers import AutoConfig
-from transformers.pipelines.base import infer_framework_load_model
 
+DATASET_REPO_URL = "https://huggingface.co/datasets/safetensors/conversions"
+DATA_FILENAME = "data.csv"
+DATA_FILE = os.path.join("data", DATA_FILENAME)
 
-class AlreadyExists(Exception):
-    pass
+HF_TOKEN = os.environ.get("HF_TOKEN")
 
+repo: Optional[Repository] = None
+if HF_TOKEN:
+    repo = Repository(local_dir="data", clone_from=DATASET_REPO_URL, token=HF_TOKEN)
 
-def shared_pointers(tensors):
-    ptrs = defaultdict(list)
-    for k, v in tensors.items():
-        ptrs[v.data_ptr()].append(k)
-    failing = []
-    for ptr, names in ptrs.items():
-        if len(names) > 1:
-            failing.append(names)
-    return failing
 
-
-def check_file_size(sf_filename: str, pt_filename: str):
-    sf_size = os.stat(sf_filename).st_size
-    pt_size = os.stat(pt_filename).st_size
-
-    if (sf_size - pt_size) / pt_size > 0.01:
-        raise RuntimeError(
-            f"""The file size different is more than 1%:
-         - {sf_filename}: {sf_size}
-         - {pt_filename}: {pt_size}
-         """
-        )
-
-
-def rename(pt_filename: str) -> str:
-    filename, ext = os.path.splitext(pt_filename)
-    local = f"{filename}.safetensors"
-    local = local.replace("pytorch_model", "model")
-    return local
-
-
-def convert_multi(model_id: str, folder: str) -> List["CommitOperationAdd"]:
-    filename = hf_hub_download(repo_id=model_id, filename="pytorch_model.bin.index.json")
-    with open(filename, "r") as f:
-        data = json.load(f)
-
-    filenames = set(data["weight_map"].values())
-    local_filenames = []
-    for filename in filenames:
-        pt_filename = hf_hub_download(repo_id=model_id, filename=filename)
-
-        sf_filename = rename(pt_filename)
-        sf_filename = os.path.join(folder, sf_filename)
-        convert_file(pt_filename, sf_filename)
-        local_filenames.append(sf_filename)
-
-    index = os.path.join(folder, "model.safetensors.index.json")
-    with open(index, "w") as f:
-        newdata = {k: v for k, v in data.items()}
-        newmap = {k: rename(v) for k, v in data["weight_map"].items()}
-        newdata["weight_map"] = newmap
-        json.dump(newdata, f, indent=4)
-    local_filenames.append(index)
-
-    operations = [
-        CommitOperationAdd(path_in_repo=local.split("/")[-1], path_or_fileobj=local) for local in local_filenames
-    ]
-
-    return operations
-
-
-def convert_single(model_id: str, folder: str) -> List["CommitOperationAdd"]:
-    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
-
-
-def convert_file(
-    pt_filename: str,
-    sf_filename: str,
-):
-    loaded = torch.load(pt_filename, map_location="cpu")
-    if "state_dict" in loaded:
-        loaded = loaded["state_dict"]
-    shared = shared_pointers(loaded)
-    for shared_weights in shared:
-        for name in shared_weights[1:]:
-            loaded.pop(name)
-
-    # For tensors to be contiguous
-    loaded = {k: v.contiguous() for k, v in loaded.items()}
-
-    dirname = os.path.dirname(sf_filename)
-    os.makedirs(dirname, exist_ok=True)
-    save_file(loaded, sf_filename, metadata={"format": "pt"})
-    check_file_size(sf_filename, pt_filename)
-    reloaded = load_file(sf_filename)
-    for k in loaded:
-        pt_tensor = loaded[k]
-        sf_tensor = reloaded[k]
-        if not torch.equal(pt_tensor, sf_tensor):
-            raise RuntimeError(f"The output tensors do not match for key {k}")
-
-
-def create_diff(pt_infos: Dict[str, List[str]], sf_infos: Dict[str, List[str]]) -> str:
-    errors = []
-    for key in ["missing_keys", "mismatched_keys", "unexpected_keys"]:
-        pt_set = set(pt_infos[key])
-        sf_set = set(sf_infos[key])
-
-        pt_only = pt_set - sf_set
-        sf_only = sf_set - pt_set
-
-        if pt_only:
-            errors.append(f"{key} : PT warnings contain {pt_only} which are not present in SF warnings")
-        if sf_only:
-            errors.append(f"{key} : SF warnings contain {sf_only} which are not present in PT warnings")
-    return "\n".join(errors)
-
-
-def check_final_model(model_id: str, folder: str):
-    config = hf_hub_download(repo_id=model_id, filename="config.json")
-    shutil.copy(config, os.path.join(folder, "config.json"))
-    config = AutoConfig.from_pretrained(folder)
-
-    _, (pt_model, pt_infos) = infer_framework_load_model(model_id, config, output_loading_info=True)
-    _, (sf_model, sf_infos) = infer_framework_load_model(folder, config, output_loading_info=True)
-
-    if pt_infos != sf_infos:
-        error_string = create_diff(pt_infos, sf_infos)
-        raise ValueError(f"Different infos when reloading the model: {error_string}")
-
-    pt_params = pt_model.state_dict()
-    sf_params = sf_model.state_dict()
-
-    pt_shared = shared_pointers(pt_params)
-    sf_shared = shared_pointers(sf_params)
-    if pt_shared != sf_shared:
-        raise RuntimeError("The reconstructed model is wrong, shared tensors are different {shared_pt} != {shared_tf}")
-
-    sig = signature(pt_model.forward)
-    input_ids = torch.arange(10).unsqueeze(0)
-    pixel_values = torch.randn(1, 3, 224, 224)
-    input_values = torch.arange(1000).float().unsqueeze(0)
-    kwargs = {}
-    if "input_ids" in sig.parameters:
-        kwargs["input_ids"] = input_ids
-    if "decoder_input_ids" in sig.parameters:
-        kwargs["decoder_input_ids"] = input_ids
-    if "pixel_values" in sig.parameters:
-        kwargs["pixel_values"] = pixel_values
-    if "input_values" in sig.parameters:
-        kwargs["input_values"] = input_values
-    if "bbox" in sig.parameters:
-        kwargs["bbox"] = torch.zeros((1, 10, 4)).long()
-    if "image" in sig.parameters:
-        kwargs["image"] = pixel_values
-
-    if torch.cuda.is_available():
-        pt_model = pt_model.cuda()
-        sf_model = sf_model.cuda()
-        kwargs = {k: v.cuda() for k, v in kwargs.items()}
-
-    pt_logits = pt_model(**kwargs)[0]
-    sf_logits = sf_model(**kwargs)[0]
-
-    torch.testing.assert_close(sf_logits, pt_logits)
-    print(f"Model {model_id} is ok !")
-
-
-def previous_pr(api: "HfApi", model_id: str, pr_title: str) -> Optional["Discussion"]:
+def run(token: str, model_id: str) -> str:
+    if token == "" or model_id == "":
+        return """
+        ### Invalid input 🐞
+        
+        Please fill a token and model_id.
+        """
     try:
-        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:
-            return discussion
-
-
-def convert_generic(model_id: str, folder: str, filenames: Set[str]) -> List["CommitOperationAdd"]:
-    operations = []
-
-    extensions = set([".bin", ".ckpt"])
-    for filename in filenames:
-        prefix, ext = os.path.splitext(filename)
-        if ext in extensions:
-            pt_filename = hf_hub_download(model_id, filename=filename)
-            _, raw_filename = os.path.split(filename)
-            if raw_filename == "pytorch_model.bin":
-                # XXX: This is a special case to handle `transformers` and the
-                # `transformers` part of the model which is actually loaded by `transformers`.
-                sf_in_repo = "model.safetensors"
-            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
-
-
-def convert(api: "HfApi", model_id: str, force: bool = False) -> Optional["CommitInfo"]:
-    pr_title = "Adding `safetensors` variant of this model"
-    info = api.model_info(model_id)
-    filenames = set(s.rfilename for s in info.siblings)
-
-    with TemporaryDirectory() as d:
-        folder = os.path.join(d, repo_folder_name(repo_id=model_id, repo_type="models"))
-        os.makedirs(folder)
-        new_pr = None
-        try:
-            operations = None
-            pr = previous_pr(api, model_id, pr_title)
-
-            library_name = getattr(info, "library_name", None)
-            if any(filename.endswith(".safetensors") for filename in filenames) and not force:
-                raise AlreadyExists(f"Model {model_id} is already converted, skipping..")
-            elif pr is not None and not force:
-                url = f"https://huggingface.co/{model_id}/discussions/{pr.num}"
-                new_pr = pr
-                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)
-                elif "pytorch_model.bin.index.json" in filenames:
-                    operations = 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)
-
-            if operations:
-                new_pr = api.create_commit(
-                    repo_id=model_id,
-                    operations=operations,
-                    commit_message=pr_title,
-                    create_pr=True,
+        api = HfApi(token=token)
+        is_private = api.model_info(repo_id=model_id).private
+        print("is_private", is_private)
+
+        commit_info = convert(api=api, model_id=model_id)
+        print("[commit_info]", commit_info)
+
+        # save in a (public) dataset:
+        if repo is not None and not is_private:
+            repo.git_pull(rebase=True)
+            print("pulled")
+            with open(DATA_FILE, "a") as csvfile:
+                writer = csv.DictWriter(
+                    csvfile, fieldnames=["model_id", "pr_url", "time"]
                 )
-                print(f"Pr created at {new_pr.pr_url}")
-            else:
-                print("No files to convert")
-        finally:
-            shutil.rmtree(folder)
-        return new_pr
-
-
-if __name__ == "__main__":
-    DESCRIPTION = """
-    Simple utility tool to convert automatically some weights on the hub to `safetensors` format.
-    It is PyTorch exclusive for now.
-    It works by downloading the weights (PT), converting them locally, and uploading them back
-    as a PR on the hub.
-    """
-    parser = argparse.ArgumentParser(description=DESCRIPTION)
-    parser.add_argument(
-        "model_id",
-        type=str,
-        help="The name of the model on the hub to convert. E.g. `gpt2` or `facebook/wav2vec2-base-960h`",
-    )
-    parser.add_argument(
-        "--force",
-        action="store_true",
-        help="Create the PR even if it already exists of if the model was already converted.",
-    )
-    args = parser.parse_args()
-    model_id = args.model_id
-    api = HfApi()
-    convert(api, model_id, force=args.force)
+                writer.writerow(
+                    {
+                        "model_id": model_id,
+                        "pr_url": commit_info.pr_url,
+                        "time": str(datetime.now()),
+                    }
+                )
+            commit_url = repo.push_to_hub()
+            print("[dataset]", commit_url)
+
+        return f"""
+        ### Success 🔥
+
+        Yay! This model was successfully converted and a PR was open using your token, here:
+
+        [{commit_info.pr_url}]({commit_info.pr_url})
+        """
+    except Exception as e:
+        return f"""
+        ### Error 😢😢😢
+        
+        {e}
+        """
+
+
+DESCRIPTION = """
+The steps are the following:
+
+- Paste a read-access token from hf.co/settings/tokens. Read access is enough given that we will open a PR against the source repo.
+- Input a model id from the Hub
+- Click "Submit"
+- That's it! You'll get feedback if it works or not, and if it worked, you'll get the URL of the opened PR 🔥
+
+⚠️ For now only `pytorch_model.bin` files are supported but we'll extend in the future.
+"""
+
+demo = gr.Interface(
+    title="Convert any model to Safetensors and open a PR",
+    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="your_hf_token"),
+        gr.Text(max_lines=1, label="model_id"),
+    ],
+    outputs=[gr.Markdown(label="output")],
+    fn=run,
+)
+
+demo.launch()