protein_language_modeling_prostT5.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b4b220d6",
   "metadata": {},
   "source": [
    "### This notebook used for finetuning the ProstT5 model. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c5bf8d4",
   "metadata": {
    "id": "4c5bf8d4"
   },
   "outputs": [],
   "source": [
    "! pip install transformers[torch] evaluate datasets requests pandas scikit-learn peft bitsandbytes matplotlib sentencepiece accelerate \n",
    "#deepspeed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19e8f77c",
   "metadata": {
    "id": "19e8f77c"
   },
   "outputs": [],
   "source": [
    "# !apt install git-lfs\n",
    "# pip install wandb\n",
    "# wandb login"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c0749e1",
   "metadata": {
    "id": "5c0749e1"
   },
   "source": [
    "# Fine-Tuning Protein Language Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "32e605a2",
   "metadata": {
    "id": "32e605a2"
   },
   "outputs": [],
   "source": [
    "model_checkpoint = \"Rostlab/ProstT5\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5bc122f",
   "metadata": {
    "id": "c5bc122f"
   },
   "source": [
    "## Data preparation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7560840d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['.git', 'README.md', 'cath_domain_list.list', 'domain_list_nonredundant_s40_v4_3.list', 'eda.ipynb', 'foldseek_seq3d.csv', 'queryDB_ss.fasta', 'sequences_nonredundant_s40_v4_3.fa', 'train.csv', 'val.csv', 'protein_language_modeling_prostT5.ipynb', '.ipynb_checkpoints']\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "\n",
    "# Change to the desired directory\n",
    "# os.chdir(\"/root\")\n",
    "# Verify the change\n",
    "print(os.listdir(\"./\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e4408970-e9d6-4c5c-8460-2f282af68ef2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "train_df = pd.read_csv(\"./train.csv\")\n",
    "val_df = pd.read_csv(\"./val.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "95f26046-e096-4995-80f3-4a5819fba45b",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_sequences = train_df[\"sequences\"].tolist()\n",
    "train_sequences_3d = train_df[\"seq3d\"].tolist()\n",
    "train_labels = train_df[\"label\"].tolist()\n",
    "\n",
    "test_sequences = val_df[\"sequences\"].tolist()\n",
    "test_sequences_3d = val_df[\"seq3d\"].tolist()\n",
    "test_labels = val_df[\"label\"].tolist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ddbe2b2d",
   "metadata": {
    "colab": {
     "referenced_widgets": [
      "cc61f599adc641da8d40eefac0179aa3",
      "4e3d35bc47c74852a5256f5b34312030",
      "8b7f3b4a47d8437f9e9cc6648cfc8984"
     ]
    },
    "id": "ddbe2b2d",
    "outputId": "131761f1-5609-4fc4-f4cc-5ed87c662774"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "You are using the default legacy behaviour of the <class 'transformers.models.t5.tokenization_t5.T5Tokenizer'>. This is expected, and simply means that the `legacy` (previous) behavior will be used so nothing changes for you. If you want to use the new behaviour, set `legacy=False`. This should only be set if you understand what it means, and thoroughly read the reason why this was added as explained in https://github.com/huggingface/transformers/pull/24565\n",
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    }
   ],
   "source": [
    "from transformers import T5Tokenizer\n",
    "\n",
    "tokenizer = T5Tokenizer.from_pretrained(model_checkpoint, do_lower_case=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "4e8f3ddf",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import T5Tokenizer\n",
    "from datasets import Dataset\n",
    "import re\n",
    "\n",
    "\n",
    "def preprocess_data(sequences, structures, tokenizer):\n",
    "    tokenized_sequences = []\n",
    "    tokenized_structures = []\n",
    "\n",
    "    for sequence, structure in zip(sequences, structures):\n",
    "        # Preprocess sequences\n",
    "        sequence = \" \".join(list(re.sub(r\"[UZOB]\", \"X\", sequence)))\n",
    "        structure = \" \".join(list(structure))\n",
    "\n",
    "        sequence = \"<AA2fold> \" + sequence if sequence.isupper() else sequence\n",
    "        structure = \"<fold2AA> \" + structure\n",
    "\n",
    "        # Tokenize sequences and structures\n",
    "        sequence_inputs = tokenizer(\n",
    "            sequence, add_special_tokens=True, padding=\"longest\", return_tensors=\"pt\"\n",
    "        )\n",
    "        structure_inputs = tokenizer(\n",
    "            structure, add_special_tokens=True, padding=\"longest\", return_tensors=\"pt\"\n",
    "        )\n",
    "\n",
    "        tokenized_sequences.append(sequence_inputs)\n",
    "        tokenized_structures.append(structure_inputs)\n",
    "\n",
    "    return tokenized_sequences, tokenized_structures\n",
    "\n",
    "\n",
    "# Preprocess and tokenize the data\n",
    "train_tokenized_sequences, train_tokenized_structures = preprocess_data(\n",
    "    train_sequences, train_sequences_3d, tokenizer\n",
    ")\n",
    "test_tokenized_sequences, test_tokenized_structures = preprocess_data(\n",
    "    test_sequences, test_sequences_3d, tokenizer\n",
    ")\n",
    "\n",
    "\n",
    "import datasets\n",
    "\n",
    "\n",
    "def create_dataset(tokenized_sequences, tokenized_structures, labels):\n",
    "    input_ids_sequence = [item[\"input_ids\"].squeeze() for item in tokenized_sequences]\n",
    "    attention_mask_sequence = [\n",
    "        item[\"attention_mask\"].squeeze() for item in tokenized_sequences\n",
    "    ]\n",
    "    input_ids_structure = [item[\"input_ids\"].squeeze() for item in tokenized_structures]\n",
    "    attention_mask_structure = [\n",
    "        item[\"attention_mask\"].squeeze() for item in tokenized_structures\n",
    "    ]\n",
    "\n",
    "    dataset_dict = {\n",
    "        \"input_ids_sequence\": input_ids_sequence,\n",
    "        \"attention_mask_sequence\": attention_mask_sequence,\n",
    "        \"input_ids_structure\": input_ids_structure,\n",
    "        \"attention_mask_structure\": attention_mask_structure,\n",
    "        \"labels\": labels,\n",
    "    }\n",
    "\n",
    "    return Dataset.from_dict(dataset_dict)\n",
    "\n",
    "\n",
    "# Create Dataset objects\n",
    "train_dataset = create_dataset(\n",
    "    train_tokenized_sequences, train_tokenized_structures, train_labels\n",
    ")\n",
    "test_dataset = create_dataset(\n",
    "    test_tokenized_sequences, test_tokenized_structures, test_labels\n",
    ")\n",
    "\n",
    "from torch.utils.data import DataLoader\n",
    "import torch\n",
    "\n",
    "\n",
    "class ProteinDataset(torch.utils.data.Dataset):\n",
    "    def __init__(self, hf_dataset):\n",
    "        self.dataset = hf_dataset\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.dataset)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        item = self.dataset[idx]\n",
    "        return {\n",
    "            \"input_ids_sequence\": torch.tensor(item[\"input_ids_sequence\"]),\n",
    "            \"attention_mask_sequence\": torch.tensor(item[\"attention_mask_sequence\"]),\n",
    "            \"input_ids_structure\": torch.tensor(item[\"input_ids_structure\"]),\n",
    "            \"attention_mask_structure\": torch.tensor(item[\"attention_mask_structure\"]),\n",
    "            \"labels\": torch.tensor(item[\"labels\"]),\n",
    "        }\n",
    "\n",
    "\n",
    "# Create custom dataset\n",
    "train_dataset = ProteinDataset(train_dataset)\n",
    "test_dataset = ProteinDataset(test_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af074a5c",
   "metadata": {
    "id": "af074a5c"
   },
   "source": [
    "## Model loading\n",
    "\n",
    "I create a customized new model class that combines the prostT5 model outputs and a classification layer, in such a way that we can optionally choose in the finetuning to use either the sequence embeddings, the structure embeddings or an average/concat of both to feed to the classifier."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "79d6bd60",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "You are using a model of type t5 to instantiate a model of type . This is not supported for all configurations of models and can yield errors.\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from transformers import T5EncoderModel, T5ForSequenceClassification, T5PreTrainedModel\n",
    "from transformers import PretrainedConfig, T5Config\n",
    "\n",
    "\n",
    "class T5ClassificationHead(nn.Module):\n",
    "    \"\"\"Head for sentence-level classification tasks.\"\"\"\n",
    "\n",
    "    def __init__(self, config: T5Config):\n",
    "        super().__init__()\n",
    "        self.dense = nn.Linear(config.d_model, config.d_model)\n",
    "        self.dropout = nn.Dropout(p=config.classifier_dropout)\n",
    "        self.out_proj = nn.Linear(config.d_model, config.num_labels)\n",
    "\n",
    "    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:\n",
    "        hidden_states = self.dropout(hidden_states)\n",
    "        hidden_states = self.dense(hidden_states)\n",
    "        hidden_states = torch.tanh(hidden_states)\n",
    "        hidden_states = self.dropout(hidden_states)\n",
    "        hidden_states = self.out_proj(hidden_states)\n",
    "        return hidden_states\n",
    "\n",
    "\n",
    "class CustomT5ForSequenceClassification(T5PreTrainedModel):\n",
    "    def __init__(self, model_checkpoint, config):\n",
    "\n",
    "        super().__init__(config)\n",
    "        self.transformer = T5EncoderModel.from_pretrained(model_checkpoint)\n",
    "        # self.classifier = nn.Linear(1024, num_labels).to(device)\n",
    "        self.classification_head = T5ClassificationHead(config)\n",
    "\n",
    "    def forward(\n",
    "        self,\n",
    "        input_ids_sequence,\n",
    "        input_ids_structure,\n",
    "        attention_mask_sequence=None,\n",
    "        attention_mask_structure=None,\n",
    "        labels=None,\n",
    "    ):\n",
    "        # here we could add an augmentation step in which we choose x% of the time the sequence only, the structure only, or both\n",
    "\n",
    "        # Get embeddings for the sequence\n",
    "        sequence_outputs = self.transformer(\n",
    "            input_ids_sequence, attention_mask=attention_mask_sequence\n",
    "        )\n",
    "        sequence_embeddings = sequence_outputs.last_hidden_state.mean(dim=1)\n",
    "\n",
    "        # Get embeddings for the structure\n",
    "        structure_outputs = self.transformer(\n",
    "            input_ids_structure, attention_mask=attention_mask_structure\n",
    "        )\n",
    "        structure_embeddings = structure_outputs.last_hidden_state.mean(dim=1)\n",
    "\n",
    "        # Combine the embeddings\n",
    "        combined_embeddings = (\n",
    "            sequence_embeddings + structure_embeddings\n",
    "        ) / 2.0  # can be changed to concatenation but the embedding d_model in the config should be adjusted accordingly\n",
    "\n",
    "        # combined_embeddings = sequence_embeddings\n",
    "\n",
    "        # Feed to classifier head\n",
    "        logits = self.classification_head(combined_embeddings)\n",
    "\n",
    "        loss = None\n",
    "        if labels is not None:\n",
    "            loss_fct = nn.CrossEntropyLoss()\n",
    "            loss = loss_fct(logits.view(-1, self.config.num_labels), labels.view(-1))\n",
    "\n",
    "        return (loss, logits) if loss is not None else logits\n",
    "\n",
    "\n",
    "# Example usage\n",
    "preconfig = PretrainedConfig.from_pretrained(model_checkpoint)\n",
    "num_labels = max(train_labels + test_labels) + 1\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "preconfig.update({\"num_labels\": num_labels, \"classifier_dropout\": 0.1})\n",
    "# preconfig.update({\"num_labels\": num_labels, \"classifier_dropout\": 0.3, \"dropout_rate\": 0.3}) if more dropout desired\n",
    "\n",
    "model = CustomT5ForSequenceClassification(model_checkpoint, preconfig).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "738b8e06",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "env: WANDB_WATCH=all\n",
      "env: WANDB_SILENT=true\n",
      "env: WANDB_LOG_MODEL=end\n",
      "env: WANDB_PROJECT=protein cath classification\n"
     ]
    }
   ],
   "source": [
    "%env WANDB_WATCH=all\n",
    "%env WANDB_SILENT=true\n",
    "%env WANDB_LOG_MODEL=end\n",
    "%env WANDB_PROJECT=protein cath classification\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "version = 1\n",
    "batch_size = 32\n",
    "train_epochs = 100\n",
    "num_workers = 8\n",
    "lr = 1e-5 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "775cb3e7",
   "metadata": {
    "id": "775cb3e7"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.10/site-packages/transformers/training_args.py:1494: FutureWarning: `evaluation_strategy` is deprecated and will be removed in version 4.46 of 🤗 Transformers. Use `eval_strategy` instead\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "from transformers import EarlyStoppingCallback\n",
    "from transformers import TrainingArguments, Trainer\n",
    "\n",
    "early_stopping = EarlyStoppingCallback(early_stopping_patience=5)\n",
    "\n",
    "model_name = model_checkpoint.split(\"/\")[-1]\n",
    "\n",
    "args = TrainingArguments(\n",
    "    f\"{model_name}-finetuned\",\n",
    "    evaluation_strategy=\"epoch\",\n",
    "    eval_strategy=\"epoch\",\n",
    "    save_strategy=\"epoch\",\n",
    "    learning_rate=lr,\n",
    "    per_device_train_batch_size=batch_size,\n",
    "    per_device_eval_batch_size=batch_size,\n",
    "    num_train_epochs=train_epochs,\n",
    "    weight_decay=0.01,\n",
    "    load_best_model_at_end=True,\n",
    "    metric_for_best_model=\"f1\",\n",
    "    greater_is_better=True,\n",
    "    push_to_hub=False,\n",
    "    fp16=True,\n",
    "    fp16_full_eval=True,\n",
    "    # bf16_full_eval=True,\n",
    "    # bf16=True,\n",
    "    save_total_limit=1,\n",
    "    gradient_checkpointing=True,\n",
    "    optim=\"adamw_torch\",\n",
    "    report_to=\"wandb\",\n",
    "    lr_scheduler_type=\"cosine\",\n",
    "    warmup_ratio=0.01,\n",
    "    logging_strategy=\"epoch\",\n",
    "    run_name=f\"{model_checkpoint.split('/')[-1]}-v-{version}\",\n",
    "    dataloader_num_workers=num_workers,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "471cef9f",
   "metadata": {
    "colab": {
     "referenced_widgets": [
      "55d5a21f0be54ea1b33c552fed2c3bd3"
     ]
    },
    "id": "471cef9f",
    "outputId": "b096a8bf-35db-4707-8538-d067865e87a4"
   },
   "outputs": [],
   "source": [
    "import evaluate\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "def compute_metrics(eval_preds):\n",
    "    metric = evaluate.combine([\"f1\", \"precision\", \"recall\"])\n",
    "    logits, labels = eval_preds\n",
    "    predictions = np.argmax(logits, axis=-1)\n",
    "    metrics = metric.compute(\n",
    "        predictions=predictions, references=labels, average=\"weighted\"\n",
    "    )\n",
    "    return metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "e212b751",
   "metadata": {
    "id": "e212b751",
    "outputId": "6bb7723b-2621-4514-ef05-e23e887d4e28"
   },
   "outputs": [],
   "source": [
    "from torch.utils.data.dataloader import default_collate\n",
    "\n",
    "\n",
    "def custom_collate_fn(batch):\n",
    "    # Extract elements\n",
    "    input_ids_sequence = [item[\"input_ids_sequence\"] for item in batch]\n",
    "    attention_mask_sequence = [item[\"attention_mask_sequence\"] for item in batch]\n",
    "    input_ids_structure = [item[\"input_ids_structure\"] for item in batch]\n",
    "    attention_mask_structure = [item[\"attention_mask_structure\"] for item in batch]\n",
    "    labels = [item[\"labels\"] for item in batch]\n",
    "\n",
    "    # Pad sequences to the maximum length in the batch\n",
    "    input_ids_sequence_padded = torch.nn.utils.rnn.pad_sequence(\n",
    "        input_ids_sequence, batch_first=True, padding_value=tokenizer.pad_token_id\n",
    "    )\n",
    "    attention_mask_sequence_padded = torch.nn.utils.rnn.pad_sequence(\n",
    "        attention_mask_sequence, batch_first=True, padding_value=0\n",
    "    )\n",
    "    input_ids_structure_padded = torch.nn.utils.rnn.pad_sequence(\n",
    "        input_ids_structure, batch_first=True, padding_value=tokenizer.pad_token_id\n",
    "    )\n",
    "    attention_mask_structure_padded = torch.nn.utils.rnn.pad_sequence(\n",
    "        attention_mask_structure, batch_first=True, padding_value=0\n",
    "    )\n",
    "\n",
    "    labels = torch.stack(labels)\n",
    "\n",
    "    return {\n",
    "        \"input_ids_sequence\": input_ids_sequence_padded,\n",
    "        \"attention_mask_sequence\": attention_mask_sequence_padded,\n",
    "        \"input_ids_structure\": input_ids_structure_padded,\n",
    "        \"attention_mask_structure\": attention_mask_structure_padded,\n",
    "        \"labels\": labels,\n",
    "    }\n",
    "\n",
    "\n",
    "trainer = Trainer(\n",
    "    model,\n",
    "    args,\n",
    "    train_dataset=train_dataset,\n",
    "    eval_dataset=test_dataset,\n",
    "    tokenizer=tokenizer,\n",
    "    compute_metrics=compute_metrics,\n",
    "    callbacks=[early_stopping],\n",
    "    data_collator=custom_collate_fn,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "9c3cf6da",
   "metadata": {
    "id": "9c3cf6da",
    "outputId": "a0780d09-fb66-4f07-b7e6-782fcc85e34b",
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.10/site-packages/torch/utils/checkpoint.py:460: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n",
      "Could not estimate the number of tokens of the input, floating-point operations will not be computed\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='3744' max='20800' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [ 3744/20800 2:01:59 < 9:16:00, 0.51 it/s, Epoch 18/100]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Epoch</th>\n",
       "      <th>Training Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "      <th>F1</th>\n",
       "      <th>Precision</th>\n",
       "      <th>Recall</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>2.095700</td>\n",
       "      <td>1.547083</td>\n",
       "      <td>0.548896</td>\n",
       "      <td>0.614960</td>\n",
       "      <td>0.549742</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2</td>\n",
       "      <td>0.814700</td>\n",
       "      <td>0.960532</td>\n",
       "      <td>0.724462</td>\n",
       "      <td>0.734050</td>\n",
       "      <td>0.721444</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3</td>\n",
       "      <td>0.312800</td>\n",
       "      <td>0.934696</td>\n",
       "      <td>0.736411</td>\n",
       "      <td>0.743351</td>\n",
       "      <td>0.736183</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4</td>\n",
       "      <td>0.152000</td>\n",
       "      <td>1.045662</td>\n",
       "      <td>0.739156</td>\n",
       "      <td>0.745034</td>\n",
       "      <td>0.740604</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>5</td>\n",
       "      <td>0.085400</td>\n",
       "      <td>1.132233</td>\n",
       "      <td>0.747653</td>\n",
       "      <td>0.762495</td>\n",
       "      <td>0.748710</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>6</td>\n",
       "      <td>0.052200</td>\n",
       "      <td>1.254073</td>\n",
       "      <td>0.740434</td>\n",
       "      <td>0.759642</td>\n",
       "      <td>0.742815</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>7</td>\n",
       "      <td>0.034000</td>\n",
       "      <td>1.310171</td>\n",
       "      <td>0.749363</td>\n",
       "      <td>0.757534</td>\n",
       "      <td>0.747237</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>8</td>\n",
       "      <td>0.027200</td>\n",
       "      <td>1.367738</td>\n",
       "      <td>0.751142</td>\n",
       "      <td>0.761050</td>\n",
       "      <td>0.747973</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>9</td>\n",
       "      <td>0.017200</td>\n",
       "      <td>1.401283</td>\n",
       "      <td>0.743830</td>\n",
       "      <td>0.748560</td>\n",
       "      <td>0.747973</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>10</td>\n",
       "      <td>0.011600</td>\n",
       "      <td>1.526976</td>\n",
       "      <td>0.753323</td>\n",
       "      <td>0.765194</td>\n",
       "      <td>0.750921</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>11</td>\n",
       "      <td>0.015000</td>\n",
       "      <td>1.524332</td>\n",
       "      <td>0.745480</td>\n",
       "      <td>0.756784</td>\n",
       "      <td>0.747973</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>12</td>\n",
       "      <td>0.010600</td>\n",
       "      <td>1.558307</td>\n",
       "      <td>0.744336</td>\n",
       "      <td>0.752671</td>\n",
       "      <td>0.747973</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>13</td>\n",
       "      <td>0.008900</td>\n",
       "      <td>1.569710</td>\n",
       "      <td>0.753776</td>\n",
       "      <td>0.763895</td>\n",
       "      <td>0.753869</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>14</td>\n",
       "      <td>0.007200</td>\n",
       "      <td>1.658265</td>\n",
       "      <td>0.751295</td>\n",
       "      <td>0.759923</td>\n",
       "      <td>0.752395</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>15</td>\n",
       "      <td>0.006400</td>\n",
       "      <td>1.742706</td>\n",
       "      <td>0.746066</td>\n",
       "      <td>0.762146</td>\n",
       "      <td>0.748710</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>16</td>\n",
       "      <td>0.005200</td>\n",
       "      <td>1.755780</td>\n",
       "      <td>0.739074</td>\n",
       "      <td>0.747839</td>\n",
       "      <td>0.739867</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>17</td>\n",
       "      <td>0.003400</td>\n",
       "      <td>1.753823</td>\n",
       "      <td>0.746907</td>\n",
       "      <td>0.755695</td>\n",
       "      <td>0.750184</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>18</td>\n",
       "      <td>0.004200</td>\n",
       "      <td>1.828261</td>\n",
       "      <td>0.745064</td>\n",
       "      <td>0.757839</td>\n",
       "      <td>0.745763</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.10/site-packages/torch/utils/checkpoint.py:460: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n",
      "/opt/conda/lib/python3.10/site-packages/torch/utils/checkpoint.py:460: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n",
      "/opt/conda/lib/python3.10/site-packages/torch/utils/checkpoint.py:460: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n",
      "/opt/conda/lib/python3.10/site-packages/torch/utils/checkpoint.py:460: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n",
      "/opt/conda/lib/python3.10/site-packages/torch/utils/checkpoint.py:460: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n",
      "Using the latest cached version of the module from /root/.cache/huggingface/modules/evaluate_modules/metrics/evaluate-metric--f1/0ca73f6cf92ef5a268320c697f7b940d1030f8471714bffdb6856c641b818974 (last modified on Fri Jul 12 15:11:03 2024) since it couldn't be found locally at evaluate-metric--f1, or remotely on the Hugging Face Hub.\n",
      "/opt/conda/lib/python3.10/site-packages/torch/utils/checkpoint.py:460: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n",
      "/opt/conda/lib/python3.10/site-packages/torch/utils/checkpoint.py:460: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n",
      "/opt/conda/lib/python3.10/site-packages/torch/utils/checkpoint.py:460: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n",
      "/opt/conda/lib/python3.10/site-packages/torch/utils/checkpoint.py:460: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n",
      "/opt/conda/lib/python3.10/site-packages/torch/utils/checkpoint.py:460: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n",
      "/opt/conda/lib/python3.10/site-packages/torch/utils/checkpoint.py:460: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n",
      "/opt/conda/lib/python3.10/site-packages/torch/utils/checkpoint.py:460: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n",
      "/opt/conda/lib/python3.10/site-packages/torch/utils/checkpoint.py:460: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n",
      "/opt/conda/lib/python3.10/site-packages/torch/utils/checkpoint.py:460: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n",
      "Using the latest cached version of the module from /root/.cache/huggingface/modules/evaluate_modules/metrics/evaluate-metric--f1/0ca73f6cf92ef5a268320c697f7b940d1030f8471714bffdb6856c641b818974 (last modified on Fri Jul 12 15:11:03 2024) since it couldn't be found locally at evaluate-metric--f1, or remotely on the Hugging Face Hub.\n",
      "/opt/conda/lib/python3.10/site-packages/torch/utils/checkpoint.py:460: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n",
      "/opt/conda/lib/python3.10/site-packages/torch/utils/checkpoint.py:460: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n",
      "/opt/conda/lib/python3.10/site-packages/torch/utils/checkpoint.py:460: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n",
      "There were missing keys in the checkpoint model loaded: ['transformer.encoder.embed_tokens.weight'].\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=3744, training_loss=0.20354116903856778, metrics={'train_runtime': 7311.208, 'train_samples_per_second': 90.86, 'train_steps_per_second': 2.845, 'total_flos': 0.0, 'train_loss': 0.20354116903856778, 'epoch': 18.0})"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c67b81d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import wandb\n",
    "\n",
    "wandb.finish()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6508b178",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1e2baf8-beca-47ee-a23c-823503b885a6",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}