notebooks/arxiv_categorization.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "48d6bd46-1b31-44bc-b35f-06d1fbaf9caa",
      "metadata": {},
      "outputs": [],
      "source": [
        "# https://www.kaggle.com/datasets/spsayakpaul/arxiv-paper-abstracts/data\n",
        "# predict category from title/abstract"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ab12b65b",
      "metadata": {},
      "outputs": [],
      "source": [
        "import typing as t\n",
        "from ast import literal_eval\n",
        "\n",
        "from transformer.models.classifier import ClassifierLM\n",
        "from transformer.dataloaders.inference import InferenceDataModule\n",
        "from transformer.params.transformer import TransformerParams\n",
        "\n",
        "import torch\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "from sklearn.preprocessing import LabelEncoder\n",
        "from lightning import Trainer\n",
        "from lightning.pytorch.callbacks.early_stopping import EarlyStopping\n",
        "from transformers import LlamaTokenizer"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "7e78e563",
      "metadata": {},
      "outputs": [],
      "source": [
        "# load and preview data\n",
        "data = pd.read_csv(\"data/arxiv.csv\")\n",
        "data.titles = data.titles.str.replace(\"\\n\", \" \")\n",
        "data.abstracts = data.abstracts.str.replace(\"\\n\", \" \")\n",
        "data.tail()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "535a4465",
      "metadata": {},
      "outputs": [],
      "source": [
        "# get titles and primary category\n",
        "X = data.titles.to_list()\n",
        "y = data.terms.apply(literal_eval).str[0]\n",
        "print(y.value_counts())\n",
        "y = y.to_numpy()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "69825956",
      "metadata": {},
      "outputs": [],
      "source": [
        "# encode categories\n",
        "label_encoder = LabelEncoder()\n",
        "y = torch.from_numpy(label_encoder.fit_transform(y))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "d6594f87",
      "metadata": {},
      "outputs": [],
      "source": [
        "# create data module\n",
        "class ArxivCategorizationDataModule(InferenceDataModule):\n",
        "    def setup(self: t.Self, stage: str) -> None:\n",
        "        self.X, self.y = X, y\n",
        "        super().setup(stage=stage)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "92febe62",
      "metadata": {},
      "outputs": [],
      "source": [
        "# initialize pretrained tokenizer\n",
        "# - llama does not add an EOS token by default, so override this\n",
        "# - llama also does not use a padding token, so this needs to be added\n",
        "tokenizer = LlamaTokenizer.from_pretrained(\n",
        "    \"huggyllama/llama-7b\", add_eos_token=True, legacy=False\n",
        ")\n",
        "tokenizer.add_special_tokens({\"pad_token\": \"<pad>\"})"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "374b7145",
      "metadata": {},
      "outputs": [],
      "source": [
        "# initialize the transformer\n",
        "context_length = 64\n",
        "model = ClassifierLM(\n",
        "    params=TransformerParams(context_length=context_length),\n",
        "    tokenizer=tokenizer,\n",
        "    num_classes=len(y.unique())\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "b20b6f8a",
      "metadata": {},
      "outputs": [],
      "source": [
        "# tokenize & encode data and prepare train/test splits\n",
        "datamodule = ArxivCategorizationDataModule(\n",
        "    tokenizer=tokenizer,\n",
        "    context_length=context_length,\n",
        "    batch_size=32,\n",
        "    val_size=0.2,\n",
        "    test_size=0.1,\n",
        "    num_workers=9,\n",
        "    persistent_workers=True,\n",
        "    limit=None,\n",
        "    random_state=1,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "4fc4cb8f",
      "metadata": {},
      "outputs": [],
      "source": [
        "%%time\n",
        "# train the model\n",
        "trainer = Trainer(\n",
        "    max_epochs=50,\n",
        "    callbacks=EarlyStopping(monitor=\"val_loss\", mode=\"min\", patience=5),\n",
        "    accelerator=\"cpu\",\n",
        ")\n",
        "trainer.fit(model=model, datamodule=datamodule)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "fa4cad49",
      "metadata": {},
      "outputs": [],
      "source": [
        "# calculate test metrics\n",
        "trainer.test(model=model, datamodule=datamodule)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "d4b525af",
      "metadata": {},
      "outputs": [],
      "source": [
        "# view first batch of test set predictions\n",
        "pred = trainer.predict(model=model, datamodule=datamodule)\n",
        "pred[:10]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "da5b7de2",
      "metadata": {},
      "outputs": [],
      "source": [
        "# calculate accuracy\n",
        "torch.tensor([x[1] == x[2] for batch in pred for x in batch]).float().mean()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "3b34a2a8",
      "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.11.8"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}