finetune_multi_lingual_MT5_training_&_evaluation_for_extended_paraphrasing.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "finetune multi-lingual MT5 training & evaluation for extended paraphrasing.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "private_outputs": true,
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Nid989/Isometric-Multi-task-NMT/blob/main/finetune_multi_lingual_MT5_training_%26_evaluation_for_extended_paraphrasing.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yPpIQCUTu9pf"
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "# un-comment below, while working on colab.\n",
        "!pip install datasets transformers sacrebleu torch sentencepiece transformers[sentencepiece] wandb boto3 --quiet "
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%%capture\n",
        "!pip install nltk -U"
      ],
      "metadata": {
        "id": "5VfBoYpqOCa0"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, DataCollatorForSeq2Seq, Seq2SeqTrainingArguments, Seq2SeqTrainer, MarianMTModel, MarianTokenizer\n",
        "from torch.utils.data import DataLoader\n",
        "from datasets import load_dataset, load_metric\n",
        "import torch\n",
        "import numpy as np\n",
        "import datasets\n",
        "import boto3\n",
        "import shutil\n",
        "import os\n",
        "import random\n",
        "from tqdm.notebook import tqdm\n",
        "# from tqdm import tqdm \n",
        "import wandb\n",
        "import logging\n",
        "import pandas as pd\n",
        "\n",
        "tqdm.pandas()"
      ],
      "metadata": {
        "id": "1sNQTEXAwBB-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "current_directory = os.getcwd()"
      ],
      "metadata": {
        "id": "putna8smDCEJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# for logging loss to wandb.ai\n",
        "access_key = \"c7deb1bb77ce9433eb246d460385f363659145a8\" # enter wandb secret_accces_key\n",
        "wandb.login(key=access_key)"
      ],
      "metadata": {
        "id": "3f0qpSQnwM1v"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "source_language = [\"de\", \"fr\", \"ru\"]\n",
        "target_languages = [\"de\", \"fr\", \"ru\"]"
      ],
      "metadata": {
        "id": "urpPTrWYIWuJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "datasets_repository = {\n",
        "    'de': ['opusparcus'],\n",
        "    'fr': ['opusparcus'],\n",
        "    'ru': ['opusparcus']\n",
        "} "
      ],
      "metadata": {
        "id": "Ex8cL0oEZv45"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# data processing\n",
        "data_types = [\"train\", \"test\", \"validation\"]"
      ],
      "metadata": {
        "id": "ltat8fVUdZ5b"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def process_opusparcus_data(data, data_type, lang):\n",
        "  path_to_data_type = os.path.join(current_directory, f\"opusparcus_{data_type}_{lang}.csv\")\n",
        "  data.to_csv(path_to_data_type, index=False)\n",
        "  df = pd.read_csv(path_to_data_type)\n",
        "  df.drop([\"lang\", \"gem_id\", \"references\", \"annot_score\"], axis=1, inplace=True)\n",
        "  df.rename(columns={\n",
        "      'input': 'input_text',\n",
        "      'target': 'target_text'\n",
        "  }, inplace=True)\n",
        "  df.to_csv(path_to_data_type, index=False)"
      ],
      "metadata": {
        "id": "r4fqkPTLdUOt"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# load dataset\n",
        "dataset_name = \"GEM/opusparcus\"\n",
        "for language in target_languages:\n",
        "  print(language)\n",
        "  globals()[f\"{language}_datasets\"] = load_dataset(dataset_name, lang=f\"{language}\", quality=95)"
      ],
      "metadata": {
        "id": "vF_Txs4faSAl"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# data processing\n",
        "def process_opusparcus_data(data, data_type, language):\n",
        "  path_to_data_type = os.path.join(current_directory, f\"opusparcus_{data_type}_{language}.csv\")\n",
        "  data.to_csv(path_to_data_type, index=False)\n",
        "  df = pd.read_csv(path_to_data_type)\n",
        "  df.drop([\"lang\", \"gem_id\", \"references\", \"annot_score\"], axis=1, inplace=True)\n",
        "  df.rename(columns={\n",
        "      'input': 'input_text',\n",
        "      'target': 'target_text'\n",
        "  }, inplace=True)\n",
        "  df.to_csv(path_to_data_type, index=False)"
      ],
      "metadata": {
        "id": "Wxp7cDYngm0R"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "for language in tqdm(target_languages, total=len(target_languages)):\n",
        "  for data_type in data_types:\n",
        "    process_opusparcus_data(globals()[f\"{language}_datasets\"][data_type], data_type=data_type, language=language)"
      ],
      "metadata": {
        "id": "wg4ZxFpagzeN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "data_directories = [file for file in os.listdir(current_directory) if file.endswith('.csv')]\n",
        "train_data_directories = [data_directory for data_directory in data_directories if \"train\" in data_directory]\n",
        "test_data_directories = [data_directory for data_directory in data_directories if \"test\" in data_directory]\n",
        "validation_data_directories = [data_directory for data_directory in data_directories if \"validation\" in data_directory]\n",
        "\n",
        "print(train_data_directories, test_data_directories, validation_data_directories)"
      ],
      "metadata": {
        "id": "JP1TDOGrHTZw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# process data\n",
        "for data_type in tqdm(data_types, total=len(data_types)):\n",
        "  path_to_data_directory = [os.path.join(current_directory, file) for file in globals()[f\"{data_type}_data_directories\"]]\n",
        "  repository = {language: file for file in globals()[f\"{data_type}_data_directories\"] for language in target_languages if language in file}\n",
        "  for language in target_languages:\n",
        "    globals()[f\"df_{language}\"] = pd.read_csv(repository[language])\n",
        "    globals()[f\"df_{language}\"]['language'] = language\n",
        "  dataframes = [globals()[f\"df_{language}\"] for language in target_languages]\n",
        "  df = pd.concat(dataframes)\n",
        "  path_to_data_file = os.path.join(current_directory, f\"{data_type}.csv\")\n",
        "  df.to_csv(path_to_data_file, index=False)"
      ],
      "metadata": {
        "id": "tQocRekWG_xz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# prepare singleton data\n",
        "data_types = [\"train\", \"test\", \"validation\"]\n",
        "for data_type in tqdm(data_types, total=len(data_types)):\n",
        "  path_to_datafiles = os.path.join(os.path.join(current_directory, f\"{data_type}.csv\"))\n",
        "  globals()[f\"{data_type}_datasets\"] = load_dataset(\"csv\", data_files={data_type: path_to_datafiles})"
      ],
      "metadata": {
        "id": "rZ90_YdiFVqz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(f\"train: {train_datasets}\\ntest: {test_datasets}\\nvalidation: {validation_datasets}\")"
      ],
      "metadata": {
        "id": "_ifjqQK6IQdC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "epoch = 1\n",
        "model_checkpoints = \"mt5-base-finetuned-multilingual-singleton-for-en\"\n",
        "session = boto3.Session(\n",
        "    aws_access_key_id='AKIA4QB2WTN5YQGLD77G',\n",
        "    aws_secret_access_key='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',\n",
        ")\n",
        "s3 = session.resource('s3')\n",
        "key = f\"{epoch}_{model_checkpoints}\"\n",
        "filename = f\"{model_checkpoints}.zip\"\n",
        "print(key)\n",
        "s3.meta.client.download_file(Bucket='tsd2022', Key=key, Filename=filename)"
      ],
      "metadata": {
        "id": "bVXAdQXfoISw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "shutil.unpack_archive(f\"{model_checkpoints}.zip\", model_checkpoints, \"zip\")"
      ],
      "metadata": {
        "id": "kuxVhoJPovp_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# pre-trained model checkpoints\n",
        "path_to_model_checkpoints = os.path.join(current_directory, model_checkpoints)\n",
        "checkpoint_folder = os.listdir(path_to_model_checkpoints)[0]\n",
        "train_model_checkpoints = os.path.join(path_to_model_checkpoints, checkpoint_folder)"
      ],
      "metadata": {
        "id": "03uF-V5DwfBJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# load the MarianMT tokenizer\n",
        "tokenizer = AutoTokenizer.from_pretrained(train_model_checkpoints)"
      ],
      "metadata": {
        "id": "SA8Wp9-xwiZi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def add_verbosity(input_list, target_list, language_list):\n",
        "  \"\"\"\n",
        "  input: list of source & target sequences\n",
        "  output: processed source sequence based on the calculated length ratios \n",
        "  \"\"\"\n",
        "  processed_input = []\n",
        "  for input, target, language in zip(input_list, target_list, language_list):\n",
        "    ts_ratio = len(target)/len(input)\n",
        "    if ts_ratio < 0.90:\n",
        "      prefix = f\"{language} para short\"\n",
        "    elif ts_ratio >= 0.90 and ts_ratio <= 1.10:\n",
        "      prefix = f\"{language} para normal\"\n",
        "    else:\n",
        "      prefix = f\"{language} para long\"\n",
        "    input = prefix + \" \" + input\n",
        "    processed_input.append(input)\n",
        "  return processed_input"
      ],
      "metadata": {
        "id": "_HRpnd8vxUGJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# preprocess MUST-C dataset\n",
        "max_input_length = 128 \n",
        "max_target_length = 128\n",
        "def preprocess_function(examples):\n",
        "    inputs = examples[\"input_text\"]\n",
        "    targets = examples[\"target_text\"]\n",
        "    languages = examples[\"language\"]\n",
        "    inputs = add_verbosity(inputs, targets, languages) # append appropriate prompts \n",
        "    model_inputs = tokenizer(inputs, max_length=max_input_length, truncation=True)\n",
        "    # setup the tokenizer for targets\n",
        "    with tokenizer.as_target_tokenizer():\n",
        "        labels = tokenizer(targets, max_length=max_target_length, truncation=True)\n",
        "    model_inputs[\"labels\"] = labels[\"input_ids\"]\n",
        "    return model_inputs"
      ],
      "metadata": {
        "id": "UGTxkmTjxXzy"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# tokenize raw data\n",
        "tokenized_train_datasets = train_datasets['train'].map(preprocess_function, batched=True)\n",
        "tokenized_validation_datasets = validation_datasets['validation'].map(preprocess_function, batched=True)"
      ],
      "metadata": {
        "id": "6_dRmTpvxv8h"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# training procedure\n",
        "model = AutoModelForSeq2SeqLM.from_pretrained(train_model_checkpoints)"
      ],
      "metadata": {
        "id": "WkukCDlFx_T5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "batch_size = 2 # change batch-size according to GPU availability \n",
        "model_name = train_model_checkpoints.split(\"/\")[-1]\n",
        "epoch = 1\n",
        "\n",
        "# define training model arguments\n",
        "args = Seq2SeqTrainingArguments(\n",
        "    f\"{model_name}-finetuned-multilingual-singleton-extended-for-paraphrasing\",\n",
        "    learning_rate=5e-5, \n",
        "    logging_strategy=\"steps\",\n",
        "    logging_steps=500,\n",
        "    evaluation_strategy=\"epoch\",\n",
        "    optim=\"adafactor\",\n",
        "    per_device_train_batch_size=batch_size,\n",
        "    per_device_eval_batch_size=batch_size,\n",
        "    weight_decay=0.01,\n",
        "    save_strategy=\"steps\",\n",
        "    save_steps=1000,\n",
        "    save_total_limit=1,\n",
        "    num_train_epochs=epoch,\n",
        "    report_to=\"wandb\",\n",
        "    predict_with_generate=True    \n",
        ")\n",
        "\n",
        "# initialize data-collator\n",
        "data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)"
      ],
      "metadata": {
        "id": "ChyA6ZvNN6C-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "sacrebleu = load_metric(\"sacrebleu\")\n",
        "meteor = load_metric(\"meteor\")\n",
        "\n",
        "def postprocess_text(preds, labels):\n",
        "    preds = [pred.strip() for pred in preds]\n",
        "    labels = [[label.strip()] for label in labels]\n",
        "    return preds, labels\n",
        "    \n",
        "def compute_metrics(eval_preds):\n",
        "    preds, labels = eval_preds\n",
        "    if isinstance(preds, tuple):\n",
        "        preds = preds[0]\n",
        "    decoded_preds = tokenizer.batch_decode(preds, 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",
        "    # Some simple post-processing\n",
        "    decoded_preds, decoded_labels = postprocess_text(decoded_preds, decoded_labels)\n",
        "    sacrebleu_result = sacrebleu.compute(predictions=decoded_preds, references=decoded_labels)\n",
        "    meteor_result = meteor.compute(predictions=decoded_preds, references=decoded_labels)\n",
        "    result = {\n",
        "        \"bleu\": sacrebleu_result[\"score\"],\n",
        "        \"meteor\": meteor_result['meteor']\n",
        "    }\n",
        "    prediction_lens = [np.count_nonzero(pred != tokenizer.pad_token_id) for pred in preds]\n",
        "    result[\"gen_len\"] = np.mean(prediction_lens)\n",
        "    result = {k: round(v, 4) for k, v in result.items()}\n",
        "    print(result)\n",
        "    return result"
      ],
      "metadata": {
        "id": "Bwk2gFWdOAJA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# initialize the trainer module\n",
        "trainer = Seq2SeqTrainer(\n",
        "    model,\n",
        "    args,\n",
        "    train_dataset=tokenized_train_datasets,\n",
        "    eval_dataset=tokenized_validation_datasets,\n",
        "    data_collator=data_collator,\n",
        "    tokenizer=tokenizer,\n",
        "    compute_metrics=compute_metrics,\n",
        ")"
      ],
      "metadata": {
        "id": "MmYgfosoywyZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# train the model\n",
        "trainer.train()"
      ],
      "metadata": {
        "id": "05-KNvDrzGOk"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# compress model checkpoint directory\n",
        "model_checkpoints = f\"{model_name}-finetuned-multilingual-singleton-extended-for-paraphrasing\"\n",
        "model_checkpoint_directory = os.path.join(current_directory, model_checkpoints)\n",
        "print(model_checkpoint_directory)\n",
        "shutil.make_archive(model_checkpoint_directory, \"zip\", model_checkpoint_directory.split('/')[-1])"
      ],
      "metadata": {
        "id": "gEnUDt7HH7wT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "session = boto3.Session(\n",
        "    aws_access_key_id='AKIA4QB2WTN5YQGLD77G',\n",
        "    aws_secret_access_key='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',\n",
        ")\n",
        "s3 = session.resource('s3')\n",
        "key = f\"{epoch}_{model_checkpoints}\"\n",
        "filename = f\"{model_checkpoints}.zip\"\n",
        "print(key)\n",
        "s3.meta.client.upload_file(Bucket='tsd2022', Key=key, Filename=filename)"
      ],
      "metadata": {
        "id": "_hhthammILOD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# delete checkpoint directory\n",
        "current_directory = os.getcwd()\n",
        "path_to_directory = os.path.join(current_directory, model_checkpoints)\n",
        "shutil.rmtree(path_to_directory)"
      ],
      "metadata": {
        "id": "xDVfaecBdtDo"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# delete zip file\n",
        "current_directory = os.getcwd()\n",
        "path_to_zip_file = os.path.join(current_directory, filename)\n",
        "os.remove(path_to_zip_file)\n",
        "\n",
        "# delete all data directories and data files\n",
        "data_directory = os.listdir(current_directory)\n",
        "for file in data_directory:\n",
        "  if '.csv' in file:\n",
        "    path_to_data_file = os.path.join(current_directory, file)\n",
        "    os.remove(path_to_data_file)"
      ],
      "metadata": {
        "id": "XHZDzdqRd8uv"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "---"
      ],
      "metadata": {
        "id": "Isy9bAIXFEk3"
      }
    }
  ]
}