ILSUM.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_DEVICE_ORDER\"]=\"PCI_BUS_ID\"\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"1\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install huggingface_hub ipywidgets \n",
    "! pip install -U \"huggingface_hub[cli]\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import notebook_login\n",
    "\n",
    "notebook_login()    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "record = []\n",
    "x = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install rouge_score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install indic-nlp-library"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install datasets transformers accelerate evaluate sentencepiece"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install sentencepiece"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of available GPUs: 1\n",
      "GPU 0 : Quadro GP100\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "# Check if CUDA (GPU support) is available\n",
    "if torch.cuda.is_available():\n",
    "    # Get the number of available GPUs\n",
    "    num_gpus = torch.cuda.device_count()\n",
    "    print(\"Number of available GPUs:\", num_gpus)\n",
    "\n",
    "    # Get the name of each GPU\n",
    "    for gpu_id in range(num_gpus):\n",
    "        gpu_name = torch.cuda.get_device_name(gpu_id)\n",
    "        print(\"GPU\", gpu_id, \":\", gpu_name)\n",
    "else:\n",
    "    print(\"CUDA is not available. Make sure you have a GPU and PyTorch with CUDA support installed.\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "!export CUDA_VISIBLE_DEVICES=0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch import cuda\n",
    "import torch\n",
    "device = 'cuda' if cuda.is_available() else 'cpu'\n",
    "torch.cuda.set_device(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GPU number: 0\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "# Check if CUDA (GPU support) is available\n",
    "if torch.cuda.is_available():\n",
    "    # Get the index of the currently selected device\n",
    "    gpu_number = torch.cuda.current_device()\n",
    "    print(\"GPU number:\", gpu_number)\n",
    "else:\n",
    "    print(\"CUDA is not available. Make sure you have a GPU and PyTorch with CUDA support installed.\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Getting the model and datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "dataset = load_dataset(\"ILSUM/ILSUM-1.0\", 'Hindi')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    }
   ],
   "source": [
    "from transformers import MBartForConditionalGeneration, AutoModelForSeq2SeqLM\n",
    "from transformers import AlbertTokenizer, AutoTokenizer\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"ai4bharat/IndicBART\", do_lower_case=False, use_fast=False, keep_accents=True)\n",
    "\n",
    "# Or use tokenizer = AlbertTokenizer.from_pretrained(\"ai4bharat/IndicBART\", do_lower_case=False, use_fast=False, keep_accents=True)\n",
    "\n",
    "model = AutoModelForSeq2SeqLM.from_pretrained(\"ai4bharat/IndicBART\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[64006,   942,    43, 32720,  8384, 64001]])\n",
      "tensor([[  942,    43, 32720,  8384, 64001]])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[64006,   942,    43, 32720,  8384]])"
      ]
     },
     "execution_count": 77,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "# Or use model = MBartForConditionalGeneration.from_pretrained(\"ai4bharat/IndicBART\")\n",
    "\n",
    "# Some initial mapping\n",
    "bos_id = tokenizer._convert_token_to_id_with_added_voc(\"<s>\")\n",
    "eos_id = tokenizer._convert_token_to_id_with_added_voc(\"</s>\")\n",
    "pad_id = tokenizer._convert_token_to_id_with_added_voc(\"<pad>\")\n",
    "# To get lang_id use any of ['<2as>', '<2bn>', '<2en>', '<2gu>', '<2hi>', '<2kn>', '<2ml>', '<2mr>', '<2or>', '<2pa>', '<2ta>', '<2te>']\n",
    "\n",
    "# First tokenize the input and outputs. The format below is how IndicBART was trained so the input should be \"Sentence </s> <2xx>\" where xx is the language code. Similarly, the output should be \"<2yy> Sentence </s>\".\n",
    "inp = tokenizer(\"I am a boy </s> <2en>\", add_special_tokens=False, return_tensors=\"pt\", padding=True).input_ids # tensor([[  466,  1981,    80, 25573, 64001, 64004]])\n",
    "\n",
    "out = tokenizer(\"<2hi> मैं  एक लड़का हूँ </s>\", add_special_tokens=False, return_tensors=\"pt\", padding=True).input_ids # tensor([[64006,   942,    43, 32720,  8384, 64001]])\n",
    "# Note that if you use any language other than Hindi or Marathi, you should convert its script to Devanagari using the Indic NLP Library.\n",
    "model_outputs=model(input_ids=inp, decoder_input_ids=out[:,0:-1], labels=out[:,1:])\n",
    "\n",
    "# For loss\n",
    "print(out)\n",
    "print(out[:,1:])\n",
    "out[:,0:-1] ## This is not label smoothed.\n",
    "\n",
    "# For logits\n",
    "#model_outputs.logits\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n"
     ]
    }
   ],
   "source": [
    "print(pad_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "I am a boy\n"
     ]
    }
   ],
   "source": [
    "\n",
    "# For generation. Pardon the messiness. Note the decoder_start_token_id.\n",
    "\n",
    "model.eval() # Set dropouts to zero\n",
    "\n",
    "model_output=model.generate(inp, use_cache=True, num_beams=4, max_length=20, min_length=1, early_stopping=True, pad_token_id=pad_id, bos_token_id=bos_id, eos_token_id=eos_id, decoder_start_token_id=tokenizer._convert_token_to_id_with_added_voc(\"<2en>\"))\n",
    "\n",
    "# Decode to get output strings\n",
    "\n",
    "decoded_output=tokenizer.decode(model_output[0], skip_special_tokens\n",
    "=True, clean_up_tokenization_spaces=False)\n",
    "\n",
    "print(decoded_output) # I am a boy\n",
    "# Note that if your output language is not Hindi or Marathi, you should convert its script from Devanagari to the desired language using the Indic NLP Library.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Preprocess the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = dataset.remove_columns(['id', 'Heading'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Maximum sequence length: 1024\n"
     ]
    }
   ],
   "source": [
    "# Context length of the model\n",
    "from transformers import AutoConfig\n",
    "\n",
    "# Load the model configuration\n",
    "config = AutoConfig.from_pretrained(\"ai4bharat/IndicBART\")\n",
    "\n",
    "# Get the maximum sequence length\n",
    "max_length = config.max_position_embeddings\n",
    "\n",
    "print(\"Maximum sequence length:\", max_length)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_words(example):\n",
    "    # Modify the column 'column_name' by adding words to each entry\n",
    "    example['input'] = [entry + \"  </s> <2hi>\" for entry in example['Article']]\n",
    "    example['target'] = [\"<2hi> \" + entry + \" </s>\" for entry in example['Summary']]\n",
    "    return example\n",
    "\n",
    "# Apply the function to the dataset using map\n",
    "mod_dataset = dataset.map(add_words, batched=True, remove_columns=[\"Article\", \"Summary\"])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_input_length = 1024\n",
    "\n",
    "\n",
    "def preprocess_function(examples):\n",
    "    model_inputs = tokenizer(\n",
    "        examples[\"input\"],\n",
    "        max_length=max_input_length,\n",
    "        add_special_tokens=True,\n",
    "        return_tensors=\"pt\",\n",
    "        padding=True,\n",
    "        truncation=True,\n",
    "    )\n",
    "    labels = tokenizer(\n",
    "        examples[\"target\"], truncation=True, add_special_tokens=True, return_tensors=\"pt\", padding=True, max_length=max_input_length, \n",
    "    )\n",
    "    model_inputs[\"labels\"] = labels[\"input_ids\"]\n",
    "    return model_inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'mod_dataset' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[92], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m tokenized_dataset \u001b[39m=\u001b[39m mod_dataset\u001b[39m.\u001b[39mmap(preprocess_function, batched\u001b[39m=\u001b[39m\u001b[39mTrue\u001b[39;00m, batch_size\u001b[39m=\u001b[39m\u001b[39m10\u001b[39m)\n",
      "\u001b[0;31mNameError\u001b[0m: name 'mod_dataset' is not defined"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "wandb: Network error (ReadTimeout), entering retry loop.\n"
     ]
    }
   ],
   "source": [
    "tokenized_dataset = mod_dataset.map(preprocess_function, batched=True, batch_size=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenized_dataset = tokenized_dataset.remove_columns(\n",
    "    mod_dataset[\"train\"].column_names + [\"token_type_ids\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "metadata": {},
   "outputs": [],
   "source": [
    "del dataset\n",
    "del mod_dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Evalutaion with ROUGE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [],
   "source": [
    "import evaluate\n",
    "\n",
    "rouge_score = evaluate.load(\"rouge\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import indicnlp.tokenize.sentence_tokenize as tok\n",
    "\n",
    "def compute_metrics(eval_pred):\n",
    "    predictions, labels = eval_pred\n",
    "    # Decode generated summaries into text\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",
    "    # Decode reference summaries into text\n",
    "    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)\n",
    "    # ROUGE expects a newline after each sentence\n",
    "    decoded_preds = [\"\\n\".join(tok.sentence_split(pred.strip(), lang='hi', delim_pat='auto')) for pred in decoded_preds]\n",
    "    decoded_labels = [\"\\n\".join(tok.sentence_split(label.strip(), lang='hi', delim_pat='auto')) for label in decoded_labels]\n",
    "    # Compute ROUGE scores\n",
    "    result = rouge_score.compute(\n",
    "        predictions=decoded_preds, references=decoded_labels\n",
    "    )\n",
    "    # Extract the median scores\n",
    "    print(result)\n",
    "    result = {key: value * 100 for key, value in result.items()}\n",
    "    return {k: round(v, 4) for k, v in result.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/shavak/anaconda3/envs/bart/lib/python3.11/site-packages/datasets/load.py:756: FutureWarning: The repository for rouge contains custom code which must be executed to correctly load the metric. You can inspect the repository content at https://raw.githubusercontent.com/huggingface/datasets/2.18.0/metrics/rouge/rouge.py\n",
      "You can avoid this message in future by passing the argument `trust_remote_code=True`.\n",
      "Passing `trust_remote_code=True` will be mandatory to load this metric from the next major release of `datasets`.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'rouge1': 0.6666666666666666, 'rouge2': 0.5, 'rougeL': 0.6666666666666666, 'rougeLsum': 0.6666666666666666}\n"
     ]
    }
   ],
   "source": [
    "from datasets import load_metric\n",
    "\n",
    "# Load the ROUGE metric\n",
    "rouge_metric = load_metric(\"rouge\")\n",
    "\n",
    "# Example data: reference summaries and generated summaries\n",
    "references = [\"Reference summary 1\", \"Reference summary 2\"]\n",
    "hypotheses = [\"Generated summary 1\", \"Generated summary 2\"]\n",
    "\n",
    "# Compute ROUGE scores\n",
    "rouge_scores = rouge_score.compute(predictions=hypotheses, references=references)\n",
    "\n",
    "# Print the scores\n",
    "print(rouge_scores)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating a subset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import DatasetDict, Dataset\n",
    "\n",
    "split_size = 0.1\n",
    "sub_train_size = int(tokenized_dataset[\"train\"].num_rows * split_size)\n",
    "sub_test_size = int(tokenized_dataset[\"test\"].num_rows * split_size)\n",
    "sub_val_size = int(tokenized_dataset[\"validation\"].num_rows * split_size)\n",
    "\n",
    "train_subset = Dataset.from_dict(tokenized_dataset[\"train\"].shuffle(seed=42)[:sub_train_size])\n",
    "test_subset = Dataset.from_dict(tokenized_dataset[\"test\"].shuffle(seed=42)[:sub_test_size])\n",
    "val_subset = Dataset.from_dict(tokenized_dataset[\"validation\"].shuffle(seed=42)[:sub_val_size])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset({\n",
      "    features: ['input_ids', 'attention_mask', 'labels'],\n",
      "    num_rows: 795\n",
      "})\n"
     ]
    }
   ],
   "source": [
    "print(train_subset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['input_ids', 'attention_mask', 'labels'],\n",
       "        num_rows: 795\n",
       "    })\n",
       "    test: Dataset({\n",
       "        features: ['input_ids', 'attention_mask', 'labels'],\n",
       "        num_rows: 284\n",
       "    })\n",
       "    validation: Dataset({\n",
       "        features: ['input_ids', 'attention_mask', 'labels'],\n",
       "        num_rows: 56\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 65,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "subset_dataset = DatasetDict()\n",
    "subset_dataset[\"train\"] = train_subset\n",
    "subset_dataset[\"test\"] = test_subset\n",
    "subset_dataset[\"validation\"] = val_subset\n",
    "subset_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import DataCollatorForSeq2Seq\n",
    "\n",
    "data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Wandb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install wandb\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "import wandb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "ERROR:wandb.jupyter:Failed to detect the name of this notebook, you can set it manually with the WANDB_NOTEBOOK_NAME environment variable to enable code saving.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mparvp\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.16.6"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Run data is saved locally in <code>/home/shavak/Desktop/Parv/wandb/run-20240409_192156-yqp4xfih</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Syncing run <strong><a href='https://wandb.ai/parvp/ILSUM/runs/yqp4xfih' target=\"_blank\">quiet-elevator-4</a></strong> to <a href='https://wandb.ai/parvp/ILSUM' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/run' target=\"_blank\">docs</a>)<br/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View project at <a href='https://wandb.ai/parvp/ILSUM' target=\"_blank\">https://wandb.ai/parvp/ILSUM</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run at <a href='https://wandb.ai/parvp/ILSUM/runs/yqp4xfih' target=\"_blank\">https://wandb.ai/parvp/ILSUM/runs/yqp4xfih</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "run = wandb.init(project=\"ILSUM\", config={\n",
    "    'learning_rate': 5e-5,\n",
    "    'architecture': 'IndicBART',\n",
    "    \"epochs\": \"8\"\n",
    "},\n",
    "entity=\"parvp\"\n",
    "           )\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generation Config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "({'activation_dropout': 0.1,\n",
       "  'activation_function': 'gelu',\n",
       "  'architectures': ['MBartForConditionalGeneration'],\n",
       "  'attention_dropout': 0.1,\n",
       "  'bos_token_id': 64000,\n",
       "  'd_model': 1024,\n",
       "  'classifier_dropout': 0.0,\n",
       "  'decoder_attention_heads': 16,\n",
       "  'decoder_ffn_dim': 4096,\n",
       "  'decoder_layerdrop': 0.0,\n",
       "  'decoder_layers': 6,\n",
       "  'dropout': 0.1,\n",
       "  'encoder_attention_heads': 16,\n",
       "  'encoder_ffn_dim': 4096,\n",
       "  'encoder_layerdrop': 0.0,\n",
       "  'encoder_layers': 6,\n",
       "  'eos_token_id': 64001,\n",
       "  'gradient_checkpointing': False,\n",
       "  'init_std': 0.02,\n",
       "  'is_encoder_decoder': True,\n",
       "  'max_position_embeddings': 1024,\n",
       "  'model_type': 'mbart',\n",
       "  'num_hidden_layers': 6,\n",
       "  'pad_token_id': 0,\n",
       "  'scale_embedding': False,\n",
       "  'transformers_version': '4.3.2',\n",
       "  'use_cache': True,\n",
       "  'vocab_size': 64014,\n",
       "  'tokenizer_class': 'AlbertTokenizer',\n",
       "  '_commit_hash': '78466a0c0e29f9229f7005623ecd6bc4243c0ae0'},\n",
       " {})"
      ]
     },
     "execution_count": 67,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.config.get_config_dict(\"ai4bharat/IndicBART\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import GenerationConfig\n",
    "generation_config = GenerationConfig(\n",
    "    max_new_tokens=512, \n",
    "    do_sample=True, \n",
    "    top_k=50,\n",
    "    use_cache=True,\n",
    "    num_beams=4,\n",
    "    min_length=5,\n",
    "    pad_token_id=pad_id,\n",
    "    bos_token_id=bos_id, \n",
    "    eos_token_id=eos_id, \n",
    "    decoder_start_token_id=tokenizer._convert_token_to_id_with_added_voc(\"<2hi>\")   \n",
    ")\n",
    "\n",
    "generation_config.save_pretrained(\"config.json\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GenerationConfig {\n",
      "  \"bos_token_id\": 64000,\n",
      "  \"decoder_start_token_id\": 64006,\n",
      "  \"do_sample\": true,\n",
      "  \"eos_token_id\": 64001,\n",
      "  \"max_new_tokens\": 512,\n",
      "  \"min_length\": 5,\n",
      "  \"num_beams\": 4,\n",
      "  \"pad_token_id\": 0\n",
      "}\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Load the configuration from the saved file\n",
    "loaded_config = GenerationConfig.from_pretrained(\"config.json\")\n",
    "print(loaded_config)\n",
    "model.generation_config = loaded_config\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GenerationConfig {\n",
      "  \"bos_token_id\": 64000,\n",
      "  \"decoder_start_token_id\": 64006,\n",
      "  \"do_sample\": true,\n",
      "  \"eos_token_id\": 64001,\n",
      "  \"max_new_tokens\": 512,\n",
      "  \"min_length\": 5,\n",
      "  \"num_beams\": 4,\n",
      "  \"pad_token_id\": 0\n",
      "}\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(model.generation_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "bruh = AutoConfig.from_pretrained(\"ai4bharat/IndicBART\")\n",
    "print(config == bruh)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import Seq2SeqTrainingArguments\n",
    "\n",
    "batch_size = 8\n",
    "num_train_epochs = 8\n",
    "# Show the training loss with every epoch\n",
    "logging_steps = len(tokenized_dataset[\"train\"])\n",
    "model_name = \"IndicBART\"\n",
    "\n",
    "args = Seq2SeqTrainingArguments(\n",
    "    output_dir=f\"{model_name}-finetuned-IndicSentenceSummarisation\",\n",
    "    evaluation_strategy=\"epoch\",\n",
    "    learning_rate=5e-5,\n",
    "    per_device_train_batch_size=batch_size,\n",
    "    per_device_eval_batch_size=batch_size,\n",
    "    weight_decay=0.01,\n",
    "    save_total_limit=3,\n",
    "    num_train_epochs=num_train_epochs,\n",
    "    predict_with_generate=True,\n",
    "    logging_steps=logging_steps,\n",
    "    push_to_hub=True,\n",
    "    report_to=\"wandb\",\n",
    "    run_name=\"first\"\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/shavak/anaconda3/envs/bart/lib/python3.11/site-packages/accelerate/accelerator.py:436: FutureWarning: Passing the following arguments to `Accelerator` is deprecated and will be removed in version 1.0 of Accelerate: dict_keys(['dispatch_batches', 'split_batches', 'even_batches', 'use_seedable_sampler']). Please pass an `accelerate.DataLoaderConfiguration` instead: \n",
      "dataloader_config = DataLoaderConfiguration(dispatch_batches=None, split_batches=False, even_batches=True, use_seedable_sampler=True)\n",
      "  warnings.warn(\n",
      "WARNING:accelerate.utils.other:Detected kernel version 3.10.0, which is below the recommended minimum of 5.5.0; this can cause the process to hang. It is recommended to upgrade the kernel to the minimum version or higher.\n"
     ]
    }
   ],
   "source": [
    "from transformers import Seq2SeqTrainer\n",
    "\n",
    "trainer = Seq2SeqTrainer(\n",
    "    model,\n",
    "    args,\n",
    "    train_dataset=tokenized_dataset[\"train\"],\n",
    "    eval_dataset=tokenized_dataset[\"validation\"],\n",
    "    data_collator=data_collator,\n",
    "    tokenizer=tokenizer,\n",
    "    compute_metrics=compute_metrics,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f2976b7949da4ef2a34a5cbcc906388e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/7960 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some non-default generation parameters are set in the model config. These should go into a GenerationConfig file (https://huggingface.co/docs/transformers/generation_strategies#save-a-custom-decoding-strategy-with-your-model) instead. This warning will be raised to an exception in v4.41.\n",
      "Non-default generation parameters: {'forced_eos_token_id': 2}\n",
      "/home/shavak/anaconda3/envs/bart/lib/python3.11/site-packages/transformers/generation/utils.py:1132: UserWarning: Using the model-agnostic default `max_length` (=20) to control the generation length. We recommend setting `max_new_tokens` to control the maximum length of the generation.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5259b594f0d5412895953549557e4166",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/72 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'rouge1': 0.11210098239975212, 'rouge2': 0.033911801925861684, 'rougeL': 0.11146325847556075, 'rougeLsum': 0.11116409464388373}\n",
      "{'eval_loss': 0.40127116441726685, 'eval_rouge1': 11.2101, 'eval_rouge2': 3.3912, 'eval_rougeL': 11.1463, 'eval_rougeLsum': 11.1164, 'eval_runtime': 33.3499, 'eval_samples_per_second': 17.062, 'eval_steps_per_second': 2.159, 'epoch': 1.0}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some non-default generation parameters are set in the model config. These should go into a GenerationConfig file (https://huggingface.co/docs/transformers/generation_strategies#save-a-custom-decoding-strategy-with-your-model) instead. This warning will be raised to an exception in v4.41.\n",
      "Non-default generation parameters: {'forced_eos_token_id': 2}\n",
      "Some non-default generation parameters are set in the model config. These should go into a GenerationConfig file (https://huggingface.co/docs/transformers/generation_strategies#save-a-custom-decoding-strategy-with-your-model) instead. This warning will be raised to an exception in v4.41.\n",
      "Non-default generation parameters: {'forced_eos_token_id': 2}\n",
      "/home/shavak/anaconda3/envs/bart/lib/python3.11/site-packages/transformers/generation/utils.py:1132: UserWarning: Using the model-agnostic default `max_length` (=20) to control the generation length. We recommend setting `max_new_tokens` to control the maximum length of the generation.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b12416ed590445628e00639b02251f0c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/72 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'rouge1': 0.11046308212389053, 'rouge2': 0.034465671021031305, 'rougeL': 0.10912697437302002, 'rougeLsum': 0.11009944458977845}\n",
      "{'eval_loss': 0.3646354675292969, 'eval_rouge1': 11.0463, 'eval_rouge2': 3.4466, 'eval_rougeL': 10.9127, 'eval_rougeLsum': 11.0099, 'eval_runtime': 30.7788, 'eval_samples_per_second': 18.487, 'eval_steps_per_second': 2.339, 'epoch': 2.0}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some non-default generation parameters are set in the model config. These should go into a GenerationConfig file (https://huggingface.co/docs/transformers/generation_strategies#save-a-custom-decoding-strategy-with-your-model) instead. This warning will be raised to an exception in v4.41.\n",
      "Non-default generation parameters: {'forced_eos_token_id': 2}\n",
      "Some non-default generation parameters are set in the model config. These should go into a GenerationConfig file (https://huggingface.co/docs/transformers/generation_strategies#save-a-custom-decoding-strategy-with-your-model) instead. This warning will be raised to an exception in v4.41.\n",
      "Non-default generation parameters: {'forced_eos_token_id': 2}\n",
      "/home/shavak/anaconda3/envs/bart/lib/python3.11/site-packages/transformers/generation/utils.py:1132: UserWarning: Using the model-agnostic default `max_length` (=20) to control the generation length. We recommend setting `max_new_tokens` to control the maximum length of the generation.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a2ac528abf8c4875b5fc060d689571e0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/72 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'rouge1': 0.11057236066023407, 'rouge2': 0.03206221136801101, 'rougeL': 0.11024502250337047, 'rougeLsum': 0.11010944240293973}\n",
      "{'eval_loss': 0.35485556721687317, 'eval_rouge1': 11.0572, 'eval_rouge2': 3.2062, 'eval_rougeL': 11.0245, 'eval_rougeLsum': 11.0109, 'eval_runtime': 30.9382, 'eval_samples_per_second': 18.392, 'eval_steps_per_second': 2.327, 'epoch': 3.0}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some non-default generation parameters are set in the model config. These should go into a GenerationConfig file (https://huggingface.co/docs/transformers/generation_strategies#save-a-custom-decoding-strategy-with-your-model) instead. This warning will be raised to an exception in v4.41.\n",
      "Non-default generation parameters: {'forced_eos_token_id': 2}\n",
      "Some non-default generation parameters are set in the model config. These should go into a GenerationConfig file (https://huggingface.co/docs/transformers/generation_strategies#save-a-custom-decoding-strategy-with-your-model) instead. This warning will be raised to an exception in v4.41.\n",
      "Non-default generation parameters: {'forced_eos_token_id': 2}\n",
      "/home/shavak/anaconda3/envs/bart/lib/python3.11/site-packages/transformers/generation/utils.py:1132: UserWarning: Using the model-agnostic default `max_length` (=20) to control the generation length. We recommend setting `max_new_tokens` to control the maximum length of the generation.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3e7456c453bd49fe881fcbeeccdef9bd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/72 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'rouge1': 0.10923125400278297, 'rouge2': 0.035528775071833066, 'rougeL': 0.10834481374903165, 'rougeLsum': 0.10929586913769683}\n",
      "{'eval_loss': 0.35208743810653687, 'eval_rouge1': 10.9231, 'eval_rouge2': 3.5529, 'eval_rougeL': 10.8345, 'eval_rougeLsum': 10.9296, 'eval_runtime': 30.8587, 'eval_samples_per_second': 18.439, 'eval_steps_per_second': 2.333, 'epoch': 4.0}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some non-default generation parameters are set in the model config. These should go into a GenerationConfig file (https://huggingface.co/docs/transformers/generation_strategies#save-a-custom-decoding-strategy-with-your-model) instead. This warning will be raised to an exception in v4.41.\n",
      "Non-default generation parameters: {'forced_eos_token_id': 2}\n",
      "Some non-default generation parameters are set in the model config. These should go into a GenerationConfig file (https://huggingface.co/docs/transformers/generation_strategies#save-a-custom-decoding-strategy-with-your-model) instead. This warning will be raised to an exception in v4.41.\n",
      "Non-default generation parameters: {'forced_eos_token_id': 2}\n",
      "/home/shavak/anaconda3/envs/bart/lib/python3.11/site-packages/transformers/generation/utils.py:1132: UserWarning: Using the model-agnostic default `max_length` (=20) to control the generation length. We recommend setting `max_new_tokens` to control the maximum length of the generation.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "92a430ce694140c8b62a4526e8224b6c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/72 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'rouge1': 0.10833174343719155, 'rouge2': 0.03367233723121245, 'rougeL': 0.10694103787249654, 'rougeLsum': 0.10815335806548457}\n",
      "{'eval_loss': 0.3489197790622711, 'eval_rouge1': 10.8332, 'eval_rouge2': 3.3672, 'eval_rougeL': 10.6941, 'eval_rougeLsum': 10.8153, 'eval_runtime': 30.8233, 'eval_samples_per_second': 18.46, 'eval_steps_per_second': 2.336, 'epoch': 5.0}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some non-default generation parameters are set in the model config. These should go into a GenerationConfig file (https://huggingface.co/docs/transformers/generation_strategies#save-a-custom-decoding-strategy-with-your-model) instead. This warning will be raised to an exception in v4.41.\n",
      "Non-default generation parameters: {'forced_eos_token_id': 2}\n",
      "Some non-default generation parameters are set in the model config. These should go into a GenerationConfig file (https://huggingface.co/docs/transformers/generation_strategies#save-a-custom-decoding-strategy-with-your-model) instead. This warning will be raised to an exception in v4.41.\n",
      "Non-default generation parameters: {'forced_eos_token_id': 2}\n",
      "/home/shavak/anaconda3/envs/bart/lib/python3.11/site-packages/transformers/generation/utils.py:1132: UserWarning: Using the model-agnostic default `max_length` (=20) to control the generation length. We recommend setting `max_new_tokens` to control the maximum length of the generation.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f0955b852c0247b380ef0f71f28ab452",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/72 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'rouge1': 0.10941469680485491, 'rouge2': 0.034690997344775897, 'rougeL': 0.10885279532028208, 'rougeLsum': 0.10933296833999817}\n",
      "{'eval_loss': 0.3471415042877197, 'eval_rouge1': 10.9415, 'eval_rouge2': 3.4691, 'eval_rougeL': 10.8853, 'eval_rougeLsum': 10.9333, 'eval_runtime': 30.7086, 'eval_samples_per_second': 18.529, 'eval_steps_per_second': 2.345, 'epoch': 6.0}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some non-default generation parameters are set in the model config. These should go into a GenerationConfig file (https://huggingface.co/docs/transformers/generation_strategies#save-a-custom-decoding-strategy-with-your-model) instead. This warning will be raised to an exception in v4.41.\n",
      "Non-default generation parameters: {'forced_eos_token_id': 2}\n",
      "Some non-default generation parameters are set in the model config. These should go into a GenerationConfig file (https://huggingface.co/docs/transformers/generation_strategies#save-a-custom-decoding-strategy-with-your-model) instead. This warning will be raised to an exception in v4.41.\n",
      "Non-default generation parameters: {'forced_eos_token_id': 2}\n",
      "/home/shavak/anaconda3/envs/bart/lib/python3.11/site-packages/transformers/generation/utils.py:1132: UserWarning: Using the model-agnostic default `max_length` (=20) to control the generation length. We recommend setting `max_new_tokens` to control the maximum length of the generation.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "15a69596df8e49e3890844bac96e9f09",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/72 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'rouge1': 0.10990522853440249, 'rouge2': 0.03524175980239248, 'rougeL': 0.1087365914694561, 'rougeLsum': 0.10954110330208744}\n",
      "{'eval_loss': 0.3471783399581909, 'eval_rouge1': 10.9905, 'eval_rouge2': 3.5242, 'eval_rougeL': 10.8737, 'eval_rougeLsum': 10.9541, 'eval_runtime': 30.8279, 'eval_samples_per_second': 18.457, 'eval_steps_per_second': 2.336, 'epoch': 7.0}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some non-default generation parameters are set in the model config. These should go into a GenerationConfig file (https://huggingface.co/docs/transformers/generation_strategies#save-a-custom-decoding-strategy-with-your-model) instead. This warning will be raised to an exception in v4.41.\n",
      "Non-default generation parameters: {'forced_eos_token_id': 2}\n",
      "Some non-default generation parameters are set in the model config. These should go into a GenerationConfig file (https://huggingface.co/docs/transformers/generation_strategies#save-a-custom-decoding-strategy-with-your-model) instead. This warning will be raised to an exception in v4.41.\n",
      "Non-default generation parameters: {'forced_eos_token_id': 2}\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.3884, 'grad_norm': 1.9990270137786865, 'learning_rate': 1.8844221105527637e-08, 'epoch': 8.0}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/shavak/anaconda3/envs/bart/lib/python3.11/site-packages/transformers/generation/utils.py:1132: UserWarning: Using the model-agnostic default `max_length` (=20) to control the generation length. We recommend setting `max_new_tokens` to control the maximum length of the generation.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "875b7692c6c54c4e8bf2ed24931c5eea",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/72 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'rouge1': 0.10831759043534087, 'rouge2': 0.033586619395055244, 'rougeL': 0.1069870424439844, 'rougeLsum': 0.10780701430965048}\n",
      "{'eval_loss': 0.34586578607559204, 'eval_rouge1': 10.8318, 'eval_rouge2': 3.3587, 'eval_rougeL': 10.6987, 'eval_rougeLsum': 10.7807, 'eval_runtime': 30.9598, 'eval_samples_per_second': 18.379, 'eval_steps_per_second': 2.326, 'epoch': 8.0}\n",
      "{'train_runtime': 4788.0587, 'train_samples_per_second': 13.295, 'train_steps_per_second': 1.662, 'train_loss': 0.3884411848969196, 'epoch': 8.0}\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=7960, training_loss=0.3884411848969196, metrics={'train_runtime': 4788.0587, 'train_samples_per_second': 13.295, 'train_steps_per_second': 1.662, 'train_loss': 0.3884411848969196, 'epoch': 8.0})"
      ]
     },
     "execution_count": 90,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = subset_dataset[\"train\"]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'input_ids': [2638, 677, 6, 45805, 82, 563, 384, 85, 3140, 115, 4808, 142, 2725, 2605, 518, 33, 4295, 3188, 2775, 33, 415, 25253, 1958, 13, 577, 5620, 44, 803, 1361, 8, 687, 822, 10, 2807, 31, 836, 12462, 13, 1635, 214, 3472, 34, 1164, 2223, 213, 384, 22332, 2832, 13, 2716, 31, 1113, 11487, 59, 3191, 34154, 45, 18217, 8618, 8, 45805, 81, 984, 53388, 3857, 47, 115, 460, 36785, 55, 34, 2775, 31, 687, 822, 10, 2807, 45, 1488, 1849, 23072, 16, 346, 11369, 40, 1794, 373, 82, 563, 384, 85, 3140, 115, 4808, 142, 33, 19578, 33, 15, 8, 115, 460, 36785, 13, 633, 7, 8449, 430, 462, 56337, 34, 295, 7, 131, 8574, 16, 9021, 1657, 45, 2725, 518, 33, 1074, 33, 14460, 391, 20, 1958, 13, 577, 34154, 1768, 16424, 801, 33, 691, 11587, 2564, 8, 199, 5755, 72, 38105, 13, 186, 7, 285, 1113, 2547, 485, 33, 22489, 44, 4985, 15, 7, 3043, 285, 12462, 13, 17134, 13, 186, 43, 36180, 7913, 4830, 735, 257, 15, 8, 518, 5487, 7202, 13, 323, 14460, 391, 20, 31484, 45, 3863, 304, 13, 186, 45805, 81, 4584, 23153, 81, 2538, 1855, 13, 814, 1112, 217, 122, 8, 2002, 31, 7, 10663, 1708, 7551, 7, 3249, 518, 143, 1980, 122, 7, 34, 9524, 31, 45805, 3140, 5502, 31, 8438, 8252, 31, 464, 1422, 243, 8, 70, 127, 11464, 34, 295, 7, 131, 8862, 7337, 34, 846, 14460, 391, 20, 1958, 1850, 13, 577, 5487, 45805, 45, 9597, 25364, 49, 8, 285, 7462, 1290, 33, 23153, 45, 883, 605, 1133, 13, 186, 15, 7, 146, 5487, 45805, 29423, 45, 1236, 381, 42, 45805, 13, 6633, 379, 4373, 10628, 13, 577, 3887, 13, 15834, 45, 2716, 5150, 3651, 582, 8, 199, 23119, 381, 42, 10668, 245, 257, 82, 563, 384, 85, 3140, 5502, 142, 822, 6649, 395, 25812, 1583, 45805, 21752, 42, 2340, 15, 8, 32681, 213, 384, 85, 3140, 5502, 45, 8912, 16351, 7, 48177, 38196, 23812, 46, 415, 24952, 81, 15721, 27, 1467, 43, 40397, 31964, 5502, 13, 323, 10668, 245, 257, 15, 8, 21752, 13, 6587, 13, 3610, 3755, 45, 2064, 31, 13542, 463, 7, 708, 3140, 5502, 184, 13, 186, 328, 128, 16, 20, 242, 143, 82, 953, 408, 2135, 142, 14278, 13, 323, 6772, 3596, 211, 21752, 42, 3315, 122, 8, 142, 563, 384, 85, 3140, 5502, 142, 12462, 31, 3496, 381, 42, 4672, 15, 81, 285, 409, 16546, 33, 660, 15, 37, 2725, 518, 33, 12462, 1958, 13, 577, 45805, 13, 21752, 42, 3315, 708, 3140, 5502, 184, 45, 44726, 245, 2236, 7, 146, 45805, 13, 1590, 3454, 18718, 7997, 434, 70, 2409, 122, 8, 616, 295, 7, 131, 953, 408, 42, 3315, 213, 384, 85, 3140, 5502, 73, 5329, 13, 3496, 81, 2538, 1855, 184, 31, 4672, 5618, 4123, 8, 285, 45805, 45, 14007, 81, 14460, 391, 20, 7102, 46, 143, 3454, 31584, 46, 24856, 26086, 8, 199, 64001, 64006, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 'labels': [64006, 3343, 36785, 13, 633, 7, 8449, 430, 462, 56337, 34, 295, 7, 131, 8574, 16, 9021, 1657, 45, 2725, 518, 33, 1074, 33, 14460, 391, 20, 1958, 13, 577, 34154, 1768, 16424, 801, 33, 691, 11587, 2564, 8, 199, 64001, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]}\n",
      "1024\n"
     ]
    }
   ],
   "source": [
    "i = 5\n",
    "print(x[i])\n",
    "print(min(len(x[z][\"input_ids\"]) for z in range(0, 700)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.evaluate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.cuda.empty_cache()\n",
    "print(torch.cuda.memory_allocated())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inp = tokenizer(\"I am a boy </s> <2en>\", add_special_tokens=False, return_tensors=\"pt\", padding=True).input_ids # tensor([[  466,  1981,    80, 25573, 64001, 64004]])\n",
    "\n",
    "out = tokenizer(\"<2hi> मैं  एक लड़का हूँ </s>\", add_special_tokens=False, return_tensors=\"pt\", padding=True).input_ids # tensor([[64006,   942,    43, 32720,  8384, 64001]])\n",
    "# Note that if you use any language other than Hindi or Marathi, you should convert its script to Devanagari using the Indic NLP Library.\n",
    "inp.to(device)\n",
    "out.to(device)\n",
    "\n",
    "model_outputs=model(input_ids=inp, decoder_input_ids=out[:,0:-1], labels=out[:,1:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "string = \"मुंबई इंडियंस ने आईपीएल 2024 के 14वें मैच में पहले बैटिंग करते हुए 20 ओवर में 9 विकेट खोकर 125 रन बनाए। राजस्थान रॉयल्स के गेंदबाजों के सामने मुंबई का टॉप-ऑर्डर फ्लॉप रहा। मुंबई के तीन बल्लेबाजों को ट्रेंट बोल्ट ने शून्य पर पवेलियन भेजा। हार्दिक और तिलक वर्मा के बीच अर्धशतकीय साझेदारी बनी। तिलक के बल्ले से 32 रन और हार्दिक ने 34 रन की पारी खेली। बोल्ट के अलावा राजस्थान के लिए चहल ने भी 3 विकेट चटकाए। इसके जवाब में 126 रन का पीछा करते हुए राजस्थान की टीम ने 15.3 ओवर में ही लक्ष्य हासिल कर लिया। राजस्थान की टीम की तरफ से रियाग पराग ने नाबाद 54 रन की पारी खेली, जिसमें 5 चौके और 3 छक्के शामिल रहे।\"\n",
    "inp = tokenizer(f\"{string} </s> <2hi>\", add_special_tokens=False, return_tensors=\"pt\", padding=True).input_ids\n",
    "inp = inp.to(model.device)\n",
    "model_output=model.generate(inp, use_cache=True,\n",
    "    min_length=30, #increase min length,\n",
    "    num_return_sequences=1,\n",
    "    do_sample=True,\n",
    "    top_k=50,\n",
    "    top_p=0.95,\n",
    "    temperature=1.0,\n",
    "    num_beams=5,  # Use beam search\n",
    "    repetition_penalty=2.0,  # Increase repetition penalty\n",
    "    length_penalty=1.0,  # Adjust length penalty\n",
    "    early_stopping=True,\n",
    "    pad_token_id=pad_id,\n",
    "    bos_token_id=bos_id,\n",
    "    eos_token_id=eos_id,\n",
    "    decoder_start_token_id=tokenizer._convert_token_to_id_with_added_voc(\"<2hi>\"))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "decoded_output=tokenizer.decode(model_output[0], skip_special_tokens=True, clean_up_tokenization_spaces=False)\n",
    "decoded_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = tokenized_dataset[\"train\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x[\"input_ids\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using Accelerator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "subset_dataset.set_format(\"torch\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "\n",
    "batch_size = 8\n",
    "train_dataloader = DataLoader(\n",
    "    subset_dataset[\"train\"],\n",
    "    shuffle=True,\n",
    "    collate_fn=data_collator,\n",
    "    batch_size=batch_size,\n",
    ")\n",
    "eval_dataloader = DataLoader(\n",
    "    subset_dataset[\"validation\"], collate_fn=data_collator, batch_size=batch_size\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "\n",
    "batch_size = 8\n",
    "train_dataloader = DataLoader(\n",
    "    subset_dataset[\"train\"],\n",
    "    shuffle=True,\n",
    "    collate_fn=data_collator,\n",
    "    batch_size=batch_size,\n",
    ")\n",
    "eval_dataloader = DataLoader(\n",
    "    subset_dataset[\"validation\"], collate_fn=data_collator, batch_size=batch_size\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.optim import AdamW\n",
    "\n",
    "optimizer = AdamW(model.parameters(), lr=1e-5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from accelerate import Accelerator\n",
    "\n",
    "accelerator = Accelerator()\n",
    "model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(\n",
    "    model, optimizer, train_dataloader, eval_dataloader\n",
    ")\n",
    "print(optimizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import get_scheduler\n",
    "\n",
    "num_train_epochs = 20\n",
    "num_update_steps_per_epoch = len(train_dataloader)\n",
    "num_training_steps = num_train_epochs * num_update_steps_per_epoch\n",
    "\n",
    "lr_scheduler = get_scheduler(\n",
    "    \"linear\",\n",
    "    optimizer=optimizer,\n",
    "    num_warmup_steps=0,\n",
    "    num_training_steps=num_training_steps,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import indicnlp.tokenize.sentence_tokenize as tok\n",
    "\n",
    "def postprocess_text(preds, labels):\n",
    "    preds = [pred.strip() for pred in preds]\n",
    "    labels = [label.strip() for label in labels]\n",
    "\n",
    "    preds = [\"\\n\".join(tok.sentence_split(pred.strip(), lang='hi', delim_pat='auto')) for pred in decoded_preds]\n",
    "    labels = [\"\\n\".join(tok.sentence_split(label.strip(), lang='hi', delim_pat='auto')) for label in decoded_labels]\n",
    "\n",
    "    return decoded_preds, decoded_labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm.auto import tqdm\n",
    "import torch\n",
    "import numpy as np\n",
    "\n",
    "progress_bar = tqdm(range(num_training_steps))\n",
    "\n",
    "for epoch in range(num_train_epochs):\n",
    "    # Training\n",
    "    model.train()\n",
    "    for step, batch in enumerate(train_dataloader):\n",
    "        outputs = model(**batch)\n",
    "        loss = outputs.loss\n",
    "        record.append(loss)\n",
    "        accelerator.backward(loss)\n",
    "        optimizer.step()\n",
    "        lr_scheduler.step()\n",
    "        optimizer.zero_grad()\n",
    "        progress_bar.update(1)\n",
    "\n",
    "    # Evaluation\n",
    "    model.eval()\n",
    "    for step, batch in enumerate(eval_dataloader):\n",
    "        with torch.no_grad():\n",
    "            generated_tokens = accelerator.unwrap_model(model).generate(\n",
    "                batch[\"input_ids\"],\n",
    "                generation_config=generation_config\n",
    "            )\n",
    "\n",
    "            generated_tokens = accelerator.pad_across_processes(\n",
    "                generated_tokens, dim=1, pad_index=tokenizer.pad_token_id\n",
    "            )\n",
    "            labels = batch[\"labels\"]\n",
    "\n",
    "            # If we did not pad to max length, we need to pad the labels too\n",
    "            labels = accelerator.pad_across_processes(\n",
    "                batch[\"labels\"], dim=1, pad_index=tokenizer.pad_token_id\n",
    "            )\n",
    "\n",
    "            generated_tokens = accelerator.gather(generated_tokens).cpu().numpy()\n",
    "            labels = accelerator.gather(labels).cpu().numpy()\n",
    "\n",
    "            # Replace -100 in the labels as we can't decode them\n",
    "            labels = np.where(labels != -100, labels, tokenizer.pad_token_id)\n",
    "            if isinstance(generated_tokens, tuple):\n",
    "                generated_tokens = generated_tokens[0]\n",
    "            decoded_preds = tokenizer.batch_decode(\n",
    "                generated_tokens, skip_special_tokens=True\n",
    "            )\n",
    "            decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True, clean_up_tokenization_spaces=False)\n",
    "\n",
    "            decoded_preds, decoded_labels = postprocess_text(\n",
    "                decoded_preds, decoded_labels\n",
    "            )\n",
    "\n",
    "            rouge_score.add_batch(predictions=decoded_preds, references=decoded_labels)\n",
    "\n",
    "    # Compute metrics\n",
    "    result = rouge_score.compute()\n",
    "    # Extract the median ROUGE scores\n",
    "    wandb.log(result)\n",
    "    result = {key: value * 100 for key, value in result.items()}\n",
    "    result = {k: round(v, 4) for k, v in result.items()}\n",
    "    print(f\"Epoch {epoch}:\", result, f\"Loss: {loss}\")\n",
    "\n",
    "    # Save and upload\n",
    "    \"\"\"accelerator.wait_for_everyone()\n",
    "    unwrapped_model = accelerator.unwrap_model(model)\n",
    "    unwrapped_model.save_pretrained(output_dir, save_function=accelerator.save)\n",
    "    if accelerator.is_main_process:\n",
    "        tokenizer.save_pretrained(output_dir)\n",
    "        repo.push_to_hub(\n",
    "            commit_message=f\"Training in progress epoch {epoch}\", blocking=False\n",
    "        )\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x.append([y.cpu().detach().item() for y in record])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "z = np.array(x[-1])\n",
    "z = np.transpose(z)\n",
    "z.size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "\n",
    "# Create x-axis values (epochs or steps)\n",
    "epochs = range(1, len(record) + 1)\n",
    "\n",
    "# Plot the loss\n",
    "plt.plot(epochs, z, 'b', label='Training loss')\n",
    "plt.title('Training Loss')\n",
    "plt.xlabel('Steps')\n",
    "plt.ylabel('Loss')\n",
    "plt.legend()\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "def sizeof_fmt(num, suffix='B'):\n",
    "    ''' by Fred Cirera,  https://stackoverflow.com/a/1094933/1870254, modified'''\n",
    "    for unit in ['','Ki','Mi','Gi','Ti','Pi','Ei','Zi']:\n",
    "        if abs(num) < 1024.0:\n",
    "            return \"%3.1f %s%s\" % (num, unit, suffix)\n",
    "        num /= 1024.0\n",
    "    return \"%.1f %s%s\" % (num, 'Yi', suffix)\n",
    "\n",
    "for name, size in sorted(((name, sys.getsizeof(value)) for name, value in list(\n",
    "                          locals().items())), key= lambda x: -x[1]):\n",
    "    print(\"{:>30}: {:>8}\".format(name, sizeof_fmt(size)))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "\n",
    "local_vars = list(locals().items())\n",
    "for var, obj in local_vars:\n",
    "    print(var, sys.getsizeof(obj))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "bart",
   "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.11.8"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}