BALM-paired.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# training BALM-paired\n",
    "\n",
    "#### architecture\n",
    "BALM-unpaired is built on the RoBERTa-large architecture, with the following hyperparameter modifications:\n",
    "* max input length of **512**, which was selected to be suitable for paired sequences and to be double the size of BALM-unpaired\n",
    "* per-GPU batch size of **32**, which is 0.5x the batch size of BALM-unpaired to equalize training steps between the two models\n",
    "* **500k** training steps, which should be roughly 200 epochs when trained on eight GPUs for a total batch size of **256**\n",
    "\n",
    "#### dataset\n",
    "We used all unique, productive paired sequences reported in [_Functional antibodies exhibit light chain coherence_](https://www.nature.com/articles/s41586-022-05371-z) (Jaffe et al, Nature 2022)\n",
    "* [dataset DOI](https://plus.figshare.com/articles/dataset/Dataset_supporting_Functional_antibodies_exhibit_light_chain_coherence_/20338177) (figshare)\n",
    "* the dataset was split in to train/eval/test subsets at a ratio of 90:5:5, which produced the following dataset sizes:\n",
    "    * **train**: `1,202,269` paired sequences\n",
    "    * **eval**: `66,791` paired sequences\n",
    "    * **test**: `66,971` paired sequences  \n",
    "    \n",
    "Each input file contains a single paired antibody amino acid sequence per line, with the heavy and light chain sequences concatenated and separated by a `</s>` token, like so:  \n",
    "  > `HEAVY_CHAIN_AA_SEQUENCE</s>LIGHT_CHAIN_AA_SEQUENCE`\n",
    "\n",
    "#### training\n",
    "Training BALM-paired on eight NVIDIA A100 GPUs took approximately 5 days."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>  \n",
    "  \n",
    "## setup  \n",
    "  \n",
    "Training BALM-paired requires several [huggingface](https://huggingface.co/) libraries. If they're not already installed, you can install them by uncommenting and running the following code block:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install transformers\n",
    "# !pip install datasets\n",
    "# !pip install accelerate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from datetime import date\n",
    "import os\n",
    "import warnings\n",
    "warnings.simplefilter('ignore')\n",
    "\n",
    "from transformers import (\n",
    "    RobertaConfig,\n",
    "    RobertaTokenizer,\n",
    "    RobertaForMaskedLM,\n",
    "    DataCollatorForLanguageModeling,\n",
    "    TrainingArguments,\n",
    "    Trainer,\n",
    ")\n",
    "\n",
    "from datasets import load_dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## BALM config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run_name = f\"BALM-paired_lc-coherence-data_90-5-5-split_{date.today().isoformat()}\"\n",
    "print(f\"Run name: {run_name}\")\n",
    "\n",
    "balm_config = {\n",
    "    \"run_name\": run_name,\n",
    "    \n",
    "    # model architecture\n",
    "    \"num_hidden_layers\": 24,\n",
    "    \"num_attention_heads\": 16,\n",
    "    \"hidden_size\": 1024,\n",
    "    \"intermediate_size\": 4096,\n",
    "    \"vocab_size\": 25,\n",
    "    \"max_len\": 512,\n",
    "    \"max_position_embeddings\": 514,\n",
    "    \n",
    "    # tokenizer\n",
    "    \"padding\": \"max_length\",\n",
    "    \"truncate\": True,\n",
    "    \"return_special_tokens_mask\": True,\n",
    "    \n",
    "    # training parameters\n",
    "    \"batch_size\": 32,\n",
    "    \"max_steps\": 500000,\n",
    "    \"warmup_steps\": 30000,\n",
    "    \"weight_decay\": 0.01,\n",
    "    \"peak_learning_rate\": 4e-4,\n",
    "    \"adam_epsilon\": 1e-6,\n",
    "    \"adam_beta1\": 0.9,\n",
    "    \"adam_beta2\": 0.98,\n",
    "    \"type_vocab_size\": 2,  # this should be 2 for paired/mixed models, 1 for unpaired models\n",
    "    \"fp16\": True,\n",
    "    \"evaluation_strategy\": \"steps\",\n",
    "    \"seed\": 42,\n",
    "    \n",
    "    # outputs and logging\n",
    "    \"save_steps\": 100000,\n",
    "    \"eval_steps\": 25000,\n",
    "    \"output_dir\": f\"./checkpoints/{run_name}\",  # where the checkpoint data will be written\n",
    "    \"logging_dir\": f\"./logs/{run_name}\",\n",
    "    \"logging_steps\": 100,\n",
    "    \"overwrite_output_dir\": True,\n",
    "    \"logging_first_step\": True,\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>  \n",
    "  \n",
    "If you'd like to use [weights and biases](https://wandb.ai) for logging, uncomment and run the following code block:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# os.environ[\"WANDB_PROJECT\"] = run_name\n",
    "# balm_config[\"report_to\"] = \"wandb\"\n",
    "\n",
    "# import wandb\n",
    "# wandb.login()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# initialize the model using the BALM config dictionary\n",
    "# defaults are based on what was used in the paper\n",
    "model_config = RobertaConfig(\n",
    "    vocab_size=balm_config.get(\"vocab_size\", 25),\n",
    "    hidden_size=balm_config.get(\"hidden_size\", 1024),\n",
    "    intermediate_size=balm_config.get(\"intermediate_size\", 4096),\n",
    "    max_position_embeddings=balm_config.get(\"max_position_embeddings\", 512),\n",
    "    num_hidden_layers=balm_config.get(\"num_hidden_layers\", 24),\n",
    "    num_attention_heads=balm_config.get(\"num_attention_heads\", 16),\n",
    "    type_vocab_size=balm_config.get(\"type_vocab_size\", 2),\n",
    ")\n",
    "    \n",
    "model = RobertaForMaskedLM(model_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_size = sum(p.numel() for p in model.parameters())\n",
    "print(f\"Model size: {model_size/1e6:.2f}M\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## load + tokenize data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "# download the train/eval/test data if it doesn't exist\n",
    "if [ ! -d \"./data/train-test-eval_paired\" ]; then\n",
    "    curl -o 'train-test-eval_paired.tar.gz' -L 'https://zenodo.org/record/8253367/files/train-test-eval_paired.tar.gz?download=1'\n",
    "    tar xzvf 'train-test-eval_paired.tar.gz' -C ./data\n",
    "    rm 'train-test-eval_paired.tar.gz'\n",
    "fi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# load the tran, eval, and test data\n",
    "data_files = {\n",
    "    \"train\": ['./data/train-test-eval_paired/train.txt'],\n",
    "    \"eval\": ['./data/train-test-eval_paired/eval.txt'],\n",
    "    \"test\": ['./data/train-test-eval_paired/test.txt']\n",
    "}\n",
    "\n",
    "dataset = load_dataset(\"text\", data_files=data_files)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "tokenizer = RobertaTokenizer.from_pretrained(\"tokenizer\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "tokenized_dataset = dataset.map(\n",
    "    lambda x: tokenizer(\n",
    "        x[\"text\"],\n",
    "        padding=balm_config.get(\"padding\", \"max_length\"),\n",
    "        truncation=balm_config.get(\"truncation\", True),\n",
    "        max_length=balm_config.get(\"max_len\", 512),\n",
    "        return_special_tokens_mask=balm_config.get(\"return_special_tokens_mask\", True),\n",
    "    ),\n",
    "    remove_columns=[\"text\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## data collator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "collator = DataCollatorForLanguageModeling(\n",
    "    tokenizer=tokenizer, mlm=True, mlm_probability=0.15\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## trainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "training_args = TrainingArguments(\n",
    "    fp16=balm_config.get(\"fp16\", True),\n",
    "    evaluation_strategy=balm_config.get(\"evaluation_strategy\", \"steps\"),\n",
    "    seed=balm_config.get(\"seed\", 42),\n",
    "    per_device_train_batch_size=balm_config.get(\"batch_size\", 32),\n",
    "    per_device_eval_batch_size=balm_config.get(\"batch_size\", 32),\n",
    "    max_steps=balm_config.get(\"max_steps\", 500000),\n",
    "    save_steps=balm_config.get(\"save_steps\", 100000),\n",
    "    logging_steps=balm_config.get(\"logging_steps\", 100),\n",
    "    eval_steps=balm_config.get(\"eval_steps\", 25000),\n",
    "    adam_beta1=balm_config.get(\"adam_beta1\", 0.9),\n",
    "    adam_beta2=balm_config.get(\"adam_beta2\", 0.98),\n",
    "    adam_epsilon=balm_config.get(\"adam_epsilon\", 1e-6),\n",
    "    weight_decay=balm_config.get(\"weight_decay\", 0.01),\n",
    "    warmup_steps=balm_config.get(\"warmup_steps\", 30000),\n",
    "    learning_rate=balm_config.get(\"peak_learning_rate\", 4e-4),\n",
    "    gradient_accumulation_steps=balm_config.get(\"gradient_accumulation_steps\", 1),\n",
    "    \n",
    "    # output and logging\n",
    "    run_name=balm_config.get(\"run_name\", None),\n",
    "    output_dir=balm_config.get(\"output_dir\", f\"./checkpoints/{run_name}\"),\n",
    "    overwrite_output_dir=balm_config.get(\"overwrite_output_dir\", True),\n",
    "    logging_dir=balm_config.get(\"logging_dir\", f\"./logs/{run_name}\"),\n",
    "    report_to=balm_config.get(\"report_to\", None),\n",
    "    logging_first_step=balm_config.get(\"logging_first_step\", True),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    data_collator=collator,\n",
    "    train_dataset=tokenized_dataset[\"train\"],\n",
    "    eval_dataset=tokenized_dataset[\"eval\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.save_model(f\"../models/{run_name}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# wandb.finish()"
   ]
  }
 ],
 "metadata": {
  "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.9.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}