Rename tank tflite/torch model dir (huggingface#219)

Author: AmosLewis

generate_sharktank.py

@@ -61,7 +61,9 @@ def save_torch_model(torch_model_list):
                 model, input, _ = get_hf_model(torch_model_name)
 
             torch_model_name = torch_model_name.replace("/", "_")
-            torch_model_dir = os.path.join(WORKDIR, str(torch_model_name))
+            torch_model_dir = os.path.join(
+                WORKDIR, str(torch_model_name) + "_torch"
+            )
             os.makedirs(torch_model_dir, exist_ok=True)
 
             mlir_importer = SharkImporter(
@@ -136,7 +138,7 @@ def save_tflite_model(tflite_model_list):
             print("tflite_model_name", tflite_model_name)
             print("tflite_model_link", tflite_model_link)
             tflite_model_name_dir = os.path.join(
-                WORKDIR, str(tflite_model_name)
+                WORKDIR, str(tflite_model_name) + "_tflite"
             )
             os.makedirs(tflite_model_name_dir, exist_ok=True)
             print(f"TMP_TFLITE_MODELNAME_DIR = {tflite_model_name_dir}")

shark/shark_downloader.py

@@ -35,18 +35,26 @@
 WORKDIR = os.path.join(home, ".local/shark_tank/")
 print(WORKDIR)
 
+
 # Checks whether the directory and files exists.
 def check_dir_exists(model_name, frontend="torch", dynamic=""):
     model_dir = os.path.join(WORKDIR, model_name)
 
     # Remove the _tf keyword from end.
     if frontend in ["tf", "tensorflow"]:
         model_name = model_name[:-3]
+    elif frontend in ["tflite"]:
+        model_name = model_name[:-7]
+    elif frontend in ["torch", "pytorch"]:
+        model_name = model_name[:-6]
 
     if os.path.isdir(model_dir):
         if (
             os.path.isfile(
-                os.path.join(model_dir, model_name + dynamic + ".mlir")
+                os.path.join(
+                    model_dir,
+                    model_name + dynamic + "_" + str(frontend) + ".mlir",
+                )
             )
             and os.path.isfile(os.path.join(model_dir, "function_name.npy"))
             and os.path.isfile(os.path.join(model_dir, "inputs.npz"))
@@ -65,27 +73,28 @@ def download_torch_model(model_name, dynamic=False):
     model_name = model_name.replace("/", "_")
     dyn_str = "_dynamic" if dynamic else ""
     os.makedirs(WORKDIR, exist_ok=True)
+    model_dir_name = model_name + "_torch"
 
     def gs_download_model():
         gs_command = (
             'gsutil -o "GSUtil:parallel_process_count=1" cp -r gs://shark_tank'
             + "/"
-            + model_name
+            + model_dir_name
             + " "
             + WORKDIR
         )
         if os.system(gs_command) != 0:
             raise Exception("model not present in the tank. Contact Nod Admin")
 
-    if not check_dir_exists(model_name, dyn_str):
+    if not check_dir_exists(model_dir_name, frontend="torch", dynamic=dyn_str):
         gs_download_model()
     else:
-        model_dir = os.path.join(WORKDIR, model_name)
+        model_dir = os.path.join(WORKDIR, model_dir_name)
         local_hash = str(np.load(os.path.join(model_dir, "hash.npy")))
         gs_hash = (
             'gsutil -o "GSUtil:parallel_process_count=1" cp gs://shark_tank'
             + "/"
-            + model_name
+            + model_dir_name
             + "/hash.npy"
             + " "
             + os.path.join(model_dir, "upstream_hash.npy")
@@ -98,8 +107,10 @@ def gs_download_model():
         if local_hash != upstream_hash:
             gs_download_model()
 
-    model_dir = os.path.join(WORKDIR, model_name)
-    with open(os.path.join(model_dir, model_name + dyn_str + ".mlir")) as f:
+    model_dir = os.path.join(WORKDIR, model_dir_name)
+    with open(
+        os.path.join(model_dir, model_name + dyn_str + "_torch.mlir")
+    ) as f:
         mlir_file = f.read()
 
     function_name = str(np.load(os.path.join(model_dir, "function_name.npy")))
@@ -115,18 +126,21 @@ def gs_download_model():
 def download_tflite_model(model_name, dynamic=False):
     dyn_str = "_dynamic" if dynamic else ""
     os.makedirs(WORKDIR, exist_ok=True)
-    if not check_dir_exists(model_name, dyn_str):
+    model_dir_name = model_name + "_tflite"
+    if not check_dir_exists(
+        model_dir_name, frontend="tflite", dynamic=dyn_str
+    ):
         gs_command = (
             'gsutil -o "GSUtil:parallel_process_count=1" cp -r gs://shark_tank'
             + "/"
-            + model_name
+            + model_dir_name
             + " "
             + WORKDIR
         )
         if os.system(gs_command) != 0:
             raise Exception("model not present in the tank. Contact Nod Admin")
 
-    model_dir = os.path.join(WORKDIR, model_name)
+    model_dir = os.path.join(WORKDIR, model_dir_name)
     with open(
         os.path.join(model_dir, model_name + dyn_str + "_tflite.mlir")
     ) as f: