arxiv_0_downloader.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c4149fb-5e49-4fc5-a2cc-88f11034e476",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import boto3\n",
    "\n",
    "# Initialize a session using Amazon S3\n",
    "session = boto3.session.Session(\n",
    "    aws_access_key_id='...',\n",
    "    aws_secret_access_key='...',\n",
    ")\n",
    "\n",
    "# Create an S3 client\n",
    "s3 = session.client('s3')\n",
    "\n",
    "# Define the bucket name and the key of the manifest file\n",
    "bucket_name = 'arxiv'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "605ebf64-9f7a-4476-b95b-1d3590e9db1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def download_arxiv_file(origin_path, target_path=None, verbose=True):\n",
    "    if target_path is None:\n",
    "        target_path = origin_path\n",
    "    target_path = './downloaded/' + target_path\n",
    "    # Ensure the directory exists\n",
    "    directory = os.path.dirname(target_path)\n",
    "    if not os.path.exists(directory):\n",
    "        os.makedirs(directory)\n",
    "    if os.path.exists(target_path):\n",
    "        if verbose:\n",
    "            print(f'🐤 {origin_path} already downloaded to {target_path}')\n",
    "        return\n",
    "    try:\n",
    "        s3.download_file(bucket_name, origin_path, target_path, ExtraArgs={'RequestPayer': 'requester'})\n",
    "        if verbose:\n",
    "            print(f'🥳 Download ({origin_path}) successfully.')\n",
    "    except Exception as e:\n",
    "        print(f\"❗️ An error occurred: {e}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9cf784ff-132f-45c2-a12a-7cdd6ec64714",
   "metadata": {},
   "outputs": [],
   "source": [
    "download_arxiv_file('pdf/arXiv_pdf_manifest.xml')\n",
    "download_arxiv_file('src/arXiv_src_manifest.xml')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a05a7254-c855-46cd-8618-24ccbed1b6cb",
   "metadata": {},
   "source": [
    "# Process PDF & SRC manifest files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c00d276e-248d-4119-a308-7ae395bc4e81",
   "metadata": {},
   "outputs": [],
   "source": [
    "import xmltodict\n",
    "\n",
    "with open('./downloaded/pdf/arXiv_pdf_manifest.xml', 'r') as manifest_pdf_file:\n",
    "    manifest_pdf_dict = xmltodict.parse(manifest_pdf_file.read())\n",
    "\n",
    "with open('./downloaded/src/arXiv_src_manifest.xml', 'r') as manifest_src_file:\n",
    "    manifest_src_dict = xmltodict.parse(manifest_src_file.read())\n",
    "\n",
    "manifest_pdf_files = manifest_pdf_dict['arXivPDF']['file']\n",
    "manifest_src_files = manifest_src_dict['arXivSRC']['file']\n",
    "\n",
    "len(manifest_pdf_files), len(manifest_src_files)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f89ef974-ed54-47cb-85dc-fbcc6b3f603f",
   "metadata": {},
   "outputs": [],
   "source": [
    "manifest_files = manifest_pdf_files + manifest_src_files\n",
    "len(manifest_files)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c42c2b3-6d4a-4f8b-81dd-719e1e7463c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_dates = [\n",
    "    '0001',\n",
    "    '0002',\n",
    "    '0003',\n",
    "    '0004',\n",
    "    '0005',\n",
    "    '0006',\n",
    "    '0007',\n",
    "    '0008',\n",
    "    '0009',\n",
    "    '0010',\n",
    "    '0011',\n",
    "    '0012',\n",
    "    '0501',\n",
    "    '1001',\n",
    "    '1501',\n",
    "    '2001',\n",
    "]\n",
    "\n",
    "sample_files = { date: [item for item in manifest_files if date in item['filename']] for date in sample_dates }\n",
    "sample_filenames = [item['filename'] for sample_date in sample_files.keys() for item in sample_files[sample_date]]\n",
    "sample_filenames.sort()\n",
    "sample_filenames"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a059abca-3693-4742-b1b9-989490789733",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The following files might be missing from the manifest file. So download them manually.\n",
    "# Comment out this block if you don't want to override the download list.\n",
    "sample_filenames = [\n",
    "    'pdf/arXiv_pdf_0009_001.tar',\n",
    "    'pdf/arXiv_pdf_0009_002.tar',\n",
    "    'pdf/arXiv_pdf_0010_001.tar',\n",
    "    'pdf/arXiv_pdf_0010_002.tar',\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a79bf5ef-dc6a-4d31-bbe7-bb28bef9f083",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm.notebook import tqdm\n",
    "import time\n",
    "\n",
    "pbar = tqdm(sample_filenames, desc=\"Downloading files\")\n",
    "for filename in pbar:\n",
    "    pbar.set_description(f\"Downloading {filename}\")\n",
    "    download_arxiv_file(filename, verbose=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d10ee41e-148d-49a4-a2b1-a759bc3faba3",
   "metadata": {},
   "source": [
    "# Download data pre-processing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "603effd7-c9b5-4d08-b86f-f2404e7b1a85",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tarfile\n",
    "import os\n",
    "import re\n",
    "import shutil\n",
    "\n",
    "from tqdm.notebook import tqdm\n",
    "import time\n",
    "\n",
    "arxiv_output_dir = './arxiv'\n",
    "temp_extraction_dir = './arxiv/temp_extracted_files'\n",
    "\n",
    "# Function to extract tar files\n",
    "def extract_tar(file_path, extraction_path=temp_extraction_dir):\n",
    "    temp_output_path = os.path.join(extraction_path, os.path.basename(file_path).split('.')[0])\n",
    "    with tarfile.open(file_path, 'r') as tar:\n",
    "        tar.extractall(path=temp_output_path)\n",
    "    return temp_output_path\n",
    "\n",
    "# Function to organize files based on dates\n",
    "def organize_by_date_and_type(filename, extracted_path):\n",
    "    match = re.match(r\"(pdf|src)/arXiv_.*_(\\d{4})_.*\\.tar\", filename)\n",
    "    if match:\n",
    "        type_dir = match.group(1)\n",
    "        date_dir = match.group(2)\n",
    "        target_dir = os.path.join(arxiv_output_dir, date_dir, type_dir)\n",
    "\n",
    "        if not os.path.exists(target_dir):\n",
    "            os.makedirs(target_dir)\n",
    "\n",
    "        temp_dir = os.path.join(extracted_path, date_dir)\n",
    "        # Move the extracted files to the date and type directory\n",
    "        for file in os.listdir(temp_dir):\n",
    "            shutil.move(os.path.join(temp_dir, file), os.path.join(target_dir, file))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79873c6c-6cde-4cd4-a145-01dc93fc44ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import subprocess\n",
    "from joblib import Parallel, delayed\n",
    "from joblib_progress import joblib_progress\n",
    "\n",
    "print(\"CPU:\", os.cpu_count())\n",
    "\n",
    "if not os.path.exists(arxiv_output_dir):\n",
    "    os.makedirs(arxiv_output_dir)\n",
    "if os.path.exists(temp_extraction_dir):\n",
    "    shutil.rmtree(temp_extraction_dir)\n",
    "os.makedirs(temp_extraction_dir)\n",
    "\n",
    "def inner_func(filename):\n",
    "    temp_output_path = extract_tar('./downloaded/' + filename)\n",
    "    organize_by_date_and_type(filename, temp_output_path)\n",
    "    shutil.rmtree(temp_output_path)\n",
    "    return filename\n",
    "\n",
    "with joblib_progress(\"Processing...\", total=len(sample_filenames)):\n",
    "    Parallel(n_jobs=os.cpu_count(), pre_dispatch='1*n_jobs')(\n",
    "        delayed(inner_func)(f) for f in sample_filenames\n",
    "    )\n",
    "\n",
    "print('Done!')\n",
    "\n",
    "shutil.rmtree(temp_extraction_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6082f4b2-0c1a-49c8-aebc-6ecd48a15546",
   "metadata": {},
   "outputs": [],
   "source": [
    "!du -sh ./arxiv"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}