Update src/saving_utils.py

src/saving_utils.py

@@ -6,6 +6,55 @@ from huggingface_hub import HfApi
 script_dir = os.path.dirname(os.path.abspath(__file__))  # Directory of the running script
 
 
+
+def download_files_from_hub(benchmark_types, repo_id="mgyigit/probe-data", repo_type="space"):
+    api = HfApi(token=os.getenv("api-key")) #load api-key secret
+
+    benchmark_types += "leaderboard"
+    for benchmark in benchmark_types:
+        file_name = f"{benchmark}_results.csv"
+        local_path = f"/tmp/{file_name}"
+        
+        try:
+            # Download the file from the specified repo
+            api.download_file(
+                repo_id=repo_id,
+                path_in_repo=file_name,
+                local_dir="/tmp",
+                repo_type=repo_type,
+            )
+            print(f"Downloaded {file_name} from {repo_id} to {local_path}")
+
+        except Exception as e:
+            print(f"Failed to download {file_name}: {e}")
+
+
+    return 0
+
+
+def upload_to_hub(benchmark_types, repo_id="mgyigit/probe-data", repo_type="space"):
+    api = HfApi(token=os.getenv("api_key"))  # Requires authentication via HF_TOKEN
+    benchmark_types += "leaderboard"
+
+    for benchmark in benchmark_types:
+        file_name = f"{benchmark}_results.csv"
+        local_path = f"/tmp/{file_name}"
+
+        api.upload_file(
+            path_or_fileobj=local_path,
+            path_in_repo=file_name,
+            repo_id=repo_id,
+            repo_type=repo_type,
+            commit_message=f"Updating {file_name}"
+        )
+        print(f"Uploaded {local_path} to {repo_id}/{file_name}")
+
+        os.remove(local_path)
+        print(f"Removed local file: {file_path}")
+
+    return 0
+
+
 def save_csv_locally(dataframe, file_name, save_dir="/tmp"):
     # Ensure the save directory exists
     os.makedirs(save_dir, exist_ok=True)
@@ -19,28 +68,13 @@ def save_csv_locally(dataframe, file_name, save_dir="/tmp"):
 
     return file_path
 
-def upload_to_hub(local_path, remote_path, repo_id, repo_type="space"):
-    api = HfApi(token=os.getenv("api_key"))  # Requires authentication via HF_TOKEN
-    api.upload_file(
-        path_or_fileobj=local_path,
-        path_in_repo=remote_path,
-        repo_id=repo_id,
-        repo_type=repo_type,
-        commit_message=f"Updating {os.path.basename(remote_path)}"
-    )
-    print(f"Uploaded {local_path} to {repo_id}/{remote_path}")
-
-def cleanup_local_file(file_path):
-    if os.path.exists(file_path):
-        os.remove(file_path)
-        print(f"Removed local file: {file_path}")
 
 def save_similarity_output(
     output_dict,
     method_name,
-    leaderboard_path="/home/user/app/src/data/leaderboard_results.csv",
-    similarity_path="/home/user/app/src/data/similarity_results.csv",
-    repo_id="mgyigit/probe3",
+    leaderboard_path="/tmp/leaderboard_results.csv",
+    similarity_path="/tmp/similarity_results.csv",
+    repo_id="mgyigit/probe-data",
 ):
     # Load or initialize the DataFrames
     if os.path.exists(leaderboard_path):
@@ -61,16 +95,12 @@ def save_similarity_output(
         new_row['Method'] = method_name
         similarity_df = pd.concat([similarity_df, pd.DataFrame([new_row])], ignore_index=True)
 
-    # Same for the leaderboard DataFrame
     if method_name not in leaderboard_df['Method'].values:
         new_row = {col: None for col in leaderboard_df.columns}
         new_row['Method'] = method_name
         leaderboard_df = pd.concat([leaderboard_df, pd.DataFrame([new_row])], ignore_index=True)
 
-    # Initialize storage for averages
     averages = {}
-
-    # Iterate through the datasets and calculate averages
     for dataset in ['sparse', '200', '500']:
         correlation_values = []
         pvalue_values = []
@@ -104,22 +134,9 @@ def save_similarity_output(
             similarity_df.loc[similarity_df['Method'] == method_name, f"{dataset}_Ave_pvalue"] = averages[f"{dataset}_Ave_pvalue"]
             leaderboard_df.loc[leaderboard_df['Method'] == method_name, f"sim_{dataset}_Ave_pvalue"] = averages[f"{dataset}_Ave_pvalue"]
 
-    # Save locally to a temporary directory
     leaderboard_file = save_csv_locally(leaderboard_df, "leaderboard_results.csv")
     similarity_file = save_csv_locally(similarity_df, "similarity_results.csv")
 
-    # Upload to Hugging Face Hub
-    try:
-        upload_to_hub(leaderboard_file, "leaderboard_results.csv", repo_id)
-        upload_to_hub(similarity_file, "similarity_results.csv", repo_id)
-    except Exception as e:
-        print(f"Failed to upload files: {e}")
-        return -1
-
-    # Clean up local files
-    cleanup_local_file(leaderboard_file)
-    cleanup_local_file(similarity_file)
-
     return 0
 
 def save_function_output(model_output, method_name, func_results_path="/home/user/app/src/data/function_results.csv", leaderboard_path="/home/user/app/src/data/leaderboard_results.csv"):