Training_QA_models.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Training QA models.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "FqIWc-eDnXnB"
      },
      "source": [
        "%%capture\r\n",
        "!pip install transformers\r\n",
        "!pip install constants\r\n",
        "!pip install --upgrade torch\r\n",
        "!wget https://raw.githubusercontent.com/huggingface/transformers/09a2f40684f77e62d0fd8485fe9d2d610390453f/examples/multiple-choice/utils_multiple_choice.py"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "v5oN3BL8nmVC"
      },
      "source": [
        "import numpy as np\r\n",
        "import pandas as pd\r\n",
        "import tensorflow as tf\r\n",
        "import torch\r\n",
        "import tqdm\r\n",
        "import utils_multiple_choice\r\n",
        "from constants import *\r\n",
        "from google.colab import auth, drive\r\n",
        "from transformers import (\r\n",
        "    AutoConfig,\r\n",
        "    AutoModelForMultipleChoice,\r\n",
        "    AutoTokenizer,\r\n",
        "    EvalPrediction,\r\n",
        "    HfArgumentParser,\r\n",
        "    Trainer,\r\n",
        "    TrainingArguments,\r\n",
        "    TFAutoModelForMultipleChoice,\r\n",
        "    TFTrainer,\r\n",
        "    TFTrainingArguments,\r\n",
        "    set_seed,\r\n",
        ")\r\n",
        "from utils_multiple_choice import MultipleChoiceDataset, RaceProcessor, Split, TFMultipleChoiceDataset, processors"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8I7PT-VnnwGc",
        "outputId": "22d26fe1-1107-4fe9-df9a-5eea7a0a6a99"
      },
      "source": [
        "from google.colab import auth, drive\r\n",
        "drive.mount('/content/drive')"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Mounted at /content/drive\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XsHNyHHYnxIO"
      },
      "source": [
        "def simple_accuracy(preds, labels):\r\n",
        "    return (preds == labels).mean()\r\n",
        "\r\n",
        "def compute_metrics(p: EvalPrediction) -> Dict:\r\n",
        "    preds = np.argmax(p.predictions, axis=1)\r\n",
        "    return {\"acc\": simple_accuracy(preds, p.label_ids)}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Lfct6Uh9nzLs"
      },
      "source": [
        "SCORE_A = 'score A'\r\n",
        "SCORE_B = 'score B'\r\n",
        "SCORE_C = 'score C'\r\n",
        "SCORE_D = 'score D'\r\n",
        "SCORES = 'scores'\r\n",
        "CORRECT = 'correct'\r\n",
        "SCORE_LABEL = 'score_label'\r\n",
        "SCORES_WRONG = 'score_wrong'\r\n",
        "LABEL = 'label'\r\n",
        "PREDICTION = 'prediction'\r\n",
        "LEVEL = 'level'\r\n",
        "SCORE_VAR = 'score variance'\r\n",
        "LIST_SCORES = [SCORE_A, SCORE_B, SCORE_C, SCORE_D]\r\n",
        "MAX_SEQ_LENGTH = 512\r\n",
        "RANDOM_SEED = 3 # 0, 42, 1, 2, 3\r\n",
        "MODEL_NAME = 'xlnet-base-cased'\r\n",
        "DATA_DIR = \"data/\""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9EtFZggFnyI2"
      },
      "source": [
        "try:\r\n",
        "    processor = processors['race']()\r\n",
        "    label_list = processor.get_labels()\r\n",
        "    num_labels = len(label_list)\r\n",
        "except KeyError:\r\n",
        "    raise ValueError(\"Task not found: %s\" % ('race'))\r\n",
        "config = AutoConfig.from_pretrained(MODEL_NAME, num_labels=num_labels, finetuning_task='race')\r\n",
        "tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mK0lxVW4n2lb"
      },
      "source": [
        "lines = []\r\n",
        "for level in ['high', 'middle']:\r\n",
        "    input_dir = os.path.join(DATA_DIR, \"train/\" + level)\r\n",
        "    files = glob.glob(input_dir + \"/*txt\")\r\n",
        "    for file in tqdm.tqdm(files, desc=\"read files\"):\r\n",
        "        with open(file, \"r\", encoding=\"utf-8\") as fin:\r\n",
        "            data_raw = json.load(fin)\r\n",
        "            data_raw[\"race_id\"] = file\r\n",
        "            lines.append(data_raw)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yDJhfzVin4zg"
      },
      "source": [
        "set_type = 'train'\r\n",
        "examples = []\r\n",
        "for (_, data_raw) in enumerate(lines):\r\n",
        "    race_id = \"%s-%s\" % (set_type, data_raw[\"race_id\"])\r\n",
        "    article = data_raw[\"article\"]\r\n",
        "    for i in range(len(data_raw[\"answers\"])):\r\n",
        "        truth = str(ord(data_raw[\"answers\"][i]) - ord(\"A\"))\r\n",
        "        question = data_raw[\"questions\"][i]\r\n",
        "        options = data_raw[\"options\"][i]\r\n",
        "\r\n",
        "        examples.append(\r\n",
        "            utils_multiple_choice.InputExample(\r\n",
        "                example_id=race_id,\r\n",
        "                question=question,\r\n",
        "                contexts=[article, article, article, article],  # this is not efficient but convenient\r\n",
        "                endings=[options[0], options[1], options[2], options[3]],\r\n",
        "                label=truth,\r\n",
        "            )\r\n",
        "        )\r\n",
        "examples[0]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cE6JHJYUn8BI"
      },
      "source": [
        "# make sure to use modified utils_multiple_choice.py to allow examples=examples otherwise it always cut the training dataset, not loading it properly\r\n",
        "import importlib\r\n",
        "importlib.reload(utils_multiple_choice)\r\n",
        "from utils_multiple_choice import MultipleChoiceDataset"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fe1KodL5n8nK"
      },
      "source": [
        "train_dataset = MultipleChoiceDataset(\r\n",
        "        data_dir=DATA_DIR,\r\n",
        "        tokenizer=tokenizer,\r\n",
        "        task='race',\r\n",
        "        max_seq_length=MAX_SEQ_LENGTH,\r\n",
        "        overwrite_cache=True,\r\n",
        "        mode=Split.train, examples=examples)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0jw1FGkAn9kT"
      },
      "source": [
        "train_dataset = MultipleChoiceDataset(\r\n",
        "        data_dir=DATA_DIR,\r\n",
        "        tokenizer=tokenizer,\r\n",
        "        task='race',\r\n",
        "        max_seq_length=MAX_SEQ_LENGTH,\r\n",
        "        overwrite_cache=False,\r\n",
        "        mode=Split.train)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4ylgmSjpn-6D"
      },
      "source": [
        "eval_dataset = MultipleChoiceDataset(\r\n",
        "        data_dir=DATA_DIR,\r\n",
        "        tokenizer=tokenizer,\r\n",
        "        task='race',\r\n",
        "        max_seq_length=MAX_SEQ_LENGTH,\r\n",
        "        overwrite_cache=False,\r\n",
        "        mode=Split.dev)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "65t3bV89oA60"
      },
      "source": [
        "torch.manual_seed(RANDOM_SEED)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZY7FV1yJoB57"
      },
      "source": [
        "model = AutoModelForMultipleChoice.from_pretrained(MODEL_NAME, config=config)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "v6tI6uMEySCO"
      },
      "source": [
        "def init_training_args(\r\n",
        "    adam_epsilon=1e-8, \r\n",
        "    learning_rate=5e-5, \r\n",
        "    num_train_epochs=3.0, \r\n",
        "    weight_decay=0, \r\n",
        "    max_steps=-1,\r\n",
        "    output_dir='drive/My Drive/Colab Data/race_results',          # output directory\r\n",
        "    logging_dir='drive/My Drive/Colab Data/race_logs',            # directory for storing logs\r\n",
        "    ):\r\n",
        "  return TrainingArguments(\r\n",
        "    do_train=True,\r\n",
        "    do_eval=True,\r\n",
        "    evaluate_during_training=True,\r\n",
        "    output_dir=output_dir,          # output directory\r\n",
        "    logging_dir=logging_dir,            # directory for storing logs\r\n",
        "    save_steps=5000,\r\n",
        "    save_total_limit=5, \r\n",
        "    per_device_train_batch_size=4,  # batch size per device during training\r\n",
        "    per_device_eval_batch_size=4,   # batch size for evaluation\r\n",
        "    adam_epsilon=adam_epsilon,\r\n",
        "    learning_rate=learning_rate,\r\n",
        "    num_train_epochs=num_train_epochs,\r\n",
        "    max_steps=max_steps,\r\n",
        "    weight_decay=weight_decay\r\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ga9NgVsIV6es"
      },
      "source": [
        "configs = {\r\n",
        "    'xlnet-base-cased': {\r\n",
        "        'adam_epsilon': 1e-6,\r\n",
        "        'learning_rate': 2e-5,\r\n",
        "        'num_train_epochs': 2.0,\r\n",
        "        'weight_decay': 0.01\r\n",
        "    },\r\n",
        "    'bert-base-cased': {\r\n",
        "        'adam_epsilon': 1e-6,\r\n",
        "        'learning_rate': 2e-5,\r\n",
        "        'num_train_epochs': 2.0,\r\n",
        "        'weight_decay': 0.05\r\n",
        "    },\r\n",
        "    'distilber-base-cased': {\r\n",
        "        'adam_epsilon': 1e-8,\r\n",
        "        'learning_rate': 5e-5,\r\n",
        "        'num_train_epochs': 3,\r\n",
        "        'weight_decay': 0.01\r\n",
        "    }\r\n",
        "}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cv4JguBwyTbY"
      },
      "source": [
        "training_args = init_training_args(\r\n",
        "    adam_epsilon=configs[MODEL_NAME]['adam_epsilon'], \r\n",
        "    learning_rate=configs[MODEL_NAME]['learning_rate'], \r\n",
        "    num_train_epochs=configs[MODEL_NAME]['num_train_epochs'], \r\n",
        "    weight_decay=configs[MODEL_NAME]['weight_decay'],\r\n",
        "    logging_dir=os.path.join(DATA_DIR, 'race_logs_{}_seed_{}'.format(MODEL_NAME, RANDOM_SEED)),            # directory for storing logs\r\n",
        "    output_dir=os.path.join(DATA_DIR, 'race_results_{}_seed_{}'.format(MODEL_NAME, RANDOM_SEED)),          # output directory \r\n",
        "    )"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FVA4Wqu_oSCR"
      },
      "source": [
        "trainer = Trainer(\r\n",
        "    model=model,\r\n",
        "    args=training_args,\r\n",
        "    train_dataset=train_dataset,\r\n",
        "    eval_dataset=eval_dataset,\r\n",
        "    compute_metrics=compute_metrics,\r\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nsbfnIK5oaR2"
      },
      "source": [
        "trainer.train()\r\n",
        "trainer.save_model()\r\n",
        "tokenizer.save_pretrained(training_args.output_dir)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xutqUnmMob-6"
      },
      "source": [
        "trainer.evaluate()\r\n",
        "\r\n",
        "# Previous (for XLNet)\r\n",
        "#   {'eval_loss': 1.0538021816397252, 'eval_acc': 0.6355637405361162, 'step': 0}\r\n",
        "#   {'eval_acc': 0.6355637405361162, 'eval_loss': 1.0538021816397252}\r\n",
        "\r\n",
        "# seed = 0\r\n",
        "# {'eval_loss': 1.039906019572468, 'eval_acc': 0.6402701043585022, 'epoch': 2.0, 'total_flos': 233886300450963456, 'step': 43934}\r\n",
        "# {'epoch': 2.0,\r\n",
        "#  'eval_acc': 0.6402701043585022,\r\n",
        "#  'eval_loss': 1.039906019572468,\r\n",
        "#  'total_flos': 233886300450963456}\r\n",
        "\r\n",
        "# seed = 42\r\n",
        "# \r\n",
        "# {'eval_loss': 1.048840732447283, 'eval_acc': 0.6351544915080827, 'epoch': 2.0, 'total_flos': 233886300450963456, 'step': 43934}\r\n",
        "# {'epoch': 2.0,\r\n",
        "#  'eval_acc': 0.6351544915080827,\r\n",
        "#  'eval_loss': 1.048840732447283,\r\n",
        "#  'total_flos': 233886300450963456}\r\n",
        "\r\n",
        "# seed = 1\r\n",
        "# \r\n",
        "# {'eval_loss': 1.0707706642666928, 'eval_acc': 0.6159197871905054, 'epoch': 2.0, 'total_flos': 233886300450963456, 'step': 43934}\r\n",
        "# {'epoch': 2.0,\r\n",
        "#  'eval_acc': 0.6159197871905054,\r\n",
        "#  'eval_loss': 1.0707706642666928,\r\n",
        "#  'total_flos': 233886300450963456}\r\n",
        "\r\n",
        "# seed = 2\r\n",
        "# \r\n",
        "# {'eval_loss': 1.0715846247221732, 'eval_acc': 0.6044608144055658, 'epoch': 2.0, 'total_flos': 233886300450963456, 'step': 43934}\r\n",
        "# {'epoch': 2.0,\r\n",
        "#  'eval_acc': 0.6044608144055658,\r\n",
        "#  'eval_loss': 1.0715846247221732,\r\n",
        "#  'total_flos': 233886300450963456}\r\n",
        "\r\n",
        "# seed = 3\r\n",
        "# \r\n",
        "# {'eval_loss': 1.049155939741198, 'eval_acc': 0.6396562308164518, 'epoch': 2.0, 'total_flos': 233886300450963456, 'step': 43934}\r\n",
        "\r\n",
        "# {'epoch': 2.0,\r\n",
        "#  'eval_acc': 0.6396562308164518,\r\n",
        "#  'eval_loss': 1.049155939741198,\r\n",
        "#  'total_flos': 233886300450963456}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9DPT32N5odAn"
      },
      "source": [
        "test_dataset = MultipleChoiceDataset(\r\n",
        "        data_dir=DATA_DIR,\r\n",
        "        tokenizer=tokenizer,\r\n",
        "        task='race',\r\n",
        "        max_seq_length=MAX_SEQ_LENGTH,\r\n",
        "        overwrite_cache=True,\r\n",
        "        mode=Split.test)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sw-UTSKzoemh"
      },
      "source": [
        "def transform_pandas(trainer, dataset):\r\n",
        "    pred = trainer.predict(dataset)\r\n",
        "    vars = []\r\n",
        "    for x in pred.predictions:\r\n",
        "        vars.append(np.var(x))\r\n",
        "    print(np.mean(vars))\r\n",
        "    df = pd.DataFrame(columns=['idx', 'level', 'document_id', 'label', 'prediction', 'score A', 'score B', 'score C', 'score D', 'score variance'])\r\n",
        "    for idx in range((len(dataset))):\r\n",
        "        sample = [idx]\r\n",
        "        result = pred.predictions[idx]\r\n",
        "        df = df.append(\r\n",
        "            pd.Series(\r\n",
        "                [\r\n",
        "                    idx, \r\n",
        "                    dataset[idx].example_id.split(\"/\")[-2], \r\n",
        "                    dataset[idx].example_id.split(\"/\")[-1],\r\n",
        "                    dataset[idx].label,\r\n",
        "                    int(np.argmax(result)),\r\n",
        "                    result[0],\r\n",
        "                    result[1],\r\n",
        "                    result[2],\r\n",
        "                    result[3],\r\n",
        "                    np.std(result)\r\n",
        "                ],\r\n",
        "                index = ['idx', 'level', 'document_id', 'label', 'prediction', 'score A', 'score B', 'score C', 'score D', 'score variance']), \r\n",
        "                ignore_index=True\r\n",
        "            )\r\n",
        "    return df, pred"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JK-TySKQofGl"
      },
      "source": [
        "df_train, pred_train = transform_pandas(trainer, train_dataset)\r\n",
        "df_train.to_csv(os.path.join(DATA_DIR, 'output_{}_seed_{}_train.csv'.format(MODEL_NAME, RANDOM_SEED)), index=False)\r\n",
        "df_train.sample(2)\r\n",
        "pickle.dump(pred_train, open(os.path.join(DATA_DIR, 'output_{}_seed_{}_pred_train.p'.format(MODEL_NAME, RANDOM_SEED)), 'wb'))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BjSaKF4_ohGd"
      },
      "source": [
        "df_test, pred_test = transform_pandas(trainer, test_dataset)\r\n",
        "df_test.to_csv(os.path.join(DATA_DIR, 'output_{}_seed_{}_test.csv'.format(MODEL_NAME, RANDOM_SEED)), index=False)\r\n",
        "df_test.head()\r\n",
        "pickle.dump(pred_test, open(os.path.join(DATA_DIR, 'output_{}_seed{}_pred_test.p'.format(MODEL_NAME, RANDOM_SEED)), 'wb'))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1quGWom7ojCy"
      },
      "source": [
        "df_eval, pred_eval = transform_pandas(trainer, eval_dataset)\r\n",
        "df_eval.to_csv(os.path.join(DATA_DIR, 'output_{}_seed_{}_eval.csv'.format(MODEL_NAME, RANDOM_SEED)), index=False)\r\n",
        "df_eval.head()\r\n",
        "pickle.dump(pred_eval, open(os.path.join(DATA_DIR, 'output_{}_seed_{}_pred_eval.p'.format(MODEL_NAME, RANDOM_SEED)), 'wb'))"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}