Update app.py

app.py

@@ -366,6 +366,25 @@ def clear_data():
         print(f"The directory '{dir_path}' does not exist.")
 
 
+def get_size(start_path):
+    total_size = 0
+    for dirpath, dirnames, filenames in os.walk(start_path):
+        for f in filenames:
+            fp = os.path.join(dirpath, f)
+            # skip if it is symbolic link
+            if not os.path.islink(fp):
+                total_size += os.path.getsize(fp)
+    return total_size
+
+
+def human_readable_size(size, decimal_places=2):
+    for unit in ['B', 'KB', 'MB', 'GB', 'TB', 'PB']:
+        if size < 1024.0:
+            break
+        size /= 1024.0
+    return f"{size:.{decimal_places}f} {unit}"
+
+
 def merge_loop():
     """
     Main function that orchestrates the merge.
@@ -381,7 +400,7 @@ def merge_loop():
     df = make_df("open-llm-leaderboard.csv", N_ROWS)
 
     # Sample two models
-    subprocess.run(["du", "-sh"], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
+    print(f"Total size: {human_readable_size(get_size('.'))}")
     sample = df.sample(n=2)
     models = [sample.iloc[i] for i in range(2)]
 
@@ -397,31 +416,30 @@ def merge_loop():
     # Merge configs
     yaml_config = create_config(models)
     print(f"YAML config:{yaml_config}")
-    subprocess.run(["du", "-sh"], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
+    print(f"Total size: {human_readable_size(get_size('.'))}")
 
     # Merge models
     merge_models()
     print("Model merged!")
-    subprocess.run(["du", "-sh"], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
-    subprocess.run(["ls", "-larth"], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
+    print(f"Total size: {human_readable_size(get_size('.'))}")
 
     # Create model card
     print("Create model card")
     create_model_card(yaml_config, model_name, USERNAME, license)
-    subprocess.run(["du", "-sh"], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
+    print(f"Total size: {human_readable_size(get_size('.'))}")
 
     # Upload model
     print("Upload model")
     upload_model(api, USERNAME, model_name)
-    subprocess.run(["du", "-sh"], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
+    print(f"Total size: {human_readable_size(get_size('.'))}")
 
     # Clear data
     clear_data()
-    subprocess.run(["du", "-sh"], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
+    print(f"Total size: {human_readable_size(get_size('.'))}")
     
     # Evaluate model on Runpod
     create_pod(model_name, USERNAME)
-    subprocess.run(["du", "-sh"], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
+    print(f"Total size: {human_readable_size(get_size('.'))}")
 
     
 # Install scrape-open-llm-leaderboard and mergekit