update

HowieHwong · Jan 31, 2024 · 24fe16f · 24fe16f
1 parent db747de
commit 24fe16f
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 0 deletions.
diff --git a/dataset/dataset.zip b/dataset/dataset.zip
diff --git a/trustllm_pkg/trustllm/dataset_download.py b/trustllm_pkg/trustllm/dataset_download.py
@@ -1,5 +1,6 @@
 import requests
 import os
+import zipfile
 def download_dataset(save_path=None):
     """
     Download a dataset from Hugging Face and save it locally.
@@ -36,6 +37,18 @@ def download_dataset(save_path=None):
                 print(f"Skipping {item['name']}, as it's not a file.")
     else:
         print("Failed to fetch repository data.")
+
+
+    zip_path = os.path.join(save_path, "dataset.zip")
+
+    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
+        zip_ref.extractall(save_path)
+
+    # Delete the ZIP file after extraction
+    os.remove(zip_path)
+
+
+
 
 
 if __name__ == "__main__":