notebooks/commonlit_readability.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "c5e4b774-8bee-4e3b-9de7-70d4675dced7",
      "metadata": {},
      "outputs": [],
      "source": [
        "# https://www.kaggle.com/competitions/commonlitreadabilityprize\n",
        "# predict ease of readability of passage"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "d3b6f9f6-b5e4-43c1-93ab-0f3a25ac7ab7",
      "metadata": {},
      "outputs": [],
      "source": [
        "import typing as t\n",
        "from ast import literal_eval\n",
        "\n",
        "from transformer.models.regressor import RegressorLM\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",
        "import matplotlib.pyplot as plt\n",
        "from sklearn.preprocessing import MinMaxScaler\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": "b1b7d8a0",
      "metadata": {},
      "outputs": [],
      "source": [
        "# load and preview data\n",
        "data = pd.read_csv(\"data/commonlit.csv\")\n",
        "# exclude:\n",
        "# - excerpts outside [-3, 1] as these are unreliable\n",
        "# - single outlier with zero easiness rating standard error\n",
        "data = data.loc[\n",
        "    data[\"BT Easiness\"].between(-3, 1) & (data[\"BT s.e.\"] > 0), \n",
        "    [\"Excerpt\", \"BT Easiness\", \"BT s.e.\"]\n",
        "]\n",
        "# convert newlines to spaces\n",
        "data.Excerpt = data.Excerpt.str.replace(\"\\n\", \" \")\n",
        "data.head()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "52993746",
      "metadata": {},
      "outputs": [],
      "source": [
        "fig, axs = plt.subplots(1, 2, figsize=(10, 3))\n",
        "data[\"BT Easiness\"].plot.hist(ax=axs[0], title=\"Easiness Rating\")\n",
        "data[\"BT s.e.\"].plot.hist(ax=axs[1], title=\"Std. Error\")\n",
        "plt.tight_layout()\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "17d32078",
      "metadata": {},
      "outputs": [],
      "source": [
        "# convert text to list of strings\n",
        "X = data[\"Excerpt\"].to_list()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ef49e7cc",
      "metadata": {},
      "outputs": [],
      "source": [
        "# use min-max scaled readability score as output\n",
        "scaler = MinMaxScaler()\n",
        "scaler.fit(np.array([[-3.0], [1.0]]))\n",
        "y = torch.from_numpy(scaler.transform(data[[\"BT Easiness\"]])).float()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "7e754c84",
      "metadata": {},
      "outputs": [],
      "source": [
        "# use reciprocal of standard error as loss function sample weights\n",
        "weights = torch.from_numpy(data[[\"BT s.e.\"]].pow(-1).to_numpy()).float()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "cafb791f",
      "metadata": {},
      "outputs": [],
      "source": [
        "# create data module\n",
        "class CommonlitReadabilityDataModule(InferenceDataModule):\n",
        "    def setup(self: t.Self, stage: str) -> None:\n",
        "        self.X, self.y, self.weights = X, y, weights\n",
        "        super().setup(stage=stage)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "34547a89",
      "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": "ab6a019e",
      "metadata": {},
      "outputs": [],
      "source": [
        "# view token sequence length distribution\n",
        "data[\"Excerpt\"].apply(tokenizer.tokenize).str.len().plot.hist(bins=50)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "3c85fd40",
      "metadata": {},
      "outputs": [],
      "source": [
        "# initialize the transformer\n",
        "context_length = 300\n",
        "model = RegressorLM(\n",
        "    params=TransformerParams(context_length=context_length),\n",
        "    tokenizer=tokenizer,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "5ecd7354",
      "metadata": {},
      "outputs": [],
      "source": [
        "# tokenize & encode data and prepare train/test splits\n",
        "datamodule = CommonlitReadabilityDataModule(\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": "c0a662a1",
      "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": "f031ef53",
      "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": "62ec61b1",
      "metadata": {},
      "outputs": [],
      "source": [
        "# calculate accuracy\n",
        "torch.tensor([x[1] == x[2] for batch in pred for x in batch]).float().mean()"
      ]
    }
  ],
  "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
}