{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "73b9d94f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b90d0ab3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>id</th>\n",
       "      <th>title</th>\n",
       "      <th>abstract</th>\n",
       "      <th>submit_date</th>\n",
       "      <th>n_citation</th>\n",
       "      <th>orgs</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1001.0279</td>\n",
       "      <td>Regularization for Matrix Completion</td>\n",
       "      <td>We consider the problem of reconstructing a lo...</td>\n",
       "      <td>2010-01-02 04:14:50</td>\n",
       "      <td>0</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1001.0405</td>\n",
       "      <td>Optimal Query Complexity for Reconstructing Hy...</td>\n",
       "      <td>In this paper we consider the problem of recon...</td>\n",
       "      <td>2010-01-03 19:54:40</td>\n",
       "      <td>0</td>\n",
       "      <td>Cs.technion.ac.il</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "          id                                              title  \\\n",
       "0  1001.0279               Regularization for Matrix Completion   \n",
       "1  1001.0405  Optimal Query Complexity for Reconstructing Hy...   \n",
       "\n",
       "                                            abstract         submit_date  \\\n",
       "0  We consider the problem of reconstructing a lo... 2010-01-02 04:14:50   \n",
       "1  In this paper we consider the problem of recon... 2010-01-03 19:54:40   \n",
       "\n",
       "   n_citation               orgs  \n",
       "0           0                     \n",
       "1           0  Cs.technion.ac.il  "
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ml_papers_df = pd.read_json(\n",
    "    \"/mnt/data-ssd-2/data/arxiv/arxiv_ml_meta.json\", \n",
    "    dtype={\"id\": str},\n",
    "    lines=True,\n",
    ")\n",
    "ml_papers_df[\"submit_date\"] = pd.to_datetime(ml_papers_df[\"submit_date\"], unit=\"ms\")\n",
    "ml_papers_df.head(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "9397fe3e",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2670 speech AI papers found\n"
     ]
    }
   ],
   "source": [
    "min_year = 2015\n",
    "search_ptns = [\n",
    "    r\"[Ss]peech [Rr]ecognition\",\n",
    "    r\"[Tt]ext [Tt]o [Ss]peech\",\n",
    "]\n",
    "search_ptn = r\"\\b(?:\" + r\"|\".join(search_ptns) + r\")\\b\"\n",
    "\n",
    "df = ml_papers_df.copy()\n",
    "df = df[df[\"submit_date\"].dt.year >= min_year]\n",
    "df = df[(df[\"title\"] + \" \" + df[\"abstract\"]).str.contains(search_ptn, regex=True)]\n",
    "df = df.reset_index(drop=True)\n",
    "\n",
    "speech_axiv_ids = set(df[\"id\"])\n",
    "print(len(speech_axiv_ids), \"speech AI papers found\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9483116",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "807c749d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tarfile\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "3b8c93f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "t0 = time.time()\n",
    "with tarfile.open(\"/mnt/data-ssd-2/data/arxiv/tarpdfs/arXiv_pdf_2102_011.tar.gz\", \"r:gz\") as f:\n",
    "    for member in f.getmembers():\n",
    "        if member.name == \"2102/2102.02382.pdf\":\n",
    "            t1 = time.time()\n",
    "            f.extract(member, \"/home/georg/notebooks/arxiv/test/\")\n",
    "            t2 = time.time()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "f0eb17e2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.007845640182495117\n"
     ]
    }
   ],
   "source": [
    "print(t2 - t1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "146303e3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.836005449295044\n"
     ]
    }
   ],
   "source": [
    "print(t1 - t0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "b44214d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import tempfile\n",
    "import shutil\n",
    "import subprocess\n",
    "from bs4 import BeautifulSoup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "cbc4a2b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "CERMINE_CMD_PTN = (\n",
    "    \"java -cp {cermine_path} pl.edu.icm.cermine.ContentExtractor \"\n",
    "    \"-path {pdf_dir} -timeout {timeout_s} -outputs jats\"\n",
    ")\n",
    "\n",
    "def extract_pdf_metadata(filepath, cermine_path, timeout_s=60):\n",
    "    if filepath.split(\".\")[-1] != \"pdf\":\n",
    "        raise ValueError(\"input needs to be of .pdf type\")\n",
    "    with tempfile.TemporaryDirectory() as tmp_dir:\n",
    "        shutil.copy(filepath, tmp_dir)\n",
    "        filename = filepath.split(\"/\")[-1][:-4]\n",
    "        out_filepath = os.path.join(tmp_dir, f\"{filename}.cermxml\")\n",
    "        cmd = CERMINE_CMD_PTN.format(cermine_path=cermine_path, pdf_dir=tmp_dir, timeout_s=timeout_s)\n",
    "        out_code = subprocess.call(\n",
    "            cmd, shell=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL\n",
    "        )\n",
    "        if out_code != 0:\n",
    "            raise ValueError(\"cermine error\")\n",
    "        with open(out_filepath) as f:\n",
    "            xml_str = f.read()\n",
    "    return xml_str\n",
    "\n",
    "def get_affiliation_info(xml_str):\n",
    "    soup = BeautifulSoup(xml_str)\n",
    "    author_infos = []\n",
    "    for e in soup.find_all(\"contrib\", attrs={\"contrib-type\" : \"author\"}):\n",
    "        author_info = {}\n",
    "        if (ee := e.find(\"string-name\")):\n",
    "            author_info[\"name\"] = ee.text\n",
    "        if (ee := e.find(\"email\")):\n",
    "            author_info[\"email\"] = ee.text\n",
    "        for ee in e.find_all(\"xref\", attrs={\"ref-type\" : \"aff\"}):\n",
    "            # TODO: this overrides multiple affiliations\n",
    "            author_info[\"affiliation\"] = ee[\"rid\"]\n",
    "        author_infos.append(author_info)\n",
    "    org_infos = {}\n",
    "    for e in soup.find_all(\"aff\"):\n",
    "        org_id = e[\"id\"]\n",
    "        org_info = {}\n",
    "        if (ee := e.find(\"institution\")):\n",
    "            org_info[\"name\"] = ee.text\n",
    "        if (ee := e.find(\"addr-line\")):\n",
    "            org_info[\"address\"] = ee.text\n",
    "        if (ee := e.find(\"country\")):\n",
    "            org_info[\"country\"] = ee.text\n",
    "        org_infos[org_id] = org_info\n",
    "    for e in author_infos:\n",
    "        if \"affiliation\" in e and e[\"affiliation\"] in org_infos and \"name\" in org_infos[e[\"affiliation\"]]:\n",
    "            e[\"affiliation\"] = org_infos[e[\"affiliation\"]][\"name\"]\n",
    "    return author_infos"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "307d65bf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 29.6 ms, sys: 12 ms, total: 41.6 ms\n",
      "Wall time: 11.3 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "CERMINE_PATH = \"/home/georg/tools/cermine/cermine-impl-1.13-jar-with-dependencies.jar\"\n",
    "\n",
    "filepath = \"/home/georg/tools/cermine/test_pdfs/wavlm.pdf\"\n",
    "# filepath = \"/home/georg/tools/cermine/test_pdfs/spgispeech.pdf\"\n",
    "xml_str = extract_pdf_metadata(filepath, CERMINE_PATH)\n",
    "author_affiliations = get_affiliation_info(xml_str)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "b782899f",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "46defc65",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "78b661d8",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "4b613455",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "38.486207485198975\n"
     ]
    }
   ],
   "source": [
    "print(t1-t0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "224ef72b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14ef8e1b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "103c477e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}