proteinbind_model.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82940d47-8ec8-498e-b5ce-9a0bca790cf3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import pandas as pd\n",
    "import torch\n",
    "from tqdm import tqdm\n",
    "from torch.utils.data import Subset\n",
    "from torch.utils.data import DataLoader\n",
    "from pytorch_metric_learning.losses import SelfSupervisedLoss, NTXentLoss\n",
    "\n",
    "from proteinbind_new import EmbeddingDataset, DualEmbeddingDataset\n",
    "from proteinbind_new import create_proteinbind\n",
    "from transformers import pipeline\n",
    "import torch.nn.functional as F\n",
    "from sklearn.linear_model import LogisticRegression\n",
    "device = torch.device(\"cpu\")\n",
    "if torch.cuda.is_available():\n",
    "    device = torch.device(\"cuda\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76cca8b4-8814-465a-907d-6f291325b79d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "batch_size = 1024\n",
    "test_num =10\n",
    "\n",
    "# GO\n",
    "go_dataset = DualEmbeddingDataset('GO_files/go_AA_embeddings.pt', 'GO_files/GO_embeddings.pt', \"go\")\n",
    "\n",
    "indices = list(range(len(go_dataset)))\n",
    "train_indices = indices[:-test_num]\n",
    "test_indices = indices[-test_num:]\n",
    "# test_indices = indices[:test_num]\n",
    "\n",
    "go_train_dataloader = DataLoader(Subset(go_dataset, train_indices), batch_size=batch_size)\n",
    "go_test_dataloader = DataLoader(Subset(go_dataset, test_indices), batch_size=batch_size)\n",
    "\n",
    "# DNA\n",
    "\n",
    "dna_dataset = DualEmbeddingDataset('DNA_files/dna_AA_embeddings.pt', 'DNA_files/uniprot_nuc_embeddings.pt', \"dna\")\n",
    "\n",
    "indices = list(range(len(dna_dataset)))\n",
    "train_indices = indices[:-test_num]\n",
    "test_indices = indices[-test_num:]\n",
    "# test_indices = indices[:test_num]\n",
    "\n",
    "dna_train_dataloader = DataLoader(Subset(dna_dataset, train_indices), batch_size=batch_size)\n",
    "dna_test_dataloader = DataLoader(Subset(dna_dataset, test_indices), batch_size=batch_size)\n",
    "\n",
    "\n",
    "# Text\n",
    "text_dataset = DualEmbeddingDataset('Text_files/text_AA_embeddings.pt', 'Text_files/protein_whole_text_embedding.pt', \"text\")\n",
    "\n",
    "indices = list(range(len(text_dataset)))\n",
    "train_indices = indices[:-test_num]\n",
    "test_indices = indices[-test_num:]\n",
    "# test_indices = indices[:test_num]\n",
    "\n",
    "text_train_dataloader = DataLoader(Subset(text_dataset, train_indices), batch_size=batch_size)\n",
    "text_test_dataloader = DataLoader(Subset(text_dataset, test_indices), batch_size=batch_size)\n",
    "\n",
    "\n",
    "# MSA\n",
    "\n",
    "msa_dataset = DualEmbeddingDataset('MSA_files/msa_AA-embeddings.pt', 'MSA_files/msa_embeddings.pt', \"msa\")\n",
    "\n",
    "indices = list(range(len(msa_dataset)))\n",
    "train_indices = indices[:-test_num]\n",
    "test_indices = indices[-test_num:]\n",
    "\n",
    "msa_train_dataloader = DataLoader(Subset(msa_dataset, train_indices), batch_size=batch_size)\n",
    "msa_test_dataloader = DataLoader(Subset(msa_dataset, test_indices), batch_size=batch_size)\n",
    "\n",
    "# PDB\n",
    "\n",
    "pdb_dataset = DualEmbeddingDataset('Struct_files/pdb_AAs-embedding.pt', 'Struct_files/pdb_embeddings.pt', \"pdb\")\n",
    "\n",
    "indices = list(range(len(pdb_dataset)))\n",
    "train_indices = indices[:-test_num]\n",
    "test_indices = indices[-test_num:]\n",
    "\n",
    "pdb_train_dataloader = DataLoader(Subset(pdb_dataset, train_indices), batch_size=batch_size)\n",
    "pdb_test_dataloader = DataLoader(Subset(pdb_dataset, test_indices), batch_size=batch_size)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c12f9c0-4c90-4fb7-9842-2dd39c10f065",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = create_proteinbind()\n",
    "loss_func = SelfSupervisedLoss(NTXentLoss())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85d598a4-269c-48b2-9737-64e9853a52f5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "num_epochs = 1000\n",
    "model.to(device)\n",
    "model.train()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.0001)\n",
    "losses = []\n",
    "\n",
    "# Get the total number of mini-batches per epoch based on the smallest data loader\n",
    "num_batches = min(len(go_train_dataloader), len(text_train_dataloader), len(dna_train_dataloader), len(msa_train_dataloader), len(pdb_train_dataloader))\n",
    "\n",
    "for epoch in tqdm(range(num_epochs)):\n",
    "    \n",
    "    go_loader_iter = iter(go_train_dataloader)\n",
    "    text_loader_iter = iter(text_train_dataloader)\n",
    "    dna_loader_iter = iter(dna_train_dataloader)\n",
    "    msa_loader_iter = iter(msa_train_dataloader)\n",
    "    pdb_loader_iter = iter(pdb_train_dataloader)\n",
    "    \n",
    "    for batch_idx in range(num_batches):\n",
    "\n",
    "        ####################\n",
    "        ### GO\n",
    "        ####################\n",
    "        go_batch = next(go_loader_iter)\n",
    "\n",
    "        \n",
    "        input_data = {\n",
    "            'go': go_batch['go'].type(torch.float32).to(device),\n",
    "            'aa': go_batch['aa'].type(torch.float32).to(device)\n",
    "        }\n",
    "\n",
    "        output = model(input_data)\n",
    "        loss = loss_func(output[\"aa\"], output[\"go\"]) + loss_func(output[\"go\"], output[\"aa\"]) \n",
    "        losses.append(loss.item())\n",
    "\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        ####################\n",
    "        ### DNA\n",
    "        ####################\n",
    "        dna_batch = next(dna_loader_iter)\n",
    "        \n",
    "        input_data = {\n",
    "            'dna': dna_batch['dna'].type(torch.float32).to(device),\n",
    "            'aa': dna_batch['aa'].type(torch.float32).to(device)\n",
    "        }\n",
    "\n",
    "        output = model(input_data)\n",
    "        loss = loss_func(output[\"aa\"], output[\"dna\"]) + loss_func(output[\"dna\"], output[\"aa\"]) \n",
    "        losses.append(loss.item())\n",
    "\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        ####################\n",
    "        ### TEXT\n",
    "        ####################\n",
    "        \n",
    "        text_batch = next(text_loader_iter)\n",
    "        \n",
    "        input_data = {\n",
    "            'text': text_batch['text'].type(torch.float32).to(device),\n",
    "            'aa': text_batch['aa'].type(torch.float32).to(device)\n",
    "        }\n",
    "\n",
    "        output = model(input_data)\n",
    "        loss = loss_func(output[\"aa\"], output[\"text\"]) + loss_func(output[\"text\"], output[\"aa\"]) \n",
    "        losses.append(loss.item())\n",
    "\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        \n",
    "        ####################\n",
    "        ### MSA\n",
    "        ####################\n",
    "        \n",
    "        msa_batch = next(msa_loader_iter)\n",
    "        \n",
    "        input_data = {\n",
    "            'msa': msa_batch['msa'].type(torch.float32).to(device),\n",
    "            'aa': msa_batch['aa'].type(torch.float32).to(device)\n",
    "        }\n",
    "\n",
    "        output = model(input_data)\n",
    "        loss = loss_func(output[\"aa\"], output[\"msa\"]) + loss_func(output[\"msa\"], output[\"aa\"]) \n",
    "        losses.append(loss.item())\n",
    "\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        ####################\n",
    "        ### PDB\n",
    "        ####################\n",
    "        \n",
    "        pdb_batch = next(pdb_loader_iter)\n",
    "        \n",
    "        input_data = {\n",
    "            'pdb': pdb_batch['pdb'].type(torch.float32).to(device),\n",
    "            'aa': pdb_batch['aa'].type(torch.float32).to(device)\n",
    "        }\n",
    "\n",
    "        output = model(input_data)\n",
    "        loss = loss_func(output[\"aa\"], output[\"pdb\"]) + loss_func(output[\"pdb\"], output[\"aa\"]) \n",
    "        losses.append(loss.item())\n",
    "\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "\n",
    "    if epoch != 0 and (epoch + 1) % 1 == 0:\n",
    "        print(f\"Epoch {epoch+1}: joint train loss {sum(losses) / num_batches}\")\n",
    "        losses = []\n",
    "\n",
    "        model.eval()\n",
    "        \n",
    "        go_test_loader_iter = iter(go_test_dataloader)\n",
    "        text_test_loader_iter = iter(text_test_dataloader)\n",
    "        dna_test_loader_iter = iter(dna_test_dataloader)\n",
    "        msa_test_loader_iter = iter(msa_test_dataloader)\n",
    "        pdb_test_loader_iter = iter(pdb_test_dataloader)\n",
    "        \n",
    "\n",
    "\n",
    "        ##############\n",
    "        ##############\n",
    "        ##############\n",
    "        go_batch = next(go_test_loader_iter)\n",
    "        \n",
    "        input_data = {\n",
    "            'go': go_batch['go'].type(torch.float32).to(device),\n",
    "            'aa': go_batch['aa'].type(torch.float32).to(device)\n",
    "        }\n",
    "\n",
    "        output = model(input_data)\n",
    "        loss = loss_func(output[\"aa\"], output[\"go\"]) + loss_func(output[\"go\"], output[\"aa\"]) \n",
    "        losses.append(loss.item())\n",
    "\n",
    "        #assigned_val = torch.matmul(output[\"aa\"], output[\"go\"].T)\n",
    "        #print(F.softmax(assigned_val, dim = 1)\n",
    "        ##############\n",
    "        ##############\n",
    "        ##############\n",
    "        dna_batch = next(dna_test_loader_iter)\n",
    "        \n",
    "        input_data = {\n",
    "            'dna': dna_batch['dna'].type(torch.float32).to(device),\n",
    "            'aa': go_batch['aa'].type(torch.float32).to(device)\n",
    "        }\n",
    "\n",
    "\n",
    "        \n",
    "        output = model(input_data)\n",
    "        #assigned_val = torch.matmul(output[\"aa\"], output[\"dna\"].T)\n",
    "        #print(assigned_val)\n",
    "        loss = loss_func(output[\"aa\"], output[\"dna\"]) + loss_func(output[\"dna\"], output[\"aa\"]) \n",
    "        losses.append(loss.item())\n",
    "        \n",
    "        ##############\n",
    "        ##############\n",
    "        ##############\n",
    "        \n",
    "        text_batch = next(text_test_loader_iter)\n",
    "        \n",
    "        input_data = {\n",
    "            'text': text_batch['text'].type(torch.float32).to(device),\n",
    "            'aa': text_batch['aa'].type(torch.float32).to(device)\n",
    "        }\n",
    "\n",
    "        output = model(input_data)\n",
    "        #assigned_val = torch.matmul(output[\"aa\"], output[\"text\"].T)\n",
    "        #print(assigned_val)\n",
    "        loss = loss_func(output[\"aa\"], output[\"text\"]) + loss_func(output[\"text\"], output[\"aa\"]) \n",
    "        losses.append(loss.item())\n",
    "\n",
    "\n",
    "        \n",
    "        ##############\n",
    "        ##############\n",
    "        ##############\n",
    "        msa_batch = next(msa_test_loader_iter)\n",
    "\n",
    "        input_data = {\n",
    "            'msa': msa_batch['msa'].type(torch.float32).to(device),\n",
    "            'aa':msa_batch['aa'].type(torch.float32).to(device)\n",
    "        }\n",
    "\n",
    "        output = model(input_data)\n",
    "        loss = loss_func(output[\"aa\"], output[\"msa\"]) + loss_func(output[\"msa\"], output[\"aa\"]) \n",
    "        losses.append(loss.item())\n",
    "\n",
    "        ##############\n",
    "        ##############\n",
    "        ##############\n",
    "        pdb_batch = next(pdb_test_loader_iter)\n",
    "        \n",
    "        input_data = {\n",
    "            'pdb': pdb_batch['pdb'].type(torch.float32).to(device),\n",
    "            'aa': pdb_batch['aa'].type(torch.float32).to(device)\n",
    "        }\n",
    "    \n",
    "        output = model(input_data)\n",
    "        loss = loss_func(output[\"aa\"], output[\"pdb\"]) + loss_func(output[\"pdb\"], output[\"aa\"]) \n",
    "        losses.append(loss.item())\n",
    "        \n",
    "        print(f\"Epoch {epoch+1}: joint test loss {sum(losses)}\")\n",
    "        losses = []\n",
    "\n",
    "        #print(f\"Epoch {epoch+1}: Linear Cls {eval_linear_task(model)} \")\n",
    "        #print(f\"Epoch {epoch+1}: DMS Correlation Cls {dms_eval_task(model)} \")\n",
    "        \n",
    "        \n",
    "        model.train()\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "efd0f78c-3105-49e4-91ff-faf886a8b043",
   "metadata": {},
   "outputs": [],
   "source": [
    "DOWNSTREAM_FILE_PATHS = 'downstream_tasks/downstream_files/'\n",
    "\n",
    "def eval_linear_task(model):\n",
    "\n",
    "    model.eval()\n",
    "    ##### Pathogenicity Task\n",
    "    \n",
    "    ################     #####################\n",
    "    ################ DATA#####################\n",
    "    ################     #####################\n",
    "    \n",
    "    # opening the file in read mode\n",
    "    my_file = open(f\"{DOWNSTREAM_FILE_PATHS}/train_mt/seq-cleaned-train_mt_labels.txt\", \"r\")\n",
    "    \n",
    "    # reading the file\n",
    "    data = my_file.read()\n",
    "\n",
    "    # replacing end splitting the text \n",
    "    # when newline ('\\n') is seen.\n",
    "    data_into_list = data.split(\"\\n\")\n",
    "    print(len(data_into_list))\n",
    "    train_labels = []\n",
    "    for entry in data_into_list:\n",
    "        if entry == 'PLP':\n",
    "            train_labels.append(1)\n",
    "        elif entry == 'BLB':\n",
    "            train_labels.append(0)\n",
    "    \n",
    "    my_file.close()\n",
    "    \n",
    "    \n",
    "    # opening the file in read mode\n",
    "    my_file = open(f\"{DOWNSTREAM_FILE_PATHS}/test_mt/seq-cleaned-test_mt-labels.txt\", \"r\")\n",
    "      \n",
    "    # reading the file\n",
    "    data = my_file.read()\n",
    "      \n",
    "    # replacing end splitting the text \n",
    "    # when newline ('\\n') is seen.\n",
    "    data_into_list = data.split(\"\\n\")\n",
    "    print(len(data_into_list))\n",
    "    test_labels = []\n",
    "    for entry in data_into_list:\n",
    "        if entry == 'PLP':\n",
    "            test_labels.append(1)\n",
    "        elif entry == 'BLB':\n",
    "            test_labels.append(0)\n",
    "    \n",
    "    my_file.close()\n",
    "    \n",
    "    ################  #####################  #####################\n",
    "    ################ Pretrained Embedded AAs #####################\n",
    "    ################ #####################   #####################\n",
    "    train_wt_aa =  torch.load(f'{DOWNSTREAM_FILE_PATHS}/train_wt/seq-cleaned-train_wt-embeddings.pt')\n",
    "    train_mt_aa =  torch.load(f'{DOWNSTREAM_FILE_PATHS}/train_mt/seq-cleaned-train_mt_embeddings.pt')\n",
    "    \n",
    "    test_wt_aa =  torch.load(f'{DOWNSTREAM_FILE_PATHS}/test_wt/seq-cleaned-test_wt-embeddings.pt')\n",
    "    test_mt_aa =  torch.load(f'{DOWNSTREAM_FILE_PATHS}/test_mt/seq-cleaned-test_mt-embeddings.pt')\n",
    "    \n",
    "    print(train_wt_aa.size(), train_mt_aa.size(), test_wt_aa.size(),test_mt_aa.size() )\n",
    "    output_train_wt = model({\n",
    "                'aa': train_wt_aa.type(torch.float32).to(device)\n",
    "            })\n",
    "    \n",
    "    output_train_mt = model({\n",
    "                'aa': train_mt_aa.type(torch.float32).to(device)\n",
    "            })\n",
    "    \n",
    "    output_test_wt = model({\n",
    "                'aa': test_wt_aa.type(torch.float32).to(device)\n",
    "            })\n",
    "    \n",
    "    output_test_mt = model({\n",
    "                'aa': test_mt_aa.type(torch.float32).to(device)\n",
    "            })\n",
    "    \n",
    "    train_concat = torch.cat((output_train_wt['aa'], output_train_mt['aa']), 1)\n",
    "    test_concat = torch.cat((output_test_wt['aa'], output_test_mt['aa']), 1)\n",
    "    \n",
    "    \n",
    "    ################  #####################  #####################\n",
    "    ################ Enrich Embedded AAs #####################\n",
    "    ################ #####################   #####################\n",
    "    \n",
    "    \n",
    "    ###### Linear Classifier ######\n",
    "    print(train_concat.size(), test_concat.size())\n",
    "    clf = LogisticRegression(random_state=0).fit(train_concat.detach().numpy(), train_labels)\n",
    "    \n",
    "    predictions = clf.predict(test_concat.detach().numpy())\n",
    "    \n",
    "    score = clf.score(x_test, test_labels)\n",
    "\n",
    "    \n",
    "    from sklearn import metrics\n",
    "    cm = metrics.confusion_matrix(test_labels, predictions)\n",
    "    print(cm)\n",
    "    return score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "827a13c9-24af-445d-b85c-f52463c85290",
   "metadata": {},
   "outputs": [],
   "source": [
    "def dms_eval_task(model):\n",
    "    model.eval()\n",
    "    ##### DMS Task\n",
    "    \n",
    "    dms_wt_aa =  torch.load(f'{DOWNSTREAM_FILE_PATHS}/DMS/data_wt/data-cleaned-wt-embeddings.pt')\n",
    "    output_wt = model({\n",
    "                'aa': dms_wt_aa.type(torch.float32).to(device)\n",
    "            })\n",
    "    \n",
    "    dms_mt_aa =  torch.load(f'{DOWNSTREAM_FILE_PATHS}/DMS/data_mt/data-cleaned-mt-embeddings.pt')\n",
    "    output_mt = model({\n",
    "                'aa': dms_mt_aa.type(torch.float32).to(device)\n",
    "            })\n",
    "    output_wt['aa'].size(), output_mt['aa'].size()\n",
    "    cos = torch.nn.CosineSimilarity(dim=1, eps=1e-6)\n",
    "    output = cos(output_wt['aa'], output_mt['aa'])\n",
    "    return "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7052289c-6c2c-4032-b45c-d50b1c974055",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}