Claim_Extraction_Training.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/biodatlab/score-claim-extraction/blob/main/Claim_Extraction_Training.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Claim Extraction Training Notebook\n",
        "This notebook preprocesses the datasets and trains the models followed by performing evaluations."
      ],
      "metadata": {
        "id": "1Pp0ygOAg2vO"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Setting Up"
      ],
      "metadata": {
        "id": "hZv-zSIDg9Fp"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "owCIugN9vov8"
      },
      "outputs": [],
      "source": [
        "# Will need to restart runtime after running this cell as numpy version is changed\n",
        "!pip install transformers==4.28.0\n",
        "!pip install -U sentence-transformers\n",
        "!pip install datasets\n",
        "!pip install evaluate\n",
        "!pip install accelerate -U\n",
        "!pip install -U tensorflow-text\n",
        "!pip install tf-models-official\n",
        "!pip install -U spacy\n",
        "!python -m spacy download en_core_web_sm\n",
        "!pip install fuzzywuzzy\n",
        "!pip install python-Levenshtein\n",
        "!pip install tabulate\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "t2F-ZzQ60ofZ"
      },
      "outputs": [],
      "source": [
        "# Log into hugging face with your token\n",
        "# from huggingface_hub import notebook_login\n",
        "# notebook_login()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "U_JI1QO70z47"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import shutil\n",
        "\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "import tensorflow as tf\n",
        "import tensorflow_hub as hub\n",
        "import tensorflow_text as text\n",
        "import matplotlib.pyplot as plt\n",
        "import spacy\n",
        "from official.nlp import optimization\n",
        "\n",
        "np.random.seed(54)\n",
        "nlp = spacy.load(\"en_core_web_sm\")\n",
        "tf.get_logger().setLevel('ERROR')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IWJfInq0UXP-"
      },
      "outputs": [],
      "source": [
        "import locale\n",
        "locale.getpreferredencoding = lambda: \"UTF-8\""
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Type 1 Preprocessing\n",
        "\n",
        "This part preprocesses Type 1 Data. It is designed for a Pandas Dataframe.\n",
        "Essentially the dataframe needs 2 columns - 'abstract' and 'claim'\n",
        "\n",
        "1. abstract (string) : This is the original abstract from the paper\n",
        "2. claim (string) : Human-coded claim for the abstract. We'll fuzzy compare this to sentences in the abstract to find the claim\n",
        "\n",
        "If you do not have human-coded data, you can skip this part"
      ],
      "metadata": {
        "id": "AJwsOc3rhDxC"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QgT_Rv6T14HN"
      },
      "outputs": [],
      "source": [
        "# Change 'train_data.csv' to the name of your csv file\n",
        "data = pd.read_csv('train_data.csv')\n",
        "data = data[['abstract','claim']]\n",
        "data.head()"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Fuzzy comparision to match sentences to human-coded claim"
      ],
      "metadata": {
        "id": "Jo05GVDRhQVy"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from fuzzywuzzy import fuzz\n",
        "\n",
        "datasets = data.to_dict(\"records\")\n",
        "type1_data = []\n",
        "for r in datasets:\n",
        "  annotations = []\n",
        "  for sent in nlp(str(r[\"abstract\"])).sents:\n",
        "    label = int(fuzz.ratio(sent.text, r[\"claim\"]) > 60)\n",
        "    annotations.append({\n",
        "      \"text\": sent.text,\n",
        "      \"label\": label,\n",
        "    })\n",
        "  type1_data.append({\n",
        "      \"abstract\": r[\"abstract\"],\n",
        "      \"claim\": r[\"claim\"],\n",
        "      \"annotations\": annotations,\n",
        "  })"
      ],
      "metadata": {
        "id": "GFRj1C1E13FN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Train-Test Split"
      ],
      "metadata": {
        "id": "xNrfuW_1hSOo"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from itertools import chain\n",
        "from sklearn.model_selection import train_test_split\n",
        "\n",
        "abstracts = [row['abstract'] for row in type1_data]\n",
        "train_ids, val_ids = train_test_split(abstracts, test_size=0.2, random_state=54)\n",
        "val_ids, test_ids = train_test_split(abstracts, test_size=0.5, random_state=54)\n",
        "type1_train = list(chain.from_iterable([r[\"annotations\"] for r in type1_data if r[\"abstract\"] in train_ids]))\n",
        "type1_val = list(chain.from_iterable([r[\"annotations\"] for r in type1_data if r[\"abstract\"] in val_ids]))\n",
        "type1_test = list(chain.from_iterable([r[\"annotations\"] for r in type1_data if r[\"abstract\"] in test_ids]))"
      ],
      "metadata": {
        "id": "kR2Zv2jR5R2P"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QSzTxliQFFnP"
      },
      "source": [
        "### Create datasets.Dataset() object and exporting"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "M3eL7OXhFHIZ"
      },
      "outputs": [],
      "source": [
        "from datasets import load_dataset, Dataset\n",
        "\n",
        "type1_train = Dataset.from_list(type1_train)\n",
        "type1_val = Dataset.from_list(type1_val)\n",
        "type1_test = Dataset.from_list(type1_test)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Hf1aU2fyDpso"
      },
      "source": [
        "## Type 2 Data Loading\n",
        "\n",
        "If your data is already preprocessed you can use this part of the code. It is designed for datasets.dataset() objects stored in csv files. Each file should have 2 columns - 'text' and 'label', in one split 'train'\n",
        "\n",
        "- text (string) : A sentence from the abstract.\n",
        "- label (int) : 1 if the text is a claim, 0 otherwise\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "W-UEmObRL7JX"
      },
      "outputs": [],
      "source": [
        "type2_train = load_dataset('csv', data_files='type2_train.csv')\n",
        "type2_val = load_dataset('csv', data_files='type2_val.csv')\n",
        "type2_test = load_dataset('csv', data_files='type2_test.csv')"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Merging and Shuffling Datasets"
      ],
      "metadata": {
        "id": "Lo9PFNIqhlxt"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oO_TExjgGOd5"
      },
      "outputs": [],
      "source": [
        "##### Uncomment to concatenate type1 and type2 data\n",
        "# from datasets import concatenate_datasets\n",
        "\n",
        "# comb_train = concatenate_datasets([type1_train, type2_train['train']])\n",
        "# comb_val = concatenate_datasets([type1_val, type2_val['train']])\n",
        "# comb_test = concatenate_datasets([type1_test, type2_test['train']])\n",
        "\n",
        "# If you are concatenating, comment these next 3 lines out\n",
        "comb_train = type2_train\n",
        "comb_val = type2_val\n",
        "comb_test = type2_test"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "h_Ao83hHH2mD"
      },
      "outputs": [],
      "source": [
        "comb_train = comb_train.shuffle(seed=54)\n",
        "comb_val = comb_val.shuffle(seed=54)\n",
        "comb_test = comb_test.shuffle(seed=54)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Model"
      ],
      "metadata": {
        "id": "kfpS3UldhvGH"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Model Initiation"
      ],
      "metadata": {
        "id": "63v83ztAvKkw"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yhiaVTlGPaVw"
      },
      "outputs": [],
      "source": [
        "from transformers import AutoTokenizer\n",
        "\n",
        "tokenizer = AutoTokenizer.from_pretrained(\"allenai/scibert_scivocab_uncased\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lFqMoKCXPtJO"
      },
      "outputs": [],
      "source": [
        "from transformers import DataCollatorWithPadding\n",
        "\n",
        "def preprocessor(batch):\n",
        "    return tokenizer(batch['text'], truncation=True)\n",
        "\n",
        "\n",
        "def get_collator(tokenizer):\n",
        "  data_collator = DataCollatorWithPadding(tokenizer=tokenizer)\n",
        "  return data_collator"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "j0Go07PbP0pO"
      },
      "outputs": [],
      "source": [
        "import evaluate\n",
        "\n",
        "clf_metrics = evaluate.combine([\"accuracy\", \"f1\", \"precision\", \"recall\"])\n",
        "\n",
        "def compute_metrics(eval_pred):\n",
        "  preds, labels = eval_pred\n",
        "  preds = np.argmax(preds,axis=1)\n",
        "  return clf_metrics.compute(predictions=preds,references=labels)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Select Hyperparameters\n",
        "\n",
        "learning_rates = [2e-05,3e-05] #@param {type:'raw'}\n",
        "num_epochs = 6 #@param {type:'integer'}\n",
        "batch_size = 32 #@param {type:'integer'}\n",
        "weight_decay=0.01 #@param {type:'number'}"
      ],
      "metadata": {
        "id": "qqNVastgrqzD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Fine-tuning"
      ],
      "metadata": {
        "id": "gTSl2iwZhzGc"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "s7Oiu_uGQI9F"
      },
      "outputs": [],
      "source": [
        "from transformers import TrainingArguments,Trainer\n",
        "from transformers import BertForSequenceClassification\n",
        "\n",
        "id2label = {0: \"Null\", 1: \"Claim\"}\n",
        "label2id = {\"Null\": 0, \"Claim\": 1}\n",
        "\n",
        "\n",
        "for learning_rate in learning_rates:\n",
        "  model = BertForSequenceClassification.from_pretrained(\n",
        "      \"allenai/scibert_scivocab_uncased\",\n",
        "      num_labels=2,\n",
        "      id2label=id2label,\n",
        "      label2id=label2id,\n",
        "      output_attentions=False,\n",
        "      output_hidden_states=False,\n",
        "  )\n",
        "  \n",
        "  model.cuda()\n",
        "\n",
        "  model_name = f\"scibert_claim_id_{learning_rate}\"\n",
        "  tokenizer = tokenizer\n",
        "  preprocessor = preprocessor\n",
        "  tokenized_train = comb_train.map(preprocessor, batched=True)\n",
        "  tokenized_val = comb_val.map(preprocessor,batched=True)\n",
        "  collator = get_collator(tokenizer)\n",
        "\n",
        "  training_args = TrainingArguments(\n",
        "      output_dir=model_name,\n",
        "      learning_rate=learning_rate,\n",
        "      per_device_train_batch_size=batch_size,\n",
        "      per_device_eval_batch_size=batch_size,\n",
        "      num_train_epochs=num_epochs,\n",
        "      weight_decay=weight_decay,\n",
        "      evaluation_strategy=\"epoch\",\n",
        "      save_strategy=\"epoch\",\n",
        "      load_best_model_at_end=True,\n",
        "      push_to_hub=False,\n",
        "  )\n",
        "\n",
        "  trainer = Trainer(\n",
        "      model=model,\n",
        "      args=training_args,\n",
        "      train_dataset=tokenized_train,\n",
        "      eval_dataset=tokenized_val,\n",
        "      tokenizer=tokenizer,\n",
        "      data_collator=collator,\n",
        "      compute_metrics=compute_metrics,\n",
        "  )\n",
        "\n",
        "  print(f'|--------------------------Now Training: {model_name} with Learning Rate = {learning_rate}------------------------------|')\n",
        "  trainer.train()\n",
        "  # trainer.push_to_hub()\n",
        "  print(f'|-----------------------------------------------------------------------------------------------------------------------|')"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Evaluation"
      ],
      "metadata": {
        "id": "WAg6OaQqieYB"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from evaluate import evaluator\n",
        "\n",
        "task_evaluator = evaluator(\"text-classification\")"
      ],
      "metadata": {
        "id": "F8rG59S3ifsp"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Change the next line to the model names that you trained\n",
        "models = [\"scibert_claim_id_2e-05\",\"scibert_claim_id_3e_05\"]\n",
        "\n",
        "dataset = comb_test\n",
        "evalres = []\n",
        "\n",
        "for model_name in models:\n",
        "  model = BertForSequenceClassification.from_pretrained(f\"/content/{model_name}/\")\n",
        "  tokenizer = tokenizer\n",
        "  preprocessor = preprocessor\n",
        "  collator = get_collator(tokenizer)\n",
        "  eval_results = task_evaluator.compute(\n",
        "    model_or_pipeline=model,\n",
        "    tokenizer=tokenizer,\n",
        "    data=dataset,\n",
        "    metric=evaluate.combine([\"accuracy\", \"recall\", \"precision\", \"f1\"]),\n",
        "    label_mapping=label2id,\n",
        "    strategy=\"simple\"\n",
        "  )\n",
        "  evalres.append([model_name,data[0],eval_results['accuracy'],eval_results['recall'],eval_results['precision'],eval_results['f1']])"
      ],
      "metadata": {
        "id": "vLAW8rVwinpM"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from tabulate import tabulate\n",
        "print(tabulate(evalres,headers = ['Model Name','Dataset Name','Accuracy','Precision','Recall','F1 Score'],tablefmt='github'))"
      ],
      "metadata": {
        "id": "E0iRrpdQmppV"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "include_colab_link": true
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}