UnderwriteMe_NER_ModelTraining.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "06z1Vx68EKVX"
      },
      "outputs": [],
      "source": [
        "!pip install datasets\n",
        "!pip install accelerate -U\n",
        "!pip install transformers[torch] -U\n",
        "!pip install seqeval"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zs-aC9WsDJ3G"
      },
      "outputs": [],
      "source": [
        "import utils\n",
        "from utils import compute_metrics\n",
        "import datasets\n",
        "from datasets import load_from_disk\n",
        "from transformers import DistilBertTokenizerFast, DataCollatorForTokenClassification, AutoModelForTokenClassification\n",
        "from transformers import TrainingArguments, Trainer\n",
        "import json\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DC-yXh_rDU3L"
      },
      "outputs": [],
      "source": [
        "# Defining model\n",
        "tokenizer = DistilBertTokenizerFast.from_pretrained(\"distilbert-base-cased\")\n",
        "model = AutoModelForTokenClassification.from_pretrained(\"distilbert-base-cased\", num_labels=25)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Oy3cTo89EXGG",
        "outputId": "5b718956-446a-4a8e-a50a-efec12bb9f11"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Dataset({\n",
            "    features: ['id', 'tokens', 'ner_tags', 'input_ids', 'attention_mask', 'labels'],\n",
            "    num_rows: 27454\n",
            "})\n",
            "Dataset({\n",
            "    features: ['id', 'tokens', 'ner_tags', 'input_ids', 'attention_mask', 'labels'],\n",
            "    num_rows: 6223\n",
            "})\n"
          ]
        }
      ],
      "source": [
        "train_dataset = load_from_disk('/content/drive/MyDrive/Colab Notebooks/underwriteme-data/combined_train_dataset')\n",
        "print(train_dataset)\n",
        "test_dataset = load_from_disk('/content/drive/MyDrive/Colab Notebooks/underwriteme-data/test_dataset')\n",
        "print(test_dataset)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Define the mapping from NER label to integer\n",
        "label_to_id = {\n",
        "    'O': 0,\n",
        "    'B-Investigation': 1,\n",
        "    'I-Investigation': 2,\n",
        "    'B-Treatment': 3,\n",
        "    'I-Treatment': 4,\n",
        "    'B-Condition': 5,\n",
        "    'I-Condition': 6,\n",
        "    'B-Behaviour': 7,\n",
        "    'I-Behaviour': 8,\n",
        "    'B-Date': 9,\n",
        "    'I-Date': 10,\n",
        "    'B-Unit': 11,\n",
        "    'I-Unit': 12,\n",
        "    'B-Sign': 13,\n",
        "    'I-Sign': 14,\n",
        "    'B-Symptom': 15,\n",
        "    'I-Symptom': 16,\n",
        "    'B-Drug': 17,\n",
        "    'I-Drug': 18,\n",
        "    'B-LabResult': 19,\n",
        "    'I-LabResult': 20,\n",
        "    'B-HealthcareProvider': 21,\n",
        "    'I-HealthcareProvider': 22,\n",
        "    'B-LabTest': 23,\n",
        "    'I-LabTest': 24,\n",
        "}\n",
        "\n",
        "id_to_label = {id: label for label, id in label_to_id.items()}"
      ],
      "metadata": {
        "id": "Qk22DbF8pcDQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d83_EPQpFLO1",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "498aca7b-5f45-4cee-c527-1c16095fd69e"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "<ipython-input-5-70a41a89b401>:14: FutureWarning: load_metric is deprecated and will be removed in the next major version of datasets. Use 'evaluate.load' instead, from the new library 🤗 Evaluate: https://huggingface.co/docs/evaluate\n",
            "  metric = datasets.load_metric(\"seqeval\")\n",
            "/usr/local/lib/python3.10/dist-packages/datasets/load.py:756: FutureWarning: The repository for seqeval contains custom code which must be executed to correctly load the metric. You can inspect the repository content at https://raw.githubusercontent.com/huggingface/datasets/2.18.0/metrics/seqeval/seqeval.py\n",
            "You can avoid this message in future by passing the argument `trust_remote_code=True`.\n",
            "Passing `trust_remote_code=True` will be mandatory to load this metric from the next major release of `datasets`.\n",
            "  warnings.warn(\n",
            "/usr/local/lib/python3.10/dist-packages/accelerate/accelerator.py:436: FutureWarning: Passing the following arguments to `Accelerator` is deprecated and will be removed in version 1.0 of Accelerate: dict_keys(['dispatch_batches', 'split_batches', 'even_batches', 'use_seedable_sampler']). Please pass an `accelerate.DataLoaderConfiguration` instead: \n",
            "dataloader_config = DataLoaderConfiguration(dispatch_batches=None, split_batches=False, even_batches=True, use_seedable_sampler=True)\n",
            "  warnings.warn(\n"
          ]
        }
      ],
      "source": [
        "# Define training args\n",
        "\n",
        "args = TrainingArguments(\n",
        "\"test-ner2\",\n",
        "evaluation_strategy = \"epoch\",\n",
        "learning_rate=2e-5,\n",
        "per_device_train_batch_size=16,\n",
        "per_device_eval_batch_size=16,\n",
        "num_train_epochs=3,\n",
        "weight_decay=0.01,\n",
        ")\n",
        "\n",
        "data_collator = DataCollatorForTokenClassification(tokenizer)\n",
        "metric = datasets.load_metric(\"seqeval\")\n",
        "\n",
        "trainer = Trainer(\n",
        "   model,\n",
        "   args,\n",
        "   train_dataset=train_dataset,\n",
        "   eval_dataset=test_dataset,\n",
        "   data_collator=data_collator,\n",
        "   tokenizer=tokenizer,\n",
        "   compute_metrics=compute_metrics\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "trainer.train()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 471
        },
        "id": "dqXmbG50_nrJ",
        "outputId": "753f5f24-9eae-4e82-f44b-d433bdc7442f"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "\n",
              "    <div>\n",
              "      \n",
              "      <progress value='4290' max='4290' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
              "      [4290/4290 11:34, Epoch 5/5]\n",
              "    </div>\n",
              "    <table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              " <tr style=\"text-align: left;\">\n",
              "      <th>Epoch</th>\n",
              "      <th>Training Loss</th>\n",
              "      <th>Validation Loss</th>\n",
              "      <th>Precision</th>\n",
              "      <th>Recall</th>\n",
              "      <th>F1</th>\n",
              "      <th>Accuracy</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <td>1</td>\n",
              "      <td>2.522200</td>\n",
              "      <td>0.981448</td>\n",
              "      <td>0.541667</td>\n",
              "      <td>0.657444</td>\n",
              "      <td>0.593966</td>\n",
              "      <td>0.758797</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>2</td>\n",
              "      <td>0.659400</td>\n",
              "      <td>0.287756</td>\n",
              "      <td>0.752183</td>\n",
              "      <td>0.840118</td>\n",
              "      <td>0.793722</td>\n",
              "      <td>0.945016</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>3</td>\n",
              "      <td>0.232800</td>\n",
              "      <td>0.114534</td>\n",
              "      <td>0.912801</td>\n",
              "      <td>0.948407</td>\n",
              "      <td>0.930264</td>\n",
              "      <td>0.982460</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>4</td>\n",
              "      <td>0.160600</td>\n",
              "      <td>0.066146</td>\n",
              "      <td>0.974733</td>\n",
              "      <td>0.984239</td>\n",
              "      <td>0.979463</td>\n",
              "      <td>0.991464</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>5</td>\n",
              "      <td>0.106900</td>\n",
              "      <td>0.055316</td>\n",
              "      <td>0.978960</td>\n",
              "      <td>0.986620</td>\n",
              "      <td>0.982775</td>\n",
              "      <td>0.992379</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table><p>"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/seqeval/metrics/v1.py:57: UndefinedMetricWarning: Precision and F-score are ill-defined and being set to 0.0 in labels with no predicted samples. Use `zero_division` parameter to control this behavior.\n",
            "  _warn_prf(average, modifier, msg_start, len(result))\n",
            "/usr/local/lib/python3.10/dist-packages/seqeval/metrics/v1.py:57: UndefinedMetricWarning: Precision and F-score are ill-defined and being set to 0.0 in labels with no predicted samples. Use `zero_division` parameter to control this behavior.\n",
            "  _warn_prf(average, modifier, msg_start, len(result))\n",
            "/usr/local/lib/python3.10/dist-packages/seqeval/metrics/v1.py:57: UndefinedMetricWarning: Precision and F-score are ill-defined and being set to 0.0 in labels with no predicted samples. Use `zero_division` parameter to control this behavior.\n",
            "  _warn_prf(average, modifier, msg_start, len(result))\n",
            "/usr/local/lib/python3.10/dist-packages/seqeval/metrics/v1.py:57: UndefinedMetricWarning: Precision and F-score are ill-defined and being set to 0.0 in labels with no predicted samples. Use `zero_division` parameter to control this behavior.\n",
            "  _warn_prf(average, modifier, msg_start, len(result))\n",
            "/usr/local/lib/python3.10/dist-packages/seqeval/metrics/v1.py:57: UndefinedMetricWarning: Precision and F-score are ill-defined and being set to 0.0 in labels with no predicted samples. Use `zero_division` parameter to control this behavior.\n",
            "  _warn_prf(average, modifier, msg_start, len(result))\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "TrainOutput(global_step=4290, training_loss=0.6399908239191229, metrics={'train_runtime': 694.8297, 'train_samples_per_second': 197.559, 'train_steps_per_second': 6.174, 'total_flos': 972931510765212.0, 'train_loss': 0.6399908239191229, 'epoch': 5.0})"
            ]
          },
          "metadata": {},
          "execution_count": 202
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "## Save model\n",
        "model.save_pretrained(\"ner_model2\")\n",
        "\n",
        "## Save tokenizer\n",
        "tokenizer.save_pretrained(\"ner_tokenizer2\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "cB_li0nWcuQi",
        "outputId": "ed6a663f-928a-414d-dfd0-2db96022644c"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "('ner_tokenizer_final/tokenizer_config.json',\n",
              " 'ner_tokenizer_final/special_tokens_map.json',\n",
              " 'ner_tokenizer_final/vocab.txt',\n",
              " 'ner_tokenizer_final/added_tokens.json',\n",
              " 'ner_tokenizer_final/tokenizer.json')"
            ]
          },
          "metadata": {},
          "execution_count": 203
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Write label-id into config file\n",
        "config = json.load(open(\"ner_model2/config.json\"))\n",
        "config[\"id2label\"] = id_to_label\n",
        "config[\"label2id\"] = label_to_id\n",
        "json.dump(config, open(\"ner_model2/config.json\",\"w\"))"
      ],
      "metadata": {
        "id": "tc9VOr9xdQ6X"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "RS0wBFyKYcJH"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}