example_notebooks/example_contact.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "499b1524-a7f3-4bc5-bff7-96aaf762236e",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/state/partition1/slurm_tmp/24926014.0.0/ipykernel_2608326/2538044701.py:20: DeprecationWarning: Please use `pearsonr` from the `scipy.stats` namespace, the `scipy.stats.stats` namespace is deprecated.\n",
      "  from scipy.stats.stats import pearsonr\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "os.chdir('..')\n",
    "import src.models_and_optimizers as model_utils\n",
    "import yaml\n",
    "from types import SimpleNamespace\n",
    "from clip_main import get_wds_loaders\n",
    "from transformers import EsmTokenizer\n",
    "import src.data_utils as data_utils\n",
    "from sklearn import preprocessing\n",
    "from sklearn.linear_model import LogisticRegression\n",
    "import torch.nn as nn\n",
    "import torch\n",
    "import sys\n",
    "import pickle\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import json\n",
    "from torch.cuda.amp import autocast\n",
    "from scipy.stats.stats import pearsonr \n",
    "import matplotlib.pyplot as plt\n",
    "import webdataset as wds\n",
    "import copy \n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e751b13a-3228-4b87-9f47-83f92fdd8dfc",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Load model parameters\n",
    "\n",
    "model_dir = \"/c/example/path\"\n",
    "path = os.path.join(model_dir, \"checkpoints/checkpoint_latest.pt\")\n",
    "args_path = os.path.join(model_dir, \n",
    "                         [u for u in os.listdir(model_dir) if u.endswith('.pt')][0])\n",
    "\n",
    "hparams = torch.load(args_path)\n",
    "hparams['args']['blacklist_file'] = \"pdb_blacklist.txt\"\n",
    "hparams['args']['coordinator_hparams'] = \"terminator_configs/standard.json\"\n",
    "args_dict = hparams['args']\n",
    "args_dict['batch_size'] = 1\n",
    "args_dict['distributed'] = 0\n",
    "args_dict['train_wds_path'] = 'wds_rocklin.tar'\n",
    "args_dict['val_wds_path'] = 'wds_rocklin.tar'\n",
    "args_dict['data_root'] = \"example_data/wds\"\n",
    "args_dict['train_wds_path'] = 'wds_contact_test.tar'\n",
    "args_dict['val_wds_path'] = 'wds_contact_test.tar'\n",
    "args_dict['arch'] = '../esm_model_150'\n",
    "args_dict['num_workers'] = 1\n",
    "\n",
    "args = SimpleNamespace(**args_dict)\n",
    "coordinator_params = data_utils.get_coordinator_params(args.coordinator_hparams)\n",
    "coordinator_params['num_positional_embeddings'] = args.gnn_num_pos_embs\n",
    "coordinator_params['zero_out_pos_embs']= args.gnn_zero_out_pos_embs\n",
    "coordinator_params['clip_mode'] = True\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7efb1d26-93e6-47f3-9402-db110ecb4ed5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading state dict from /data1/groups/keating_madry/runs/new_blacklist/version_0/checkpoints/checkpoint_latest.pt\n",
      "building model based on path\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at /data1/groups/keating_madry/huggingface/esm2_t30_150M_UR50D were not used when initializing EsmModel: ['lm_head.dense.weight', 'lm_head.dense.bias', 'lm_head.bias', 'lm_head.layer_norm.weight', 'lm_head.decoder.weight', 'lm_head.layer_norm.bias']\n",
      "- This IS expected if you are initializing EsmModel 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 EsmModel 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 EsmModel were not initialized from the model checkpoint at /data1/groups/keating_madry/huggingface/esm2_t30_150M_UR50D and are newly initialized: ['esm.pooler.dense.bias', 'esm.pooler.dense.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ZERO OUT POS EMB False\n",
      "GNN Potts Model Encoder hidden dimensionality is 128\n",
      "freeze_llm False\n",
      "{'training_args': {'epochs': 10, 'lr': 0.001, 'weight_decay': 0.001, 'momentum': 0.9, 'label_smoothing': 0.0, 'lr_peak_epoch': 2, 'eval_epochs': 2, 'only_non_bn_weight_decay': False, 'opt': 'ADAM', 'lr_scheduler': 'cosine', 'burn_in': 0, 'mixed_precision': True, 'max_len': 1024, 'self_supervised': True}, 'epoch': 9, 'training_metrics': {'loss': 0.6333152562079322, 'acc': 0.6356185298161776}, 'val_metrics': {'loss': 1.6227540006807477, 'acc': 0.5308778800395427}, 'model_building_args': {'esm_arch': 'facebook/esm2_t30_150M_UR50D', 'terminator_hparams': {'cov_features': 'all_raw', 'cov_compress': 'project', 'term_use_mpnn': True, 'matches': 'transformer', 'energies_style': 'mpnn', 'energies_use_mpnn': True, 'energies_full_graph': True, 'contact_idx': True, 'energies_encoder_layers': 3, 'energies_hidden_dim': 128, 'resnet_linear': True, 'matches_linear': True, 'transformer_linear': True, 'term_mpnn_linear': True, 'use_terms': False, 'res_embed_linear': True, 'model': 'multichain', 'term_hidden_dim': 32, 'gradient_checkpointing': True, 'num_pair_stats': 28, 'num_sing_stats': 0, 'resnet_blocks': 4, 'term_layers': 4, 'term_heads': 4, 'conv_filter': 3, 'matches_layers': 4, 'matches_num_heads': 4, 'k_neighbors': 30, 'k_cutoff': None, 'cie_dropout': 0.1, 'cie_scaling': 500, 'cie_offset': 0, 'transformer_dropout': 0.1, 'energies_protein_features': 'full', 'energies_augment_eps': 0, 'energies_dropout': 0.1, 'energies_output_dim': 400, 'energies_gvp': False, 'struct2seq_linear': False, 'use_coords': True, 'num_features': 9, 'condense_options': 'residue_local_reduce_edges_expand_transformer_reduce_log_probs', 'chain_handle': 'replace', 'merge_dups': 'broken', 'num_positional_embeddings': 16, 'zero_out_pos_embs': False, 'energies_type': 'MPNN', 'energies_geometric': False, 'graphformer_config': {'num_layers': 5, 'num_heads': 4, 'embed_per_head': 64, 'mlp_multiplier': 4}, 'clip_mode': True, 'energies_input_dim': 0}, 'self_supervised': True, 'gnn_checkpoint': None, 'freeze_llm': False, 'freeze_text_proj': False, 'language_head': False, 'language_head_type': 'MLP'}}\n"
     ]
    }
   ],
   "source": [
    "## Load model\n",
    "\n",
    "trained_model = model_utils.load_model(path, args_dict['arch'], 'cuda')\n",
    "tokenizer = EsmTokenizer.from_pretrained(args_dict['arch'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f63e40c3-a5cd-4860-9cc1-460a4ccee71d",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Functions to extract attention \"contact\" values from ESM-RLA\n",
    "\n",
    "import torch\n",
    "from torch.cuda.amp import GradScaler, autocast\n",
    "import torch\n",
    "import numpy as np\n",
    "from torch.nn.utils.rnn import pad_sequence\n",
    "\n",
    "def get_formatted_inp(batch):\n",
    "    seq_batch, coords_batch = batch\n",
    "    seqs = seq_batch['string_sequence']\n",
    "    text_inp = tokenizer(seqs, return_tensors='pt', padding=True, \n",
    "                              truncation=True, max_length=1024+2)\n",
    "    text_inp['position_ids'] = seq_batch['pos_embs'][0]\n",
    "    text_inp = {k: v.to('cuda') for k, v in text_inp.items()}\n",
    "    return text_inp, seq_batch, coords_batch\n",
    "\n",
    "def postprocess_attention_features(attention_features, inp_dict, tokenizer, placeholder_mask):\n",
    "    # remove class token, eos token\n",
    "    attention_features, text_mask = _adjust_text_features(attention_features, inp_dict, tokenizer)\n",
    "    # remove placeholders\n",
    "    attention_features, text_mask = _remove_placeholders(attention_features, text_mask, placeholder_mask)\n",
    "    return attention_features, text_mask\n",
    "\n",
    "\n",
    "def _adjust_text_features(attention_features, inp_dict, tokenizer):\n",
    "    mask = inp_dict['attention_mask'].clone()\n",
    "    toks = inp_dict['input_ids']\n",
    "    eos_token = tokenizer.convert_tokens_to_ids(tokenizer.eos_token)\n",
    "    mask[toks == eos_token] = 0\n",
    "    # ignore class token and last eos token\n",
    "    mask = mask[:, 1:-1]\n",
    "    attention_features = attention_features[:, :, 1:-1, 1:-1]\n",
    "    return attention_features, mask\n",
    "\n",
    "def _remove_placeholders(attention_features, text_mask, placeholder_mask):\n",
    "    B = placeholder_mask.shape[0]\n",
    "    filtered = []\n",
    "    new_masks = []\n",
    "    for b in range(B):\n",
    "        p_m = placeholder_mask[b]\n",
    "        new_text_feat = attention_features[b][:, p_m]\n",
    "        new_text_feat = new_text_feat[:, :, p_m]\n",
    "        filtered.append(new_text_feat)\n",
    "        new_masks.append(text_mask[b][p_m])\n",
    "    new_masks = pad_sequence(new_masks, batch_first=True)\n",
    "    dest_shape = 1024\n",
    "    padded_filtered = []\n",
    "    for f in filtered:\n",
    "        amt_to_pad = dest_shape - f.shape[-1]\n",
    "        padded = torch.nn.functional.pad(f, (0, amt_to_pad, 0, amt_to_pad))\n",
    "        padded_filtered.append(padded)\n",
    "    return torch.stack(padded_filtered), new_masks\n",
    "\n",
    "def get_attentions(model, text_inp_):\n",
    "    with torch.no_grad():\n",
    "        out = model.text_model(**text_inp_, output_attentions=True)\n",
    "        all_attns = torch.cat(out['attentions'], 1)\n",
    "        num_heads = len(out['attentions'])\n",
    "    attns = all_attns.cpu()\n",
    "    return attns, num_heads\n",
    "\n",
    "def APC_correction(attn_):\n",
    "    F_i = attn_.sum(keepdims=True, dim=2)\n",
    "    F_j = attn_.sum(keepdims=True, dim=3)\n",
    "    F = attn_.sum(dim=2, keepdims=True).sum(dim=3, keepdims=True)\n",
    "    apc_corr = (F_i * F_j)/F\n",
    "    attn_corr = attn_ - apc_corr\n",
    "    return attn_corr\n",
    "\n",
    "def pred_and_contact_gt(model, batch_):\n",
    "    text_inp, seq_batch, coord_batch = get_formatted_inp(batch_)\n",
    "    mask = seq_batch['seq_loss_mask'][0]\n",
    "    mask = torch.nn.functional.pad(mask, (0, 1024 - mask.shape[-1]))\n",
    "    B = mask.shape[0]\n",
    "    T = mask.shape[-1]\n",
    "    # get 2D mask\n",
    "    M = mask.int()\n",
    "    \n",
    "    M = M.unsqueeze(-1) @ M.unsqueeze(1)\n",
    "    local_mask = torch.ones(T, T)\n",
    "    for i in range(6):\n",
    "        torch.diagonal(local_mask, i).zero_()\n",
    "        torch.diagonal(local_mask, i*-1).zero_()\n",
    "    \n",
    "    # M[:, local_mask == 0] = 0 \n",
    "    M = M==1\n",
    "    # get index distance\n",
    "    U = torch.arange(T).float().unsqueeze(1)\n",
    "    U = torch.cdist(U, U).expand((B, T, T))\n",
    "    \n",
    "    # get indices\n",
    "    U1 = torch.arange(T).unsqueeze(1)\n",
    "    U1 = U1.expand((B, T, T))\n",
    "    \n",
    "    U2 = torch.arange(T).unsqueeze(0)\n",
    "    U2 = U2.expand((B, T, T))\n",
    "    \n",
    "    # get distances and contacts\n",
    "    coords = coord_batch['coords'][0][:, :, 1]\n",
    "    dists = torch.cdist(coords, coords)\n",
    "    contacts = dists < 8\n",
    "    amt_to_pad = 1024 - dists.shape[-1] \n",
    "    dists = torch.nn.functional.pad(dists, (0, amt_to_pad, 0, amt_to_pad))\n",
    "    contacts = torch.nn.functional.pad(contacts, (0, amt_to_pad, 0, amt_to_pad))\n",
    "    \n",
    "    \n",
    "    # get attn\n",
    "    attns, num_heads = get_attentions(model, text_inp)\n",
    "    attns, _ = postprocess_attention_features(attns, text_inp, tokenizer, seq_batch['placeholder_mask'][0])\n",
    "    attns.permute(0, 2, 3, 1)[~M] = 0\n",
    "    attns = APC_correction(attns)\n",
    "    attns = attns.permute(0, 2, 3, 1)\n",
    "    return {\n",
    "        'X_': attns,\n",
    "        'Y_': contacts,\n",
    "        'num_heads': num_heads,\n",
    "        'dists': dists,\n",
    "        'M': M,\n",
    "        'U': U,\n",
    "        'U1': U1,\n",
    "        'U2': U2,\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "78d13f66-adeb-4c72-8b88-ea65d8eb836f",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Pass data through ESM_RLA to be able to get attention \"contact\" scores\n",
    "def get_X_and_Y(models, loader, num_iters=5):\n",
    "    g_cpu = torch.Generator()\n",
    "    g_cpu.manual_seed(10)\n",
    "    OUTPUT_KEYS = ['X_', 'Y_', 'dists', 'U', 'U1', 'U2']\n",
    "    output = {\n",
    "        o: {k: [] for k in models.keys()} \n",
    "        for o in OUTPUT_KEYS\n",
    "    }\n",
    "    for batch_idx, batch_ in enumerate(tqdm(loader)): \n",
    "        train_indices = None\n",
    "        if torch.sum(batch_[0]['placeholder_mask'][0]) < 30:\n",
    "            continue\n",
    "        for k in models.keys():\n",
    "            model = models[k]\n",
    "            result = pred_and_contact_gt(model, batch_)\n",
    "            t2 = time.time()\n",
    "            result_subset = {o: result[o][result['M']] for o in OUTPUT_KEYS}\n",
    "            if train_indices is None:\n",
    "                pos_contacts = torch.where(result_subset['Y_'] == 1)[0]\n",
    "                neg_contacts = torch.where(result_subset['Y_'] == 0)[0]\n",
    "                neg_contacts = neg_contacts[torch.randperm(len(neg_contacts), generator=g_cpu)[:len(pos_contacts)]]\n",
    "                train_indices = torch.cat([pos_contacts, neg_contacts])\n",
    "            for o in OUTPUT_KEYS:\n",
    "                output[o][k].append(result_subset[o][train_indices])\n",
    "            t3 = time.time()\n",
    "        if batch_idx == num_iters:\n",
    "            break\n",
    "            \n",
    "            \n",
    "    output = {\n",
    "        o: {k: torch.cat(v) for k, v in output_sub.items()}\n",
    "        for o, output_sub in output.items()\n",
    "    }\n",
    "    \n",
    "    \n",
    "    \n",
    "    return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "44fa0174-9527-4f3f-bd7f-c67b0ee46897",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'max_coords_len': 2000, 'shuffle_coords': False, 'max_seq_len': 1024, 'pos_offset': 128, 'burn_in': 0, 'k_neighbors': 30, 'crop_type': 'absolute', 'shuffle_chains': False, 'num_mutations': -1, 'masked_rate': -1.0, 'masked_mode': 'MASK'}\n",
      "added select filtering... 30\n",
      "{'max_coords_len': 2000, 'shuffle_coords': False, 'max_seq_len': 1024, 'pos_offset': 128, 'burn_in': 0, 'k_neighbors': 30, 'crop_type': 'absolute', 'shuffle_chains': False, 'num_mutations': -1, 'masked_rate': -1.0, 'masked_mode': 'MASK'}\n",
      "added select filtering... 30\n",
      "0 0\n"
     ]
    }
   ],
   "source": [
    "## Load sample data\n",
    "_, val_loader, _, _ = get_wds_loaders(args, coordinator_params, gpu=None, shuffle_train=False, val_only=False, return_count=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "12d199c1-96f2-4a85-a311-da5985648f2e",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "20it [00:29,  1.50s/it]\n"
     ]
    }
   ],
   "source": [
    "models = {\n",
    "    'trained': trained_model,\n",
    "}\n",
    "val_output = get_X_and_Y(models, val_loader, num_iters=20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "b59bcffe-e6ad-4a75-bf22-aef8054900e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "val_X, val_Y, val_dists, val_U = val_output['X_'], val_output['Y_'], val_output['dists'], val_output['U']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "5242367e-1355-40a8-9362-1887ace6f59c",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Create top layer to predict contacts from attention\n",
    "def gelu(x):\n",
    "    \"\"\"\n",
    "    This is the gelu implementation from the original ESM repo. Using F.gelu yields subtly wrong results.\n",
    "    \"\"\"\n",
    "    return x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0)))\n",
    "\n",
    "class MultiLayerLinear(nn.Module):\n",
    "    def __init__(self, in_features, out_features, num_layers, activation_layers='relu', dropout=0):\n",
    "        super(MultiLayerLinear, self).__init__()\n",
    "        self.activation_layers = activation_layers\n",
    "        self.layers = nn.ModuleList()\n",
    "        for i in range(num_layers):\n",
    "            self.layers.append(nn.Linear(in_features[i], out_features[i]))\n",
    "        self.dropout_prob = dropout\n",
    "        if dropout > 0:\n",
    "            self.dropout = nn.Dropout(dropout)\n",
    "            \n",
    "    def forward(self, x):\n",
    "        for layer in self.layers:  \n",
    "            x = layer(x)\n",
    "            if self.activation_layers == 'relu':\n",
    "                x = nn.functional.relu(x)\n",
    "            else:\n",
    "                x = gelu(x)\n",
    "        if self.dropout_prob > 0:\n",
    "            x = self.dropout(x)\n",
    "        return x\n",
    "    \n",
    "    \n",
    "def train_finetune_model(X_, Y_, U_, D_, crit='BCE'):\n",
    "    inp_dim = 600\n",
    "\n",
    "    finetune_model = MultiLayerLinear(in_features=[600,100], out_features=[100,1], num_layers=2, activation_layers='gelu', dropout=0).cuda()\n",
    "\n",
    "    # finetune_model = nn.Linear(inp_dim, 1).cuda()\n",
    "    \n",
    "    if crit == 'BCE':\n",
    "        criterion = nn.BCEWithLogitsLoss()\n",
    "    else:\n",
    "        criterion = nn.MSELoss()\n",
    "    \n",
    "    optimizer = FullBatchLBFGS(finetune_model.parameters(), lr=0.01)\n",
    "    print('loaded data')\n",
    "    optimizer.zero_grad()\n",
    "    out = finetune_model(X_).squeeze(1)\n",
    "    loss = criterion(out, Y_)\n",
    "    loss.backward()\n",
    "\n",
    "    epochs=300\n",
    "    pbar = tqdm(range(epochs))\n",
    "    for epoch in pbar:\n",
    "        def closure():\n",
    "            optimizer.zero_grad()\n",
    "            out = finetune_model(X_).squeeze(1)\n",
    "            if crit=='BCE':\n",
    "                loss = criterion(out, Y_)\n",
    "            else:\n",
    "                loss = criterion(out, D_)\n",
    "            return loss\n",
    "        options = {'closure': closure, 'current_loss': loss}\n",
    "        loss = optimizer.step(options)[0]\n",
    "        pbar.set_postfix({'loss': loss.item()})\n",
    "    return finetune_model\n",
    "\n",
    "def metric_fn(pred, gt, crit='BCE'):\n",
    "    frac_neg = (gt[0] == 0).float().mean()\n",
    "    frac_pos = (gt[0] == 1).float().mean()\n",
    "    if crit == 'BCE':\n",
    "        neg = (pred[gt[0] == 0] == 0).float().mean().item()\n",
    "        pos = (pred[gt[0] == 1] == 1).float().mean().item()\n",
    "        total = (gt[0] == pred).float().mean().item()\n",
    "    else:\n",
    "        pred = pred.squeeze(1)\n",
    "        neg = ((pred[gt[0] == 0] - gt[1][gt[0] == 0])**2).float().mean().item()\n",
    "        pos = ((pred[gt[0] == 1] - gt[1][gt[0] == 1])**2).float().mean().item()\n",
    "        total = ((pred - gt[1])**2).float().mean().item()\n",
    "    print(f\"All -> {total:0.4f} \\t Neg ({frac_neg*100:0.2f}%) -> {neg:0.4f} \\t Pos({frac_pos*100:0.2f}%) -> {pos:0.4f}\")\n",
    "    \n",
    "def get_U_mask(U, min_, max_):\n",
    "    return (U > min_) & (U <= max_)\n",
    "\n",
    "def evaluate(finetune_model, X_, Y_, U, D, crit='BCE'):\n",
    "    with torch.no_grad():\n",
    "        out = finetune_model(X_)\n",
    "        if crit == 'BCE':\n",
    "            preds = (nn.Sigmoid()(out.squeeze(1)) > 0.5).float()\n",
    "        else:\n",
    "            preds = out.float()\n",
    "        print(\"overall\")\n",
    "        metric_fn(pred=preds, gt=(Y_, D), crit=crit)\n",
    "        \n",
    "        mask = get_U_mask(U, 6, 12)\n",
    "        print(\"short\", mask.sum().item())\n",
    "        metric_fn(pred=preds[mask], gt=(Y_[mask], D[mask]), crit=crit)\n",
    "        \n",
    "        mask = get_U_mask(U, 12, 24)\n",
    "        print(\"medium\", mask.sum().item())\n",
    "        metric_fn(pred=preds[mask], gt=(Y_[mask], D[mask]), crit=crit)\n",
    "        \n",
    "        mask = get_U_mask(U, 24, np.inf)\n",
    "        print(\"long\", mask.sum().item())\n",
    "        metric_fn(pred=preds[mask], gt=(Y_[mask], D[mask]), crit=crit)\n",
    "    return preds.cpu()\n",
    "\n",
    "def evaluate_ft(finetune_model, scaler, MODEL_TO_ANALYZE):\n",
    "    val_X_, val_Y_, val_D_ = val_X[MODEL_TO_ANALYZE], val_Y[MODEL_TO_ANALYZE], val_dists[MODEL_TO_ANALYZE]\n",
    "    val_X_norm = torch.tensor(scaler.transform(val_X_.numpy()))\n",
    "    val_X_norm, val_Y_norm = val_X_norm.cuda(), val_Y_.cuda().float()\n",
    "    U = val_output['U'][MODEL_TO_ANALYZE]\n",
    "    val_preds = evaluate(finetune_model, val_X_norm, val_Y_norm, U, val_D_)\n",
    "    return val_preds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "d72b8481-4c4a-4af0-a7c7-e7d29d584e7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from joblib import load"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "edc3618e-c924-4cfa-9aaf-421a8c95d881",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/gridsan/fbirnbaum/.local/lib/python3.9/site-packages/sklearn/base.py:288: UserWarning: Trying to unpickle estimator StandardScaler from version 1.2.2 when using version 1.2.0. This might lead to breaking code or invalid results. Use at your own risk. For more info please refer to:\n",
      "https://scikit-learn.org/stable/model_persistence.html#security-maintainability-limitations\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "# evaluate\n",
    "contact_model_ = nn.Linear(600, 1).cuda()\n",
    "contact_model_.load_state_dict(torch.load(os.path.join(os.getcwd(), \"contact_head/finetune.pt\")))\n",
    "contact_scaler = load(os.path.join(os.getcwd(), \"contact_head/finetune_scaler.bin\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "e938fe87-069c-40e9-a71c-3daa0fd46aa8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "overall\n",
      "All -> 0.7318 \t Neg (50.00%) -> 0.9314 \t Pos(50.00%) -> 0.5323\n",
      "short 7388\n",
      "All -> 0.8499 \t Neg (36.55%) -> 0.6352 \t Pos(63.45%) -> 0.9735\n",
      "medium 14651\n",
      "All -> 0.9183 \t Neg (34.12%) -> 0.8062 \t Pos(65.88%) -> 0.9764\n",
      "long 124991\n",
      "All -> 0.9142 \t Neg (80.23%) -> 0.9505 \t Pos(19.77%) -> 0.7668\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([0., 0., 0.,  ..., 1., 1., 0.])"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "evaluate_ft(contact_model_, contact_scaler, 'trained')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:.conda-rla-public]",
   "language": "python",
   "name": "conda-env-.conda-rla-public-py"
  },
  "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.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}