dataset.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61f2c04d-d309-405c-899f-e666dc8516f1",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import json\n",
    "import pathlib\n",
    "from itertools import chain\n",
    "from collections import Counter\n",
    "\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from arxiv2text import arxiv_to_text\n",
    "from arxiv.taxonomy import definitions\n",
    "\n",
    "path_root = pathlib.Path(\".\").parent\n",
    "path_dataset = path_root / \"dataset\"\n",
    "path_database = path_root / \"database\"\n",
    "path_dataset.mkdir(exist_ok=True)\n",
    "path_database.mkdir(exist_ok=True)\n",
    "\n",
    "CATEGORIES = {\n",
    "    cat: meta\n",
    "    for cat, meta in definitions.CATEGORIES_ACTIVE.items()\n",
    "    if cat.split(\".\")[0] == \"cs\"\n",
    "}\n",
    "\n",
    "cat2idx = {cat: idx for idx, cat in enumerate(CATEGORIES)}\n",
    "idx2cat = {idx: cat for idx, cat in enumerate(CATEGORIES)}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9ca893f-fe2f-4def-80e3-cd8f564966e0",
   "metadata": {},
   "source": [
    "## Build Database\n",
    "\n",
    "First, we need to discard all metadata that is unrelated to Computer Science papers. This step is taken to restrict the size of the dataset, making it more manageable. By doing so, we also limit the number of categories in the classification task, enabling it to comfortably fit into the LLM prompt without muddling the context. If your prompt contains numerous categories, you can refer to [this notebook](https://colab.research.google.com/drive/1CpsOiLiLYKeGrhmq579_FmtGsD5uZ3Qe); dspy-arxiv heavily relies on the techniques discussed in this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "faec2f98-a934-4497-b572-47af2f7b6ea0",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(path_database / \"arxiv.json\", \"w\") as outfile:\n",
    "    with open(path_root / \"arxiv.json\") as infile:\n",
    "        for line in infile:\n",
    "            paper = json.loads(line.strip())\n",
    "            cats = {cat for cat in paper[\"categories\"].split(\" \")}\n",
    "            if cats & set(CATEGORIES.keys()):\n",
    "                outfile.write(line)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8937de4-b085-4e14-b212-1ed24f1716e5",
   "metadata": {},
   "source": [
    "## Build Dataset\n",
    "\n",
    "In the following, we will perform a brief data analysis to better understand the data and determine which papers to select. The goal is to choose papers that belong to a single category and also include a few papers that have multiple categories."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f756444-11db-4979-9aa2-06bdca396f65",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def reader(path_file):\n",
    "    with open(path_file, \"r\") as f:\n",
    "        for line in f:\n",
    "            paper = json.loads(line)\n",
    "            yield {\n",
    "                \"id\": paper[\"id\"],\n",
    "                \"categories\": paper[\"categories\"].split(\" \"),\n",
    "            }\n",
    "\n",
    "\n",
    "df = pd.DataFrame(reader(path_database / \"arxiv.json\"))\n",
    "df.set_index(\"id\", inplace=True)\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "355c3539-347f-4b74-92b9-74957fb42c80",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "occurences = Counter(\n",
    "    s for s in chain.from_iterable(df.categories) if s.startswith(\"cs\")\n",
    ")\n",
    "occurences.update(CATEGORIES.keys())\n",
    "assert set(occurences.keys()) == set(CATEGORIES.keys())\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(10, 10))\n",
    "numbers = list(occurences.keys())\n",
    "counts = list(occurences.values())\n",
    "ax.barh(numbers, counts, height=0.8)\n",
    "ax.set_ylabel(\"Number\")\n",
    "ax.set_xlabel(\"Frequency\")\n",
    "ax.set_xscale(\"log\")\n",
    "ax.set_title(\"Number of papers\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6dde45e4-1305-4cb3-99d2-5636fa31d5f6",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "df_single_cat = df[df.categories.apply(len) == 1].categories.apply(lambda x: x[0])\n",
    "occurences = Counter(s for s in df_single_cat if s.startswith(\"cs\"))\n",
    "occurences.update(CATEGORIES.keys())\n",
    "assert set(occurences.keys()) == set(CATEGORIES.keys())\n",
    "# The category cs.IT always come with math.IT\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(10, 10))\n",
    "numbers = list(occurences.keys())\n",
    "counts = list(occurences.values())\n",
    "ax.barh(numbers, counts, height=0.8)\n",
    "ax.set_ylabel(\"Number\")\n",
    "ax.set_xlabel(\"Frequency\")\n",
    "ax.set_xscale(\"log\")\n",
    "ax.set_title(\"Number of paper with single category\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb351d6b-3d11-4006-8769-f3af7ec5eb1a",
   "metadata": {},
   "source": [
    "We create three splits:\n",
    "\n",
    "- `trainset` used for \"training\" the pipeline\n",
    "- `valset` used to evaluate the performace during training\n",
    "- `testset` used to evaluate the preformace after training "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30584475-5761-4430-bbcc-eaca1aabf953",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# splits\n",
    "trainset, valset, testset = set(), set(), set()\n",
    "\n",
    "# Single category\n",
    "for cat in CATEGORIES:\n",
    "    if cat == \"cs.IT\":  # The category cs.IT always come with math.IT\n",
    "        df_cat = df[df.categories.apply(lambda x: x == [\"cs.IT\", \"math.IT\"])]\n",
    "    else:\n",
    "        df_cat = df_single_cat[df_single_cat == cat]\n",
    "    sample = df_cat.sample(n=3, random_state=1).index\n",
    "    trainset.add(sample[0])\n",
    "    valset.add(sample[1])\n",
    "    testset.add(sample[2])\n",
    "\n",
    "# Multiple categories: add multi-categories paper to reach 50 papers in each split\n",
    "# random_state 1 sample pdf that are not withdrawn\n",
    "num = 50 - len(CATEGORIES)\n",
    "sample = df[df.categories.apply(len) > 2].sample(n=3 * num, random_state=1).index\n",
    "trainset |= set(sample[:num])\n",
    "valset |= set(sample[num : num * 2])\n",
    "testset |= set(sample[num * 2 :])\n",
    "\n",
    "dataset = trainset | valset | testset\n",
    "\n",
    "# ensure no overlapping\n",
    "assert not (trainset & valset)\n",
    "assert not (trainset & testset)\n",
    "assert not (valset & testset)\n",
    "assert len(trainset) == len(valset) == len(testset)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ba397a1-b64a-4903-a022-0f4c917ff498",
   "metadata": {},
   "source": [
    "Now, we want to download the PDFs of the selected papers (50 per split) and extract the full text body. \n",
    "For this task, we utilize [arxiv2text](https://github.com/dsdanielpark/arxiv2text).\n",
    "\n",
    "> In the simple example of the pipeline proposed in *features.ipynb*, we only utilize the title and abstract of the paper, without using the full text body. However, we have decided to include the full PDF text in the dataset to facilitate future experimentation with more complex pipelines that can process documents in chunks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "688f5c98-3702-47e5-9a3d-d25a1f957353",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "(path_dataset / \"trainset\").mkdir(exist_ok=True)\n",
    "(path_dataset / \"valset\").mkdir(exist_ok=True)\n",
    "(path_dataset / \"testset\").mkdir(exist_ok=True)\n",
    "\n",
    "i = 0\n",
    "with open(path_database / \"arxiv.json\") as file:\n",
    "    for line in file:\n",
    "        paper = json.loads(line.strip())\n",
    "        if paper[\"id\"] in dataset:\n",
    "            if paper[\"id\"] in trainset:\n",
    "                path_paper = path_dataset / \"trainset\"\n",
    "            if paper[\"id\"] in valset:\n",
    "                path_paper = path_dataset / \"valset\"\n",
    "            if paper[\"id\"] in testset:\n",
    "                path_paper = path_dataset / \"testset\"\n",
    "\n",
    "            i += 1\n",
    "            path_paper = path_paper / f'{paper[\"id\"].replace(\"/\", \"-\")}.json'\n",
    "            url = f\"https://arxiv.org/pdf/{paper['id']}.pdf\"\n",
    "            print(f\"[{i:>3}/{len(dataset)}] Processing {path_paper.stem}\")\n",
    "            if not path_paper.exists():\n",
    "                paper[\"text\"] = arxiv_to_text(url)\n",
    "                with open(path_paper, \"w\") as outfile:\n",
    "                    json.dump(paper, outfile, indent=4)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dspy-arxiv",
   "language": "python",
   "name": "dspy-arxiv"
  },
  "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}