notebooks/pretraining.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "19b9940d-ef44-4f95-92b0-b0cb3316c747",
   "metadata": {},
   "source": [
    "# Pretrain METL on Rosetta data\n",
    "\n",
    "This notebook demonstrates how to pretrain a METL source model on Rosetta data.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c67101d4-81c1-4092-b39a-53d74ad234ca",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-16T22:51:36.573559Z",
     "start_time": "2024-02-16T22:51:36.569490Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ca743fb9-7041-4688-91bb-9cd800c6dfdd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "\n",
    "# define the name of the project root directory\n",
    "project_root_dir_name = \"metl\"\n",
    "\n",
    "# find the project root by checking each parent directory\n",
    "current_dir = os.getcwd()\n",
    "while os.path.basename(current_dir) != project_root_dir_name and current_dir != os.path.dirname(current_dir):\n",
    "    current_dir = os.path.dirname(current_dir)\n",
    "\n",
    "# change the current working directory to the project root directory\n",
    "if os.path.basename(current_dir) == project_root_dir_name:\n",
    "    os.chdir(current_dir)\n",
    "else:\n",
    "    print(\"project root directory not found\")\n",
    "    \n",
    "# add the project code folder to the system path so imports work\n",
    "module_path = os.path.abspath(\"code\")\n",
    "if module_path not in sys.path:\n",
    "    sys.path.append(module_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4fd163d-007e-44ca-8e5b-0049b0a09213",
   "metadata": {},
   "source": [
    "# Training arguments\n",
    "\n",
    "The script for pretraining on Rosetta data is [train_source_model.py](../code/train_source_model.py).\n",
    "This script has a number of arguments to specify various aspects of the model, training, and logging.\n",
    "You can view the arguments by uncommenting and running the following line."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5df63031-b4d5-4d44-94aa-a9a3ec16f791",
   "metadata": {},
   "outputs": [],
   "source": [
    "# %run code/train_source_model.py -h"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64e777ce-a24e-4338-a5a7-23f593c09a58",
   "metadata": {},
   "source": [
    "Note this won't show model-specific arguments. For information on what arguments to use to set up the model, see the function in [code/models.py](../code/models.py) that corresponds to the `--model_name` you want to use.\n",
    "\n",
    "We will set up arguments to pretrain a toy METL-Local model on the sample avGFP Rosetta dataset located in [data/rosetta_data](data/rosetta_data). See the README in that directory for more information about this sample dataset. \n",
    "\n",
    "The arguments are contained in the file [pretrain_avgfp_local.txt](../args/pretrain_avgfp_local.txt) in the [args](../args) directory. Uncomment and run the cell below to view the contents of the argument file. The sections below will walk through and explain the key arguments."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "30600142-23a7-4fb2-8442-a7ad6ccac8bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# with open(\"args/pretrain_avgfp_local.txt\", \"r\") as file:\n",
    "#     contents = file.read()\n",
    "#     print(contents)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef0fb348-195a-476d-bd66-327357708902",
   "metadata": {},
   "source": [
    "## Dataset arguments"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "447d57c3-f687-495b-8422-a1fd4de3e283",
   "metadata": {},
   "source": [
    "Define the path to the Rosetta dataset (the .db file).\n",
    "\n",
    "```\n",
    "--ds_fn\n",
    "data/rosetta_data/avgfp/avgfp.db\n",
    "```\n",
    "\n",
    "Define the path to the train/val/test split to use for this run.\n",
    "\n",
    "```\n",
    "--split_dir\n",
    "data/rosetta_data/avgfp/splits/standard_tr0.8_tu0.1_te0.1_w1aea30517f4f_r4991\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a3e1530-e3d3-4773-8af3-bac3637c1862",
   "metadata": {},
   "source": [
    "## Optimizer arguments\n",
    "Source models are trained with the **AdamW** optimizer by default. See `RosettaTask.configure_optimizers()` in [tasks.py](../code/tasks.py) if you want to change the optimizer."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a1b44f9-2ee1-4219-9296-16211fcafe66",
   "metadata": {},
   "source": [
    "Basic optimizer arguments include the batch size, learning rate, and maximum number of epochs to train for. Unless early stopping is enabled, the model will train for the given number of epochs. \n",
    "\n",
    "```\n",
    "--batch_size\n",
    "128\n",
    "--learning_rate\n",
    "0.001\n",
    "--max_epochs\n",
    "30\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45b4cdaf-f521-4705-ba1a-cacb75b30ffd",
   "metadata": {},
   "source": [
    "Learning rate schedule determines how the learning rate changes over the course of training. We implemented a few different schedules and settled on a constant learning rate with a linear warmup period for our models. The number of warmup steps can be specified as an integer or a fraction of total training steps (determined by number of epochs and batch size).\n",
    "\n",
    "```\n",
    "--lr_scheduler\n",
    "warmup_constant\n",
    "--warmup_steps\n",
    ".02\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce91ed10-a1ee-4080-9ff8-3e70dfdf5763",
   "metadata": {},
   "source": [
    "Gradient clipping can help prevent exploding gradients and make training smoother. We used gradient norm clipping with a threshold of 0.5.\n",
    "```\n",
    "--gradient_clip_val\n",
    "0.5\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbb79830-fc79-4baa-a1b6-5e6ce3832679",
   "metadata": {},
   "source": [
    "## Model architecture arguments\n",
    "\n",
    "METL models are transformer-based models. Our implementation supports a number of options including structure-based (3D) or sequence-based (1D) relative position embeddings, different numbers of layers, etc. \n",
    "\n",
    "The arguments below are what we used for our 2M parameter METL-Local models."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "953f935c-1ba2-437b-9cb7-75e75079d53a",
   "metadata": {},
   "source": [
    "```\n",
    "--model_name\n",
    "transformer_encoder\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d1dd4ae-be6f-4bd4-90c1-0b00c111410a",
   "metadata": {},
   "source": [
    "The 3D structure-based relative position embeddings need a specified contact threshold and clipping threshold. The contact threshold determines the cutoff distance for contact map / structure graph, which the 3D RPE uses to compute neighboring residues. The clipping threshold determines the maximum relative distance cutoff. With a clipping threshold of 3, a relative distance of 0 represents a node with itself, 1 signifies direct neighbors, 2 signifies second degree neighbors, and 3 encapsulates any other node not covered by the previous categories.\n",
    "\n",
    "```\n",
    "--pos_encoding\n",
    "relative_3D\n",
    "--contact_threshold\n",
    "8\n",
    "--clipping_threshold\n",
    "3\n",
    "```\n",
    "\n",
    "To use a 1D sequence-based embedding, you would instead specify `relative` for the `pos_encoding` and a `clipping_threshold` determining the maximum sequence-based relative distance cutoff. We used `clipping_threshold` of `8` for our 1D-based models."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a044eea-c863-47f7-a380-c0021dae1308",
   "metadata": {},
   "source": [
    "The following arguments determine the transformer encoder architecture.\n",
    "```\n",
    "--embedding_len\n",
    "256\n",
    "--num_hidden\n",
    "1024\n",
    "--num_heads\n",
    "4\n",
    "--num_enc_layers\n",
    "3\n",
    "--enc_layer_dropout\n",
    "0.1\n",
    "--use_final_encoder_norm\n",
    "--global_average_pooling\n",
    "--use_final_hidden_layer\n",
    "--final_hidden_size\n",
    "256\n",
    "--use_final_hidden_layer_norm\n",
    "--final_hidden_layer_norm_before_activation\n",
    "--use_final_hidden_layer_dropout\n",
    "--final_hidden_layer_dropout_rate\n",
    "0.1\n",
    "--activation\n",
    "relu\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d058b2b-0ea5-4768-a2d9-d2660a5177fc",
   "metadata": {},
   "source": [
    "## Logging arguments\n",
    "We have built in functionality for tracking model training with Weights & Biases. If you have a Weights and Biases account, set the environment variable `WANDB_API_KEY` to your API key and set the flag `--use_wandb` instead of `--no_use_wandb` below.\n",
    "\n",
    "```\n",
    "--no-use_wandb\n",
    "--wandb_project\n",
    "metl-source\n",
    "--wandb_online\n",
    "--experiment\n",
    "default\n",
    "```\n",
    "\n",
    "Another flag that may be of interest is `--wandb_log_grad` which will additionally log gradients."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1251ac44-6040-4363-95f2-6237f0d40ee5",
   "metadata": {},
   "source": [
    "The below argument determines where to place the log directory locally.\n",
    "```\n",
    "--log_dir_base\n",
    "output/training_logs\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b115337-7f38-44f7-a7b4-ec10fbaa2c56",
   "metadata": {},
   "source": [
    "# Running training\n",
    "In addition to the arguments explained above and defined the in the arguments file, we are going to limit the amount of training and testing data using `--limit_(train,val,test)_batches` to speed up training for this example. We also overwrite the number of training epochs from 30 to 5 to speed up this example.\n",
    "\n",
    "The training script will do the following:\n",
    "1. Assign a UUID for this model\n",
    "2. Create a log directory with the UUID name in the base log directory [output/training_logs](../output/training_logs)\n",
    "3. Print out various information including a summary of the model architecture (the output may not fit in this Jupyter Notebook)\n",
    "4. Initiate training, which will be tracked via a progress bar printed in this notebook\n",
    "5. Evaluate the final model on the test set and print metrics\n",
    "\n",
    "Running the training script will generate output in the cell below and place training logs and checkpoints in the log directory in `output/training_logs`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c7df5c76-276c-4cf5-bc52-b3827d60975a",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Global seed set to 1\n",
      "Created model UUID: En8Ux9wH\n",
      "Created log directory: output/training_logs/En8Ux9wH\n",
      "Final UUID: En8Ux9wH\n",
      "Final log directory: output/training_logs/En8Ux9wH\n",
      "This is version: 0\n",
      "Version-specific logs will be saved to: output/training_logs/En8Ux9wH/version_0\n",
      "No checkpoint found, training from scratch\n",
      "Using example_input_array with pdb_fn='1gfl_cm.pdb' and aa_seq_len=237\n",
      "Trainer already configured with model summary callbacks: [<class 'pytorch_lightning.callbacks.rich_model_summary.RichModelSummary'>]. Skipping setting a default `ModelSummary` callback.\n",
      "GPU available: True (mps), used: False\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n",
      "/opt/homebrew/Caskroom/mambaforge/base/envs/metl/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py:1789: UserWarning: MPS available but not used. Set `accelerator` and `devices` using `Trainer(accelerator='mps', devices=1)`.\n",
      "  rank_zero_warn(\n",
      "Loading `train_dataloader` to estimate number of stepping batches.\n",
      "/opt/homebrew/Caskroom/mambaforge/base/envs/metl/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py:1892: PossibleUserWarning: The number of training batches (5) is smaller than the logging interval Trainer(log_every_n_steps=50). Set a lower value for log_every_n_steps if you want to see logs for the training epoch.\n",
      "  rank_zero_warn(\n",
      "Number of training steps is 25\n",
      "Number of warmup steps is 0.5\n",
      "┏┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳┳┳┳┓\n",
      "┃┃\u001b[1;35m \u001b[0m\u001b[1;35mName                                                                   \u001b[0m\u001b[1;35m \u001b[0m┃┃┃┃┃\n",
      "┡╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇╇╇╇┩\n",
      "││ model                                                                   │││││\n",
      "││ model.model                                                             │││││\n",
      "││ model.model.embedder                                                    │││││\n",
      "││ model.model.embedder.embedding                                          │││││\n",
      "││ model.model.tr_encoder                                                  │││││\n",
      "││ model.model.tr_encoder.layers                                           │││││\n",
      "││ model.model.tr_encoder.layers.0                                         │││││\n",
      "││ model.model.tr_encoder.layers.0.self_attn                               │││││\n",
      "││ model.model.tr_encoder.layers.0.self_attn.relative_position_k           │││││\n",
      "││ model.model.tr_encoder.layers.0.self_attn.relative_position_k.embeddin… │││││\n",
      "││ model.model.tr_encoder.layers.0.self_attn.relative_position_v           │││││\n",
      "││ model.model.tr_encoder.layers.0.self_attn.relative_position_v.embeddin… │││││\n",
      "││ model.model.tr_encoder.layers.0.self_attn.q_proj                        │││││\n",
      "││ model.model.tr_encoder.layers.0.self_attn.k_proj                        │││││\n",
      "││ model.model.tr_encoder.layers.0.self_attn.v_proj                        │││││\n",
      "││ model.model.tr_encoder.layers.0.self_attn.out_proj                      │││││\n",
      "││ model.model.tr_encoder.layers.0.self_attn.dropout                       │││││\n",
      "││ model.model.tr_encoder.layers.0.linear1                                 │││││\n",
      "││ model.model.tr_encoder.layers.0.dropout                                 │││││\n",
      "││ model.model.tr_encoder.layers.0.linear2                                 │││││\n",
      "││ model.model.tr_encoder.layers.0.norm1                                   │││││\n",
      "││ model.model.tr_encoder.layers.0.norm2                                   │││││\n",
      "││ model.model.tr_encoder.layers.0.dropout1                                │││││\n",
      "││ model.model.tr_encoder.layers.0.dropout2                                │││││\n",
      "││ model.model.tr_encoder.layers.1                                         │││││\n",
      "││ model.model.tr_encoder.layers.1.self_attn                               │││││\n",
      "││ model.model.tr_encoder.layers.1.self_attn.relative_position_k           │││││\n",
      "││ model.model.tr_encoder.layers.1.self_attn.relative_position_k.embeddin… │││││\n",
      "││ model.model.tr_encoder.layers.1.self_attn.relative_position_v           │││││\n",
      "││ model.model.tr_encoder.layers.1.self_attn.relative_position_v.embeddin… │││││\n",
      "││ model.model.tr_encoder.layers.1.self_attn.q_proj                        │││││\n",
      "││ model.model.tr_encoder.layers.1.self_attn.k_proj                        │││││\n",
      "││ model.model.tr_encoder.layers.1.self_attn.v_proj                        │││││\n",
      "││ model.model.tr_encoder.layers.1.self_attn.out_proj                      │││││\n",
      "││ model.model.tr_encoder.layers.1.self_attn.dropout                       │││││\n",
      "││ model.model.tr_encoder.layers.1.linear1                                 │││││\n",
      "││ model.model.tr_encoder.layers.1.dropout                                 │││││\n",
      "││ model.model.tr_encoder.layers.1.linear2                                 │││││\n",
      "││ model.model.tr_encoder.layers.1.norm1                                   │││││\n",
      "││ model.model.tr_encoder.layers.1.norm2                                   │││││\n",
      "││ model.model.tr_encoder.layers.1.dropout1                                │││││\n",
      "││ model.model.tr_encoder.layers.1.dropout2                                │││││\n",
      "││ model.model.tr_encoder.layers.2                                         │││││\n",
      "││ model.model.tr_encoder.layers.2.self_attn                               │││││\n",
      "││ model.model.tr_encoder.layers.2.self_attn.relative_position_k           │││││\n",
      "││ model.model.tr_encoder.layers.2.self_attn.relative_position_k.embeddin… │││││\n",
      "││ model.model.tr_encoder.layers.2.self_attn.relative_position_v           │││││\n",
      "││ model.model.tr_encoder.layers.2.self_attn.relative_position_v.embeddin… │││││\n",
      "││ model.model.tr_encoder.layers.2.self_attn.q_proj                        │││││\n",
      "││ model.model.tr_encoder.layers.2.self_attn.k_proj                        │││││\n",
      "││ model.model.tr_encoder.layers.2.self_attn.v_proj                        │││││\n",
      "││ model.model.tr_encoder.layers.2.self_attn.out_proj                      │││││\n",
      "││ model.model.tr_encoder.layers.2.self_attn.dropout                       │││││\n",
      "││ model.model.tr_encoder.layers.2.linear1                                 │││││\n",
      "││ model.model.tr_encoder.layers.2.dropout                                 │││││\n",
      "││ model.model.tr_encoder.layers.2.linear2                                 │││││\n",
      "││ model.model.tr_encoder.layers.2.norm1                                   │││││\n",
      "││ model.model.tr_encoder.layers.2.norm2                                   │││││\n",
      "││ model.model.tr_encoder.layers.2.dropout1                                │││││\n",
      "││ model.model.tr_encoder.layers.2.dropout2                                │││││\n",
      "││ model.model.tr_encoder.norm                                             │││││\n",
      "││ model.model.avg_pooling                                                 │││││\n",
      "││ model.model.fc1                                                         │││││\n",
      "││ model.model.fc1.fc                                                      │││││\n",
      "││ model.model.fc1.activation                                              │││││\n",
      "││ model.model.fc1.norm                                                    │││││\n",
      "││ model.model.fc1.dropout                                                 │││││\n",
      "││ model.model.prediction                                                  │││││\n",
      "││ test_pearson_collection                                                 │││││\n",
      "││ test_pearson_collection.0                                               │││││\n",
      "││ test_pearson_collection.1                                               │││││\n",
      "││ test_pearson_collection.2                                               │││││\n",
      "││ test_pearson_collection.3                                               │││││\n",
      "││ test_pearson_collection.4                                               │││││\n",
      "││ test_pearson_collection.5                                               │││││\n",
      "││ test_pearson_collection.6                                               │││││\n",
      "││ test_pearson_collection.7                                               │││││\n",
      "││ test_pearson_collection.8                                               │││││\n",
      "││ test_pearson_collection.9                                               │││││\n",
      "││ test_pearson_collection.10                                              │││││\n",
      "││ test_pearson_collection.11                                              │││││\n",
      "││ test_pearson_collection.12                                              │││││\n",
      "││ test_pearson_collection.13                                              │││││\n",
      "││ test_pearson_collection.14                                              │││││\n",
      "││ test_pearson_collection.15                                              │││││\n",
      "││ test_pearson_collection.16                                              │││││\n",
      "││ test_pearson_collection.17                                              │││││\n",
      "││ test_pearson_collection.18                                              │││││\n",
      "││ test_pearson_collection.19                                              │││││\n",
      "││ test_pearson_collection.20                                              │││││\n",
      "││ test_pearson_collection.21                                              │││││\n",
      "││ test_pearson_collection.22                                              │││││\n",
      "││ test_pearson_collection.23                                              │││││\n",
      "││ test_pearson_collection.24                                              │││││\n",
      "││ test_pearson_collection.25                                              │││││\n",
      "││ test_pearson_collection.26                                              │││││\n",
      "││ test_pearson_collection.27                                              │││││\n",
      "││ test_pearson_collection.28                                              │││││\n",
      "││ test_pearson_collection.29                                              │││││\n",
      "││ test_pearson_collection.30                                              │││││\n",
      "││ test_pearson_collection.31                                              │││││\n",
      "││ test_pearson_collection.32                                              │││││\n",
      "││ test_pearson_collection.33                                              │││││\n",
      "││ test_pearson_collection.34                                              │││││\n",
      "││ test_pearson_collection.35                                              │││││\n",
      "││ test_pearson_collection.36                                              │││││\n",
      "││ test_pearson_collection.37                                              │││││\n",
      "││ test_pearson_collection.38                                              │││││\n",
      "││ test_pearson_collection.39                                              │││││\n",
      "││ test_pearson_collection.40                                              │││││\n",
      "││ test_pearson_collection.41                                              │││││\n",
      "││ test_pearson_collection.42                                              │││││\n",
      "││ test_pearson_collection.43                                              │││││\n",
      "││ test_pearson_collection.44                                              │││││\n",
      "││ test_pearson_collection.45                                              │││││\n",
      "││ test_pearson_collection.46                                              │││││\n",
      "││ test_pearson_collection.47                                              │││││\n",
      "││ test_pearson_collection.48                                              │││││\n",
      "││ test_pearson_collection.49                                              │││││\n",
      "││ test_pearson_collection.50                                              │││││\n",
      "││ test_pearson_collection.51                                              │││││\n",
      "││ test_pearson_collection.52                                              │││││\n",
      "││ test_pearson_collection.53                                              │││││\n",
      "││ test_pearson_collection.54                                              │││││\n",
      "└┴─────────────────────────────────────────────────────────────────────────┴┴┴┴┘\n",
      "\u001b[1mTrainable params\u001b[0m: 2.5 M                                                         \n",
      "\u001b[1mNon-trainable params\u001b[0m: 0                                                         \n",
      "\u001b[1mTotal params\u001b[0m: 2.5 M                                                             \n",
      "\u001b[1mTotal estimated model params size (MB)\u001b[0m: 9                                       \n",
      "/opt/homebrew/Caskroom/mambaforge/base/envs/metl/lib/python3.9/site-packages/pytorch_lightning/loggers/csv_logs.py:57: UserWarning: Experiment logs directory output/training_logs/En8Ux9wH/version_0 exists and is not empty. Previous log files in this directory will be deleted when the new ones are saved!\n",
      "  rank_zero_warn(\n",
      "Sanity Checking DataLoader 0:   0%|                       | 0/2 [00:00<?, ?it/s]/opt/homebrew/Caskroom/mambaforge/base/envs/metl/lib/python3.9/site-packages/pytorch_lightning/utilities/data.py:98: UserWarning: Trying to infer the `batch_size` from an ambiguous collection. The batch size we found is 128. To avoid any miscalculations, use `self.log(..., batch_size=batch_size)`.\n",
      "  warning_cache.warn(\n",
      "Epoch 0:  50%|███████▌       | 5/10 [00:12<00:12,  2.47s/it, loss=63.6, v_num=0]\n",
      "Validation: 0it [00:00, ?it/s]\u001b[A\n",
      "Validation:   0%|                                         | 0/5 [00:00<?, ?it/s]\u001b[A\n",
      "Validation DataLoader 0:   0%|                            | 0/5 [00:00<?, ?it/s]\u001b[A\n",
      "Epoch 0:  60%|█████████      | 6/10 [00:12<00:08,  2.17s/it, loss=63.6, v_num=0]\u001b[A\n",
      "Epoch 0:  70%|██████████▌    | 7/10 [00:13<00:05,  1.93s/it, loss=63.6, v_num=0]\u001b[A\n",
      "Epoch 0:  80%|████████████   | 8/10 [00:13<00:03,  1.75s/it, loss=63.6, v_num=0]\u001b[A\n",
      "Epoch 0:  90%|█████████████▌ | 9/10 [00:14<00:01,  1.61s/it, loss=63.6, v_num=0]\u001b[A\n",
      "Epoch 0: 100%|█| 10/10 [00:14<00:00,  1.49s/it, loss=63.6, v_num=0, val_loss=56.\u001b[A\n",
      "Epoch 1:  50%|▌| 5/10 [00:10<00:10,  2.05s/it, loss=59.7, v_num=0, val_loss=56.8\u001b[A\n",
      "Validation: 0it [00:00, ?it/s]\u001b[A\n",
      "Validation:   0%|                                         | 0/5 [00:00<?, ?it/s]\u001b[A\n",
      "Validation DataLoader 0:   0%|                            | 0/5 [00:00<?, ?it/s]\u001b[A\n",
      "Epoch 1:  60%|▌| 6/10 [00:10<00:07,  1.81s/it, loss=59.7, v_num=0, val_loss=56.8\u001b[A\n",
      "Epoch 1:  70%|▋| 7/10 [00:11<00:04,  1.62s/it, loss=59.7, v_num=0, val_loss=56.8\u001b[A\n",
      "Epoch 1:  80%|▊| 8/10 [00:11<00:02,  1.48s/it, loss=59.7, v_num=0, val_loss=56.8\u001b[A\n",
      "Epoch 1:  90%|▉| 9/10 [00:12<00:01,  1.36s/it, loss=59.7, v_num=0, val_loss=56.8\u001b[A\n",
      "Epoch 1: 100%|█| 10/10 [00:12<00:00,  1.27s/it, loss=59.7, v_num=0, val_loss=56.\u001b[A\n",
      "Epoch 2:  50%|▌| 5/10 [00:10<00:10,  2.06s/it, loss=58.7, v_num=0, val_loss=56.5\u001b[A\n",
      "Validation: 0it [00:00, ?it/s]\u001b[A\n",
      "Validation:   0%|                                         | 0/5 [00:00<?, ?it/s]\u001b[A\n",
      "Validation DataLoader 0:   0%|                            | 0/5 [00:00<?, ?it/s]\u001b[A\n",
      "Epoch 2:  60%|▌| 6/10 [00:10<00:07,  1.81s/it, loss=58.7, v_num=0, val_loss=56.5\u001b[A\n",
      "Epoch 2:  70%|▋| 7/10 [00:11<00:04,  1.62s/it, loss=58.7, v_num=0, val_loss=56.5\u001b[A\n",
      "Epoch 2:  80%|▊| 8/10 [00:11<00:02,  1.48s/it, loss=58.7, v_num=0, val_loss=56.5\u001b[A\n",
      "Epoch 2:  90%|▉| 9/10 [00:12<00:01,  1.37s/it, loss=58.7, v_num=0, val_loss=56.5\u001b[A\n",
      "Epoch 2: 100%|█| 10/10 [00:12<00:00,  1.28s/it, loss=58.7, v_num=0, val_loss=56.\u001b[A\n",
      "Epoch 3:  50%|▌| 5/10 [00:10<00:10,  2.05s/it, loss=57.6, v_num=0, val_loss=56.3\u001b[A\n",
      "Validation: 0it [00:00, ?it/s]\u001b[A\n",
      "Validation:   0%|                                         | 0/5 [00:00<?, ?it/s]\u001b[A\n",
      "Validation DataLoader 0:   0%|                            | 0/5 [00:00<?, ?it/s]\u001b[A\n",
      "Epoch 3:  60%|▌| 6/10 [00:10<00:07,  1.81s/it, loss=57.6, v_num=0, val_loss=56.3\u001b[A\n",
      "Epoch 3:  70%|▋| 7/10 [00:11<00:04,  1.62s/it, loss=57.6, v_num=0, val_loss=56.3\u001b[A\n",
      "Epoch 3:  80%|▊| 8/10 [00:11<00:02,  1.48s/it, loss=57.6, v_num=0, val_loss=56.3\u001b[A\n",
      "Epoch 3:  90%|▉| 9/10 [00:12<00:01,  1.36s/it, loss=57.6, v_num=0, val_loss=56.3\u001b[A\n",
      "Epoch 3: 100%|█| 10/10 [00:12<00:00,  1.28s/it, loss=57.6, v_num=0, val_loss=56.\u001b[A\n",
      "Epoch 4:  50%|▌| 5/10 [00:10<00:10,  2.06s/it, loss=56, v_num=0, val_loss=56.40,\u001b[A\n",
      "Validation: 0it [00:00, ?it/s]\u001b[A\n",
      "Validation:   0%|                                         | 0/5 [00:00<?, ?it/s]\u001b[A\n",
      "Validation DataLoader 0:   0%|                            | 0/5 [00:00<?, ?it/s]\u001b[A\n",
      "Epoch 4:  60%|▌| 6/10 [00:10<00:07,  1.82s/it, loss=56, v_num=0, val_loss=56.40,\u001b[A\n",
      "Epoch 4:  70%|▋| 7/10 [00:11<00:04,  1.62s/it, loss=56, v_num=0, val_loss=56.40,\u001b[A\n",
      "Epoch 4:  80%|▊| 8/10 [00:11<00:02,  1.48s/it, loss=56, v_num=0, val_loss=56.40,\u001b[A\n",
      "Epoch 4:  90%|▉| 9/10 [00:12<00:01,  1.37s/it, loss=56, v_num=0, val_loss=56.40,\u001b[A\n",
      "Epoch 4: 100%|█| 10/10 [00:12<00:00,  1.28s/it, loss=56, v_num=0, val_loss=56.20\u001b[A\n",
      "Epoch 4: 100%|█| 10/10 [00:12<00:00,  1.28s/it, loss=56, v_num=0, val_loss=56.20\u001b[A`Trainer.fit` stopped: `max_epochs=5` reached.\n",
      "Epoch 4: 100%|█| 10/10 [00:12<00:00,  1.28s/it, loss=56, v_num=0, val_loss=56.20\n",
      "output/training_logs/En8Ux9wH/checkpoints/epoch=4-step=25-val_loss=56.23.ckpt\n",
      "tensor(56.2257)\n",
      "/opt/homebrew/Caskroom/mambaforge/base/envs/metl/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py:1397: UserWarning: `.test(ckpt_path=\"best\")` is called with Trainer configured with multiple `ModelCheckpoint` callbacks. It will use the best checkpoint path from first checkpoint callback.\n",
      "  rank_zero_warn(\n",
      "Restoring states from the checkpoint path at output/training_logs/En8Ux9wH/checkpoints/epoch=4-step=25-val_loss=56.23.ckpt\n",
      "Loaded model weights from checkpoint at output/training_logs/En8Ux9wH/checkpoints/epoch=4-step=25-val_loss=56.23.ckpt\n",
      "Testing DataLoader 0: 100%|███████████████████████| 5/5 [00:02<00:00,  1.89it/s]\n",
      "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
      "┃\u001b[1m \u001b[0m\u001b[1m            Test metric             \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m            DataLoader 0             \u001b[0m\u001b[1m \u001b[0m┃\n",
      "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
      "│\u001b[36m \u001b[0m\u001b[36m      test/pearson_buried_all       \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.22604259848594666         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m       test/pearson_buried_np       \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.3186895549297333          \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m         test/pearson_cbeta         \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        0.0027122749015688896        \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m        test/pearson_cenpack        \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.03428901731967926         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m test/pearson_centroid_total_score  \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m          0.031411312520504          \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m      test/pearson_contact_all      \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.1516418159008026          \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m  test/pearson_contact_buried_core  \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        -0.025072049349546432        \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36mtest/pearson_contact_buried_core_bo…\u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        -0.047540564090013504        \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m        test/pearson_degree         \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m          0.142898291349411          \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m      test/pearson_degree_core      \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        0.050124455243349075         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m test/pearson_degree_core_boundary  \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        -0.043538663536310196        \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m          test/pearson_env          \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        0.0012075577396899462        \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m test/pearson_exposed_hydrophobics  \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.05842389166355133         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m  test/pearson_exposed_np_AFIMLWVY  \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        -0.03353390470147133         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m    test/pearson_exposed_polars     \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         -0.0300059262663126         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m     test/pearson_exposed_total     \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.06565481424331665         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m        test/pearson_fa_atr         \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        0.0019589329604059458        \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m        test/pearson_fa_dun         \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.13203604519367218         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m        test/pearson_fa_elec        \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.16673614084720612         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m     test/pearson_fa_intra_rep      \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.2192520648241043          \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m  test/pearson_fa_intra_sol_xover4  \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         -0.0497041717171669         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m        test/pearson_fa_rep         \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         -0.0482688769698143         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m        test/pearson_fa_sol         \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.3524876534938812          \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m      test/pearson_hbond_bb_sc      \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.06194213777780533         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m      test/pearson_hbond_lr_bb      \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.1017257496714592          \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m       test/pearson_hbond_sc        \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        0.0021758440416306257        \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m      test/pearson_hbond_sr_bb      \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        -0.03285767510533333         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m        test/pearson_hs_pair        \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        0.026854708790779114         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m      test/pearson_lk_ball_wtd      \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        -0.004749468993395567        \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m         test/pearson_omega         \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.10179334133863449         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m     test/pearson_one_core_each     \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        0.0043275924399495125        \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m        test/pearson_p_aa_pp        \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        -0.003584893187507987        \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m         test/pearson_pack          \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        -0.01860794425010681         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m         test/pearson_pair          \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        0.059420570731163025         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m       test/pearson_pro_close       \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.06833246350288391         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m      test/pearson_rama_prepro      \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.05540400370955467         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m          test/pearson_ref          \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        -0.008935367688536644        \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m test/pearson_res_count_buried_core \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.05195057392120361         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36mtest/pearson_res_count_buried_core_…\u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        -0.06926741451025009         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36mtest/pearson_res_count_buried_np_co…\u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.09391546249389648         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36mtest/pearson_res_count_buried_np_co…\u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.05288832634687424         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m          test/pearson_rg           \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.09583169966936111         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m        test/pearson_rsigma         \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        -0.006123807281255722        \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m         test/pearson_sheet         \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        -0.005697159096598625        \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m  test/pearson_ss_contributes_core  \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        0.048998136073350906         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m        test/pearson_ss_mis         \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.19278551638126373         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m        test/pearson_ss_pair        \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        0.016048330813646317         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m   test/pearson_total_hydrophobic   \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.48500728607177734         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36mtest/pearson_total_hydrophobic_AFIL…\u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m          0.419537216424942          \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m      test/pearson_total_sasa       \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        0.023418603464961052         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m      test/pearson_total_score      \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.24490368366241455         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m     test/pearson_two_core_each     \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.04460342973470688         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m      test/pearson_unsat_hbond      \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        -0.09420519322156906         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m          test/pearson_vdw          \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         -0.109303779900074          \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m     test/pearson_yhh_planarity     \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        -0.042719367891550064        \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m             test_loss              \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         55.401283264160156          \u001b[0m\u001b[35m \u001b[0m│\n",
      "└──────────────────────────────────────┴───────────────────────────────────────┘\n"
     ]
    }
   ],
   "source": [
    "!python code/train_source_model.py @args/pretrain_avgfp_local.txt --max_epochs 5 --limit_train_batches 5 --limit_val_batches 5 --limit_test_batches 5\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "248413d3-d43c-47ba-afc6-7769ca59f866",
   "metadata": {},
   "source": [
    "# Resuming training\n",
    "If you need to train your model for more epochs, you can adjust the max_epochs and resume training by providing the existing UUID of model. Here, we provide the previous UUID `En8Ux9wH` and change the max_epochs to 7 to train for an additional 2 epochs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "6bee988a-45f0-443e-aa1f-da9f5e3c490e",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Global seed set to 1\n",
      "User gave model UUID: En8Ux9wH\n",
      "Found existing log directory corresponding to given UUID: output/training_logs/En8Ux9wH\n",
      "Final UUID: En8Ux9wH\n",
      "Final log directory: output/training_logs/En8Ux9wH\n",
      "This is version: 1\n",
      "Version-specific logs will be saved to: output/training_logs/En8Ux9wH/version_1\n",
      "Found checkpoint, resuming training from: output/training_logs/En8Ux9wH/checkpoints/last.ckpt\n",
      "Using example_input_array with pdb_fn='1gfl_cm.pdb' and aa_seq_len=237\n",
      "Trainer already configured with model summary callbacks: [<class 'pytorch_lightning.callbacks.rich_model_summary.RichModelSummary'>]. Skipping setting a default `ModelSummary` callback.\n",
      "GPU available: True (mps), used: False\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n",
      "/opt/homebrew/Caskroom/mambaforge/base/envs/metl/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py:1789: UserWarning: MPS available but not used. Set `accelerator` and `devices` using `Trainer(accelerator='mps', devices=1)`.\n",
      "  rank_zero_warn(\n",
      "/opt/homebrew/Caskroom/mambaforge/base/envs/metl/lib/python3.9/site-packages/pytorch_lightning/callbacks/model_checkpoint.py:616: UserWarning: Checkpoint directory output/training_logs/En8Ux9wH/checkpoints exists and is not empty.\n",
      "  rank_zero_warn(f\"Checkpoint directory {dirpath} exists and is not empty.\")\n",
      "Restoring states from the checkpoint path at output/training_logs/En8Ux9wH/checkpoints/last.ckpt\n",
      "Loading `train_dataloader` to estimate number of stepping batches.\n",
      "/opt/homebrew/Caskroom/mambaforge/base/envs/metl/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py:1892: PossibleUserWarning: The number of training batches (5) is smaller than the logging interval Trainer(log_every_n_steps=50). Set a lower value for log_every_n_steps if you want to see logs for the training epoch.\n",
      "  rank_zero_warn(\n",
      "Number of training steps is 35\n",
      "Number of warmup steps is 0.7000000000000001\n",
      "┏┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳┳┳┳┓\n",
      "┃┃\u001b[1;35m \u001b[0m\u001b[1;35mName                                                                   \u001b[0m\u001b[1;35m \u001b[0m┃┃┃┃┃\n",
      "┡╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇╇╇╇┩\n",
      "││ model                                                                   │││││\n",
      "││ model.model                                                             │││││\n",
      "││ model.model.embedder                                                    │││││\n",
      "││ model.model.embedder.embedding                                          │││││\n",
      "││ model.model.tr_encoder                                                  │││││\n",
      "││ model.model.tr_encoder.layers                                           │││││\n",
      "││ model.model.tr_encoder.layers.0                                         │││││\n",
      "││ model.model.tr_encoder.layers.0.self_attn                               │││││\n",
      "││ model.model.tr_encoder.layers.0.self_attn.relative_position_k           │││││\n",
      "││ model.model.tr_encoder.layers.0.self_attn.relative_position_k.embeddin… │││││\n",
      "││ model.model.tr_encoder.layers.0.self_attn.relative_position_v           │││││\n",
      "││ model.model.tr_encoder.layers.0.self_attn.relative_position_v.embeddin… │││││\n",
      "││ model.model.tr_encoder.layers.0.self_attn.q_proj                        │││││\n",
      "││ model.model.tr_encoder.layers.0.self_attn.k_proj                        │││││\n",
      "││ model.model.tr_encoder.layers.0.self_attn.v_proj                        │││││\n",
      "││ model.model.tr_encoder.layers.0.self_attn.out_proj                      │││││\n",
      "││ model.model.tr_encoder.layers.0.self_attn.dropout                       │││││\n",
      "││ model.model.tr_encoder.layers.0.linear1                                 │││││\n",
      "││ model.model.tr_encoder.layers.0.dropout                                 │││││\n",
      "││ model.model.tr_encoder.layers.0.linear2                                 │││││\n",
      "││ model.model.tr_encoder.layers.0.norm1                                   │││││\n",
      "││ model.model.tr_encoder.layers.0.norm2                                   │││││\n",
      "││ model.model.tr_encoder.layers.0.dropout1                                │││││\n",
      "││ model.model.tr_encoder.layers.0.dropout2                                │││││\n",
      "││ model.model.tr_encoder.layers.1                                         │││││\n",
      "││ model.model.tr_encoder.layers.1.self_attn                               │││││\n",
      "││ model.model.tr_encoder.layers.1.self_attn.relative_position_k           │││││\n",
      "││ model.model.tr_encoder.layers.1.self_attn.relative_position_k.embeddin… │││││\n",
      "││ model.model.tr_encoder.layers.1.self_attn.relative_position_v           │││││\n",
      "││ model.model.tr_encoder.layers.1.self_attn.relative_position_v.embeddin… │││││\n",
      "││ model.model.tr_encoder.layers.1.self_attn.q_proj                        │││││\n",
      "││ model.model.tr_encoder.layers.1.self_attn.k_proj                        │││││\n",
      "││ model.model.tr_encoder.layers.1.self_attn.v_proj                        │││││\n",
      "││ model.model.tr_encoder.layers.1.self_attn.out_proj                      │││││\n",
      "││ model.model.tr_encoder.layers.1.self_attn.dropout                       │││││\n",
      "││ model.model.tr_encoder.layers.1.linear1                                 │││││\n",
      "││ model.model.tr_encoder.layers.1.dropout                                 │││││\n",
      "││ model.model.tr_encoder.layers.1.linear2                                 │││││\n",
      "││ model.model.tr_encoder.layers.1.norm1                                   │││││\n",
      "││ model.model.tr_encoder.layers.1.norm2                                   │││││\n",
      "││ model.model.tr_encoder.layers.1.dropout1                                │││││\n",
      "││ model.model.tr_encoder.layers.1.dropout2                                │││││\n",
      "││ model.model.tr_encoder.layers.2                                         │││││\n",
      "││ model.model.tr_encoder.layers.2.self_attn                               │││││\n",
      "││ model.model.tr_encoder.layers.2.self_attn.relative_position_k           │││││\n",
      "││ model.model.tr_encoder.layers.2.self_attn.relative_position_k.embeddin… │││││\n",
      "││ model.model.tr_encoder.layers.2.self_attn.relative_position_v           │││││\n",
      "││ model.model.tr_encoder.layers.2.self_attn.relative_position_v.embeddin… │││││\n",
      "││ model.model.tr_encoder.layers.2.self_attn.q_proj                        │││││\n",
      "││ model.model.tr_encoder.layers.2.self_attn.k_proj                        │││││\n",
      "││ model.model.tr_encoder.layers.2.self_attn.v_proj                        │││││\n",
      "││ model.model.tr_encoder.layers.2.self_attn.out_proj                      │││││\n",
      "││ model.model.tr_encoder.layers.2.self_attn.dropout                       │││││\n",
      "││ model.model.tr_encoder.layers.2.linear1                                 │││││\n",
      "││ model.model.tr_encoder.layers.2.dropout                                 │││││\n",
      "││ model.model.tr_encoder.layers.2.linear2                                 │││││\n",
      "││ model.model.tr_encoder.layers.2.norm1                                   │││││\n",
      "││ model.model.tr_encoder.layers.2.norm2                                   │││││\n",
      "││ model.model.tr_encoder.layers.2.dropout1                                │││││\n",
      "││ model.model.tr_encoder.layers.2.dropout2                                │││││\n",
      "││ model.model.tr_encoder.norm                                             │││││\n",
      "││ model.model.avg_pooling                                                 │││││\n",
      "││ model.model.fc1                                                         │││││\n",
      "││ model.model.fc1.fc                                                      │││││\n",
      "││ model.model.fc1.activation                                              │││││\n",
      "││ model.model.fc1.norm                                                    │││││\n",
      "││ model.model.fc1.dropout                                                 │││││\n",
      "││ model.model.prediction                                                  │││││\n",
      "││ test_pearson_collection                                                 │││││\n",
      "││ test_pearson_collection.0                                               │││││\n",
      "││ test_pearson_collection.1                                               │││││\n",
      "││ test_pearson_collection.2                                               │││││\n",
      "││ test_pearson_collection.3                                               │││││\n",
      "││ test_pearson_collection.4                                               │││││\n",
      "││ test_pearson_collection.5                                               │││││\n",
      "││ test_pearson_collection.6                                               │││││\n",
      "││ test_pearson_collection.7                                               │││││\n",
      "││ test_pearson_collection.8                                               │││││\n",
      "││ test_pearson_collection.9                                               │││││\n",
      "││ test_pearson_collection.10                                              │││││\n",
      "││ test_pearson_collection.11                                              │││││\n",
      "││ test_pearson_collection.12                                              │││││\n",
      "││ test_pearson_collection.13                                              │││││\n",
      "││ test_pearson_collection.14                                              │││││\n",
      "││ test_pearson_collection.15                                              │││││\n",
      "││ test_pearson_collection.16                                              │││││\n",
      "││ test_pearson_collection.17                                              │││││\n",
      "││ test_pearson_collection.18                                              │││││\n",
      "││ test_pearson_collection.19                                              │││││\n",
      "││ test_pearson_collection.20                                              │││││\n",
      "││ test_pearson_collection.21                                              │││││\n",
      "││ test_pearson_collection.22                                              │││││\n",
      "││ test_pearson_collection.23                                              │││││\n",
      "││ test_pearson_collection.24                                              │││││\n",
      "││ test_pearson_collection.25                                              │││││\n",
      "││ test_pearson_collection.26                                              │││││\n",
      "││ test_pearson_collection.27                                              │││││\n",
      "││ test_pearson_collection.28                                              │││││\n",
      "││ test_pearson_collection.29                                              │││││\n",
      "││ test_pearson_collection.30                                              │││││\n",
      "││ test_pearson_collection.31                                              │││││\n",
      "││ test_pearson_collection.32                                              │││││\n",
      "││ test_pearson_collection.33                                              │││││\n",
      "││ test_pearson_collection.34                                              │││││\n",
      "││ test_pearson_collection.35                                              │││││\n",
      "││ test_pearson_collection.36                                              │││││\n",
      "││ test_pearson_collection.37                                              │││││\n",
      "││ test_pearson_collection.38                                              │││││\n",
      "││ test_pearson_collection.39                                              │││││\n",
      "││ test_pearson_collection.40                                              │││││\n",
      "││ test_pearson_collection.41                                              │││││\n",
      "││ test_pearson_collection.42                                              │││││\n",
      "││ test_pearson_collection.43                                              │││││\n",
      "││ test_pearson_collection.44                                              │││││\n",
      "││ test_pearson_collection.45                                              │││││\n",
      "││ test_pearson_collection.46                                              │││││\n",
      "││ test_pearson_collection.47                                              │││││\n",
      "││ test_pearson_collection.48                                              │││││\n",
      "││ test_pearson_collection.49                                              │││││\n",
      "││ test_pearson_collection.50                                              │││││\n",
      "││ test_pearson_collection.51                                              │││││\n",
      "││ test_pearson_collection.52                                              │││││\n",
      "││ test_pearson_collection.53                                              │││││\n",
      "││ test_pearson_collection.54                                              │││││\n",
      "└┴─────────────────────────────────────────────────────────────────────────┴┴┴┴┘\n",
      "\u001b[1mTrainable params\u001b[0m: 2.5 M                                                         \n",
      "\u001b[1mNon-trainable params\u001b[0m: 0                                                         \n",
      "\u001b[1mTotal params\u001b[0m: 2.5 M                                                             \n",
      "\u001b[1mTotal estimated model params size (MB)\u001b[0m: 9                                       \n",
      "/opt/homebrew/Caskroom/mambaforge/base/envs/metl/lib/python3.9/site-packages/pytorch_lightning/loggers/csv_logs.py:57: UserWarning: Experiment logs directory output/training_logs/En8Ux9wH/version_1 exists and is not empty. Previous log files in this directory will be deleted when the new ones are saved!\n",
      "  rank_zero_warn(\n",
      "Restored all states from the checkpoint file at output/training_logs/En8Ux9wH/checkpoints/last.ckpt\n",
      "Sanity Checking DataLoader 0:   0%|                       | 0/2 [00:00<?, ?it/s]/opt/homebrew/Caskroom/mambaforge/base/envs/metl/lib/python3.9/site-packages/pytorch_lightning/utilities/data.py:98: UserWarning: Trying to infer the `batch_size` from an ambiguous collection. The batch size we found is 128. To avoid any miscalculations, use `self.log(..., batch_size=batch_size)`.\n",
      "  warning_cache.warn(\n",
      "Epoch 5:  50%|███████▌       | 5/10 [00:12<00:12,  2.47s/it, loss=56.4, v_num=1]\n",
      "Validation: 0it [00:00, ?it/s]\u001b[A\n",
      "Validation:   0%|                                         | 0/5 [00:00<?, ?it/s]\u001b[A\n",
      "Validation DataLoader 0:   0%|                            | 0/5 [00:00<?, ?it/s]\u001b[A\n",
      "Epoch 5:  60%|█████████      | 6/10 [00:12<00:08,  2.16s/it, loss=56.4, v_num=1]\u001b[A\n",
      "Epoch 5:  70%|██████████▌    | 7/10 [00:13<00:05,  1.92s/it, loss=56.4, v_num=1]\u001b[A\n",
      "Epoch 5:  80%|████████████   | 8/10 [00:13<00:03,  1.74s/it, loss=56.4, v_num=1]\u001b[A\n",
      "Epoch 5:  90%|█████████████▌ | 9/10 [00:14<00:01,  1.60s/it, loss=56.4, v_num=1]\u001b[A\n",
      "Epoch 5: 100%|█| 10/10 [00:14<00:00,  1.48s/it, loss=56.4, v_num=1, val_loss=56.\u001b[A\n",
      "Epoch 6:  50%|▌| 5/10 [00:10<00:10,  2.06s/it, loss=55.7, v_num=1, val_loss=56.2\u001b[A\n",
      "Validation: 0it [00:00, ?it/s]\u001b[A\n",
      "Validation:   0%|                                         | 0/5 [00:00<?, ?it/s]\u001b[A\n",
      "Validation DataLoader 0:   0%|                            | 0/5 [00:00<?, ?it/s]\u001b[A\n",
      "Epoch 6:  60%|▌| 6/10 [00:10<00:07,  1.82s/it, loss=55.7, v_num=1, val_loss=56.2\u001b[A\n",
      "Epoch 6:  70%|▋| 7/10 [00:11<00:04,  1.63s/it, loss=55.7, v_num=1, val_loss=56.2\u001b[A\n",
      "Epoch 6:  80%|▊| 8/10 [00:11<00:02,  1.49s/it, loss=55.7, v_num=1, val_loss=56.2\u001b[A\n",
      "Epoch 6:  90%|▉| 9/10 [00:12<00:01,  1.37s/it, loss=55.7, v_num=1, val_loss=56.2\u001b[A\n",
      "Epoch 6: 100%|█| 10/10 [00:12<00:00,  1.28s/it, loss=55.7, v_num=1, val_loss=56.\u001b[A\n",
      "Epoch 6: 100%|█| 10/10 [00:12<00:00,  1.28s/it, loss=55.7, v_num=1, val_loss=56.\u001b[A`Trainer.fit` stopped: `max_epochs=7` reached.\n",
      "Epoch 6: 100%|█| 10/10 [00:12<00:00,  1.29s/it, loss=55.7, v_num=1, val_loss=56.\n",
      "output/training_logs/En8Ux9wH/checkpoints/epoch=5-step=30-val_loss=56.22.ckpt\n",
      "tensor(56.2178)\n",
      "/opt/homebrew/Caskroom/mambaforge/base/envs/metl/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py:1397: UserWarning: `.test(ckpt_path=\"best\")` is called with Trainer configured with multiple `ModelCheckpoint` callbacks. It will use the best checkpoint path from first checkpoint callback.\n",
      "  rank_zero_warn(\n",
      "Restoring states from the checkpoint path at output/training_logs/En8Ux9wH/checkpoints/epoch=5-step=30-val_loss=56.22.ckpt\n",
      "Loaded model weights from checkpoint at output/training_logs/En8Ux9wH/checkpoints/epoch=5-step=30-val_loss=56.22.ckpt\n",
      "Testing DataLoader 0: 100%|███████████████████████| 5/5 [00:02<00:00,  1.97it/s]\n",
      "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
      "┃\u001b[1m \u001b[0m\u001b[1m            Test metric             \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m            DataLoader 0             \u001b[0m\u001b[1m \u001b[0m┃\n",
      "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
      "│\u001b[36m \u001b[0m\u001b[36m      test/pearson_buried_all       \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.19549381732940674         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m       test/pearson_buried_np       \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.2820279002189636          \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m         test/pearson_cbeta         \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        0.006789690814912319         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m        test/pearson_cenpack        \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.06755797564983368         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m test/pearson_centroid_total_score  \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        0.024258162826299667         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m      test/pearson_contact_all      \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.20111863315105438         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m  test/pearson_contact_buried_core  \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        -0.016893047839403152        \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36mtest/pearson_contact_buried_core_bo…\u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        -0.07590403407812119         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m        test/pearson_degree         \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.14845307171344757         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m      test/pearson_degree_core      \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        0.044732216745615005         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m test/pearson_degree_core_boundary  \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        -0.015861663967370987        \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m          test/pearson_env          \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        -0.02680979296565056         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m test/pearson_exposed_hydrophobics  \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.0865994542837143          \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m  test/pearson_exposed_np_AFIMLWVY  \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        0.061353255063295364         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m    test/pearson_exposed_polars     \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        -0.007581458427011967        \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m     test/pearson_exposed_total     \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.04606642946600914         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m        test/pearson_fa_atr         \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.06978549808263779         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m        test/pearson_fa_dun         \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.17426544427871704         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m        test/pearson_fa_elec        \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.21024872362613678         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m     test/pearson_fa_intra_rep      \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.20041094720363617         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m  test/pearson_fa_intra_sol_xover4  \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        0.021917205303907394         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m        test/pearson_fa_rep         \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        -0.009428195655345917        \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m        test/pearson_fa_sol         \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.21377339959144592         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m      test/pearson_hbond_bb_sc      \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.05525718256831169         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m      test/pearson_hbond_lr_bb      \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.09545594453811646         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m       test/pearson_hbond_sc        \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.02604163996875286         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m      test/pearson_hbond_sr_bb      \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        -0.06445325165987015         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m        test/pearson_hs_pair        \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        0.019333532080054283         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m      test/pearson_lk_ball_wtd      \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        -0.008282855153083801        \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m         test/pearson_omega         \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.15580017864704132         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m     test/pearson_one_core_each     \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        0.005425552371889353         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m        test/pearson_p_aa_pp        \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        -0.056554194539785385        \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m         test/pearson_pack          \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        0.0010201290715485811        \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m         test/pearson_pair          \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.04599135369062424         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m       test/pearson_pro_close       \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        0.048722509294748306         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m      test/pearson_rama_prepro      \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.05651899427175522         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m          test/pearson_ref          \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        0.009523695334792137         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m test/pearson_res_count_buried_core \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        0.028048822656273842         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36mtest/pearson_res_count_buried_core_…\u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        -0.05315422639250755         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36mtest/pearson_res_count_buried_np_co…\u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.07821594178676605         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36mtest/pearson_res_count_buried_np_co…\u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        -0.059419240802526474        \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m          test/pearson_rg           \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.15277768671512604         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m        test/pearson_rsigma         \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        0.018071671947836876         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m         test/pearson_sheet         \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        0.0005552759976126254        \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m  test/pearson_ss_contributes_core  \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.05250401049852371         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m        test/pearson_ss_mis         \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.16041426360607147         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m        test/pearson_ss_pair        \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        0.0022559859789907932        \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m   test/pearson_total_hydrophobic   \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.5237370133399963          \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36mtest/pearson_total_hydrophobic_AFIL…\u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.4336868226528168          \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m      test/pearson_total_sasa       \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.05494881048798561         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m      test/pearson_total_score      \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.2772155702114105          \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m     test/pearson_two_core_each     \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         0.04130518063902855         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m      test/pearson_unsat_hbond      \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        -0.07321402430534363         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m          test/pearson_vdw          \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        -0.10069911926984787         \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m     test/pearson_yhh_planarity     \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        -0.050693441182374954        \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m             test_loss              \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m          55.35341262817383          \u001b[0m\u001b[35m \u001b[0m│\n",
      "└──────────────────────────────────────┴───────────────────────────────────────┘\n"
     ]
    }
   ],
   "source": [
    "!python code/train_source_model.py @args/pretrain_avgfp_local.txt --uuid En8Ux9wH --max_epochs 7 --limit_train_batches 5 --limit_val_batches 5 --limit_test_batches 5\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8133a6e6-66a6-4149-8239-b01292d000fc",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a151fe0-2043-4d2b-8b56-a3eb69bea55d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39e7c178-e46d-4c94-a998-e005b8b537b6",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac870913-3407-4381-aed9-48c7a63268d0",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}