Add sharded support

This PR adds the sharded model checking support

app.py

@@ -1,3 +1,6 @@
+import json
+import shutil
+import gc
 import gradio as gr
 
 import torch
@@ -6,14 +9,7 @@ import safetensors
 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"
-
-
+def check_simple_file(st_weights_path, torch_weights_path):
     st_weights = safetensors.torch.load_file(st_weights_path)
     torch_weights = torch.load(torch_weights_path)
 
@@ -21,7 +17,7 @@ def run(pr_number, model_id):
     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}"
+        return f"keys are not the same ! Conversion failed - unexpected keys are: {unexpected_keys} for the file {st_weights_path}"
 
     total_errors = []
 
@@ -33,6 +29,54 @@ def run(pr_number, model_id):
         except Exception as e:
             total_errors.append(e)
 
+    del st_weights
+    del torch_weights
+    gc.collect()
+
+    return total_errors
+
+def run(pr_number, model_id):
+    is_sharded = False
+    try:
+        st_sharded_index_file = hf_hub_download(repo_id=model_id, filename="model.safetensors.index.json", revision=f"refs/pr/{pr_number}")
+        torch_sharded_index_file = hf_hub_download(repo_id=model_id, filename="pytorch_model.bin.index.json")
+
+        is_sharded = True
+    except:
+        pass
+
+    if not is_sharded:
+        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"
+
+        total_errors = check_simple_file(st_weights_path, torch_weights_path)  
+    else:
+        total_errors = []
+        total_st_files = set(json.load(open(st_sharded_index_file, "r"))["weight_map"].values())
+        total_pt_files = set(json.load(open(torch_sharded_index_file, "r"))["weight_map"].values())
+
+        if len(total_st_files) != len(total_pt_files):
+            return f"weights are not the same there are {len(total_st_files)} files in safetensors and {len(total_pt_files)} files in torch ! Conversion failed - {len(total_errors)} errors : {total_errors}"
+         
+        # check if the mapping are correct
+        if not all([pt_file.replace("pytorch_model", "model").replace(".bin", ".safetensors") in total_st_files for pt_file in total_pt_files]):
+            return f"Conversion failed! Safetensors files are not the same as torch files - make sure you have the correct files in the PR"
+
+        for pt_file in total_pt_files:
+            st_file = pt_file.replace("pytorch_model", "model").replace(".bin", ".safetensors")
+            
+            st_weights_path = hf_hub_download(repo_id=model_id, filename=st_file, revision=f"refs/pr/{pr_number}")
+            torch_weights_path = hf_hub_download(repo_id=model_id, filename=pt_file)
+
+            total_errors += check_simple_file(st_weights_path, torch_weights_path)
+
+            # remove files for memory optimization
+            shutil.rmtree(st_weights_path)
+            shutil.rmtree(torch_weights_path)
+
 
     if len(total_errors) > 0:
         return f"weights are not the same ! Conversion failed - {len(total_errors)} errors : {total_errors}"
@@ -47,7 +91,7 @@ The steps are the following:
 - 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.
+This checker also support sharded weights.
 """
 
 demo = gr.Interface(