scripts/contact-based_model.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "10df24ce-dece-47be-9d35-0cabfb02cd85",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List, Tuple, Optional, Dict, NamedTuple, Union, Callable\n",
    "import numpy as np\n",
    "import torch\n",
    "from scipy.spatial.distance import squareform, pdist, cdist\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib as mpl\n",
    "\n",
    "import biotite.structure as bs\n",
    "from biotite.structure.io.pdbx import PDBxFile, get_structure\n",
    "from biotite.database import rcsb\n",
    "from tqdm import tqdm\n",
    "import pandas as pd\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "883b7fe4-49af-43c3-8977-f275a48351e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run it in cuda or cpu\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8ef2232-8a05-4254-8cca-5555f901d8ca",
   "metadata": {},
   "source": [
    "## Get the true lbdb contacts from cif file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "54ceae83-a93f-4798-9aef-f4faa02b2d18",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"Adapted from: https://github.com/facebookresearch/esm/blob/main/examples/contact_prediction.ipynb\"\"\"\n",
    "def extend(a, b, c, L, A, D):\n",
    "    \"\"\"\n",
    "    input:  3 coords (a,b,c), (L)ength, (A)ngle, and (D)ihedral\n",
    "    output: 4th coord\n",
    "    \"\"\"\n",
    "\n",
    "    def normalize(x):\n",
    "        return x / np.linalg.norm(x, ord=2, axis=-1, keepdims=True)    #求范数\n",
    "\n",
    "    bc = normalize(b - c)\n",
    "    n = normalize(np.cross(b - a, bc))\n",
    "    m = [bc, np.cross(n, bc), n]\n",
    "    d = [L * np.cos(A), L * np.sin(A) * np.cos(D), -L * np.sin(A) * np.sin(D)]\n",
    "    return c + sum([m * d for m, d in zip(m, d)])\n",
    "\n",
    "\n",
    "def contacts_from_pdb(\n",
    "    structure: bs.AtomArray,\n",
    "    distance_threshold: float = 8.0,\n",
    "    chain: Optional[str] = None,\n",
    ") -> np.ndarray:\n",
    "    mask = ~structure.hetero\n",
    "    if chain is not None:\n",
    "        mask &= structure.chain_id == chain\n",
    "\n",
    "    N = structure.coord[mask & (structure.atom_name == \"N\")]\n",
    "    CA = structure.coord[mask & (structure.atom_name == \"CA\")]\n",
    "    C = structure.coord[mask & (structure.atom_name == \"C\")]\n",
    "\n",
    "    Cbeta = extend(C, N, CA, 1.522, 1.927, -2.143)\n",
    "    dist = squareform(pdist(Cbeta))    #矩阵对称化\n",
    "    \n",
    "    contacts = dist < distance_threshold\n",
    "    contacts = contacts.astype(np.int64)\n",
    "    contacts[np.isnan(dist)] = -1\n",
    "    return contacts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d35de7df-70e1-4ea9-b03e-d4ab728fc2c4",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/nfs/home/9403_huangjiaquan/.conda/envs/AMP/lib/python3.8/site-packages/biotite/structure/io/pdbx/convert.py:311: UserWarning: Attribute 'auth_seq_id' not found within 'atom_site' category. The fallback attribute 'label_seq_id' will be used instead\n",
      "  warnings.warn(\n",
      "/nfs/home/9403_huangjiaquan/.conda/envs/AMP/lib/python3.8/site-packages/biotite/structure/io/pdbx/convert.py:311: UserWarning: Attribute 'auth_comp_id' not found within 'atom_site' category. The fallback attribute 'label_comp_id' will be used instead\n",
      "  warnings.warn(\n",
      "/nfs/home/9403_huangjiaquan/.conda/envs/AMP/lib/python3.8/site-packages/biotite/structure/io/pdbx/convert.py:311: UserWarning: Attribute 'auth_atom_id' not found within 'atom_site' category. The fallback attribute 'label_atom_id' will be used instead\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "lbd=get_structure(PDBxFile.read('../database/lbdb.cif'))[0]\n",
    "contacts=contacts_from_pdb(lbd)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ad353bf-ddbd-4cd1-aa60-33b5d12e6caa",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Finetune logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b4e3419c-4b47-404f-8de0-623168d45133",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.optim import AdamW\n",
    "from transformers import get_scheduler\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from transformers import EsmTokenizer, EsmForSequenceClassification, EsmModel, EsmConfig,EsmForMaskedLM\n",
    "from sklearn.metrics import accuracy_score\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "8f2c0b22-8466-4445-a49f-c13f750ff324",
   "metadata": {},
   "outputs": [],
   "source": [
    "def symmetrize(x):\n",
    "    \"Make layer symmetric in final two dimensions, used for contact prediction.\"\n",
    "    return x + x.transpose(-1, -2)\n",
    "\n",
    "def apc(x):\n",
    "    \"Perform average product correct, used for contact prediction.\"\n",
    "    a1 = x.sum(-1, keepdims=True)\n",
    "    a2 = x.sum(-2, keepdims=True)\n",
    "    a12 = x.sum((-1, -2), keepdims=True)\n",
    "\n",
    "    avg = a1 * a2\n",
    "    avg.div_(a12)  # in-place to reduce memory\n",
    "    normalized = x - avg\n",
    "    return normalized"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ad7a062a-fc0c-4e68-a62a-df23cdc82d5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Adapted from https://github.com/facebookresearch/esm/blob/main/esm/modules.py\n",
    "class AttentionLogisticRegression(nn.Module):\n",
    "    \"\"\"Performs symmetrization, apc, and computes a logistic regression on the output features\"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        in_features:int,\n",
    "        prepend_bos: bool,\n",
    "        append_eos: bool,\n",
    "        bias=True,\n",
    "        eos_idx: Optional[int] = None,\n",
    "    ):\n",
    "        super().__init__()\n",
    "        self.in_features=in_features\n",
    "        self.prepend_bos = prepend_bos\n",
    "        self.append_eos = append_eos\n",
    "        if append_eos and eos_idx is None:\n",
    "            raise ValueError(\"Using an alphabet with eos token, but no eos token was passed in.\")\n",
    "        self.eos_idx = eos_idx\n",
    "        self.regression = nn.Linear(in_features, 1, bias)\n",
    "        self.activation = nn.Sigmoid()\n",
    "    \n",
    "    def forward(self, tokens, attentions):\n",
    "        # remove eos token attentions\n",
    "        if self.append_eos:\n",
    "            eos_mask = tokens.ne(self.eos_idx).to(attentions)\n",
    "            eos_mask = eos_mask.unsqueeze(1) * eos_mask.unsqueeze(2)\n",
    "            attentions = attentions * eos_mask[:, None, None, :, :]\n",
    "            attentions = attentions[..., :-1, :-1]\n",
    "        # remove cls token attentions\n",
    "        if self.prepend_bos:\n",
    "            attentions = attentions[..., 1:, 1:]\n",
    "        batch_size, layers, heads, seqlen, _ = attentions.size()\n",
    "        attentions = attentions.view(batch_size, layers * heads, seqlen, seqlen)\n",
    "        attentions = attentions.to(self.regression.weight.device)  # attentions always float32, may need to convert to float16\n",
    "        attentions= apc(symmetrize(attentions))\n",
    "        attentions = attentions.permute(0, 2, 3, 1)\n",
    "        \n",
    "        return self.activation(self.regression(attentions).squeeze(3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "46957418-c127-40c0-906f-d1ce4b6a9282",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_layers=33\n",
    "attention_heads=20\n",
    "mymodel=AttentionLogisticRegression(in_features=num_layers*attention_heads,prepend_bos=True,append_eos=True,eos_idx=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "193354a0-8cdd-4658-b194-00db78e9c21f",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at ../models/esm2_650M were not used when initializing EsmForSequenceClassification: ['lm_head.decoder.weight', 'lm_head.dense.bias', 'lm_head.dense.weight', 'lm_head.layer_norm.weight', 'lm_head.bias', 'lm_head.layer_norm.bias']\n",
      "- This IS expected if you are initializing EsmForSequenceClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing EsmForSequenceClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "Some weights of EsmForSequenceClassification were not initialized from the model checkpoint at ../models/esm2_650M and are newly initialized: ['classifier.dense.weight', 'classifier.dense.bias', 'classifier.out_proj.weight', 'classifier.out_proj.bias']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "PeftModelForSequenceClassification(\n",
       "  (base_model): LoraModel(\n",
       "    (model): EsmForSequenceClassification(\n",
       "      (esm): EsmModel(\n",
       "        (embeddings): EsmEmbeddings(\n",
       "          (word_embeddings): Embedding(33, 1280, padding_idx=1)\n",
       "          (dropout): Dropout(p=0.0, inplace=False)\n",
       "          (position_embeddings): Embedding(1026, 1280, padding_idx=1)\n",
       "        )\n",
       "        (encoder): EsmEncoder(\n",
       "          (layer): ModuleList(\n",
       "            (0-32): 33 x EsmLayer(\n",
       "              (attention): EsmAttention(\n",
       "                (self): EsmSelfAttention(\n",
       "                  (query): Linear(\n",
       "                    in_features=1280, out_features=1280, bias=True\n",
       "                    (lora_dropout): ModuleDict(\n",
       "                      (default): Dropout(p=0.6, inplace=False)\n",
       "                    )\n",
       "                    (lora_A): ModuleDict(\n",
       "                      (default): Linear(in_features=1280, out_features=48, bias=False)\n",
       "                    )\n",
       "                    (lora_B): ModuleDict(\n",
       "                      (default): Linear(in_features=48, out_features=1280, bias=False)\n",
       "                    )\n",
       "                    (lora_embedding_A): ParameterDict()\n",
       "                    (lora_embedding_B): ParameterDict()\n",
       "                  )\n",
       "                  (key): Linear(\n",
       "                    in_features=1280, out_features=1280, bias=True\n",
       "                    (lora_dropout): ModuleDict(\n",
       "                      (default): Dropout(p=0.6, inplace=False)\n",
       "                    )\n",
       "                    (lora_A): ModuleDict(\n",
       "                      (default): Linear(in_features=1280, out_features=48, bias=False)\n",
       "                    )\n",
       "                    (lora_B): ModuleDict(\n",
       "                      (default): Linear(in_features=48, out_features=1280, bias=False)\n",
       "                    )\n",
       "                    (lora_embedding_A): ParameterDict()\n",
       "                    (lora_embedding_B): ParameterDict()\n",
       "                  )\n",
       "                  (value): Linear(\n",
       "                    in_features=1280, out_features=1280, bias=True\n",
       "                    (lora_dropout): ModuleDict(\n",
       "                      (default): Dropout(p=0.6, inplace=False)\n",
       "                    )\n",
       "                    (lora_A): ModuleDict(\n",
       "                      (default): Linear(in_features=1280, out_features=48, bias=False)\n",
       "                    )\n",
       "                    (lora_B): ModuleDict(\n",
       "                      (default): Linear(in_features=48, out_features=1280, bias=False)\n",
       "                    )\n",
       "                    (lora_embedding_A): ParameterDict()\n",
       "                    (lora_embedding_B): ParameterDict()\n",
       "                  )\n",
       "                  (dropout): Dropout(p=0.0, inplace=False)\n",
       "                  (rotary_embeddings): RotaryEmbedding()\n",
       "                )\n",
       "                (output): EsmSelfOutput(\n",
       "                  (dense): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "                  (dropout): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "                (LayerNorm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "              )\n",
       "              (intermediate): EsmIntermediate(\n",
       "                (dense): Linear(in_features=1280, out_features=5120, bias=True)\n",
       "              )\n",
       "              (output): EsmOutput(\n",
       "                (dense): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "                (dropout): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "              (LayerNorm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "            )\n",
       "          )\n",
       "          (emb_layer_norm_after): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "        (contact_head): EsmContactPredictionHead(\n",
       "          (regression): Linear(in_features=660, out_features=1, bias=True)\n",
       "          (activation): Sigmoid()\n",
       "        )\n",
       "      )\n",
       "      (classifier): ModulesToSaveWrapper(\n",
       "        (original_module): EsmClassificationHead(\n",
       "          (dense): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "          (dropout): Dropout(p=0.0, inplace=False)\n",
       "          (out_proj): Linear(in_features=1280, out_features=2, bias=True)\n",
       "        )\n",
       "        (modules_to_save): ModuleDict(\n",
       "          (default): EsmClassificationHead(\n",
       "            (dense): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "            (dropout): Dropout(p=0.0, inplace=False)\n",
       "            (out_proj): Linear(in_features=1280, out_features=2, bias=True)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from peft import PeftModel, PeftConfig\n",
    "from transformers import EsmTokenizer, EsmForSequenceClassification, EsmModel, EsmConfig\n",
    "config = PeftConfig.from_pretrained('../models/esm2_650M_LORA_SEQ_CLS_0.99')\n",
    "model = EsmForSequenceClassification.from_pretrained('../models/esm2_650M', num_labels=2) \n",
    "model = PeftModel.from_pretrained(model, '../models/esm2_650M_LORA_SEQ_CLS_0.99')\n",
    "model.eval()\n",
    "model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d26b81af-3dd8-4887-9068-ecc09d85c6b7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "model_name = \"../models/esm2_650M\"\n",
    "tokenizer = EsmTokenizer.from_pretrained(model_name)\n",
    "seqs='CTYKVKPKLQRFKLYFLGTVTC'\n",
    "inputs = tokenizer(seqs, return_tensors='pt', padding=\"max_length\", truncation=True, max_length=24).to(device)\n",
    "outputs=model.esm(**inputs,output_attentions=True,output_hidden_states=True)\n",
    "lbd_lora_attention=torch.stack(outputs.attentions,1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "a310cf29-7601-4094-b853-559525d287e9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "mymodel.to(device)\n",
    "lbd_lora_attention=lbd_lora_attention.detach()\n",
    "contact_y=torch.tensor(contacts,dtype=torch.float32).to(device)\n",
    "\n",
    "# Define loss function and optimizer\n",
    "criterion = nn.BCELoss()\n",
    "optimizer = optim.Adam(mymodel.parameters(), lr=5e-3)\n",
    "\n",
    "# Training loop\n",
    "num_epochs = 6000  # Adjust the number of epochs as needed\n",
    "\n",
    "lr_scheduler = get_scheduler(\n",
    "    \"linear\", \n",
    "    optimizer=optimizer, \n",
    "    num_warmup_steps=0, \n",
    "    num_training_steps=num_epochs\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "ccba808e-de11-4c12-9448-09ade1dd9d88",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 6000/6000 [00:06<00:00, 976.43it/s] \n"
     ]
    }
   ],
   "source": [
    "loss_list=[]\n",
    "accurate_list=[]\n",
    "for epoch in tqdm(range(num_epochs)):\n",
    "    mymodel.train()\n",
    "    outputs=mymodel(inputs['input_ids'], lbd_lora_attention)\n",
    "    \n",
    "    total_loss = 0.0\n",
    "    accurate=0\n",
    "    \n",
    "    for index,one_contact_head in enumerate(outputs):\n",
    "        loss = criterion(one_contact_head, contact_y)\n",
    "        total_loss += loss.item()\n",
    "        accurate+=(torch.where(one_contact_head < 0.9, torch.tensor(0), torch.tensor(1))==contact_y).sum()/(one_contact_head.shape[0]*one_contact_head.shape[0])\n",
    "        \n",
    "        loss.backward()# calulate loss\n",
    "        optimizer.step()# update gradient\n",
    "        lr_scheduler.step()# update learn rate\n",
    "        optimizer.zero_grad()# reset gradient\n",
    "        \n",
    "        train_accuracy = accurate / len(outputs)\n",
    "        average_loss = total_loss / len(outputs)\n",
    "        \n",
    "    loss_list.append(average_loss)\n",
    "    accurate_list.append(train_accuracy)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "2b5f6289-a906-49ec-8fcd-909af9645776",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7f890003f280>]"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAXQAAAD4CAYAAAD8Zh1EAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAAsTAAALEwEAmpwYAAAdbUlEQVR4nO3de3wV9Z3/8dcnJzcgQEISEEggKHgBWgUiXtAu1WqRVm3XVrE3tW7ZR1172bpttfZnt+6vj992tw+7vdhaWy+t3dV6qS6luHi/bFuFACICAuFiCaAJtxAuuZ18f3/MJJycnJBDODlz5uT9fDzmMXNmvpnz+eLxncl35syYcw4REQm/nKALEBGR1FCgi4hkCQW6iEiWUKCLiGQJBbqISJbIDeqNy8rKXFVVVVBvLyISSitWrNjtnCtPtC2wQK+qqqKmpiaotxcRCSUze6e3bRpyERHJEgp0EZEsoUAXEckSCnQRkSyhQBcRyRJ9BrqZ3W9m9Wb2Vi/bzcx+bGa1Zvammc1MfZkiItKXZI7QHwTmHWP7ZcAUf1oI/PzEyxIRkePV53XozrlXzKzqGE2uBH7jvPvwvmZmxWY21jm3K1VFiqTD8+vfY/X2/UGXIYPAxWeM4czK4pTvNxVfLBoPbI95Xeev6xHoZrYQ7yieCRMmpOCtZTCrP9DMv/xxPS1t0RPelwOeXfceAGYnvDuRYxo9ojBjAz1pzrl7gXsBqqur9WSNQaot2kFbtOO4f+7Zde9x17Mb6Xwmy1/3HgbgpBGFFA/NO+G6po0bwbc/MpXzTik94X2JBCEVgb4DqIx5XeGvE+mhvqmZ2d97/oT28fEZ4wGYNbGEipIhfO2SUzEdVoukJNAXATeb2SPAOUCjxs8llnOOP67ZxYEj7Tz4560AzH/fSZxZUXzc+zqrsphzTtYRtEgifQa6mT0MzAXKzKwO+A6QB+CcuwdYAswHaoHDwA0DVayEw0sb6nnh7fqu12+/28SyrXu7Xo8ozOXuT83UUbVIiiVzlcu1fWx3wD+krCIJlQf+tJWn17zbbd2ybV54l/jj2g4YM6KAB66fTWlRPiOH5CnMRQZAYLfPlcx2sKWdxiNtAPxh9U5+9mItic5iNzW3A3BezDDIeSeXcsOcKi6ddlI6ShURnwJdetj0XhOX/PCVHutvmFPVY51hXDVrPNPGjUxDZSJyLAp06abxSFtXmM+bdhIXnT4agKnjRjB9vEJbJJMp0KXL2p2NLPjFa4B3NH7HR6dqrFskRBTog5RzDufg5y9v5qlV3tcGNtUfBOCSqWP4zuXTgixPRPpBgT5INB5uo26/983K/Yfb+Ox9r9MRc5Zz/vtOYsqYIs6ZVMp151cFU6SInBAFepbae6iVlzfW0+F/w/6Wx1b3aHPx6aM5s7KY+e8by+TRRWmuUERSTYGepb7x+GqeW1/fbd35pxw9+i4qyOX8U0o1Ri6SRRToWaL+QDPfevItWtq9Ow++umk3RQW5LPnyhYB3B8GKkiEKcJEspkAPqdb2DpqavS/+vFnXyA0PLgdgYulQSoflM2tiCTd/cDITSocGWaaIpJECPYSamts4+3vP0dzW/Ra0N809ha9/+DQdhYsMUgr0kHllYwNfengVzW0dvL9iJJ+YVQHAhFFDmXva6ICrE5EgKdBD4Pcr6/jL5j1EOxy/968ZP7NiJI9/8XzyIsk8FlZEBgMFegarrW/im0+sYcU7+wAYN7KQipIhfPeKaVx0+mgNrYhINwr0DLD/cCtzf/ASTc3t5MRkdFvU++bP+aeUcsulpzJr4qiAKhSRMFCgB+jVTQ0cONLOk6t2sP9wG7MnjaJ6Ykm3NlVlw7i6urKXPYiIHKVAD8irmxr47H3Luq37zednU5gXCagiEQk7BXqa/fj5TbxZ18hz698D4KEbZzNmRCGjhuUrzEXkhCjQT1C0w3GkLdpnux8/v4nFq3eys7EZgKljR3DWhGIunFI+0CWKyCAR2kBvi3awbOteTi4fxtiRQwKr4/oHlvHqpt1Jt7+6uoKb5k6mqmzYAFYlIoNRaAP9uvuX8efNe8gx2PL/PjJg77Nj/xFefLu+1+012/ZxzqRRfOiMMcfcj5l3n/GJpQpyERkYoQ30v2zZA0CH8+5rkp+b+i/YtEc7+PAPX+FgS/sx232yurLrG5siIkEJbaD/zanlvLShAYB9h1sZM6Iw5e/xmfte52BLO+dMGsVPPjUjYZvcnBxGDctP+XuLiByv0Ab6vsNtXcu7D7YcV6DX7TvMR3/yv7RHHb1917KlvYPWaAdFBbncd/3ZFBWE9p9KRAaJ0KZU4+FWyocX0NDUwq1PrOEPX7rgmO237j7EZv+Zmbc8tprGI2188LRyJpX1/qSevFzjCxeerDAXkVAIbVLtO9zGmZXFNDQ1sGZHI0daowzJ73kd96b3mli69l1+8MzGbusnjy7igRtmp6tcEZEBF8pAj3Y4DjS3cVZlMVsaDlK37wgvb6xn3vSx3dodaY1yyQ9f6Xp9/flVXDXTO3l5ymhdbSIi2SWUgX7gSBvOQcnQPJ68aQ5nf+85du5vpj169IEP33j8za5bzS44u5L/+7Hp5OpWsyKSxUIZ6PsOtwJQPDSv6wqTOxev487F67q1i+QY15xdyXcun6owF5GsF9JA965wKR6aTyTmfrO3XHJq13IkYlxdXUlZUUHa6xMRCUIoA73xiHeEXjLUOzp/9RsfZMSQPEYOyQuyLBGRQCU1DmFm88xsg5nVmtmtCbZPMLMXzWyVmb1pZvNTX+pRB1u8m2EVFXhXtVSOGqowF5FBr89AN7MIcDdwGTAVuNbMpsY1+zbwqHNuBrAA+FmqC40V7fBOfkZyNC4uItIpmUScDdQ657Y451qBR4Ar49o4YIS/PBLYmboSe+q8mCWiZ2qKiHRJJtDHA9tjXtf562L9M/AZM6sDlgBfSrQjM1toZjVmVtPQ0NCPcj0dHc7fX793ISKSdVI1ZnEt8KBzrgKYDzxkZj327Zy71zlX7ZyrLi/v/4MdOpwX6LFXuIiIDHbJBPoOIPYpxRX+ulg3Ao8COOf+AhQCZakoMJGoAl1EpIdkAn05MMXMJplZPt5Jz0Vxbf4KXAxgZmfgBXr/x1T60DnkkqMxFxGRLn0GunOuHbgZWAqsx7uaZa2Z3WlmV/jNbgG+YGargYeB653zD6MHQLRDR+giIvGS+mKRc24J3snO2HV3xCyvA+aktrTeRf1fFbrKRUTkqFBeyN015BLK6kVEBkYoI1EnRUVEegpnoOukqIhID6EM9M7zrQp0EZGjQhnoXV/915CLiEiXcAZ61xF6wIWIiGSQUAZ6R4cjx8A05CIi0iWUgR51TsMtIiJxQhno3hG6Al1EJFYoAz3aoSN0EZF4oQz0tmgHeZFQli4iMmBCmYqtUadAFxGJE8pUbI92kB/RkIuISKxQBnprtINcHaGLiHQTvlSsX8/p7y5mTFFSd/4VERk0whfom57hi/t/wMSR4StdRGQghS4V9x6JAlAxsiDgSkREMkvoAn3NzkMAnFxaGHAlIiKZJXSBPmNiKQAfmT464EpERDJL6AJ9xFBvqCVCR8CViIhkltAFOjn+1S0d0WDrEBHJMCEM9Ig372gPtg4RkQwTvkA3P9CdjtBFRGKFL9C7hlw0hi4iEiuEge6XrCEXEZFuwhfoGnIREUkofIGuq1xERBIKYaDrKhcRkUTCF+hdQy46KSoiEit8gd51hK4hFxGRWCEOdA25iIjESirQzWyemW0ws1ozu7WXNleb2TozW2tm/5XaMmPfSFe5iIgk0udjf8wsAtwNXALUAcvNbJFzbl1MmynAbcAc59w+Mxu4WyHqKhcRkYSSOUKfDdQ657Y451qBR4Ar49p8AbjbObcPwDlXn9oyY2gMXUQkoWQCfTywPeZ1nb8u1qnAqWb2JzN7zczmJdqRmS00sxozq2loaOhfxRpyERFJKFUnRXOBKcBc4Frgl2ZWHN/IOXevc67aOVddXl7ev3fSEbqISELJBPoOoDLmdYW/LlYdsMg51+ac2wpsxAv41NNVLiIiCSUT6MuBKWY2yczygQXAorg2T+EdnWNmZXhDMFtSV2YMDbmIiCTUZ6A759qBm4GlwHrgUefcWjO708yu8JstBfaY2TrgReDrzrk9A1OxrnIREUmkz8sWAZxzS4AlcevuiFl2wNf8aWBpDF1EJKHwfVPU/JI15CIi0k34Al1DLiIiCYUw0HWVi4hIIuELdF3lIiKSUPgCXQ+JFhFJKISBrodEi4gkEsJA7zxCV6CLiMQKX6BH8r15tDXYOkREMkyIA70t2DpERDJM+ALdDHLyINoSdCUiIhklfIEO3lG6jtBFRLoJZ6Dn5kO7jtBFRGKFM9Aj+TopKiISJ6SBXqBAFxGJE9JAz1Ogi4jECWeg5xZoDF1EJE44Az2Sp6tcRETihDTQC3QduohInJAGuq5DFxGJF85A13XoIiI9hDPQdR26iEgPCnQRkSwRzkDPLYD25qCrEBHJKOEM9Pxh0Ho46CpERDJKSAO9CFoPBV2FiEhGCWmgD4O2Q3pQtIhIjPAGOkCbhl1ERDqFNNCLvLmGXUREuoQ80A8GW4eISAYJaaD7Qy4KdBGRLkkFupnNM7MNZlZrZrceo91VZubMrDp1JSbQGegtCnQRkU59BrqZRYC7gcuAqcC1ZjY1QbvhwFeA11NdZA9DSrx58/4BfysRkbBI5gh9NlDrnNvinGsFHgGuTNDuX4DvAwP/Fc5h5d78UMOAv5WISFgkE+jjge0xr+v8dV3MbCZQ6Zz747F2ZGYLzazGzGoaGk4gjIeVeXMFuohIlxM+KWpmOcBdwC19tXXO3eucq3bOVZeXl/f/TXMLoGAEHNrd/32IiGSZZAJ9B1AZ87rCX9dpODAdeMnMtgHnAosG/MTosDIdoYuIxEgm0JcDU8xskpnlAwuARZ0bnXONzrky51yVc64KeA24wjlXMyAVdyoaA03vDuhbiIiESZ+B7pxrB24GlgLrgUedc2vN7E4zu2KgC+xVSRXs2xbY24uIZJrcZBo555YAS+LW3dFL27knXlYSSqpg9SPQ1gx5hWl5SxGRTBbOb4oClEwCHOz/a9CViIhkhPAG+qhJ3nzf1mDrEBHJEOEN9LIp3rx+XbB1iIhkiPAG+pASGFkJ760NuhIRkYwQ3kAHGDMN3n0r6CpERDJC+AN990Zobwm6EhGRwIU80KeDi0L9+qArEREJXLgDfdwMb75jRbB1iIhkgHAHekmVdyvduuVBVyIiErhwB7oZVMyG7cuCrkREJHDhDnSAyrNh72Y4tCfoSkREAhX+QK+Y7c3rdJQuIoNb+AN9/CzILYQtLwddiYhIoMIf6HmFMHEObH4+6EpERAIV/kAHmHyx9wWj/dv7bisikqWyI9BPucibb34h2DpERAKUHYFefrp3o64NS/puKyKSpbIj0M1g6pXeEXpzY9DViIgEIjsCHbxAj7bChv8JuhIRkUBkT6CPr4bh42DdU0FXIiISiOwJ9JwcmPZxqH0ODu8NuhoRkbTLnkAHmPFpb9jlzd8FXYmISNplV6CPmQbjZsLKh8C5oKsREUmr7Ap0gJmfhfq1sHNV0JWIiKRV9gX69Ksgbygs/1XQlYiIpFX2BXrhSDjr07DmMWh6L+hqRETSJvsCHeDcL0K0TUfpIjKoZGegl54Cp13mBXrr4aCrERFJi+wMdIDzboYje2Hlr4OuREQkLbI30KvmQNWF8OpdOkoXkUEhewMd4IO3w6F6jaWLyKCQVKCb2Twz22BmtWZ2a4LtXzOzdWb2ppk9b2YTU19qP0w8z7tX+p/+A1qagq5GRGRA9RnoZhYB7gYuA6YC15rZ1Lhmq4Bq59z7gceBf0t1of120bfh8B545QdBVyIiMqCSOUKfDdQ657Y451qBR4ArYxs45150znUOVL8GVKS2zBMwfhac+Sl47WewZ3PQ1YiIDJhkAn08EPuwzjp/XW9uBJ5OtMHMFppZjZnVNDQ0JF/lifrQdyCSD0tvT997ioikWUpPiprZZ4Bq4N8TbXfO3eucq3bOVZeXl6fyrY9t+Enwga/DxqfhbT2mTkSyUzKBvgOojHld4a/rxsw+BNwOXOGca0lNeSl07k0wZjos/qruly4iWSmZQF8OTDGzSWaWDywAFsU2MLMZwC/wwrw+9WWmQG4+fOxncGg3LP1W0NWIiKRcn4HunGsHbgaWAuuBR51za83sTjO7wm/270AR8JiZvWFmi3rZXbDGngkX3gKrH4b1i4OuRkQkpcwF9CCI6upqV1NTk/43bm+F+y+FvVvg71+Bkqr01yAi0k9mtsI5V51oW3Z/UzSR3Hz45IPggMdu8AJeRCQLDL5AB++o/GN3w86V8PTX9bg6EckKgzPQAc64HC74R1jxIPzlp0FXIyJywnKDLiBQF90Be7fCM/8HiifC1Cv6/hkRkQw1eI/QAXJy4OP3QEU1PPF3sPnFoCsSEem3wR3oAHlD4FOPQulkePha2Pa/QVckItIvCnSAoaPgc/8NJRPhP6+Gra8GXZGIyHFToHcqKvdCvbgSfvu3sPapoCsSETkuCvRYw0+CG56GcTPgseth2S+DrkhEJGkK9HhDR8Fnn4LTLoMl/wSLvgztmXevMRGReAr0RPKHwjW/9e77svLX8MB8aOxxg0kRkYyiQO9NTgQuvsML9oa34efnw9ong65KRKRXCvS+nHE5LHwZSk/xxtWf+AIc2Rd0VSIiPSjQk1E2GT7/DMz9Frz1BPz0bFj9iO4BIyIZRYGerEguzP0mLHzRu03Ak38PD34E3n0r6MpERAAF+vEbeybc+Cxc/iOoXwf3XOANw+zdGnRlIjLIKdD7IycHZl0PX1oJc74C6/8AP62Gxf/oPThDRCQACvQTMXQUXPJd+PIqmHkdrPot/GQWPPo5qFsRdHUiMsgo0FNhxFj46F3w1TXeEfvml+BXF8EvL4YVv4aWg0FXKCKDwOB7pmg6tDTByoe8h2fs3gD5RTD9Knj/NTDhXO8adxGRfjjWM0UV6APJOdi+DFb+Btb+HtoOQ9EYmHolTP2Ywl1EjpsCPRO0HoKNS71vm256BtqbYcgomHwxTL7Emw8rC7pKEclwxwr0wf0IunTKHwbT/9abWg7CpqWw8RmofQ7WPAYYjDsLJs7xpgnneiddRUSSpCP0oHV0wK43vGDf/ALsWAHRVm9b+Rkw8TwYN9ML+/LTIZIXZLUiEjANuYRJWzPsXAnv/Nmbti+D1iZvW6QAxkzzwn3smV7Al58GQ0oCLVlE0kdDLmGSVwgTz/cm8I7g927xjuJ3roJdq2HNE1Bz/9GfKRoDZaceDfjSyVBSBSMrdEQvMogo0DNdTo53c7CyyfC+T3jrOjpg/zuweyM0bPCnt70bhnUezQNYBEaO98K9cyqeCCPGwfCx3pRXGECnRGQgKNDDKCcHRk3yplM/fHS9c3BgJ+zbCvu2dZ82/A8cqu+5ryGjjgb8iLEwYjwUjYahZd5VN8PKYWgpFBZ77ysiGUuBnk3MvCPykeOh6oKe21sOQuN2L/SbdsGBXdC08+h812o41AAkOK9iES/Yh5XDsFIv8IeWQuFIbxpSfHS5MHZ5pK61F0kTBfpgUlAEo8/wpt60t8Lh3XBotz/f44V857rO9btWw+E90HIAXEcf7zvCC/aCEV4N+cP8qcif/OVjbcsrhNwhR+f6a0GkBwW6dJeb7w3BjBiXXHvnoPUgHNkPzY3+FLvceHRbywGvbUuT91dB6yHvdevBo5dqJiuSHxPwhZA3JLl5boF3ojiSHzfl+dv85Ujscv6xf87sOP+RRQZGUoFuZvOAHwER4FfOuX+N214A/AaYBewBrnHObUttqZKRzKBguDdR2f/9tLdC2yFvWKj1kD81efOWg9B+xLukM3be3gJtR7xv3cbOWw96f0n0+JlmcNGUdb1LJB9yciEnzxteysn1pkju0eVE03Fvj/jv4b+2HO8vFYv4y/7cIv76+G2RmJ/pbBeJa3eC+zMDzH+daNnilnOSW9YvzaT0GehmFgHuBi4B6oDlZrbIObcuptmNwD7n3GQzWwB8H7hmIAqWLJWb700DfU19R9T7ayDaCtE275dC53K0JW59W0zbmKk9/nULdLR7++5o96e27q+jca8717Ud8Zfbu2+Ln7ptbxvYf6NMluwvAAwMjv3LJe4XDX574hdj18W16/aLJpl1/nzuN70b9qVYMkfos4Fa59wWry57BLgSiA30K4F/9pcfB35qZuaC+taSSG9yIpAzxBuCCbOOqHfuonPu4l/HbYtd37Ut9mdcXLv+7s8Bzt9fR8xyzHr8bX0up6p9fC0J6oK4ZwQnsy5mWzLrYvdVWJzsf+njkkygjwe2x7yuA87prY1zrt3MGoFSYHdsIzNbCCwEmDBhQj9LFhHvyqGIvjgm3aT1UgHn3L3OuWrnXHV5eXk631pEJOslE+g76H62q8Jfl7CNmeUCI/FOjoqISJokE+jLgSlmNsnM8oEFwKK4NouA6/zlTwAvaPxcRCS9+hxD98fEbwaW4l22eL9zbq2Z3QnUOOcWAfcBD5lZLbAXL/RFRCSNkroO3Tm3BFgSt+6OmOVm4JOpLU1ERI6Hvj8tIpIlFOgiIllCgS4ikiUCewSdmTUA7/Tzx8uI+9JSiKkvmSlb+pIt/QD1pdNE51zCL/IEFugnwsxqenumXtioL5kpW/qSLf0A9SUZGnIREckSCnQRkSwR1kC/N+gCUkh9yUzZ0pds6QeoL30K5Ri6iIj0FNYjdBERiaNAFxHJEqELdDObZ2YbzKzWzG4Nup5EzOx+M6s3s7di1o0ys2fNbJM/L/HXm5n92O/Pm2Y2M+ZnrvPbbzKz6xK91wD3o9LMXjSzdWa21sy+EuK+FJrZMjNb7fflu/76SWb2ul/z7/w7imJmBf7rWn97Vcy+bvPXbzCzD6e7L34NETNbZWaLQ96PbWa2xszeMLMaf13oPl9+DcVm9riZvW1m683svLT3xTkXmgnvbo+bgZOBfGA1MDXouhLU+QFgJvBWzLp/A271l28Fvu8vzweexnvY4LnA6/76UcAWf17iL5ekuR9jgZn+8nBgIzA1pH0xoMhfzgNe92t8FFjgr78H+KK/fBNwj7+8APidvzzV/9wVAJP8z2MkgM/Y14D/Ahb7r8Paj21AWdy60H2+/Dp+Dfydv5wPFKe7L2ntcAr+wc4Dlsa8vg24Lei6eqm1iu6BvgEY6y+PBTb4y78Aro1vB1wL/CJmfbd2AfXpv/EeFh7qvgBDgZV4j1LcDeTGf77wbhd9nr+c67ez+M9cbLs01l8BPA9cBCz26wpdP/z33UbPQA/d5wvvoT5b8S80CaovYRtySfR80/EB1XK8xjjndvnL7wJj/OXe+pRRffX/VJ+Bd2Qbyr74wxRvAPXAs3hHpfudc+0J6ur2nFyg8zm5mdCX/wC+AXT4r0sJZz/Ae4LyM2a2wrxnDkM4P1+TgAbgAX8o7FdmNow09yVsgZ4VnPerNzTXi5pZEfAE8FXn3IHYbWHqi3Mu6pw7C+8IdzZwerAVHT8z+yhQ75xbEXQtKXKBc24mcBnwD2b2gdiNIfp85eINs/7cOTcDOIQ3xNIlHX0JW6An83zTTPWemY0F8Of1/vre+pQRfTWzPLww/0/n3O/91aHsSyfn3H7gRbyhiWLznoMbX1dvz8kNui9zgCvMbBvwCN6wy48IXz8AcM7t8Of1wJN4v2jD+PmqA+qcc6/7rx/HC/i09iVsgZ7M800zVexzV6/DG4/uXP85/6z3uUCj/yfaUuBSMyvxz4xf6q9LGzMzvMcLrnfO3RWzKYx9KTezYn95CN65gPV4wf4Jv1l8XxI9J3cRsMC/emQSMAVYlpZOAM6525xzFc65KrzP/wvOuU8Tsn4AmNkwMxveuYz3uXiLEH6+nHPvAtvN7DR/1cXAOtLdl3SfBEnByYf5eFdbbAZuD7qeXmp8GNgFtOH95r4Rb9zyeWAT8Bwwym9rwN1+f9YA1TH7+TxQ6083BNCPC/D+RHwTeMOf5oe0L+8HVvl9eQu4w19/Ml6Q1QKPAQX++kL/da2//eSYfd3u93EDcFmAn7O5HL3KJXT98Gte7U9rO/9/DuPny6/hLKDG/4w9hXeVSlr7oq/+i4hkibANuYiISC8U6CIiWUKBLiKSJRToIiJZQoEuIpIlFOgiIllCgS4ikiX+P+jIQdZXbSSCAAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.plot(torch.tensor(accurate_list))\n",
    "plt.plot(torch.tensor(loss_list))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "c35c927d-8d75-4ae7-aa79-c5c03e06c2c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "#  # Save the model\n",
    "# torch.save(mymodel, './models/contact-based model.pt')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a974057b-ee05-4fd0-93cc-6732b74bcd8b",
   "metadata": {},
   "source": [
    "## Model Evaluation  \n",
    "<font size=3>we follow the guidelence of the script shown on the github: https://github.com/facebookresearch/esm/blob/main/examples/contact_prediction.ipynb to get the plots and precision"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "203cbda3-f80e-45c5-8032-bc9371266aeb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_precisions(\n",
    "    predictions: torch.Tensor,\n",
    "    targets: torch.Tensor,\n",
    "    src_lengths: Optional[torch.Tensor] = None,\n",
    "    minsep: int = 6,\n",
    "    maxsep: Optional[int] = None,\n",
    "    override_length: Optional[int] = None,  # for casp\n",
    "):\n",
    "    if isinstance(predictions, np.ndarray):\n",
    "        predictions = torch.from_numpy(predictions)\n",
    "    if isinstance(targets, np.ndarray):\n",
    "        targets = torch.from_numpy(targets)\n",
    "    if predictions.dim() == 2:\n",
    "        predictions = predictions.unsqueeze(0)\n",
    "    if targets.dim() == 2:\n",
    "        targets = targets.unsqueeze(0)\n",
    "    override_length = (targets[0, 0] >= 0).sum()\n",
    "\n",
    "    # Check sizes\n",
    "    if predictions.size() != targets.size():\n",
    "        raise ValueError(\n",
    "            f\"Size mismatch. Received predictions of size {predictions.size()}, \"\n",
    "            f\"targets of size {targets.size()}\"\n",
    "        )\n",
    "    device = predictions.device\n",
    "\n",
    "    batch_size, seqlen, _ = predictions.size()\n",
    "    seqlen_range = torch.arange(seqlen, device=device)\n",
    "\n",
    "    sep = seqlen_range.unsqueeze(0) - seqlen_range.unsqueeze(1)\n",
    "    sep = sep.unsqueeze(0)\n",
    "    valid_mask = sep >= minsep\n",
    "    valid_mask = valid_mask & (targets >= 0)  # negative targets are invalid\n",
    "\n",
    "    if maxsep is not None:\n",
    "        valid_mask &= sep < maxsep\n",
    "\n",
    "    if src_lengths is not None:\n",
    "        valid = seqlen_range.unsqueeze(0) < src_lengths.unsqueeze(1)\n",
    "        valid_mask &= valid.unsqueeze(1) & valid.unsqueeze(2)\n",
    "    else:\n",
    "        src_lengths = torch.full([batch_size], seqlen, device=device, dtype=torch.long)\n",
    "\n",
    "    predictions = predictions.masked_fill(~valid_mask, float(\"-inf\"))\n",
    "\n",
    "    x_ind, y_ind = np.triu_indices(seqlen, minsep)\n",
    "    predictions_upper = predictions[:, x_ind, y_ind]\n",
    "    targets_upper = targets[:, x_ind, y_ind]\n",
    "\n",
    "    topk = seqlen if override_length is None else max(seqlen, override_length)\n",
    "    indices = predictions_upper.argsort(dim=-1, descending=True)[:, :topk]\n",
    "    topk_targets = targets_upper[torch.arange(batch_size).unsqueeze(1), indices]\n",
    "    if topk_targets.size(1) < topk:\n",
    "        topk_targets = F.pad(topk_targets, [0, topk - topk_targets.size(1)])\n",
    "\n",
    "    cumulative_dist = topk_targets.type_as(predictions).cumsum(-1)\n",
    "\n",
    "    gather_lengths = src_lengths.unsqueeze(1)\n",
    "    if override_length is not None:\n",
    "        gather_lengths = override_length * torch.ones_like(\n",
    "            gather_lengths, device=device\n",
    "        )\n",
    "\n",
    "    gather_indices = (\n",
    "        torch.arange(0.1, 1.1, 0.1, device=device).unsqueeze(0) * gather_lengths\n",
    "    ).type(torch.long) - 1\n",
    "\n",
    "    binned_cumulative_dist = cumulative_dist.gather(1, gather_indices)\n",
    "    binned_precisions = binned_cumulative_dist / (gather_indices + 1).type_as(\n",
    "        binned_cumulative_dist\n",
    "    )\n",
    "\n",
    "    pl5 = binned_precisions[:, 1]\n",
    "    pl2 = binned_precisions[:, 4]\n",
    "    pl = binned_precisions[:, 9]\n",
    "    auc = binned_precisions.mean(-1)\n",
    "\n",
    "    return {\"AUC\": auc, \"P@L\": pl, \"P@L2\": pl2, \"P@L5\": pl5}\n",
    "\n",
    "\n",
    "def evaluate_prediction(\n",
    "    predictions: torch.Tensor,\n",
    "    targets: torch.Tensor,\n",
    ") -> Dict[str, float]:\n",
    "    if isinstance(targets, np.ndarray):\n",
    "        targets = torch.from_numpy(targets)\n",
    "    contact_ranges = [\n",
    "        (\"local\", 3, 6),\n",
    "        (\"short\", 6, 12),\n",
    "        (\"medium\", 12, 24),\n",
    "#         (\"long\", 24, None),\n",
    "    ]\n",
    "    metrics = {}\n",
    "    targets = targets.to(predictions.device)\n",
    "    for name, minsep, maxsep in contact_ranges:\n",
    "        rangemetrics = compute_precisions(\n",
    "            predictions,\n",
    "            targets,\n",
    "            minsep=minsep,\n",
    "            maxsep=maxsep,\n",
    "        )\n",
    "        for key, val in rangemetrics.items():\n",
    "            metrics[f\"{name}_{key}\"] = val.item()\n",
    "    return metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "052c1fc8-76cf-4752-9bb2-85ab9fda0409",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"Adapted from: https://github.com/rmrao/evo/blob/main/evo/visualize.py\"\"\"\n",
    "def plot_contacts_and_predictions(\n",
    "    predictions: Union[torch.Tensor, np.ndarray],\n",
    "    contacts: Union[torch.Tensor, np.ndarray],\n",
    "    ax: Optional[mpl.axes.Axes] = None,\n",
    "    # artists: Optional[ContactAndPredictionArtists] = None,\n",
    "    cmap: str = \"Blues\",\n",
    "    ms: float = 1,\n",
    "    title: Union[bool, str, Callable[[float], str]] = True,\n",
    "    animated: bool = False,\n",
    ") -> None:\n",
    "\n",
    "    if isinstance(predictions, torch.Tensor):\n",
    "        predictions = predictions.detach().cpu().numpy()\n",
    "    if isinstance(contacts, torch.Tensor):\n",
    "        contacts = contacts.detach().cpu().numpy()\n",
    "    if ax is None:\n",
    "        ax = plt.gca()\n",
    "\n",
    "    seqlen = contacts.shape[0]\n",
    "    relative_distance = np.add.outer(-np.arange(seqlen), np.arange(seqlen))\n",
    "    bottom_mask = relative_distance < 0  #形成一个三角\n",
    "    masked_image = np.ma.masked_where(bottom_mask, predictions)\n",
    "    invalid_mask = np.abs(np.add.outer(np.arange(seqlen), -np.arange(seqlen))) < 6\n",
    "    predictions = predictions.copy()\n",
    "    predictions[invalid_mask] = float(\"-inf\")\n",
    "\n",
    "    topl_val = np.sort(predictions.reshape(-1))[-seqlen]\n",
    "    pred_contacts = predictions >= topl_val\n",
    "    true_positives = contacts & pred_contacts & ~bottom_mask\n",
    "    false_positives = ~contacts & pred_contacts & ~bottom_mask\n",
    "    other_contacts = contacts & ~pred_contacts & ~bottom_mask\n",
    "\n",
    "    if isinstance(title, str):\n",
    "        title_text: Optional[str] = title\n",
    "    elif title:\n",
    "        long_range_pl = compute_precisions(predictions, contacts, minsep=24)[\n",
    "            \"P@L\"\n",
    "        ].item()\n",
    "        if callable(title):\n",
    "            title_text = title(long_range_pl)\n",
    "        else:\n",
    "            title_text = f\"Long Range P@L: {100 * long_range_pl:0.1f}\"\n",
    "    else:\n",
    "        title_text = None\n",
    "\n",
    "    img = ax.imshow(masked_image, cmap=cmap, animated=animated)\n",
    "    oc = ax.plot(*np.where(other_contacts), \"o\", c=\"grey\", ms=ms)[0]\n",
    "    fn = ax.plot(*np.where(false_positives), \"o\", c=\"r\", ms=ms)[0]\n",
    "    tp = ax.plot(*np.where(true_positives), \"o\", c=\"b\", ms=ms)[0]\n",
    "    ti = ax.set_title(title_text) if title_text is not None else None\n",
    "    # artists = ContactAndPredictionArtists(img, oc, fn, tp, ti)\n",
    "\n",
    "    ax.axis(\"square\")\n",
    "    ax.set_xlim([-0.5, seqlen-0.5])\n",
    "    ax.set_ylim([-0.5, seqlen-0.5])\n",
    "    ax.set_xticks(np.arange(1,seqlen,2))\n",
    "    ax.set_yticks(np.arange(1,seqlen,2))\n",
    "    plt.colorbar(img, fraction=0.045)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "073739a5-7621-432b-8e10-52e42586ccb2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AttentionLogisticRegression(\n",
       "  (regression): Linear(in_features=660, out_features=1, bias=True)\n",
       "  (activation): Sigmoid()\n",
       ")"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mymodel=torch.load('../models/contact-based model.pt')\n",
    "mymodel.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "52c8c5a5-952f-4558-b012-70205f9fbb1d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAYcAAAFeCAYAAABuL9OZAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAAsTAAALEwEAmpwYAAAp+klEQVR4nO3de5xdZX3v8c83k+twEWXQwyVcrIECAQMOgRovHG+N2sKpHm2i9lWpLbU1r0pb4Wi1qHjal5ZW21pObaqAtVxULBY1iqBgVCRkgCGZJAQi5ZKIZMJFLplcJvmdP/Yau+e6Zz97r7Vn7fm+89qvtS/rWc+z96zs336eZ63fUkRgZmZWbUarG2BmZlOPg4OZmY3i4GBmZqM4OJiZ2SgODmZmNoqDg5mZjeLgYNOWpI9J+vfs/tGSnpXUkbCdv5D0+ea30Kx1HBwMAEkPSnrdiOfOlrQ/+9J8VtI2SR8fsU5Iei57/XFJ35P02yPWuVXSrmydX0haLemUIt7XZEXEwxFxYETsm2i97DPZOqLsX0fE7+fbQrNiOThYLT/LvjQPBF4BvEfS/xqxzkuz108ArgT+SdJHR6yzIlvnBcCtwJea2UhJM5u5PbPpzsHBJi0i/gu4DThpnNd3RMSXgD8CPiTp0DHW2QdcO942qmXDPtdJ+rKkZyTdJemlVa8/KOn/SFoHPCdppqSzJN0m6SlJ90g6u2r94yT9INvWTUBX1WvHZr2gmdnjF0i6QtLPJD0p6euSDgC+DRxR1Zs6onp4Kit7jqQNWRtulXTiiDZ/QNK6rBf1ZUlza30WZkVzcLBJk7QAWALcXmPV/wRmAovH2MZs4J2T2MaQc4GvUulxXA18XdKsqteXA28GDgFeBHwL+L/Z+h8AvibpsGzdq4E7qQSFTwC/O0G9XwI6gZOBFwKfiYjngDdS1ZuKiJ+NeH/HA9cAFwCHAauAb2Tve8jbgaXAccCpwLsn91GYFcfBwWo5IvsF/DRwH7AG+NFEBSJiL7CDyhf0kH+U9BTwDLAC+PgYRcdyZ0Rcl23z08Bc4Kzq7UbEIxExALwLWBURqyJif0TcBPQAb5J0NHAG8JcRsTsiVgPfGKtCSYdTCQLvjYgnI2JvRPxgku39beBbEXFT1ua/BeYBLx/R5p9FxBNZGxZNcttmhXFwsFp+FhGHRMTBVH6dDwBfnKhA9sv+MOCJqqf/JCIOofJF+RvAdZJOnUT9jwzdiYj9wFbgiLFeB44B3pYFs6eyYPQK4PCszJPZr/8hD41T53zgiYh4chLtG+mI6u1mbX4EOLJqnZ9X3d8JHJhQj1muHBxs0iLiF1SGZn6zxqrnAoPAHWNsY39E/BDYArxhEtXOH7ojaQZwFFA9lFOdVvgR4EtZMBu6HRARnwQeBZ6fzRsMOXqcOh8BXiDpkDFeq5XG+GdUgtRQm5W9h201yplNKQ4OVm2WpLlDNyrzBr8k6UBgGbBhrMLZJO47gcuAT0XE4+Os92tUJqTH3M4IL5P0lmyi+AJgN+PPV/w78JuSfl1SR/Y+zpZ0VEQ8RGWI6eOSZkt6BeMEuYh4lMrE8/+T9HxJsyS9Knv5MeBQSc8bpw1fAd4s6bVZD+rPszbfNon3ajZl+PA/q7ZqxOMfkx2Zkz0e+mJ+54j17pEUwB7gHuBPI+LqEev8k6S/z+7/HPhIRHx7Em36Tyrj+F+k0tt4SzaWP0pEPCLpXOBvqEwK76PSe/mjbJV3ZNt5AvgJ8G9UhsrG8jvAZ4B7gdnALcDqiLhX0jXAA6qcMDfsqKuI2CzpXcBnqQwl9QK/GRF7JvFezaYM+WI/NlVJ+hjwkoh4V6vbYjbdeFjJzMxG8bCStZSkbwOvHOOlvy66LWb23zysZGZmo3hYyczMRnFwMDOzUQqdc+jq6opjjz22yCrNzH5pf+Io+t133bkjIg6rvWb7KDQ4HHvssfT09BRZpZnZLz23Jy06HDhnxnipViat4+BjIgYHksrGQP+NEbG00TbUw0crmZkVIAZ3MedXlyWV3XX3Z7tqr9VcNeccJM2XdIukjVmO+vdnz78te7xfUnf+TTUzKzEBUtqtBSbTcxgE/jwi7pJ0EHBndqGUPuAtwL/k2UAzs7ah8hwDVLOlEfFoRNyV3X8G2AQcGRGbImJz3g00M2sbOfUcJC2VtFnSFkkfHOP1z0jqzW73ZensJ1RXGJN0LHAalQu+NM3OnTv58Y9/zM6dO5u5WbPkfasU++SOHXDppZWlAY39ve+4/bap/fceR5YA8jIqF6g6CVguaWRCyD+NiEURsYhKUsj/qLXdSQeHLF3z14ALIuLpOsqdL6lHUk9/f/+Y69x9993cfPPN3H333ZPdrNmkpO5bpdgnr7gCLrqosjQg/e/Wt66X1bfcTN+63nwaBoAqw0opt4ktBrZExANZ9t9rqVxTZTzLqWQtntCkjlbK8tJ/DbgqImpGnGoRsRJYCdDd3T3mcWSnnXbasKVZs6TuW6XYJ887b/jSkv9uC09dNGyZm3wml49k+BURtwJnjl29jqFy7fLv19pozeCQXcnqC8CmiPj0pJpap87OTpYsWZLHpm2aS923SrFPdnXBhRe2uhVTSiN/78Vnvbz2io0QjUxId0mqPklsZfbDu17LgOsiYl+tFSfTc1hC5cIn6yX1Zs/9BTCHytjVYcC3JPVGxK8nNNbMbBpo6LDUHREx3ikD26i6nC6VS+mOd1naZcD7JlNhzeAQET+iEvPGcv1kKjEzM/I6lHUtsEDScVSCwjIqVz0cXrX0q8DzqVwFsabyHHRrZlZ2ORzKGhGDwArgRiqnGnwlIjZIukTSOVWrLgOujUlep8HpM8zMSi4iVjHiGvARcfGIxx+rZ5sODmZWOuu3PptU7rjDOpvcknqoVGdIOziYmRVhKLdSSTg4mJkVxT0HMzMbzsNKZmY2lhnlGVaaEmGsFEnOrJRKs2+1cRK9IpMf7t61k/vWrWX3rin49x46Q7r5uZVyMSWCQymSnFkplWbfauMkekUmP3zovg30rV3NQ/dtqLeZNsKUGFYqRZIzK6XS7FttnESvyOSHxxx/8rDllFOio5U0yZPlmqK7uzt6enpqr2hmNoGiz3M4aG7HnRPkNpqUGQcfFXMWr0gqu+t7H2q4/npNiZ6Dmdm0UKKeg4ODmVlRfCirmZkNM8nrQU8V5QljZmZWGPcczMyK4mElM7PJecNlt9dd5vo/GPMSyTV1tPoM5RINKzk4mJkVwrmVzMxsLCXqOdQMY5Iul7RdUl/Vcy+V9BNJ6yV9Q9LB+TbTzKzk2jC30pXA0hHPfR74YEScAlwPXNjkdpmZtRm1V3CIiNXAEyOePh5Ynd2/CXhrk9uVnzbOfmmjFZ6VNXH/KkP22CKzqwLM3L+HI3f+FzP376mrnDVHakjaAJyb3X8bMH+8FSWdL6lHUk9/f39idU3UxtkvbbTCs7Im7l9lyB5bZHZVgBft2saxO+/jRbu21VVuShs6Ea7eWwukTkj/HvCPkv4SuAEYN7RHxEpgJVQS7yXW1zxtnP3SRis8K2vi/lWG7LFFZlcFeGzukcOWbaHdj1aKiHuBNwBIOh54czMblauuLrjQUyTTRWdnJ0uWLCmuwsT9q/B2JkhtY2q5wRmz2dZ5XN3lprQSHa2UFBwkvTAitkuaAXwE+Fxzm2Vm1mbUZuc5SLoGOBvokrQV+ChwoKT3Zav8B+ABfDOzWtqp5xARy8d56R+a3BYzs7amEgWH8vRxzMysME6fYWZNMe+0tEtgPnHHZ+suU6Zf4ENEudrt4GBmVgRlt5JwcDAzK4TcczAzs9EcHMzMbJQyBQcfrWRmZqOUOzg4w+q0kpLds+hMokWWK0MbAebNFotfMo95s8vzqzkvkpJuk9juUkmbJW2R9MFx1nm7pI2SNki6utY2yx0cnGF1WknJ7ll0JtEiy5WhjQCnHD2X1yw8kFOOnltXubajBm4TbVbqAC4D3gicBCyXdNKIdRYAHwKWRMTJwAW1mlvuOQdnWJ1WUrJ7Fp1JtMhyZWgjwPqHdw1bTlfK72ilxcCWiHgAQNK1VC6psLFqnT8ALouIJwEiYnutjSqiuCza3d3d0dPTU1h9Zlacdj4Jbt4s3RkR3Y1sY+ahL46D3viJpLJPXfWuceuX9L+BpRHx+9nj3wHOjIgVVet8HbgPWAJ0AB+LiO9M2N6klpqZWd0aCGpdkqp/Wa/MrpUzWTOBBVSSqB4FrJZ0SkQ8NVEBMzMrQAPBYccEPZdtDL8a51HZc9W2AmsiYi/wX5LuoxIs1o5XYbknpM3MbC2wQNJxkmYDy6hcobPa16n0GpDUBRwPPDDRRh0czMyKkNPRShExCKwAbgQ2AV+JiA2SLpF0TrbajcDjkjYCtwAXRsTjE23Xw0pmNsyuwbRym27626RyRU4u79/f2svY5/VeI2IVsGrEcxdX3Q/gz7LbpDg4mJkVIMdDWXPh4GBmVpC2Cg6SLgd+A9geEQuz574MnJCtcgjwVEQsyqmNZmbtoTyxYVI9hyuBfwL+beiJiPjtofuS/g74RdNbZmbWTtRmPYeIWC3p2LFeU+Wdvh14TZPbZWZmLdTooayvBB6LiPub0RizpitJ5t7UjKdF1rVz507W/KSYNrarvLKy5qHR4LAcuGaiFSSdL6lHUk9/f3+D1ZnVqSSZe1MznhZZ1/p77ubW79/M+nvyb2O7KlNwSD5aSdJM4C3AyyZaL8v/sRIqifdS6zNLUpLMvakZT4us65SXnjZsafWZToeyvg64NyK2NqsxZk3X1QUXXtjqVtTU2dnJkiVLpnRdnZ2dnPlrxbSxbZUnNtQeVpJ0DfAT4ARJWyW9J3tpGTWGlMzMLKM2G1aKiOXjPP/uprfGzMymBJ8hbWZWkOky52BmZnVwcDCzlut9+JmkcnsG9yeVO/Xo5yWVSzG4L62NHTNa/OVcntjg4GBmVhT3HMzMbJhWHnmUwsHBzKwgZQoOvkyomZmN4p6DmVlB3HMwm0hJMqWmaCTjab3lUuvavWuA+9evZfeugbrKWRMo8dYCDg5WvJJkSk2RmvE0pVxqXQ/f38fGnh/y8P19dZWzxrVV+gyzpitJptQUqRlPU8ql1nX0goXDllaQdrsSnFnTlSRTaopGMp7WWy61rjlz57HglDPqLmeNEVCi2OBhJTMzG809BzOzQvgkODMzG0OJYoODg5lZUdxzMLOmenZ3/VlITzj8wKS6iv4C27+//kvLJxQB4Onn9qYVbAa552BmZiMImNHqlOF1cHAwMytImXoONQ9llXS5pO2S+qqe+5ikbZJ6s9ub8m2mmZkVaTLnOVwJLB3j+c9ExKLstqq5zTIzaz9lSp9RMzhExGrgiQLaYlZ6RSfeW7vmtrrrshbJJqRTbq3QyBnSKySty4adnj/eSpLOl9Qjqae/v7+B6symviIT721Y38vqW77HhvW9dbbSWqGSPiOfnoOkpZI2S9oi6YNjvP5uSf1VUwG/X2ubqRPS/wx8Aohs+XfA7421YkSsBFYCdHd3Jx6AZlYORSbeO/mURcOWNtXlM0QkqQO4DHg9sBVYK+mGiNg4YtUvR8SKyW43KThExGNVDftX4Jsp2zFrN0Un3jvjzJfXXZe1Tk5DRIuBLRHxQKUOXQucC4wMDnVJGlaSdHjVw98CnBjezKyGnIaVjgQeqXq8NXtupLdmUwHXSZpfa6M1ew6SrgHOBrokbQU+CpwtaRGVYaUHgT+stR0zM0vWJamn6vHKbMh+sr4BXBMRuyX9IfBF4DUTFagZHCJi+RhPf6GORpmZWWNHHu2IiO5xXtsGVPcEjsqe+6WIeLzq4eeBv6lVoa/nYGZWgByPVloLLJB0nKTZwDLghmF1D58KOAfYVGujTp9hVqDHnk5L/HbAnI66y6QeGbNnsP4kfwCpaYNS2rk3sY0Hz2vtV14eE9IRMShpBXAj0AFcHhEbJF0C9ETEDcCfSDoHGKRy3tq7a23XwcHMrCB5ne2cZalYNeK5i6vufwj4UD3bdHAwMytIWyXeMzOz6cc9BzOzIshXgjMzsxEqRyu1uhWT52Ela2tFZkltpNzAwE56e25nYMAZVttX2mGsUzZlt1mZFZkltZFymzes4/Yffp/NG9bVVc7KpUwpuz2sZG2tyCypjZQ74eRThy2tPXnOwWyKKDJLaiPl5s3rZFH3WXWXsxJpYS8ghYeVzMxsFPcczMwKMJRbqSwcHMzMCuLgYGZmo5QoNjg4mKV4bk/a5dBTsqsCzOxo3+nB53YP1l1m3qy0zzFafBV79xzMzGy4kh2t5OBgZlYA0bqznVPU7KtKulzSdkl9Vc99IrtQda+k70o6It9mmplZkSYzkHklsHTEc5dGxKkRsQj4JnDxyEJmZjZcW6XPiIjVko4d8dzTVQ8PAFo8zWNmNvXNaKdhpfFI+itJjwDvxD2H8e3YAZdeWlm2T1UtkZLxtBVZWe+4/ba6y9n0UKaeQ3JwiIgPR8R84CpgxXjrSTpfUo+knv7+/tTqyuuKK+CiiyrL9qmqJVIynhadXbVvXS+rb7mZvnW9dZWz9qfsYj9lSdndjKOVrqJyYeuPjvViRKwEVgJ0d3dPv+Gn884bvmyPqloiJeNp0dlVF566aNjSrNqM8owqpQUHSQsi4v7s4bnAvc1rUpvp6oILL2y3qloiJeNpK7KyLj7r5XWXM5tqagYHSdcAZwNdkrZS6SG8SdIJwH7gIeC9eTbSzKwdlOk8h8kcrbR8jKe/kENbzMzaWolig8+QNjMrgqicJV0WDg5mZgVp+wlps3bxzK79SeX27ksr1zkn7b/cnsH660v9Ikottz/xWMQDEj6TSEyv2tHKb+cWHpaawsHBzKwgJYoNvoa0mZmN5p6DmVkBRLlyKzk4mJkVpESxwcHBzKwoZZqQ9pxDG0nNJFq0IjOlptY1sHMna9fcxsAU/yytPFIzspYuK6tNPamZRItWZKbU5Oyq63v54a3fo299b13lzCYyQ0q61SJpqaTNkrZI+uAE671VUkjqrrVNDyu1kdRMokUrMlNqcnbVUxYNW5o1Qx6dAEkdwGXA64GtwFpJN0TExhHrHQS8H1gzme06OLSR1EyiRSsyU2pqXfM6OznjTGdXtVJYDGyJiAcAJF1LJVv2xhHrfQL4FDCp3M0eVjIzK0gDF/vpGrpoWnY7v2qzRwKPVD3emj1XXe/pwPyI+NZk2+qeg5lZASrnOSQX3xERNecJxqxXmgF8Gnh3PeUcHMzMipBfbqVtwPyqx0dlzw05CFgI3JrV/z+AGySdExE9423UwcHMrCA5HZa6Flgg6TgqQWEZ8I6hFyPiF0DXf7dBtwIfmCgwgIODtYlnd6dlSd2fmN0zNbtqajbRmQVmEx3Ysy+pXEdHWhv3Ddb/mcyZmTZdui81dWyT5NFziIhBSSuAG4EO4PKI2CDpEqAnIm5I2a6Dg5lZyUXEKmDViOcuHmfdsyezTQcHM7MCNDghXbiafTNJl0vaLqmv6rm3Sdogaf9kzrQzM7OGDmUt3GQG7q4Elo54rg94C7C62Q0yM2tXSry1Qs3gEBGrgSdGPLcpIjbn1qp2smMHXHppZZkzJ94bu8zaNbdN+c/E2p+UX26lPOR+hrSk84fO6uvv78+7uqnniivgoosqy5w58d5oG9b3svqW77HBCfRsCihTVtbcJ6QjYiWwEqC7u7u1x5G1wnnnDV/myIn3Rjs5S5x3shPo2RRQpus5+GilvHV1wYWTynPVMCfeG7uME+iZ1c/BwcysICXqONQODpKuAc6mkhVwK/BRKhPUnwUOA74lqTcifj3PhpqZlZlo3eRyiprBISKWj/PS9U1ui5lZ+2rh5HIKDyuZmRXEE9JmZjZKma6u5uBgU87Pf7Gn7jKpWVLnzOpIKpea3bPI3Dp79qVlqp07O+0zSX1rKZlxU3+Ad7Twl7soV8+hTIHMzMwK4p6DmVlBypSV1cHBzKwgDg5mZjZMJU9SeaKD5xzaSJHZTltRbmBgJ709tzMw4AyrVk4zlHZrSVtbU63lochsp60ot3nDOtb86BY2b1hXVzmzqcJZWa0lisx22opyJ5x86rClmeXHwaGNFJnttBXl5s3rZFH3WXWXM5sKKteQLs+cg4ODmVlByjSO7+BgZlaQEnUcHBzMzIqgFl4POoWDg5lZQUoUGxwcLD8pCfQAZnXUPzI7e2baaG7RCfQG96XV15FQ4e69aYn3Uj5/SP/im5GQsi8hVx/Q+i/nMp0hXab5ETMzK4h7DmZmBSjboawN9RwkvV9Sn6QNki5oUpvMzNrStDhDWtJC4A+AxcAe4DuSvhkRW5rVODOzttHCPEkpGuk5nAisiYidETEI/AB4S3OaZWbWfpT4rxUaCQ59wCslHSqpE3gTML85zcrPjh1w6aWVZbtxdlWzqasy5zANsrJGxCbgU8B3ge8AvcC+ketJOl9Sj6Se/v7+1Oqa5oor4KKLKst24+yqZlNbmYJDQ0crRcQXgC8ASPprYOsY66wEVgJ0d3cnHp3cPOedN3zZTpxd1Wx6krQU+AegA/h8RHxyxOvvBd5H5Qf8s8D5EbFxom02FBwkvTAitks6msp8w5RPmdnVBRde2OpW5MPZVc2mtjyuBCepA7gMeD2VH+hrJd0w4sv/6oj4XLb+OcCngaUTbbfR8xy+JulQYC/wvoh4qsHtmZm1paE5hxwsBrZExAMAkq4FzgV+GRwi4umq9Q8Aao7iNDqs9MpGypuZTRuNnbPQJamn6vHKbMge4EjgkarXtgJnjqpeeh/wZ8Bs4DW1KvQZ0mZmBWngDOkdEdHdSN0RcRlwmaR3AB8Bfnei9R0czMwKkOOw0jaGn0ZwVPbceK4F/rnWRh0crKbNP087H+HI589NKjczMStokRKTuSYPKwzsGXWUeE3zZnekVZYoNVNqitTPMfXvNsWtBRZIOo5KUFgGvKN6BUkLIuL+7OGbgfupwcHBzKwgeeRJiohBSSuAG6kcynp5RGyQdAnQExE3ACskvY7KwUNPUmNICRwczMwKoqRrV0xGRKwCVo147uKq+++vd5sODmZmBRCtv9hQPRwczMyKULKsrA4OZmYFmTYX+ymj1EyiRdeXUq7o7Kq7B3ayqfcOdju7qllNQ8NKZbnYz7QLDqmZRIuuL6Vc0dlVH9jcxz1rfsADm/vqKmdmU9+0G1ZKzSRadH0p5YrOrvriExYOW5rZxMo0rDTtgkNqJtGi60spV3R21TnzOjlx0eK6y5lNVyWKDdMvOJiZtYIo1zi+g4OZWRGUz/Uc8uLgYGZWkPKEBgcHM7NCVLKylic8ODhMIz+678mkcicdeVBSudTsqnsH99ddZkbiqaeRmEo09T/5roT3BjBnVnGj1aln8aZmZU35KFPrKtMZyq3m4GBmVpAyxSYHBzOzgpRoVCn9yCpJJ0jqrbo9LemCJrbNzKyNCCnt1grJPYeI2AwsApDUQeUKRNc3p1lmZu1lup7n8FrgpxHxUJO2Z2bWdsp0nkOzAtky4JombcsoNivrnt0DPLTpTvbsHqi3mWbWphoODpJmA+cAXx3n9fMl9Ujq6e/vb7S6aaPIrKyPPrCRLb0/5tEHNtbbTDOrgxJvrdCMYaU3AndFxGNjvRgRK4GVAN3d3YlHJ08/RWZlPfzFJw1bmlkOpmH6jOV4SKnpiszKOnvOPI458WV1lzOzyZtWE9KSDgBeD/xhc5pjZta+pk3PISKeAw5tUlvMzNpaeUKDz5A2MytMiToOpRoCMzOzgrjnUFKv+Nsf1l3mm3/88qS65s7uSCqXKiXDavIPssSfcmseeCKp3Fm/8oKkcqlZSMsg5b2lZuFNyfjbLJUJ6fJ0HRwczMwKUqZhJQcHM7NCCLnnYGZmI7nnYGZmw5RtzmHaHa2Umpyu6PpSys2KPRyz5yFmxZ56m2lmNsy0Cw6pyemKri+l3BF7H+X4vT/liL2P1ttMM8ubKsNKKbdWmHbDSqnJ6YquL6Xcz2YdPmxpZlNLmeYcpl3PYSg5XWdn55SuL6XcXs3modnHsFez622mmRVAif9qbldaKmmzpC2SPjjG638maaOkdZK+J+mYWtucdsHBzKwVBMxQ2m3C7VYu03wZlcsnnAQslzQy//7dQHdEnApcB/xNrfY6OJiZFSSnnsNiYEtEPBARe4BrgXOrV4iIWyJi6OiW24Gjam102s05mJm1SgNzDl2Seqoer8wupAZwJPBI1WtbgTMn2NZ7gG/XqtDBwcxs6tsREd2NbkTSu4Bu4NW11nVwaLEnnhtMKnfVeWfUXSY1gd7gvrRkZfv2p2WLm9VR/2hnal2pv+TOfHFaAr3UdhYp9USt/YnZATsSkuil7pMzO1p7uFBO6TO2AfOrHh+VPTe8bul1wIeBV0fE7lobdXAwMyvA0IR0DtYCCyQdRyUoLAPeMaxu6TTgX4ClEbF9Mht1cDAzK0Q+ifciYlDSCuBGoAO4PCI2SLoE6ImIG4BLgQOBr2aXKn04Is6ZaLuNXkP6QeAZYB8w2IwxMTOztpTj2c4RsQpYNeK5i6vuv67ebTaj5/A/I2JHE7ZjZtbWSnSCtIeVzMyKUJlzKE94aPQkuAC+K+lOSec3o0HtJjUr68DOndx1x08YKCh7rJlZtUZ7Dq+IiG2SXgjcJOneiFhdvUIWNM4HOProoxusrnyGsqsCLFmyZNLlNvXdw49Xfx+A0xf/Wi5tM7Nilaff0GBwiIht2XK7pOupnMa9esQ6K4GVAN3d3VP/IO8mS83KeuLClw5bmlkbKFF0SA4Okg4AZkTEM9n9NwCXNK1lbWIou2q95nV2usdg1mamyzWkXwRcnx0zOxO4OiK+05RWmZm1oRLNR6cHh4h4APCYh5nZJJUoNjhlt5mZjebzHMzMilKiroODQ5M89PiupHIvPHhOUrnOOcX96VKyZkL6CT97EzJu3vXQU0l1nX7MIUnlNj/2bFK5E484KKlcSsZTJX7+qYcUJteXUGHqPjm4r3UHTIrpMyFtZmaTlWNupTw4OJiZFaREscHBwcysMCWKDg4OZmaFyOd6DnnxoaxmZjbKtAsOqVlSU8vtGthJ311r2DXg7Kpm052UdmuFaRcchrKk3n333YWU27JpPT233cqWTevrKmdm7UUN3Fph2s05pGZJTS33khNPGbY0s2msPFMO0y84pGZJTS03d14nC08/s+5yZtZ+yjQhPe2Cg5lZq/gkODMzG6VEsWH6TUibmVlt7jmYmRWhlYceJXBwGOGZXfVnBAU49MDZSeXSM1nWn11yf8EJKQcTsqtCWlbQM457flJdqZlji8yuCmn7Sco+AumfSZHvbW9idtXUbK7N4glpMzMbRnhC2szMxlCi2JAeHCTNBVYDc7LtXBcRH21Ww8zM2k6JokMjPYfdwGsi4llJs4AfSfp2RNzepLaZmVmLJAeHqMx2DV0rcVZ2a901+MzMprgyTUg3dJ6DpA5JvcB24KaIWNOUVk1SSqbU1OyqAzt3snbNbQzUWc7MbMi0ycoaEfsiYhFwFLBY0sKR60g6X1KPpJ7+/v5GqhslJVNqanbVvvW9/PDW79G3vrfOVpqZVUy7rKwR8ZSkW4ClQN+I11YCKwG6u7ubOuyUkik1NbvqwlMWDVuamdUtp296SUuBfwA6gM9HxCdHvP4q4O+BU4FlEXFdrW0m9xwkHSbpkOz+POD1wL2p20sxlCm1s7Mz1zIA8zo7OePMlzOvznJmZjDUC0j7N+F2pQ7gMuCNwEnAckknjVjtYeDdwNWTbW8jPYfDgS9mDZsBfCUivtnA9szM2ld+8weLgS0R8QCApGuBc4GNQytExIPZa5NOW9DI0UrrgPrGZszMrNmOBB6perwVaPgiMj5D2sysIA10HLok9VQ9XpnN5+bGwcHMrCjp0WFHRHSP89o2YH7V46Oy5xrS1sFh21N76i4zMzFr4/M6ZyWVS82cuWew/oynRWekTGkjwAFzitstU8eAE/9szOwo8hIqidlVE9P3pr6353YN1l3mgLlp+0jqe2uO2pPLidYCCyQdRyUoLAPe0ehGfbEfM7OC5HESXEQMAiuAG4FNVA4O2iDpEknnVOrVGZK2Am8D/kXShlptbeueg5nZVJHnCW0RsQpYNeK5i6vur6Uy3DRpDg5mZkUpT2olDyuZmdlopQ4OKUn0dg3sZN2dt7NrwAn0zKxYeZwhnZdSB4eUJHr3bVzHHT+6hfs2rsuxZWZmo5UpK2up5xxSkugdf9Kpw5ZmZkUp0ZRDuYPDUBK9esyd18mpLzsrpxaZmY2jhb2AFKUODmZm5VKe6FDqOQczM8uHew5mZgUQHlYyM7MxlCg2lCM4PLcnLVnWjqd3113mhCMOSqorVWoesNkz6x8R3LsvrbLUhH0HJiZHK4PUX4D7Ckz8VnCeRdY9/Iukcqce/by6ywzuS0vqOKPFP93dczAzs1FadUJbCgcHM7OilCc2ODiYmRWlRLGhsUNZJV0uabukvmY1yMzMWq/R8xyuBJY2oR1mZm0tNa9SqyaxGwoOEbEaeKLRRqRkVx0qd8ftt9VdzsysFZyVtYqk8yX1SOrp7+8fc52U7KoAfet6WX3LzfSt621CS83McqbEWwvkPiEdESuBlQDd3d1jHuSdkl0VYOGpi4YtzcymsjJNSE+Jo5VSsqsOlVt81stzaJGZWfP5JDgzMxuhdfMHKRo9lPUa4CfACZK2SnpPc5plZmat1FDPISKWN6shZmbtrGxZWX09BzMzG6XQOYc9g8HDj++qu9zBnbOS6is6w2qKiMQsnQk/QYrO0pmagTQlc2bRWVJTs3um/g1SdpPUXWt/YsFT5h+cVC4lw2pqpuACk+KOqUw9B09Im5kVpEwT0g4OZmZFaGEqjBSeczAzs1HcczAzK0ALM2EkcXAwMytKiaLDlBhW2jWwk/V3rWHXgLOrmln7clbWOt2/aT09t93K/ZvWt7opZma5KdP1HKbEsNKCE08ZtjQza0clGlWaGsFh7rxOTjn9zFY3w8wsXyWKDlNiWMnMzKaWKdFzMDObDnyGtJmZDVO2rKwODmZmBbjrrjtvnDdLXYnFdzS1MZOg5KygKZVJ/cBDTd5sF8V9cEXWVXR9fm/lrK+d31vR9U1U1zERcVhB7ZgSCg0OeZDUExHd7VZX0fX5vZWzvnZ+b0XXV/R7m+p8tJKZmY3i4GBmZqO0Q3BY2aZ1FV2f31s562vn91Z0fUW/tymt9HMOZmbWfO3QczAzsyYrbXCQdLmk7ZL6CqhrrqQ7JN0jaYOkj+dc34OS1kvqldSTc10nZPUM3Z6WdEGO9b1fUl/2OTa9nrH2C0lvy+rbL6mpR6OMU98nJK3LPs/vSjoix7o+Jmlb1d/vTc2oa4L6vlxV14OSenOs66WSfpL9X/iGpIObVNd8SbdI2pjtF+/Pns9tPymliCjlDXgVcDrQV0BdAg7M7s8C1gBn5Vjfg0BXCz7TDuDnVI7pzmP7C4E+oJPKCZg3Ay/Je78ATgROAG4Fuguo7+Cq+38CfC7Huj4GfCCnv9eE/8eAvwMuzvG9rQVend3/PeATTarrcOD07P5BwH3ASXnuJ2W8lbbnEBGrgScKqisi4tns4azs1o6TNa8FfhoRzT5RcciJwJqI2BkRg8APgLc0s4Kx9ouI2BQRm5tZT436nq56eABN2leK3Odr1SdJwNuBa3Ks63hgdXb/JuCtTarr0Yi4K7v/DLAJODLP/aSMShsciiapI+tCbwduiog1OVYXwHcl3Snp/BzrGWkZTfrPPo4+4JWSDpXUCbwJmJ9jfS0j6a8kPQK8E7g45+pWZMNYl0t6fs51DXkl8FhE3J9jHRuAc7P7byOHfUXSscBpVEYDrIqDwyRFxL6IWAQcBSyWtDDH6l4REacDbwTeJ+lVOdYFgKTZwDnAV/OqIyI2AZ8Cvgt8B+gF9uVVXytFxIcjYj5wFbAix6r+GfgVYBHwKJWhniIsJ98fElAZSvpjSXdSGf7Z08yNSzoQ+BpwwYjenuHgULeIeAq4BViaYx3bsuV24HpgcV51VXkjcFdEPJZnJRHxhYh4WUS8CniSynhvO7uKJg2HjCUiHst+uOwH/pUC9hVJM6kMB345z3oi4t6IeENEvIxKIPpps7YtaRaVwHBVRPxHs7bbThwcJkHSYZIOye7PA14P3JtTXQdIOmjoPvAGKsMxeSvilyCSXpgtj6byBXN13nUWTdKCqofnktO+ktV1eNXD36KYfeV1wL0RsTXPSqr2lRnAR4DPNWm7Ar4AbIqITzdjm22p1TPiqTcqX2SPAnuBrcB7cqzrVOBuYB2V/3xNOUJjnLpeDNyT3TYAHy7gszwAeBx4XgF1/RDYmL2/1xaxX1D50twK7AYeA27Mub6vZfvJOuAbVCY786rrS8D6rK4bgMPzfG/Z81cC7y3g7/Z+Kj3L+4BPkp2024S6XkFlXm8dlaHNXirzX7ntJ2W8+QxpMzMbxcNKZmY2ioODmZmN4uBgZmajODiYmdkoDg5mZjaKg4OZmY3i4GBmZqM4OJiZ2Sj/H/6/93Bm8h8kAAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 432x432 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "#contact predicts from lora_esm2\n",
    "lbd_lora_out=model.esm.contact_head(inputs['input_ids'],lbd_lora_attention)\n",
    "fig, axes = plt.subplots(figsize=(6, 6))\n",
    "plot_contacts_and_predictions(\n",
    "        lbd_lora_out[0], contacts, ax=axes, title = 'LBDB_prediction'\n",
    "    )\n",
    "# plt.savefig('./outputs/lbdb_contact_map_initial.png',dpi=800)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "0f18b7fa-3016-4845-91cc-e38b37f48670",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'id': 'lbdb',\n",
       " 'model': 'ESM-2 (Unsupervised)',\n",
       " 'local_AUC': 0.37979942560195923,\n",
       " 'local_P@L': 0.3636363744735718,\n",
       " 'local_P@L2': 0.3636363744735718,\n",
       " 'local_P@L5': 0.5,\n",
       " 'short_AUC': 0.2595936954021454,\n",
       " 'short_P@L': 0.22727273404598236,\n",
       " 'short_P@L2': 0.1818181872367859,\n",
       " 'short_P@L5': 0.25,\n",
       " 'medium_AUC': 0.004545454867184162,\n",
       " 'medium_P@L': 0.04545454680919647,\n",
       " 'medium_P@L2': 0.0,\n",
       " 'medium_P@L5': 0.0}"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "metrics = {\"id\": 'lbdb', \"model\": \"ESM-2 (Unsupervised)\"}\n",
    "metrics.update(evaluate_prediction(lbd_lora_out, contacts))\n",
    "metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "486ea8e2-d4a9-4ad3-a508-70afbb53635c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAYcAAAFeCAYAAABuL9OZAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAAsTAAALEwEAmpwYAAAkR0lEQVR4nO3de5ydVX3v8c93ktAkgFVI7OGSiD1FjxYUdA7VIsrBSwMe5ZS+tEQ950g9TdsjLVjFY6sFlNpXLcVrqXZaLtYiXkuLGhXsC0yViwwkYsLFUoqSgGRCgjdALvmdP/YzujO3PbP286y9197fd177NXv2ftaz1szszG/Wep7nuxURmJmZtRvp9QDMzKz/uDiYmdk0Lg5mZjaNi4OZmU3j4mBmZtO4OJiZ2TQuDja0JJ0t6R+q+6sl/UjSooT9/LGkv6t/hGa94+JgAEi6S9JLpjx2rKTd1S/NH0naJumdU7YJST+unr9f0r9I+s0p21wt6eFqm+9L2iDp8Bxf13xFxHcjYp+IeHyu7arvydYpbf8sIv5PsyM0y8vFwTq5p/qluQ/wAuANkv7HlG2eXT3/dOBi4K8knTVlm1OrbfYDrgY+VucgJS2uc39mw87FweYtIv4DuAZ45izP74iIjwG/B/yRpP1n2OZx4BOz7aNdtezzGUmflPRDSTdJenbb83dJ+n+SbgZ+LGmxpOdJukbSA5K+KenYtu2fKumr1b6uBFa0PXdINQtaXH2+n6SLJN0jaZekf5K0N/BF4MC22dSB7ctTVdtXStpSjeFqSc+YMua3SLq5mkV9UtLSTt8Ls9xcHGzeJB0KHA1c12HTfwYWA0fNsI+9gNfOYx+TTgQ+TWvG8XHgnyQtaXt+LfBy4InALwBfAP602v4twGclray2/ThwI62icA7wv+fo92PAcuCXgScD74uIHwPH0zabioh7pnx9TwMuBU4HVgLrgc9VX/ekVwNrgKcCzwJeP79vhVk+Lg7WyYHVX8A/AL4NXA98ba4GEfEosIPWL+hJH5T0APBD4FTgnTM0ncmNEfGZap/vBZYCz2vfb0TcHREPAa8D1kfE+ojYHRFXAuPACZJWA/8V+JOI+ElEbAA+N1OHkg6gVQR+NyJ2RcSjEfHVeY73N4EvRMSV1Zj/ElgG/OqUMd8TETurMRwxz32bZePiYJ3cExFPjIgn0Prr/CHgo3M1qP6yXwnsbHv4DyLiibR+Uf534DOSnjWP/u+evBMRu4GtwIEzPQ88BXhVVcweqIrRC4ADqja7qr/+J31nlj5XATsjYtc8xjfVge37rcZ8N3BQ2zbfa7v/ILBPQj9mjXJxsHmLiO/TWpp5RYdNTwQeA74xwz52R8S/AncAL5tHt6sm70gaAQ4G2pdy2mOF7wY+VhWzydveEfHnwL3Ak6rjBpNWz9Ln3cB+kp44w3OdYozvoVWkJses6mvY1qGdWV9xcbB2SyQtnbzROm7wU5L2AU4GtszUuDqI+1rgfOA9EXH/LNs9n9YB6Rn3M8VzJZ1UHSg+HfgJsx+v+AfgFZJ+TdKi6us4VtLBEfEdWktM75S0l6QXMEuRi4h7aR14/mtJT5K0RNILq6fvA/aX9POzjOFTwMslvbiaQb25GvM18/hazfqGT/+zduunfP51qjNzqs8nfzG/dsp235QUwCPAN4E3RcTHp2zzV5LeX93/HvCOiPjiPMb0z7TW8T9Ka7ZxUrWWP01E3C3pROAvaB0UfpzW7OX3qk1eU+1nJ3At8Pe0lspm8j+B9wG3AXsBVwEbIuI2SZcCd6p1wdweZ11FxO2SXgd8iNZS0ibgFRHxyDy+VrO+Ib/Zj/UrSWcDvxQRr+v1WMyGjZeVzMxsGi8rWU9J+iJwzAxP/VnusZjZz3hZycyscJIupHWK+PaIOGyG5wV8ADiB1unTr4+Im+bap5eVzMzKdzGtq+5nczxwaHVbB3y40w5dHMzMCldd8b9zjk1OBP4+Wq4DnlglAcwq6zGHFStWxCGHHJKzSzOzn7rplu8mtYuHJnZExMrOW/atg9gzTWBr9di9szXIWhwOOeQQxsfHc3ZpZvZTy448Nandw5vOny1qZd4WPeEpEY89lNQ2HprYAjzc9tBYRIx1O6a5+GwlM7MM4rGH+bn/cnJS24c3fujhiBjtovtttEXR0IqhmTPSpeMxB0mrJF0l6ZYqo/606vFXVZ/vltTNoM3MBp8AKe3WvcuB/6WW5wHfr2JiZjWfmcNjwJsj4iZJ+wI3Vm+Ushk4CfibbkdtZjYU1Mw5QFWky7HACrXexvYsYAlARHyEVjTOCbQiaB4ETum0z47Foaou91b3fyjpVuCgKisf1VPVzMwGX0O/LyNibYfnA3jjQva5oDIm6RDgSFpv+FKbHTvg3HNbH83qlPra8muyTA8++CBf//rXefDBBxfUbtle4qhfWsayvfzH7qR5F4cqrvmzwOkR8YMFtFsnaVzS+MTExIzbXHQRvPWtrY9mdUp9bfk1WaaNGzfyla98hY0bNy6o3eGrl3LcYftw+Oom385brWWllFsPzOtspSqX/rPAJRHxjwvpoDrdagxgdHR0xqyOU07Z86NZXVJfW35NlunII4/c4+N8feu7D+/xsTEFLcN3LA5VJscFwK0R8d4mBrFiBZxxRhN7tmGX+trya7JMy5cv5+ijj15wu4ceCb5xR9o1CPMmejYLSDGfmcPRtN745FuSNlWP/THwc7Te0GQl8AVJmyLi1xoZpZlZ8Wo7LTWL+Zyt9DVaNW8ml9U7HDOzATZgMwczM6tDQTOHcsqYmZll45mDmRUnNUCvt+RlJTMzm2IyW6kQLg5mZrl45mBmZnvyspKZmc1kpJxlpb4oYw45s6akBrFZfVJ/Bint+jpAb/IK6UKylfqiODjkzJqSGsRm9Un9GaS0yxOgNxz6YlnJIWfWlNQgNqtP6s8gpV22AL1UPltpYRxyZk1JDWKz+qT+DFLaZQnQS+YD0mZmNhPPHMzMbBrPHMzMbA8qK7K7nDJmZmbZeOZgZpaLl5XMzOYnZ8LqxHUfTGq379Lz6xlAQctKLg5mZln4VFYzM5tJQTOHjmVM0oWStkva3PbYsyVdK+lbkj4n6QnNDtPMrHADmK10MbBmymN/B7wtIg4HLgN8fbOZ2Zw0WMUhIjYAO6c8/DRgQ3X/SuA3ah6XWS1KSfwtIT02Z7oq9HnC6hBILUlbgBOr+68CVs22oaR1ksYljU9MTCR2Z5amlMTfEtJjc6arwoAmrE5eCLfQWw+kHpD+LeCDkv4EuBx4ZLYNI2IMGAMYHR2NxP7MkpSS+FtCemzOdFUoIGE1xaCfrRQRtwEvA5D0NODldQ7KrC6lJP6WkB6bM10V+j1hNVFBZyslFQdJT46I7ZJGgHcAH6l3WGZmA0YDdp2DpEuBY4EVkrYCZwH7SHpjtck/An2+omtm1gcGaeYQEWtneeoDNY/FzGygqaDiUM4cx8zMsnF8hpnVImeAHovSfnUtGundX+6irJmDi4OZWQ6qboVwcTAzy0KeOZiZ2XQuDmZmNk1JxcFnK5mZ2TQuDlaMlHTP1FTW1HY5k0tzp6Q6XbV7kpJuveDiYMVISfdMTWVNbZczuTR3SqrTVbukLm494GMOVoyUdM/UVNbUdjmTS3OnpDpdtTsq7GwlReRL0R4dHY3x8fFs/ZlZPiVcBLfz2vcltVu+18iNETGa1LiyeP9fjH2PPyep7QOXvK7r/hfKMwczs0xKmjm4OJiZZVJScfABaTMzm8YzBzOzHJytZGYly3pgOdGOa9IOLPd6WafX/S+Ei4OZWQalncrqYw5mZpk0dYW0pDWSbpd0h6S3zfD8aklXSdoo6WZJJ3TaZ8fiIOlCSdslbW577JOSNlW3uyRt6jh6M7Nh18AV0pIWAecDxwPPBNZKeuaUzd4BfCoijgROBv6601Dns6x0MfBXwN9PPhARv9k2sPOA789jP2Zmw0uNHXM4CrgjIu4EkPQJ4ETglrZtAnhCdf/ngXs67bRjcYiIDZIOmek5tb7SVwPHddqPmZk14iDg7rbPtwK/MmWbs4ErJP0+sDfwkk477faYwzHAfRHxb13ux2yopSae5uzL6ard6+KYwwpJ4223dQvsei1wcUQcDJwAfEzSnL//uz1baS1w6VwbVF/EOoDVq1d32Z3ZYJpMPAU4+uij+7KvyXRVgG/c8VAjYxt0XSwr7ZgjW2kbsKrt84Orx9q9AVgDEBHXSloKrAC2z9ZhcnGQtBg4CXjuXNtFxBgwBq3gvdT+zAZZauJpzr6crtqdBk9lvQE4VNJTaRWFk4HXTNnmu8CLgYslPQNYCkzMtdNuZg4vAW6LiK1d7MPMgOXLlzc+Y+i2r4ceCc8YutVAbYiIxySdCnwZWARcGBFbJL0LGI+Iy4E3A38r6U20Dk6/PjpEcncsDpIuBY6ltea1FTgrIi6gVZ3mXFIyM7NKc2crERHrgfVTHjuz7f4twIL+IpjP2UprZ3n89QvpyMzMyuH4DDOzTEqKz3BxMDPLxMXBzHoue7pq4i+++6/7YK6ueq+gcbs4mJll4pmDmZntYb4Jq/3CxcHMLJOSioPfz8HMzKbxzMHMLBPPHMyGVGriaUo7p6sWqIE3+2mKi4NZjSYTTzdu3Nh4u9S+JtNVD1+9dEHtrHtNvU1oE7ysZFaj1MTTlHZOVy1Mg9lKTXBxMKtRauJpSjunq5ZFlHXxnpeVzMxsGs8czMyy8EVwZmY2g4Jqg4uDmVkunjmYWa1yJqxOJKSkAqx8/mlJ7XbP/W6VM1o8kna4dPfuHr6NvTxzMDOzKQSMjJRTHVwczMwyKWnm0HFuJulCSdslbW577GxJ2yRtqm4nNDtMMzPLaT4LdxcDa2Z4/H0RcUR1W1/vsMzMBk9J8Rkdi0NEbAB2ZhiLWfFyBu85QK8w1QHplFsvdHOF9KmSbq6WnZ4020aS1kkalzQ+MTHRRXdm/S9n8J4D9MrSis8oZ+aQekD6w8A5QFQfzwN+a6YNI2IMGAMYHR3t4XlkZs3LGbznAL3SDMEV0hFx3+R9SX8LfL62EZkVLGfwngP0ylNQbUgrDpIOiIh7q09/Hdg81/ZmZjZgV0hLuhQ4FlghaStwFnCspCNoLSvdBfxOc0M0M7PcOhaHiFg7w8MXNDAWM7PB5fgMMzObavJspVK4OJhllDNAL9XKX31TWkOlnRm/8kVvW3CbHRvek9RXr381F1QbXBzMzHLxzMHMzKYpqDb4PaTNzGw6zxzMzHKQl5XMzGyK1tlKvR7F/HlZyQZazpTUbto5YXUYpIXu9W1kt1nJcqakdtPOCavDoaTIbi8r2UDLmZLaTTsnrA4HH3Mw6xM5U1K7aeeE1SFQWHyGl5XMzGwazxzMzDJwtpKZmc3IxcHMzKYpqDa4OJilyJ6umvJbZdGSpK52XnNeUrv9fvXNaf3968ITVncnvhv9yEhvfzt75mBmZnsq7GwlFwczswxE7652TtHxVFZJF0raLmlz22PnSLpZ0iZJV0g6sNlhmplZTvO5zuFiYM2Ux86NiGdFxBHA54Ezax6XmdnAGaj4jIjYIOmQKY/9oO3TvYHEw0NmZsNjZJCWlWYj6d2S7gZei2cOfWHHDjj33NbHQZSSeOp0VesnJc0ckotDRLw9IlYBlwCzntcnaZ2kcUnjExMTqd3ZPFx0Ebz1ra2Pgygl8dTpqtYvVL3ZTymR3XWcrXQJsB44a6YnI2IMGAMYHR318lODTjllz4+DJiXx1Omq1k+ausxC0hrgA8Ai4O8i4s9n2ObVwNm0DgN8MyJeM9c+k4qDpEMj4t+qT08EbkvZj9VrxQo444xej6I5KYmnTle1QSdpEXA+8FJgK3CDpMsj4pa2bQ4F/gg4OiJ2SXpyp/12LA6SLgWOBVZI2kprhnCCpKcDu4HvAL+78C/JzGy4NLREdBRwR0TcWfXxCVp/tN/Sts1vA+dHxC6AiNjeaafzOVtp7QwPXzCfEZuZ2c90URtWSBpv+3ysWrIHOAi4u+25rcCvTGn/tFb/+jqtpaezI+JLc3XoK6TNzDIQraukE+2IiNEuul8MHEprFehgYIOkwyPigbkamJlZBg0dkN4GrGr7/ODqsXZbgesj4lHgPyR9m1axuGG2nbo42FDLnq66KO2/3I5r3rfwrnInkC7eK1tXEWknPj7p+WnJsbVo7rTUG4BDJT2VVlE4GZh6JtI/AWuBiyStoLXMdOdcO3VxMDPLpInaEBGPSToV+DKt4wkXRsQWSe8CxiPi8uq5l0m6BXgcOCMi7p9rvy4OZmaFi4j1tK43a3/szLb7AfxhdZsXFwczswxEWdlKLg5mZpkUVBtcHMzMchmoN/uxcpSSypozKdXpqtYvUhNZi0tltf5TSiprzqRUp6taPxmRkm694GWlAVJKKmvOpFSnq1o/KWke6uIwQEpJZc2ZlOp0VbM0Lg5mZpmUdEDaxcHMLIPWdQ69HsX8uTiYmeXQw7f8TOHiYGaWSUG1wcXBBsOy5/x+3g5T/5ePLEpqlj1hNcXjjyY12+/otyy8UWoCbOxOa1eTkmYOvs7BzMym8czBzCyD0g5Id5w5SLpQ0nZJm9see5WkLZJ2S+rmrevMzIaGqoPSC731wnyWlS4G1kx5bDNwErCh7gGZmQ0qJd56oWNxiIgNwM4pj90aEbc3NipLkhoyl1vO4D0H6Fm/kMrKVmr8gLSkdZLGJY1PTEw03d1QSw2Zyy1n8N7hq5dy3C/v7QA96wslpbI2fkA6IsaAMYDR0dG0dwW3eUkNmcstZ/CeA/Ssn5R0KqvPVhogqSFzueUM3nOAnlkaFwczs0wKmjh0Lg6SLgWOBVZI2gqcResA9YeAlcAXJG2KiF9rcqBmZiUTvTu4nKJjcYiItbM8dVnNYzEzG1w9PLicwstKZmaZ+IC0mZlNU1KYnYuD9Z1lR56ar7PElNTU9YFd15yX1G737oWfBZ76R2ryX7dL0q4l2Xn1ny64Te6/wJct+UDX+xBlzRxKKmRmZpaJZw5mZpmUlMrq4mBmlomLg5mZ7aGVk1ROdfAxhwGSM+20F+2csGqlG1HarSdj7U231oScaae9aHf46qUcd9g+Tli1YjmV1XoiZ9ppL9o5YdUsHxeHAZIz7bQX7ZywaiVrvYd0OUuiLg5mZpmUtI7v4mBmlklBEwcXBzOzHNTD94NO4eJgZpZJQbXBxcGakzVAb/Feae1id1KzXde9P62/RCMZT3aPSHur911ffXfNIxk8JV0hXdLxETMzy8QzBzOzDEo7lbWrmYOk0yRtlrRF0uk1jcnMbCANxRXSkg4Dfhs4CngE+JKkz0fEHXUNzsxsYPQwJylFNzOHZwDXR8SDEfEY8FXgpHqGZWY2eJT4rxe6KQ6bgWMk7S9pOXACsKqeYTVnxw4499zWx0HjdFWz/tU65jAEqawRcSvwHuAK4EvAJuDxqdtJWidpXNL4xMREane1uegieOtbWx8HjdNVzfpbScWhq7OVIuIC4AIASX8GbJ1hmzFgDGB0dDTtBOoanXLKnh8HidNVzawuXRUHSU+OiO2SVtM63vC8eobVnBUr4Iwzej2KZjhd1ay/lfROcN1e5/BZSfsDjwJvjIgHuh+SmdngmTzmUIpul5WOqWsgZmYDrYfXLKTwFdJmZpkMzRXSZmY2P02eyippjaTbJd0h6W1zbPcbkkLSaKd9euZgHWVNVwUYWbTgJjuvOS+pq5IOEC5UarqqlUXSIuB84KW0zhi9QdLlEXHLlO32BU4Drp/Pfj1zMDPLpKFspaOAOyLizoh4BPgEcOIM251D69q0eZ1D7uJgZpaFGEm8ASsmLyaubuvadnwQcHfb51urx37Ws/QcYFVEfGG+o/WykplZBqKrs5V2RETH4wQz9iuNAO8FXr+Qdi4OZmY5NBeFsY09c+0Orh6btC9wGHB1dYztPwGXS3plRIzPtlMXBzOzTBo6lfUG4FBJT6VVFE4GXjP5ZER8H1gx+bmkq4G3zFUYYAiPOeROZc2ZeOp0VbP+NbmsVPcB6eotE04FvgzcCnwqIrZIepekV6aOd+hmDpOprJAnY2kyuRRYUA5RSrucfcHP0lUBZyWZ9VBErAfWT3nszFm2PXY++xy64pA7lTVn4qnTVc36W0lXSA9dccidypoz8dTpqmb9raDaMHzFwcysF0RZB3ldHMzMclBZcS0uDmZmmZRTGlwczMyyaKWyllMeXByGSPZ01dT/CAntSpqu55L6Pcmd5prSn3/ezXNxMDPLpKSS5uJgZpZJSROe5DOrJD1d0qa22w8knV7j2MzMBoiQ0m69kDxziIjbgSPgp+9EtA24rJ5hmZkNlmG9zuHFwL9HxHdq2p+Z2cAp6UB6XYXsZODSmvZl5E1ldbqqmU3VdXGQtBfwSuDTszy/bvKt7SYmJrrtbmhMJqVu3Lix0Tbws3TVw1cvXegwzWwBlHjrhTqWlY4HboqI+2Z6MiLGgDGA0dHRvCdQFyxnKqvTVc0yGML4jLV4Sal2OVNZna5q1ryhOiAtaW/gpcDv1DMcM7PBNTQzh4j4MbB/TWMxMxto5ZQGXyFtZpZNQROHopbAzMwsE88cCrXsOb/f6yE0Z9GSBTd50tGJ7/2a+U+5XV/7i6R2OZNLc6+Ll7QO343WAelyvlYXBzOzTEqqgy4OZmZZCHnmYGZmU3nmYGZmeyjtmMPQna2UGk6Xu7+Udg7QM7O6DF1xSA2ny91fSrvDVy/luF/e2wF6Zv1IrWWllFsvDN2yUmo4Xe7+Uto5QM+sv/mYQx9LDafL3V9KOwfomfU3n61kZmZ7EDBSTm1wcTAzy8UzBzMzm6akYw5Dd7aSmZl15plDjy078tReD6GzkUVJze6/9v1p3WVcmE0JtIP8YXG7E4a5KHGIu1M6I/3nljNU8LHHdye1q4uXlczMbA8+IG1mZjMYouA9SXcBPwQeBx6LiNE6BmVmNnB6eLVzijpmDv8tInbUsB8zs4FWUG3wspKZWQ6tYw7llIduT2UN4ApJN0paV8eABk1qKqsTVs2sl7qdObwgIrZJejJwpaTbImJD+wZV0VgHsHr16i67K89kuiqwoKykw1cv5bjD9gFwXpLZgCjpT72uikNEbKs+bpd0GXAUsGHKNmPAGMDo6GjaCdQFS01ldcKq2QAqqDokLytJ2lvSvpP3gZcBm+sa2KCYTFddvnz5gtpNJqw+9MjQ1VOzgaXEf73QzczhF4DLqisVFwMfj4gv1TIqM7MBVNDx6PTiEBF3As+ucSxmZgOtoNrg4D0zM5vO1zmYmeVS0NTBxaEmJaSr3n/9h5La7X/0m5Pa5UxXTZU7XTXVoozfy9w/t5w/g5zfx6mEU1nNzGyqIcxWMjOzeSioNrg4mJllU1B18NlKZmZZpF4C17miSFoj6XZJd0h62wzP/6GkWyTdLOlfJD2l0z5dHMzMCiZpEXA+cDzwTGCtpGdO2WwjMBoRzwI+A/xFp/0OXXFITUl1uqqZdUtKu3VwFHBHRNwZEY8AnwBObN8gIq6KiMlfXtcBB3fa6dAVh8mU1I0bN2ZpN5muevjqpQtqZ2aDRV3cOjgIuLvt863VY7N5A/DFTjsdugPSqSmpTlc1s66lLyCskDTe9vlYlXi9sO6l1wGjwIs6bTt0xWEyJTVXu8l0VTOzLi6C2xERo7M8tw1Y1fb5wdVje/YtvQR4O/CiiPhJpw6HrjiYmfVKQxfB3QAcKumptIrCycBr9uxXRwJ/A6yJiO3z2amLg5lZJk3Uhoh4TNKpwJeBRcCFEbFF0ruA8Yi4HDgX2Af4dBVX8t2IeOVc+3VxMDMrXESsB9ZPeezMtvsvWeg+XRzMzHKY56lH/cLFYYrc6arfu+YDSe32WpzvLOSJr/1lUruItLc43Z3QrJdpm5bf7pQXSR9wKquZme1BOJXVzMxmUFBtSC8OkpYCG4Cfq/bzmYg4q66BmZkNnIKqQzczh58Ax0XEjyQtAb4m6YsRcV1NYzMzsx5JLg7ROtr4o+rTJdWtzKNEZmYZlHRAuqtTXiQtkrQJ2A5cGRHX1zKqeUpJSnW6qpn1SkOprI3oqjhExOMRcQStLI+jJB02dRtJ6ySNSxqfmJjoprtpUpJSna5qZr3SUCprI2o5WykiHpB0FbAG2DzluTFgDGB0dLTWZaeUpFSnq5pZzxS08NDN2UorgUerwrAMeCnwntpGNg8pSalOVzWzXmjNAsqpDt3MHA4APlq9Rd0I8KmI+Hw9wzIzGzA9PH6QopuzlW4GFrY2Y2ZmRfAV0mZmmRQ0cXBxMDPLpqDqMNDFIXfCaooli9LOJt7v+W9acJud174vqa+Vx5yR1I7EVNZd1743rT8bGiNFpvBqaA5Im5nZAgzFAWkzM5u/wt7rx8XBzCybgqpDvrcTMzOzYhRdHFJC9BygZ2a9osR/vVB0cUgJ0XOAnpn1SkmprEUfc0gJ0XOAnpn1SknrFUUXh5QQPQfomVlPDEu2kpmZLVQ51aHoYw5mZtYMzxzMzDIQXlYyM7MZFFQbyigOJQTo3X/9h7L2t+u69y+4zWOP707r65rzktpZfSIhxFCJf6am9NVNfylSx5jYrDaeOZiZ2TROZTUzs+nKqQ0uDmZmuRRUG7o7lVXShZK2S9pc14DMzKz3ur3O4WJgTQ3jMDMbaKm5Sr06iN1VcYiIDcDObgeRkq4KTlg1s7I4lbWNpHWSxiWNT0xMzLhNSroqOGHVzAqjxFsPNH5AOiLGgDGA0dHRGc8yTklXBSesmllZSlrj6IuzlVLSVcEJq2ZWFl8EZ2ZmU/Tu+EGKbk9lvRS4Fni6pK2S3lDPsMzMrJe6mjlExNq6BmJmNshKS2X1+zmYmdk0WY853HTLd7MmrO7ImJQ6MpL2J8Hu3fliIhcvSvtboISUTitXyutrv1/5g9TO0trVpKT/Ej4gbWaWSUkHpF0czMxy6GEURgofczAzs2k8czAzy6CHSRhJXBzMzHIpqDr0xbKS01XNbBiUlMraFzOHyXRVwFlJZjawSjog3RfFwemqZjYMCqoN/VEcnK5qZkOhoOrQF8cczMysv/TFzMHMbBj4CmkzM9tDaamsLg5mZhncdNONX162RCsSm++odTDzoNTEzaTOpAngOzXvdgX5vnE5+8rdn7+2Mvsb5K8td39z9fWUiFiZaRx9IWtxaIKk8YgYHbS+cvfnr63M/gb5a8vdX+6vrd/5bCUzM5vGxcHMzKYZhOIwNqB95e7PX1uZ/Q3y15a7v9xfW18r/piDmZnVbxBmDmZmVrNii4OkCyVtl7Q5Q19LJX1D0jclbZH0zob7u0vStyRtkjTecF9Pr/qZvP1A0ukN9neapM3V97H2fmZ6XUh6VdXfbkm1no0yS3/nSLq5+n5eIenABvs6W9K2tp/fCXX0NUd/n2zr6y5Jmxrs69mSrq3+L3xO0hNq6muVpKsk3VK9Lk6rHm/sdVKkiCjyBrwQeA6wOUNfAvap7i8Brgee12B/dwErevA9XQR8j9Y53U3s/zBgM7Cc1gWYXwF+qenXBfAM4OnA1cBohv6e0Hb/D4CPNNjX2cBbGvp5zfl/DDgPOLPBr+0G4EXV/d8CzqmprwOA51T39wW+DTyzyddJibdiZw4RsQHYmamviIgfVZ8uqW6DeLDmxcC/R0TdFypOegZwfUQ8GBGPAV8FTqqzg5leFxFxa0TcXmc/Hfr7Qdune1PTayXna75Tf5IEvBq4tMG+ngZsqO5fCfxGTX3dGxE3Vfd/CNwKHNTk66RExRaH3CQtqqbQ24ErI+L6BrsL4ApJN0pa12A/U51MTf/ZZ7EZOEbS/pKWAycAqxrsr2ckvVvS3cBrgTMb7u7UahnrQklParivSccA90XEvzXYxxbgxOr+q2jgtSLpEOBIWqsB1sbFYZ4i4vGIOAI4GDhK0mENdveCiHgOcDzwRkkvbLAvACTtBbwS+HRTfUTErcB7gCuALwGbgMeb6q+XIuLtEbEKuAQ4tcGuPgz8Z+AI4F5aSz05rKXZPySgtZT0fyXdSGv555E6dy5pH+CzwOlTZnuGi8OCRcQDwFXAmgb72FZ93A5cBhzVVF9tjgduioj7muwkIi6IiOdGxAuBXbTWewfZJdS0HDKTiLiv+sNlN/C3ZHitSFpMaznwk032ExG3RcTLIuK5tArRv9e1b0lLaBWGSyLiH+va7yBxcZgHSSslPbG6vwx4KXBbQ33tLWnfyfvAy2gtxzQtx1+CSHpy9XE1rV8wH2+6z9wkHdr26Yk09Fqp+jqg7dNfJ89r5SXAbRGxtclO2l4rI8A7gI/UtF8BFwC3RsR769jnQOr1EfHUG61fZPcCjwJbgTc02NezgI3AzbT+89VyhsYsff0i8M3qtgV4e4bv5d7A/cDPZ+jrX4Fbqq/vxTleF7R+aW4FfgLcB3y54f4+W71ObgY+R+tgZ1N9fQz4VtXX5cABTX5t1eMXA7+b4ed2Gq2Z5beBP6e6aLeGvl5A67jezbSWNjfROv7V2OukxJuvkDYzs2m8rGRmZtO4OJiZ2TQuDmZmNo2Lg5mZTePiYGZm07g4mJnZNC4OZmY2jYuDmZlN8/8BWpW/yFoTBqAAAAAASUVORK5CYII=\n",
      "text/plain": [
       "<Figure size 432x432 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "#contact predicts from contact-based model\n",
    "new_predictions=mymodel(inputs['input_ids'], lbd_lora_attention)[0].cpu()\n",
    "fig, axes = plt.subplots(figsize=(6, 6))\n",
    "plot_contacts_and_predictions(\n",
    "        new_predictions, contacts, ax=axes, title = 'LBDB_prediction'\n",
    "    )\n",
    "# plt.savefig('./output/lbdb_contact_map_finetune.png',dpi=800)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "ce8a00a1-b9a4-4acf-9f07-d43d05ed6d02",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'id': 'lbdb',\n",
       " 'model': 'ESM-2 (Unsupervised)',\n",
       " 'local_AUC': 0.9510554075241089,\n",
       " 'local_P@L': 0.7272727489471436,\n",
       " 'local_P@L2': 1.0,\n",
       " 'local_P@L5': 1.0,\n",
       " 'short_AUC': 0.8305492401123047,\n",
       " 'short_P@L': 0.5,\n",
       " 'short_P@L2': 1.0,\n",
       " 'short_P@L5': 1.0,\n",
       " 'medium_AUC': 0.2866150736808777,\n",
       " 'medium_P@L': 0.09090909361839294,\n",
       " 'medium_P@L2': 0.1818181872367859,\n",
       " 'medium_P@L5': 0.5}"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "metrics = {\"id\": 'lbdb', \"model\": \"ESM-2 (Unsupervised)\"}\n",
    "metrics.update(evaluate_prediction(new_predictions, contacts))\n",
    "metrics"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "AMP (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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}