training_c2t.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hkrtK4vnYapr"
      },
      "outputs": [],
      "source": [
        "!pip3 install torch\n",
        "!pip3 install torchvision\n",
        "!pip3 install sentencepiece\n",
        "!pip3 install transformers\n",
        "!pip3 install datasets\n",
        "!pip install rouge-score\n",
        "#!pip3 install wandb"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iR7jHkmIYzec"
      },
      "outputs": [],
      "source": [
        "#from google.colab import drive\n",
        "#drive.mount('/content/drive')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_kRb-e4NasiF"
      },
      "outputs": [],
      "source": [
        "import datetime\n",
        "import os\n",
        "import time\n",
        "import sys\n",
        "\n",
        "import numpy as np\n",
        "import random\n",
        "import pandas as pd\n",
        "import nltk\n",
        "import numpy as np\n",
        "import json\n",
        "import wandb\n",
        "import torch\n",
        "from torch.utils.data import Dataset, DataLoader, random_split, RandomSampler, SequentialSampler\n",
        "torch.manual_seed(42)\n",
        "\n",
        "from datasets.dataset_dict import DatasetDict\n",
        "from datasets import load_metric, load_dataset \n",
        "from transformers import AutoTokenizer,  AutoModelForSeq2SeqLM, DataCollatorForSeq2Seq, Seq2SeqTrainingArguments, Seq2SeqTrainer\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "lr = 3e-4\n",
        "weight_dec = 0.01\n",
        "epochs = 12\n",
        "save_path = ''"
      ],
      "metadata": {
        "id": "T5j-16EQs9PL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aoF0J76Cp4G5"
      },
      "outputs": [],
      "source": [
        "nltk.download('punkt')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "psKbE40feuk9"
      },
      "outputs": [],
      "source": [
        "train_data = load_dataset('csv', data_files='data/c2t-big/train_c2t_big.csv', sep='\\t')\n",
        "val_data = load_dataset('csv', data_files='data/c2t-big//val_c2t_big.csv', sep='\\t')\n",
        "dataset = DatasetDict({\n",
        "    'train': train_data['train'],\n",
        "    'valid': val_data['train']})\n",
        "dataset\n",
        "'''\n",
        "train_data = load_dataset('csv', data_files='data/c2t-big/train_c2t_small.csv', sep='\\t')\n",
        "val_data = load_dataset('csv', data_files='data/c2t-big//val_c2t_small.csv', sep='\\t')\n",
        "dataset = DatasetDict({\n",
        "    'train': train_data['train'],\n",
        "    'valid': val_data['train']})\n",
        "dataset\n",
        "'''\n",
        "\n",
        "'''\n",
        "train_data = load_dataset('csv', data_files='data/c2t-big/train_c2t_small_nli.csv', sep='\\t')\n",
        "val_data = load_dataset('csv', data_files='data/c2t-big//val_c2t_small_nli.csv', sep='\\t')\n",
        "dataset = DatasetDict({\n",
        "    'train': train_data['train'],\n",
        "    'valid': val_data['train']})\n",
        "dataset\n",
        "'''"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gOyvxDbJb8g5"
      },
      "outputs": [],
      "source": [
        "metric = load_metric(\"rouge\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mq8StZpPcFct"
      },
      "outputs": [],
      "source": [
        "model_checkpoint = 't5-base'\n",
        "device = \"cuda:0\" if torch.cuda.is_available() else \"cpu\"\n",
        "tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)\n",
        "model =  AutoModelForSeq2SeqLM.from_pretrained(model_checkpoint)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XyT1uL13c8F_"
      },
      "outputs": [],
      "source": [
        "prefix = \"C2T: \""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "77nErzWIc_9L"
      },
      "outputs": [],
      "source": [
        "max_input_length = 1024\n",
        "max_target_length = 512\n",
        "\n",
        "def preprocess_function(examples):\n",
        "    inputs = [prefix + doc for doc in examples[\"Data\"]]\n",
        "    model_inputs = tokenizer(inputs, max_length=max_input_length, truncation=True, padding='max_length')\n",
        "\n",
        "    # Setup the tokenizer for targets\n",
        "    with tokenizer.as_target_tokenizer():\n",
        "        labels = tokenizer(examples[\"Summaries\"], max_length=max_target_length, truncation=True,  padding='max_length')\n",
        "\n",
        "    model_inputs[\"labels\"] = labels[\"input_ids\"]\n",
        "    return model_inputs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "k_l3_VZ0dEC1"
      },
      "outputs": [],
      "source": [
        "tokenized_datasets = dataset.map(preprocess_function, batched=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9WR5v-IYnO_m"
      },
      "outputs": [],
      "source": [
        "model =  AutoModelForSeq2SeqLM.from_pretrained(model_checkpoint)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nlqJD0XDnYYF"
      },
      "outputs": [],
      "source": [
        "batch_size = 8\n",
        "model_name = model_checkpoint.split(\"/\")[-1]\n",
        "args = Seq2SeqTrainingArguments(\n",
        "    save_path,\n",
        "    evaluation_strategy = \"epoch\",\n",
        "    learning_rate=lr,\n",
        "    per_device_train_batch_size=batch_size,\n",
        "    per_device_eval_batch_size=batch_size,\n",
        "    weight_decay=weight_dec,\n",
        "    save_total_limit=1,\n",
        "    num_train_epochs=epochs,\n",
        "    predict_with_generate=True\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DY3GBzRzntGt"
      },
      "outputs": [],
      "source": [
        "data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dvEy_tCadEIt"
      },
      "outputs": [],
      "source": [
        "def compute_metrics(eval_pred):\n",
        "    predictions, labels = eval_pred\n",
        "    decoded_preds = tokenizer.batch_decode(predictions, skip_special_tokens=True)\n",
        "    # Replace -100 in the labels as we can't decode them.\n",
        "    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)\n",
        "    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)\n",
        "    \n",
        "    # Rouge expects a newline after each sentence\n",
        "    decoded_preds = [\"\\n\".join(nltk.sent_tokenize(pred.strip())) for pred in decoded_preds]\n",
        "    decoded_labels = [\"\\n\".join(nltk.sent_tokenize(label.strip())) for label in decoded_labels]\n",
        "    \n",
        "    result = metric.compute(predictions=decoded_preds, references=decoded_labels, use_stemmer=True)\n",
        "    # Extract a few results\n",
        "    result = {key: value.mid.fmeasure * 100 for key, value in result.items()}\n",
        "    \n",
        "    # Add mean generated length\n",
        "    prediction_lens = [np.count_nonzero(pred != tokenizer.pad_token_id) for pred in predictions]\n",
        "    result[\"gen_len\"] = np.mean(prediction_lens)\n",
        "    \n",
        "    return {k: round(v, 4) for k, v in result.items()}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jgxJBeX1nzwq"
      },
      "outputs": [],
      "source": [
        "trainer = Seq2SeqTrainer(\n",
        "    model,\n",
        "    args,\n",
        "    train_dataset=tokenized_datasets[\"train\"],\n",
        "    eval_dataset=tokenized_datasets[\"valid\"],\n",
        "    data_collator=data_collator,\n",
        "    tokenizer=tokenizer,\n",
        "    compute_metrics=compute_metrics\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5UKAlDivn64h"
      },
      "outputs": [],
      "source": [
        "trainer.train()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HspjeSq-smaV"
      },
      "outputs": [],
      "source": [
        "trainer.save_model('t5base_C2T_big')"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "HekpWuqfoW-O"
      ],
      "provenance": []
    },
    "gpuClass": "premium",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}