Tachygraphy-MicroText-Transformers-GooglePegasus.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8fe131e7",
   "metadata": {
    "papermill": {
     "duration": 0.003041,
     "end_time": "2024-10-22T18:28:59.616233",
     "exception": false,
     "start_time": "2024-10-22T18:28:59.613192",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Google Pegasus\n",
    "### Tachygraphy Transformers Generative Model\n",
    "### PyTorch Distributed Data Parallel in use.\n",
    "### Dated - 25.09.2024"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c87378ae",
   "metadata": {
    "_cell_guid": "b1076dfc-b9ad-4769-8c92-a6c4dae69d19",
    "_uuid": "8f2839f25d086af736a60e9eeb907d3b93b6e0e5",
    "execution": {
     "iopub.execute_input": "2024-10-22T18:28:59.623564Z",
     "iopub.status.busy": "2024-10-22T18:28:59.623160Z",
     "iopub.status.idle": "2024-10-22T18:28:59.651415Z",
     "shell.execute_reply": "2024-10-22T18:28:59.650377Z"
    },
    "papermill": {
     "duration": 0.03504,
     "end_time": "2024-10-22T18:28:59.653827",
     "exception": false,
     "start_time": "2024-10-22T18:28:59.618787",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing ddp.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile ddp.py\n",
    "\n",
    "\n",
    "\n",
    "# This Python 3 environment comes with many helpful analytics libraries installed\n",
    "# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n",
    "# For example, here's several helpful packages to load\n",
    "\n",
    "import numpy as np # linear algebra\n",
    "import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n",
    "\n",
    "# Input data files are available in the read-only \"../input/\" directory\n",
    "# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n",
    "\n",
    "\n",
    "\n",
    "# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n",
    "# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "import torch\n",
    "from torch.utils.data import DataLoader\n",
    "from transformers import BartTokenizer, BartForConditionalGeneration, AdamW, BertModel, BertTokenizer\n",
    "# from torch.cuda.amp import autocast, GradScaler\n",
    "from datasets import load_dataset\n",
    "from sklearn.metrics.pairwise import cosine_similarity\n",
    "import ray\n",
    "from ray import tune\n",
    "import torch.nn as nn\n",
    "import torch.distributed as dist\n",
    "\n",
    "\n",
    "import os\n",
    "import gc\n",
    "import re\n",
    "import ast\n",
    "import sys\n",
    "import copy\n",
    "import json\n",
    "import time\n",
    "import math\n",
    "import string\n",
    "import pickle\n",
    "import random\n",
    "import joblib\n",
    "import itertools\n",
    "import warnings\n",
    "\n",
    "\n",
    "import argparse # CPMP\n",
    "\n",
    "import scipy as sp\n",
    "\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import Parameter\n",
    "import torch.nn.functional as F\n",
    "from torch.optim import Adam, SGD, AdamW\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "\n",
    "# CPMP imports for DDP\n",
    "from torch.nn.parallel import DistributedDataParallel\n",
    "from torch.utils.data.distributed import DistributedSampler\n",
    "\n",
    "from torch.optim import AdamW\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from transformers import RobertaTokenizer, RobertaModel\n",
    "from sklearn.model_selection import train_test_split\n",
    "import optuna\n",
    "# from optuna.integration import PyTorchLightningPruner\n",
    "from ray import tune\n",
    "import ray\n",
    "from ray import tune\n",
    "from ray.tune import CLIReporter\n",
    "from ray.tune import CLIReporter\n",
    "from ray.tune.schedulers import ASHAScheduler\n",
    "from ray.tune.schedulers import ASHAScheduler\n",
    "# from ray.tune.integration.pytorch import TuneReportCallback\n",
    "from torch.amp import GradScaler, autocast\n",
    "# from ray.tune.integration.optuna import OptunaSearch\n",
    "from ray.tune.search.optuna import OptunaSearch\n",
    "from ray import tune\n",
    "from ray.tune.search.hyperopt import HyperOptSearch\n",
    "from torch import autocast\n",
    "# from ray import tune\n",
    "# from ray.tune.integration.tensorboard import TensorBoardReporter\n",
    "from ray.tune.logger import TBXLogger\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "from ray.train import report\n",
    "# from ray.tune.integration.jupyter import JupyterNotebookReporter\n",
    "from ray.tune import JupyterNotebookReporter\n",
    "# from torch.cuda.amp import GradScaler, autocast\n",
    "from torch.optim.lr_scheduler import StepLR, CosineAnnealingLR\n",
    "from ray.tune.schedulers import HyperBandScheduler, AsyncHyperBandScheduler\n",
    "from transformers import AutoTokenizer, AutoModel\n",
    "from transformers import PegasusTokenizer, PegasusForConditionalGeneration\n",
    "from accelerate import notebook_launcher\n",
    "from accelerate import Accelerator\n",
    "# import evaluate\n",
    "from torch.utils.data import Dataset\n",
    "import torch\n",
    "\n",
    "\n",
    "import os\n",
    "import gc\n",
    "import re\n",
    "import ast\n",
    "import sys\n",
    "import copy\n",
    "import json\n",
    "import time\n",
    "import math\n",
    "import string\n",
    "import pickle\n",
    "import random\n",
    "import joblib\n",
    "import itertools\n",
    "import warnings\n",
    "\n",
    "\n",
    "import emoji\n",
    "from bs4 import BeautifulSoup\n",
    "import os\n",
    "import re\n",
    "import string\n",
    "import json\n",
    "\n",
    "\n",
    "import argparse # CPMP\n",
    "\n",
    "import scipy as sp\n",
    "\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "\n",
    "def setup(rank, world_size):\n",
    "    os.environ['MASTER_ADDR'] = 'localhost'\n",
    "    os.environ['MASTER_PORT'] = '12355'  # Use a port that's available on Kaggle\n",
    "    os.environ['RANK'] = str(rank)\n",
    "    os.environ['WORLD_SIZE'] = str(world_size)\n",
    "    # Initialize process group\n",
    "    dist.init_process_group(\"nccl\", rank=rank, world_size=world_size)\n",
    "    torch.cuda.set_device(rank)\n",
    "\n",
    "def cleanup():\n",
    "    dist.destroy_process_group()\n",
    "\n",
    "    \n",
    "    \n",
    "contraction_mapping = {\"ain't\": \"is not\", \"aren't\": \"are not\",\"can't\": \"cannot\", \"'cause\": \"because\", \"could've\": \"could have\", \"couldn't\": \"could not\", \n",
    "                       \"didn't\": \"did not\",  \"doesn't\": \"does not\", \"don't\": \"do not\", \"hadn't\": \"had not\", \"hasn't\": \"has not\", \"haven't\": \"have not\", \n",
    "                       \"he'd\": \"he would\",\"he'll\": \"he will\", \"he's\": \"he is\", \"how'd\": \"how did\", \"how'd'y\": \"how do you\", \"how'll\": \"how will\", \n",
    "                       \"how's\": \"how is\",  \"I'd\": \"I would\", \"I'd've\": \"I would have\", \"I'll\": \"I will\", \"I'll've\": \"I will have\",\"I'm\": \"I am\",\n",
    "                       \"I've\": \"I have\", \"i'd\": \"i would\", \"i'd've\": \"i would have\", \"i'll\": \"i will\",  \"i'll've\": \"i will have\",\"i'm\": \"i am\", \n",
    "                       \"i've\": \"i have\", \"isn't\": \"is not\", \"it'd\": \"it would\", \"it'd've\": \"it would have\", \"it'll\": \"it will\", \"it'll've\": \"it will have\",\n",
    "                       \"it's\": \"it is\", \"let's\": \"let us\", \"ma'am\": \"madam\", \"mayn't\": \"may not\", \"might've\": \"might have\",\"mightn't\": \"might not\",\n",
    "                       \"mightn't've\": \"might not have\", \"must've\": \"must have\", \"mustn't\": \"must not\", \"mustn't've\": \"must not have\", \"needn't\": \"need not\", \n",
    "                       \"needn't've\": \"need not have\",\"o'clock\": \"of the clock\", \"oughtn't\": \"ought not\", \"oughtn't've\": \"ought not have\", \"shan't\": \"shall not\",\n",
    "                       \"sha'n't\": \"shall not\", \"shan't've\": \"shall not have\", \"she'd\": \"she would\", \"she'd've\": \"she would have\", \"she'll\": \"she will\", \n",
    "                       \"she'll've\": \"she will have\", \"she's\": \"she is\", \"should've\": \"should have\", \"shouldn't\": \"should not\", \"shouldn't've\": \"should not have\",\n",
    "                       \"so've\": \"so have\",\"so's\": \"so as\", \"this's\": \"this is\",\"that'd\": \"that would\", \"that'd've\": \"that would have\", \"that's\": \"that is\",\n",
    "                       \"there'd\": \"there would\", \"there'd've\": \"there would have\", \"there's\": \"there is\", \"here's\": \"here is\",\"they'd\": \"they would\",\n",
    "                       \"they'd've\": \"they would have\", \"they'll\": \"they will\", \"they'll've\": \"they will have\", \"they're\": \"they are\", \"they've\": \"they have\",\n",
    "                       \"to've\": \"to have\", \"wasn't\": \"was not\", \"we'd\": \"we would\", \"we'd've\": \"we would have\", \"we'll\": \"we will\", \"we'll've\": \"we will have\",\n",
    "                       \"we're\": \"we are\", \"we've\": \"we have\", \"weren't\": \"were not\", \"what'll\": \"what will\", \"what'll've\": \"what will have\", \n",
    "                       \"what're\": \"what are\",  \"what's\": \"what is\", \"what've\": \"what have\", \"when's\": \"when is\", \"when've\": \"when have\", \"where'd\": \"where did\",\n",
    "                       \"where's\": \"where is\", \"where've\": \"where have\", \"who'll\": \"who will\", \"who'll've\": \"who will have\", \"who's\": \"who is\", \n",
    "                       \"who've\": \"who have\", \"why's\": \"why is\", \"why've\": \"why have\", \"will've\": \"will have\", \"won't\": \"will not\", \"won't've\": \"will not have\", \n",
    "                       \"would've\": \"would have\", \"wouldn't\": \"would not\", \"wouldn't've\": \"would not have\", \"y'all\": \"you all\", \"y'all'd\": \"you all would\",\n",
    "                       \"y'all'd've\": \"you all would have\",\"y'all're\": \"you all are\",\"y'all've\": \"you all have\",\"you'd\": \"you would\", \"you'd've\": \"you would have\",\n",
    "                       \"you'll\": \"you will\", \"you'll've\": \"you will have\", \"you're\": \"you are\", \"you've\": \"you have\", 'u.s':'america', 'e.g':'for example'}\n",
    "\n",
    "punct = [',', '.', '\"', ':', ')', '(', '-', '!', '?', '|', ';', \"'\", '$', '&', '/', '[', ']', '>', '%', '=', '#', '*', '+', '\\\\', '•',  '~', '@', '£', \n",
    " '·', '_', '{', '}', '©', '^', '®', '`',  '<', '→', '°', '€', '™', '›',  '♥', '←', '×', '§', '″', '′', 'Â', '█', '½', 'à', '…', \n",
    " '“', '★', '”', '–', '●', 'â', '►', '−', '¢', '²', '¬', '░', '¶', '↑', '±', '¿', '▾', '═', '¦', '║', '―', '¥', '▓', '—', '‹', '─', \n",
    " '▒', '：', '¼', '⊕', '▼', '▪', '†', '■', '’', '▀', '¨', '▄', '♫', '☆', 'é', '¯', '♦', '¤', '▲', 'è', '¸', '¾', 'Ã', '⋅', '‘', '∞', \n",
    " '∙', '）', '↓', '、', '│', '（', '»', '，', '♪', '╩', '╚', '³', '・', '╦', '╣', '╔', '╗', '▬', '❤', 'ï', 'Ø', '¹', '≤', '‡', '√', ]\n",
    "\n",
    "punct_mapping = {\"‘\": \"'\", \"₹\": \"e\", \"´\": \"'\", \"°\": \"\", \"€\": \"e\", \"™\": \"tm\", \"√\": \" sqrt \", \"×\": \"x\", \"²\": \"2\", \"—\": \"-\", \"–\": \"-\", \"’\": \"'\", \"_\": \"-\",\n",
    "                 \"`\": \"'\", '“': '\"', '”': '\"', '“': '\"', \"£\": \"e\", '∞': 'infinity', 'θ': 'theta', '÷': '/', 'α': 'alpha', '•': '.', 'à': 'a', '−': '-', \n",
    "                 'β': 'beta', '∅': '', '³': '3', 'π': 'pi', '!':' '}\n",
    "\n",
    "mispell_dict = {'colour': 'color', 'centre': 'center', 'favourite': 'favorite', 'travelling': 'traveling', 'counselling': 'counseling', 'theatre': 'theater',\n",
    "                'cancelled': 'canceled', 'labour': 'labor', 'organisation': 'organization', 'wwii': 'world war 2', 'citicise': 'criticize', 'youtu ': 'youtube ',\n",
    "                'Qoura': 'Quora', 'sallary': 'salary', 'Whta': 'What', 'narcisist': 'narcissist', 'howdo': 'how do', 'whatare': 'what are', 'howcan': 'how can',\n",
    "                'howmuch': 'how much', 'howmany': 'how many', 'whydo': 'why do', 'doI': 'do I', 'theBest': 'the best', 'howdoes': 'how does', \n",
    "                'mastrubation': 'masturbation', 'mastrubate': 'masturbate', \"mastrubating\": 'masturbating', 'pennis': 'penis', 'Etherium': 'Ethereum', \n",
    "                'narcissit': 'narcissist', 'bigdata': 'big data', '2k17': '2017', '2k18': '2018', 'qouta': 'quota', 'exboyfriend': 'ex boyfriend', \n",
    "                'airhostess': 'air hostess', \"whst\": 'what', 'watsapp': 'whatsapp', 'demonitisation': 'demonetization', 'demonitization': 'demonetization',\n",
    "                'demonetisation': 'demonetization'}\n",
    "\n",
    "def clean_text(text):\n",
    "    '''Clean emoji, Make text lowercase, remove text in square brackets,remove links,remove punctuation\n",
    "    and remove words containing numbers.'''\n",
    "#     text = emoji.demojize(text)\n",
    "#     text = re.sub(r'\\:(.*?)\\:','',text)\n",
    "#     text = str(text).lower()    #Making Text Lowercase\n",
    "#     text = re.sub('\\[.*?\\]', '', text)\n",
    "    #The next 2 lines remove html text\n",
    "    text = BeautifulSoup(text, 'lxml').get_text()\n",
    "    text = re.sub('https?://\\S+|www\\.\\S+', '', text)\n",
    "    text = re.sub('<.*?>+', '', text)\n",
    "    text = re.sub('\\n', '', text)\n",
    "#     text = re.sub('\\w*\\d\\w*', '', text)\n",
    "    # replacing everything with space except (a-z, A-Z, 0-9, \"%\", \".\", \"&\", \",\", \"'\", \"?\", \"!\", \",\", \"'\", \";\", \"-\")\n",
    "    text = re.sub(r\"[^a-zA-Z0-9?.!,¿'%&,';-]+\", \" \", text)\n",
    "    return text\n",
    "\n",
    "def clean_contractions(text, mapping):\n",
    "    '''Clean contraction using contraction mapping'''    \n",
    "    specials = [\"’\", \"‘\", \"´\", \"`\"]\n",
    "    for s in specials:\n",
    "        text = text.replace(s, \"'\")\n",
    "    for word in mapping.keys():\n",
    "        if \"\"+word+\"\" in text:\n",
    "            text = text.replace(\"\"+word+\"\", \"\"+mapping[word]+\"\")\n",
    "    #Remove Punctuations\n",
    "    text = re.sub('[%s]' % re.escape(string.punctuation), '', text)\n",
    "    # creating a space between a word and the punctuation following it\n",
    "    # eg: \"he is a boy.\" => \"he is a boy .\"\n",
    "    text = re.sub(r\"([?.!,¿])\", r\" \\1 \", text)\n",
    "    text = re.sub(r'[\" \"]+', \" \", text)\n",
    "    return text\n",
    "\n",
    "def clean_special_chars(text, punct, mapping):\n",
    "    '''Cleans special characters present(if any)'''   \n",
    "    for p in mapping:\n",
    "        text = text.replace(p, mapping[p])\n",
    "    \n",
    "    for p in punct:\n",
    "        text = text.replace(p, f' {p} ')\n",
    "    \n",
    "    specials = {'\\u200b': ' ', '…': ' ... ', '\\ufeff': '', 'करना': '', 'है': ''}  \n",
    "    for s in specials:\n",
    "        text = text.replace(s, specials[s])\n",
    "    \n",
    "    return text\n",
    "\n",
    "def correct_spelling(x, dic):\n",
    "    '''Corrects common spelling errors'''   \n",
    "    for word in dic.keys():\n",
    "        x = x.replace(word, dic[word])\n",
    "    return x\n",
    "\n",
    "def remove_space(text):\n",
    "    '''Removes awkward spaces'''   \n",
    "    #Removes awkward spaces \n",
    "    text = text.strip()\n",
    "    text = text.split()\n",
    "    return \" \".join(text)\n",
    "\n",
    "def text_preprocessing_pipeline(text):\n",
    "    '''Cleaning and parsing the text.'''\n",
    "#     text = clean_contractions(text, contraction_mapping)\n",
    "    text = clean_text(text)\n",
    "#     text = clean_contractions(text, contraction_mapping)\n",
    "#     text = clean_special_chars(text, punct, punct_mapping)\n",
    "#     text = correct_spelling(text, mispell_dict)\n",
    "    text = remove_space(text)\n",
    "    return text\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "class TextDataset(Dataset):\n",
    "    def __init__(self, tokenizer, dataset, max_length=256):\n",
    "        self.tokenizer = tokenizer\n",
    "        self.texts = dataset['input']\n",
    "        self.targets = dataset['target']\n",
    "        self.max_length = max_length\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.texts)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        text = self.texts[idx]\n",
    "        target = self.targets[idx]\n",
    "        encodings = self.tokenizer(text, truncation=True, padding='max_length', max_length=self.max_length)\n",
    "        target_encodings = self.tokenizer(target, truncation=True, padding='max_length', max_length=self.max_length)\n",
    "\n",
    "        return {\n",
    "            'input_ids': torch.tensor(encodings['input_ids'], dtype=torch.long),\n",
    "            'attention_mask': torch.tensor(encodings['attention_mask'], dtype=torch.long),\n",
    "            'labels': torch.tensor(target_encodings['input_ids'], dtype=torch.long)\n",
    "        }\n",
    "\n",
    "# %% [code] {\"jupyter\":{\"outputs_hidden\":false}}\n",
    "def test_train_split(dataset, test_ratio=0.1):\n",
    "  test_indices = np.random.rand(len(dataset)) < test_ratio\n",
    "  return dataset[~test_indices], dataset[test_indices]\n",
    "\n",
    "\n",
    "def calculate_cosine_similarity(pred_text, target_text, model, tokenizer, device, model_type=\"pegasus\"):\n",
    "    # Move model to GPU if available\n",
    "    model.to(device)\n",
    "\n",
    "    # Tokenize and encode both texts, moving inputs to the GPU\n",
    "    pred_inputs = tokenizer(pred_text, return_tensors=\"pt\", truncation=True, padding=True).to(device)\n",
    "    target_inputs = tokenizer(target_text, return_tensors=\"pt\", truncation=True, padding=True).to(device)\n",
    "\n",
    "    # Get embeddings based on the model type\n",
    "    with torch.no_grad():\n",
    "#         with torch.cuda.amp.autocast(enabled=True):  # Autocast for mixed precision training/inference\n",
    "        if model_type in [\"bart\", \"pegasus\", \"t5\"]:  # For seq2seq models like BART, Pegasus, T5\n",
    "            # Use encoder outputs\n",
    "            pred_embeddings = model.model.encoder(input_ids=pred_inputs.input_ids, attention_mask=pred_inputs.attention_mask).last_hidden_state.mean(dim=1).to(device)\n",
    "            target_embeddings = model.model.encoder(input_ids=target_inputs.input_ids, attention_mask=target_inputs.attention_mask).last_hidden_state.mean(dim=1).to(device)\n",
    "\n",
    "        elif model_type == \"llama\":  # For LLaMA models\n",
    "            pred_embeddings = model(**pred_inputs).hidden_states[-1].mean(dim=1).to(device)\n",
    "            target_embeddings = model(**target_inputs).hidden_states[-1].mean(dim=1).to(device)\n",
    "\n",
    "        else:  # For models like BERT, RoBERTa (directly exposing `last_hidden_state`)\n",
    "            pred_embeddings = model(**pred_inputs).last_hidden_state.mean(dim=1).to(device)\n",
    "            target_embeddings = model(**target_inputs).last_hidden_state.mean(dim=1).to(device)\n",
    "\n",
    "    # Calculate cosine similarity\n",
    "    similarity = cosine_similarity(pred_embeddings.cpu(), target_embeddings.cpu())[0].item()\n",
    "    return similarity\n",
    "\n",
    "\n",
    "\n",
    "class EarlyStopping:\n",
    "    def __init__(self, patience=5, verbose=False, delta=0):\n",
    "        \"\"\"\n",
    "        EarlyStopping to stop training when a metric has stopped improving.\n",
    "\n",
    "        Args:\n",
    "            patience (int): How long to wait after last time validation loss improved.\n",
    "            verbose (bool): If True, prints a message for each validation loss improvement.\n",
    "            delta (float): Minimum change to qualify as an improvement.\n",
    "        \"\"\"\n",
    "        self.patience = patience\n",
    "        self.verbose = verbose\n",
    "        self.delta = delta\n",
    "        self.counter = 0\n",
    "        self.best_loss = np.inf\n",
    "        self.early_stop = False\n",
    "\n",
    "    def __call__(self, val_loss, model):\n",
    "        if val_loss < self.best_loss - self.delta:\n",
    "            self.best_loss = val_loss\n",
    "            self.counter = 0\n",
    "#             if self.verbose:\n",
    "#                 print(f'Validation loss improved to {val_loss:.4f}')\n",
    "#             torch.save(model.state_dict(), 'checkpoint.pt')\n",
    "        else:\n",
    "            self.counter += 1\n",
    "#             if self.verbose:\n",
    "#                 print(f'Validation loss did not improve. Patience: {self.patience}, Counter: {self.counter}')\n",
    "        \n",
    "        if self.counter >= self.patience:\n",
    "            self.early_stop = True\n",
    "            \n",
    "            \n",
    "            \n",
    "            \n",
    "            \n",
    "def train_model(rank, config, train_dataset, val_dataset, world_size):\n",
    "    \n",
    "    setup(rank, world_size)\n",
    "    \n",
    "#     tokenizer = BartTokenizer.from_pretrained('facebook/bart-large')\n",
    "#     model = BartForConditionalGeneration.from_pretrained('facebook/bart-large')\n",
    "\n",
    "    tokenizer = PegasusTokenizer.from_pretrained('google/pegasus-large')\n",
    "    model = PegasusForConditionalGeneration.from_pretrained('google/pegasus-large')\n",
    "#     bert_model = BertModel.from_pretrained('bert-base-uncased')\n",
    "#     bert_tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')\n",
    "\n",
    "#     device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "    \n",
    "#     bert_model.to(device)\n",
    "\n",
    "#     model = nn.DataParallel(model, device_ids=[0, 1])\n",
    "#     model.to(device)\n",
    "    model = model.to(rank)\n",
    "    model = DistributedDataParallel(model, device_ids=[rank])\n",
    "    \n",
    "#     train_texts = train_ds_pd['input']\n",
    "#     train_labels = train_ds_pd['target']\n",
    "#     val_texts = validation_ds_pd['input']\n",
    "#     val_labels = validation_ds_pd['target']\n",
    "\n",
    "#     train_dataset = TextDataset(tokenizer, train_texts, train_labels)\n",
    "#     val_dataset = TextDataset(tokenizer, val_texts, val_labels)\n",
    "\n",
    "    train_sampler = DistributedSampler(train_dataset,\n",
    "                                       num_replicas=world_size,\n",
    "                                       rank=rank,\n",
    "                                       shuffle=False,\n",
    "                                       )\n",
    "    valid_sampler = DistributedSampler(val_dataset,\n",
    "                                       num_replicas=world_size,\n",
    "                                       rank=rank,\n",
    "                                       shuffle=False,\n",
    "                                       )\n",
    "    \n",
    "    train_loader = DataLoader(train_dataset,\n",
    "                              batch_size=int(config['batch_size']),\n",
    "                              shuffle=False,\n",
    "                              sampler=train_sampler)\n",
    "    val_loader = DataLoader(val_dataset,\n",
    "                            batch_size=int(config['batch_size']),\n",
    "                            shuffle=False,\n",
    "                            sampler=valid_sampler)\n",
    "\n",
    "    optimizer = AdamW(model.parameters(), lr=config['lr'])\n",
    "    scaler = torch.cuda.amp.GradScaler()  # Initialize the gradient scaler for mixed precision\n",
    "    \n",
    "    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=2, verbose=True)\n",
    "    \n",
    "    early_stopping = EarlyStopping(patience=3, verbose=True)\n",
    "\n",
    "    # Training loop with mixed precision\n",
    "    for epoch in range(int(config['epochs'])):\n",
    "        model.train()\n",
    "        \n",
    "        total_train_loss = 0.0\n",
    "        \n",
    "        for batch in train_loader:\n",
    "            \n",
    "            input_ids = batch['input_ids'].to(rank)\n",
    "            attention_mask = batch['attention_mask'].to(rank)\n",
    "            labels = batch['labels'].to(rank)\n",
    "            \n",
    "            optimizer.zero_grad()\n",
    "            \n",
    "#             with torch.cuda.amp.autocast(enabled=True):  # Use autocast for mixed precision\n",
    "            outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)\n",
    "            loss = outputs.loss\n",
    "\n",
    "            # Scale the loss and backpropagate\n",
    "            scaler.scale(loss).backward()\n",
    "            scaler.step(optimizer)\n",
    "            scaler.update()\n",
    "\n",
    "            total_train_loss += loss.item()\n",
    "            \n",
    "        avg_train_loss = total_train_loss / len(train_loader)\n",
    "\n",
    "        # Validation loop with mixed precision\n",
    "        model.eval()\n",
    "        \n",
    "        total_val_loss = 0.0\n",
    "        val_similarity = 0.0\n",
    "        num_val_samples = 0\n",
    "        \n",
    "        with torch.no_grad():\n",
    "            for batch in val_loader:\n",
    "                input_ids = batch['input_ids'].to(rank)\n",
    "                attention_mask = batch['attention_mask'].to(rank)\n",
    "                labels = batch['labels'].to(rank)\n",
    "                \n",
    "#                 with torch.cuda.amp.autocast(enabled=True):\n",
    "                outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)\n",
    "                loss = outputs.loss\n",
    "\n",
    "#                 preds = model.module.generate(input_ids=input_ids, attention_mask=attention_mask, max_length=512)\n",
    "                preds = model.module.generate(input_ids=input_ids, attention_mask=attention_mask)\n",
    "                pred_texts = tokenizer.batch_decode(preds, skip_special_tokens=True)\n",
    "                target_texts = tokenizer.batch_decode(labels, skip_special_tokens=True)\n",
    "\n",
    "                for pred_text, target_text in zip(pred_texts, target_texts):\n",
    "#                     similarity = calculate_cosine_similarity(pred_text, target_text, bert_model, bert_tokenizer)\n",
    "                    similarity = calculate_cosine_similarity(pred_text, target_text, model.module, tokenizer, device=rank)\n",
    "                    val_similarity += similarity\n",
    "                    num_val_samples += 1\n",
    "\n",
    "                total_val_loss += loss.item()\n",
    "                \n",
    "        avg_val_loss = total_val_loss / len(val_loader)\n",
    "#         avg_val_similarity = val_similarity / len(val_loader)\n",
    "        avg_val_similarity = val_similarity / num_val_samples\n",
    "        \n",
    "\n",
    "\n",
    "        # Report both loss and similarity\n",
    "#         tune.report(train_loss=avg_train_loss, val_loss=avg_val_loss, val_similarity=avg_val_similarity)\n",
    "        \n",
    "#         report({\n",
    "#             \"loss\": avg_val_loss,\n",
    "#             \"similarity\": avg_val_similarity,\n",
    "#             \"train_loss\": avg_train_loss,\n",
    "#             \"early_stopping_epoch\": epoch + 1,\n",
    "#         })\n",
    "\n",
    "        print(f\"\"\"\n",
    "             loss: {avg_val_loss},\n",
    "             similarity: {avg_val_similarity},\n",
    "             train_loss: {avg_train_loss},\n",
    "             early_stopping_epoch: {epoch + 1},\n",
    "            \"\"\")\n",
    "        \n",
    "        \n",
    "        # Learning rate scheduler step\n",
    "        scheduler.step(avg_val_loss)\n",
    "        \n",
    "        \n",
    "        early_stopping(avg_val_loss, model)\n",
    "        if early_stopping.early_stop:\n",
    "#             print(\"Early stopping\")\n",
    "            break\n",
    "    \n",
    "    if rank == 0:  # Only save from the main process\n",
    "        torch.save(model.module.state_dict(), \"ddp_model.pth\")\n",
    "    \n",
    "    cleanup()\n",
    "    \n",
    "    return model\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "def run_ddp(demo_fn, world_size):\n",
    "    \n",
    "    \n",
    "    dataset = pd.read_excel('/kaggle/input/dataset-tachygraphy/Tachygraphy_MicroText-AIO-V2.xlsx')\n",
    "    \n",
    "    df = dataset\n",
    "    \n",
    "    df.rename(columns={'Informal Text':'input', 'Expanded Meaning':'target'}, inplace = True)\n",
    "    \n",
    "    df['input'] = df['input'].astype(str)\n",
    "    df['target'] = df['target'].astype(str)\n",
    "    \n",
    "    df['input'] = df['input'].apply(lambda x: text_preprocessing_pipeline(x))\n",
    "    df['target'] = df['target'].apply(lambda x: text_preprocessing_pipeline(x))\n",
    "    \n",
    "    df['input'] = df['input'].astype(str)\n",
    "    df['target'] = df['target'].astype(str)\n",
    "    \n",
    "    train_ds_pd, validation_ds_pd = test_train_split(df)\n",
    "    print(\"{} examples in training, {} examples in testing.\".format(\n",
    "    len(train_ds_pd), len(validation_ds_pd)))\n",
    "    \n",
    "    train_ds_pd = train_ds_pd.reset_index(drop=True)\n",
    "    validation_ds_pd = validation_ds_pd.reset_index(drop=True)\n",
    "    \n",
    "    \n",
    "    config = {\n",
    "        \"lr\": 5e-6,\n",
    "        \"batch_size\": 2,\n",
    "        \"epochs\": 7\n",
    "    }\n",
    "    \n",
    "    tokenizer = PegasusTokenizer.from_pretrained('google/pegasus-large')\n",
    "#     bert_tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')\n",
    "    \n",
    "    train_dataset = TextDataset(tokenizer, train_ds_pd)\n",
    "    val_dataset = TextDataset(tokenizer, validation_ds_pd)\n",
    "    \n",
    "    torch.multiprocessing.spawn(demo_fn,\n",
    "             args=(config, train_dataset, val_dataset, world_size),\n",
    "             nprocs=world_size,\n",
    "             join=True)\n",
    "    \n",
    "    \n",
    "\n",
    "    \n",
    "if __name__ == \"__main__\":\n",
    "    \n",
    "    import os\n",
    "    for dirname, _, filenames in os.walk('/kaggle/input'):\n",
    "        for filename in filenames:\n",
    "            print(os.path.join(dirname, filename))\n",
    "    n_gpus = torch.cuda.device_count()\n",
    "    print(f\"total GPUs: {n_gpus}\")\n",
    "    assert n_gpus >= 2, f\"Requires at least 2 GPUs to run, but got {n_gpus}\"\n",
    "    world_size = n_gpus\n",
    "    \n",
    "    config = {\n",
    "        \"lr\": 5e-6,\n",
    "        \"batch_size\": 2,\n",
    "        \"epochs\": 7\n",
    "    }\n",
    "    \n",
    "    run_ddp(train_model, world_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "02d832ae",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-22T18:28:59.659936Z",
     "iopub.status.busy": "2024-10-22T18:28:59.659618Z",
     "iopub.status.idle": "2024-10-23T00:10:14.365475Z",
     "shell.execute_reply": "2024-10-23T00:10:14.364319Z"
    },
    "papermill": {
     "duration": 20474.711531,
     "end_time": "2024-10-23T00:10:14.367881",
     "exception": false,
     "start_time": "2024-10-22T18:28:59.656350",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/kaggle/input/dataset-tachygraphy/Tachygraphy_EmotionMoodtags_Dataset.csv\r\n",
      "/kaggle/input/dataset-tachygraphy/Tachygraphy_dataset_main.csv\r\n",
      "/kaggle/input/dataset-tachygraphy/Tachygraphy_MicroText-AIO-V2.xlsx\r\n",
      "total GPUs: 2\r\n",
      "/kaggle/working/ddp.py:205: MarkupResemblesLocatorWarning: The input looks more like a filename than markup. You may want to open this file and pass the filehandle into Beautiful Soup.\r\n",
      "  text = BeautifulSoup(text, 'lxml').get_text()\r\n",
      "9247 examples in training, 1033 examples in testing.\r\n",
      "tokenizer_config.json: 100%|██████████████████| 88.0/88.0 [00:00<00:00, 502kB/s]\r\n",
      "spiece.model: 100%|████████████████████████| 1.91M/1.91M [00:00<00:00, 7.67MB/s]\r\n",
      "special_tokens_map.json: 100%|████████████████| 65.0/65.0 [00:00<00:00, 477kB/s]\r\n",
      "config.json: 100%|█████████████████████████| 3.09k/3.09k [00:00<00:00, 25.6MB/s]\r\n",
      "/opt/conda/lib/python3.10/site-packages/transformers/tokenization_utils_base.py:1601: FutureWarning: `clean_up_tokenization_spaces` was not set. It will be set to `True` by default. This behavior will be depracted in transformers v4.45, and will be then set to `False` by default. For more details check this issue: https://github.com/huggingface/transformers/issues/31884\r\n",
      "  warnings.warn(\r\n",
      "/opt/conda/lib/python3.10/site-packages/transformers/tokenization_utils_base.py:1601: FutureWarning: `clean_up_tokenization_spaces` was not set. It will be set to `True` by default. This behavior will be depracted in transformers v4.45, and will be then set to `False` by default. For more details check this issue: https://github.com/huggingface/transformers/issues/31884\r\n",
      "  warnings.warn(\r\n",
      "/opt/conda/lib/python3.10/site-packages/transformers/tokenization_utils_base.py:1601: FutureWarning: `clean_up_tokenization_spaces` was not set. It will be set to `True` by default. This behavior will be depracted in transformers v4.45, and will be then set to `False` by default. For more details check this issue: https://github.com/huggingface/transformers/issues/31884\r\n",
      "  warnings.warn(\r\n",
      "pytorch_model.bin: 100%|████████████████████| 2.28G/2.28G [00:07<00:00, 318MB/s]\r\n",
      "Some weights of PegasusForConditionalGeneration were not initialized from the model checkpoint at google/pegasus-large and are newly initialized: ['model.decoder.embed_positions.weight', 'model.encoder.embed_positions.weight']\r\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\r\n",
      "generation_config.json: 100%|██████████████████| 260/260 [00:00<00:00, 1.52MB/s]\r\n",
      "Some weights of PegasusForConditionalGeneration were not initialized from the model checkpoint at google/pegasus-large and are newly initialized: ['model.decoder.embed_positions.weight', 'model.encoder.embed_positions.weight']\r\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\r\n",
      "/kaggle/working/ddp.py:414: FutureWarning: `torch.cuda.amp.GradScaler(args...)` is deprecated. Please use `torch.amp.GradScaler('cuda', args...)` instead.\r\n",
      "  scaler = torch.cuda.amp.GradScaler()  # Initialize the gradient scaler for mixed precision\r\n",
      "/kaggle/working/ddp.py:414: FutureWarning: `torch.cuda.amp.GradScaler(args...)` is deprecated. Please use `torch.amp.GradScaler('cuda', args...)` instead.\r\n",
      "  scaler = torch.cuda.amp.GradScaler()  # Initialize the gradient scaler for mixed precision\r\n",
      "\r\n",
      "             loss: 0.19319527156104452,\r\n",
      "             similarity: 0.9600018575874011,\r\n",
      "             train_loss: 6.714410341193313,\r\n",
      "             early_stopping_epoch: 1,\r\n",
      "            \r\n",
      "\r\n",
      "             loss: 0.18639743099399056,\r\n",
      "             similarity: 0.9564971863877612,\r\n",
      "             train_loss: 6.718675607757379,\r\n",
      "             early_stopping_epoch: 1,\r\n",
      "            \r\n",
      "\r\n",
      "             loss: 0.08569591824310512,\r\n",
      "             similarity: 0.9746533005103843,\r\n",
      "             train_loss: 0.18531837193843195,\r\n",
      "             early_stopping_epoch: 2,\r\n",
      "            \r\n",
      "\r\n",
      "             loss: 0.08148885949273885,\r\n",
      "             similarity: 0.9762756826799197,\r\n",
      "             train_loss: 0.18440464021522024,\r\n",
      "             early_stopping_epoch: 2,\r\n",
      "            \r\n",
      "\r\n",
      "             loss: 0.06843656893373687,\r\n",
      "             similarity: 0.9688450907369647,\r\n",
      "             train_loss: 0.09053845685877705,\r\n",
      "             early_stopping_epoch: 3,\r\n",
      "            \r\n",
      "\r\n",
      "             loss: 0.06383124081072718,\r\n",
      "             similarity: 0.9703054850059967,\r\n",
      "             train_loss: 0.09093190728329045,\r\n",
      "             early_stopping_epoch: 3,\r\n",
      "            \r\n",
      "\r\n",
      "             loss: 0.0563205881349915,\r\n",
      "             similarity: 0.965283863894944,\r\n",
      "             train_loss: 0.06596328444232291,\r\n",
      "             early_stopping_epoch: 4,\r\n",
      "            \r\n",
      "\r\n",
      "             loss: 0.0614965278960937,\r\n",
      "             similarity: 0.9627613374527465,\r\n",
      "             train_loss: 0.06648177200699527,\r\n",
      "             early_stopping_epoch: 4,\r\n",
      "            \r\n",
      "\r\n",
      "             loss: 0.05607606566234215,\r\n",
      "             similarity: 0.9615067079865033,\r\n",
      "             train_loss: 0.05438341679401772,\r\n",
      "             early_stopping_epoch: 5,\r\n",
      "            \r\n",
      "\r\n",
      "             loss: 0.050987673405572254,\r\n",
      "             similarity: 0.9631281760260031,\r\n",
      "             train_loss: 0.05354880031878427,\r\n",
      "             early_stopping_epoch: 5,\r\n",
      "            \r\n",
      "\r\n",
      "             loss: 0.05184301281973547,\r\n",
      "             similarity: 0.959649500002944,\r\n",
      "             train_loss: 0.04758226231797852,\r\n",
      "             early_stopping_epoch: 6,\r\n",
      "            \r\n",
      "\r\n",
      "             loss: 0.047280419288871944,\r\n",
      "             similarity: 0.9599812803807988,\r\n",
      "             train_loss: 0.04719354398965271,\r\n",
      "             early_stopping_epoch: 6,\r\n",
      "            \r\n",
      "\r\n",
      "             loss: 0.043701620276255036,\r\n",
      "             similarity: 0.9666443883335106,\r\n",
      "             train_loss: 0.04293582811239651,\r\n",
      "             early_stopping_epoch: 7,\r\n",
      "            \r\n",
      "\r\n",
      "             loss: 0.04775519826373941,\r\n",
      "             similarity: 0.962459180410411,\r\n",
      "             train_loss: 0.04310499188204159,\r\n",
      "             early_stopping_epoch: 7,\r\n",
      "            \r\n"
     ]
    }
   ],
   "source": [
    "!python /kaggle/working/ddp.py"
   ]
  }
 ],
 "metadata": {
  "kaggle": {
   "accelerator": "nvidiaTeslaT4",
   "dataSources": [
    {
     "datasetId": 5557640,
     "sourceId": 9500773,
     "sourceType": "datasetVersion"
    }
   ],
   "dockerImageVersionId": 30776,
   "isGpuEnabled": true,
   "isInternetEnabled": true,
   "language": "python",
   "sourceType": "notebook"
  },
  "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.10.14"
  },
  "papermill": {
   "default_parameters": {},
   "duration": 20478.084789,
   "end_time": "2024-10-23T00:10:14.595543",
   "environment_variables": {},
   "exception": null,
   "input_path": "__notebook__.ipynb",
   "output_path": "__notebook__.ipynb",
   "parameters": {},
   "start_time": "2024-10-22T18:28:56.510754",
   "version": "2.6.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}