1 - L1-Aware-T5 Training Pipeline.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "mysterious-holmes",
   "metadata": {},
   "outputs": [],
   "source": [
    "BASE_DIR = '' # Working dir\n",
    "DATA_DIR = f'{BASE_DIR}data/'\n",
    "MODELS_DIR = f'{BASE_DIR}models/'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "particular-auditor",
   "metadata": {},
   "source": [
    "# Load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "contained-crowd",
   "metadata": {},
   "outputs": [],
   "source": [
    "IMPROVE_TOKEN = \"improve_english: \"\n",
    "IMPROVE_TOKEN_MULTI = \"improve_english\"\n",
    "\n",
    "def get_improve_token(language=None, use_single_token=False):\n",
    "    if use_single_token:\n",
    "        return IMPROVE_TOKEN\n",
    "    return IMPROVE_TOKEN if not language else f'{IMPROVE_TOKEN_MULTI} {language}: '\n",
    "\n",
    "def get_reverse_token(language=None, use_single_token=False):\n",
    "    if use_single_token:\n",
    "        return TRANSFORM_TOKEN\n",
    "    return TRANSFORM_TOKEN if not language else f'{TRANSFORM_TOKEN_MULTI} {language}: '\n",
    "\n",
    "def prepare_sentences(df, limit=1000, language=None, use_single_token=False, reverse_model=False):    \n",
    "    sentences_inputs = []\n",
    "    sentences_outputs = []\n",
    "    counter = 0\n",
    "    for index, row in df.iterrows():\n",
    "        \n",
    "        if reverse_model:\n",
    "            sentences_inputs.append(get_reverse_token(language, use_single_token) + row['en'])        \n",
    "                \n",
    "            if language:\n",
    "                sentences_outputs.append(row[f'trans_{language}'])        \n",
    "            else:\n",
    "                sentences_outputs.append(row['trans'])                \n",
    "        else:\n",
    "            if language:\n",
    "                sentences_inputs.append(get_improve_token(language, use_single_token) + row[f'trans_{language}'])        \n",
    "            else:\n",
    "                sentences_inputs.append(get_improve_token(language, use_single_token) + row['trans'])                    \n",
    "            sentences_outputs.append(row['en'])\n",
    "            \n",
    "        counter += 1\n",
    "        if counter >=limit:\n",
    "            break\n",
    "    return sentences_inputs, sentences_outputs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "favorite-strength",
   "metadata": {},
   "source": [
    "# Multi lingual "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "auburn-cincinnati",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "physical-entertainment",
   "metadata": {},
   "outputs": [],
   "source": [
    "from random import shuffle\n",
    "\n",
    "def prepare_sentences_multilingual(df, languages, size, use_single_token=False, reverse_model=False):\n",
    "    sentences_inputs, sentences_outputs = [], []\n",
    "    for language in languages:\n",
    "        inputs, outputs = prepare_sentences(df, size*1000, language, use_single_token, reverse_model)\n",
    "        sentences_inputs.extend(inputs)\n",
    "        sentences_outputs.extend(outputs)\n",
    "    return sentences_inputs, sentences_outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "provincial-resort",
   "metadata": {},
   "outputs": [],
   "source": [
    "def prepare_data(data_frame, language, size, reverse_model=False, languages_multi=None):\n",
    "    use_single_token = language == 'all'\n",
    "    \n",
    "    if language == 'all' or language == 'multi':\n",
    "        languages = languages_multi\n",
    "    else:\n",
    "        languages = [language]\n",
    "    sentences_inputs, sentences_outputs = prepare_sentences_multilingual(data_frame, languages, size, use_single_token, reverse_model)   \n",
    "\n",
    "    grouped = list(zip(sentences_inputs, sentences_outputs))\n",
    "    shuffle(grouped)\n",
    "\n",
    "    sentences_inputs, sentences_outputs = zip(*grouped)\n",
    "\n",
    "    return sentences_inputs, sentences_outputs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "worst-method",
   "metadata": {},
   "source": [
    "# Prepare data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "binary-harvard",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers import T5Tokenizer\n",
    "from torch.utils.data import Dataset, random_split, IterableDataset\n",
    "\n",
    "class LazyLoadDataset(IterableDataset):\n",
    "    def __init__(self, tokenizer, inputs, outputs, length):\n",
    "        super().__init__()\n",
    "        self.tokenizer = tokenizer\n",
    "        self.sentences = list(zip(inputs, outputs))\n",
    "        self.length = length \n",
    "        \n",
    "    def __len__(self):\n",
    "        return self.length\n",
    "    \n",
    "    def get_next_sentence(self):    \n",
    "        for sentence in self.sentences:\n",
    "            yield sentence\n",
    "    \n",
    "    def __iter__(self):\n",
    "        return self.get_next_sentence()\n",
    "    \n",
    "class PairedDataset(Dataset):\n",
    "    def __init__(self, encodings):\n",
    "        self.encodings = encodings\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}\n",
    "        return item\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.encodings['input_ids'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "lightweight-senator",
   "metadata": {},
   "source": [
    "## Train and validation set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "signed-niger",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader, RandomSampler, SequentialSampler\n",
    "\n",
    "def get_dataloader(dataset, tokenizer, split_size=1, batch_size = 32):\n",
    "        \n",
    "    def tokenize_batch(batch):\n",
    "        sentence_inputs = [item[0] for item in batch]\n",
    "        sentence_outputs = [item[1] for item in batch]\n",
    "        encodings_input = tokenizer(sentence_inputs, truncation=True, padding='longest')\n",
    "        encodings_output = tokenizer(sentence_outputs, truncation=True, padding='longest')\n",
    "            \n",
    "        return {\n",
    "            'inputs': {key: torch.tensor(val) for key, val in encodings_input.items() },\n",
    "            'outputs': {key: torch.tensor(val) for key, val in encodings_output.items()}\n",
    "        }  \n",
    "\n",
    "    return DataLoader(\n",
    "        dataset,\n",
    "        batch_size=batch_size,\n",
    "        collate_fn=tokenize_batch\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "further-communist",
   "metadata": {},
   "source": [
    "# Prepare training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "marked-basket",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import T5ForConditionalGeneration, AdamW, BertConfig\n",
    "from transformers import get_linear_schedule_with_warmup\n",
    "\n",
    "def prepare_training(model_name, tokenizer, data_size, epochs):    \n",
    "    model = T5ForConditionalGeneration.from_pretrained(model_name).cuda()\n",
    "    \n",
    "    optimizer = AdamW(model.parameters(),\n",
    "                      lr = 5e-5, \n",
    "                      eps = 1e-8 \n",
    "                    )\n",
    "    num_warmup_steps = 0\n",
    "\n",
    "    total_steps = data_size * epochs\n",
    "    scheduler = get_linear_schedule_with_warmup(optimizer, \n",
    "                                                num_warmup_steps = num_warmup_steps, \n",
    "                                                num_training_steps = total_steps)\n",
    "    \n",
    "    return model, optimizer, scheduler"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "increased-spokesman",
   "metadata": {},
   "source": [
    "# Training loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ordinary-roman",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import datetime\n",
    "import random\n",
    "import numpy as np\n",
    "\n",
    "seed_val = 42\n",
    "\n",
    "random.seed(seed_val)\n",
    "np.random.seed(seed_val)\n",
    "torch.manual_seed(seed_val)\n",
    "torch.cuda.manual_seed_all(seed_val)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "hundred-trainer",
   "metadata": {},
   "outputs": [],
   "source": [
    "import wandb\n",
    "from tqdm import tqdm\n",
    "\n",
    "def train(run, model, optimizer, scheduler, epochs, model_name, dataloader):    \n",
    "        \n",
    "    with run:\n",
    "        for epoch_i in range(0, epochs):\n",
    "            model.train()\n",
    "\n",
    "            progress_bar = tqdm(dataloader)\n",
    "            step = 0\n",
    "            total_train_loss = 0\n",
    "\n",
    "            for batch in progress_bar:\n",
    "                inputs = batch['inputs']\n",
    "                outputs = batch['outputs']\n",
    "\n",
    "                b_input_ids = inputs['input_ids'].to(device)\n",
    "                b_masks = inputs['attention_mask'].to(device)\n",
    "                b_labels = outputs['input_ids'].to(device)\n",
    "\n",
    "                model.zero_grad()        \n",
    "                outputs = model(  b_input_ids,\n",
    "                                  labels=b_labels, \n",
    "                                  attention_mask = b_masks\n",
    "                                )\n",
    "                loss = outputs.loss\n",
    "                total_train_loss += loss.item()\n",
    "\n",
    "                loss.backward()\n",
    "\n",
    "                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",
    "\n",
    "                optimizer.step()\n",
    "                scheduler.step()\n",
    "\n",
    "                if step % 5 == 0 and step > 0:\n",
    "                    progress_bar.set_postfix({'loss': round(total_train_loss/step,2)})\n",
    "                    train_loss = total_train_loss/step\n",
    "                \n",
    "                    run.log({\n",
    "                        'epoch': epoch_i,\n",
    "                        'step': step,\n",
    "                        'train_loss': train_loss\n",
    "                    })\n",
    "                    \n",
    "                step += 1\n",
    "                \n",
    "                \n",
    "        run.finish()\n",
    "\n",
    "    model.save_pretrained(MODELS_DIR+model_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "confirmed-basis",
   "metadata": {},
   "outputs": [],
   "source": [
    "import gc\n",
    "\n",
    "def main(data, model_versions):            \n",
    "    tokenizer_cache = {}\n",
    "    \n",
    "    for train_model in model_versions:\n",
    "        # Clean up \n",
    "        torch.cuda.empty_cache()  \n",
    "                \n",
    "        model_name = f'{train_model[\"name\"]}-s{train_model[\"size\"]}-{train_model[\"version\"]}'\n",
    "        model_base = train_model['base']\n",
    "        \n",
    "        if not tokenizer_cache.get(model_base):\n",
    "            tokenizer_cache[model_base] = T5Tokenizer.from_pretrained(model_base)\n",
    "            \n",
    "        tokenizer = tokenizer_cache.get(model_base)\n",
    "        epochs = train_model['epochs']\n",
    "        size = train_model['size']\n",
    "        batch_size = train_model.get('batch_size', 32)\n",
    "        languages_multi = train_model.get('languages', ['pt','es', 'de', 'fr'])\n",
    "        \n",
    "        # Prepare data and model\n",
    "        model_language = model_name.split('-')[2] \n",
    "        model_type = model_name.split('-')[1] \n",
    "        is_reverse = model_type == 'reverse'\n",
    "        \n",
    "        sentences_inputs, sentences_outputs = prepare_data(data, model_language, size, is_reverse, languages_multi)        \n",
    "        \n",
    "        paired_dataset = LazyLoadDataset(tokenizer, sentences_inputs, sentences_outputs, len(sentences_inputs))    \n",
    "        training_dataloader = get_dataloader(paired_dataset, tokenizer, batch_size=batch_size)\n",
    "        model, optimizer, scheduler = prepare_training(model_base, tokenizer, len(training_dataloader), epochs)        \n",
    "    \n",
    "        # Print info\n",
    "        print('\\n')\n",
    "        print('Training', model_name, 'Sentences', len(sentences_inputs), 'Batches', len(training_dataloader))\n",
    "        print('=> Sample input:', sentences_inputs[0])\n",
    "        print('=> Sample output:', sentences_outputs[0])\n",
    "        time.sleep(1)\n",
    "            \n",
    "        # Init wandb\n",
    "        project_name = '-'.join(model_name.split('-')[:-1])\n",
    "        config_wandb = {\n",
    "            'epochs': epochs,\n",
    "            'size': size,\n",
    "            'batch_size': batch_size,\n",
    "            'base': model_base,\n",
    "            'dataset_size': len(sentences_inputs),\n",
    "        }\n",
    "        \n",
    "        if model_type == 'all' or model_type == 'multi':\n",
    "            config_wandb['languages'] = languages_multi\n",
    "            \n",
    "        run = wandb.init(reinit=True, project=project_name, config=config_wandb)\n",
    "        # Train model\n",
    "        train(run, model, optimizer, scheduler, epochs, model_name, training_dataloader)\n",
    "        \n",
    "        # Clean up\n",
    "        del model, optimizer, scheduler, training_dataloader, paired_dataset, sentences_inputs, sentences_outputs\n",
    "        gc.collect()\n",
    "        time.sleep(10)\n",
    "        torch.cuda.empty_cache()        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "tough-affairs",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_training = pd.read_csv(f'{DATA_DIR}pt-es-en-parallel-corpus.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "brazilian-museum",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda')\n",
    "\n",
    "BATCH_SIZE=16\n",
    "EPOCHS = 3\n",
    "\n",
    "models_config = [\n",
    "{\n",
    "    'id': 1,\n",
    "    'name': 't5sm-l1aware-multi',\n",
    "    'base': 't5-small', \n",
    "    'epochs': EPOCHS,\n",
    "    'batch_size': BATCH_SIZE,\n",
    "    'size': 260,\n",
    "    'languages': ['pt','es'],\n",
    "    'version': 'v1',\n",
    "},\n",
    "{\n",
    "    'id': 2,    \n",
    "    'name': 't5lg-l1aware-multi',\n",
    "    'base': 't5-base', \n",
    "    'epochs': EPOCHS,\n",
    "    'batch_size': BATCH_SIZE,\n",
    "    'size': 130,\n",
    "    'languages': ['pt','es'],\n",
    "    'version': 'v1',\n",
    "},\n",
    "]\n",
    "\n",
    "train_ids = [1, 2]\n",
    "\n",
    "to_train = [model for model in models_config if model['id'] in train_ids]\n",
    "main(df_training, to_train)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}