Add download_dataset.py

category_classification/dataset/download_dataset.py

@@ -0,0 +1,39 @@
+import os
+import subprocess
+import zipfile
+import glob
+import pandas as pd
+
+def download_and_extract_dataset(dataset="Cornell-University/arxiv", dest_dir="arxiv_dataset"):
+    os.makedirs(dest_dir, exist_ok=True)
+    print("Downloading dataset...")
+
+    try:
+        subprocess.run(
+            ["kaggle", "datasets", "download", "-d", dataset, "-p", dest_dir],
+            check=True,
+        )
+        print("Download finished")
+
+    except subprocess.CalledProcessError as e:
+        print("Downloading error: ", e)
+        return
+
+    zip_filename = os.path.join(dest_dir, "arxiv.zip")
+    if os.path.exists(zip_filename):
+        print("Unpacking dataset...")
+        with zipfile.ZipFile(zip_filename, "r") as zip_ref:
+            zip_ref.extractall(dest_dir)
+        print("Unpacking finished")
+    else:
+        print("Zip-file is not found")
+
+def filter_csv(dest_dir="arxiv_dataset"):
+    csv_file = glob.glob(os.path.join(dest_dir, "*.csv"))[0]
+    df = pd.read_csv(csv_file)
+    df = df[["title", "authors", "abstract", "categories"]]
+    df.to_csv(os.path.join(dest_dir, "filtered_arxiv.csv"), index=False)
+
+if __name__ == "__main__":
+    download_and_extract_dataset()
+    filter_csv()