2. create_embeddings.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Embedding generation\n",
    "This Jupyter notebook is designed to generate embeddings for each protein."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/data/nasif12/home_if12/l_schewinski/.conda/envs/l-schewinski-tp-4/lib/python3.7/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import esm\n",
    "import torch\n",
    "import pandas as pd\n",
    "import tqdm\n",
    "import tqdm\n",
    "import numpy as np\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ESM2(\n",
       "  (embed_tokens): Embedding(33, 1280, padding_idx=1)\n",
       "  (layers): ModuleList(\n",
       "    (0): TransformerLayer(\n",
       "      (self_attn): MultiheadAttention(\n",
       "        (k_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (v_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (q_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (out_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (rot_emb): RotaryEmbedding()\n",
       "      )\n",
       "      (self_attn_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "      (fc1): Linear(in_features=1280, out_features=5120, bias=True)\n",
       "      (fc2): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "      (final_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (1): TransformerLayer(\n",
       "      (self_attn): MultiheadAttention(\n",
       "        (k_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (v_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (q_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (out_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (rot_emb): RotaryEmbedding()\n",
       "      )\n",
       "      (self_attn_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "      (fc1): Linear(in_features=1280, out_features=5120, bias=True)\n",
       "      (fc2): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "      (final_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (2): TransformerLayer(\n",
       "      (self_attn): MultiheadAttention(\n",
       "        (k_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (v_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (q_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (out_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (rot_emb): RotaryEmbedding()\n",
       "      )\n",
       "      (self_attn_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "      (fc1): Linear(in_features=1280, out_features=5120, bias=True)\n",
       "      (fc2): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "      (final_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (3): TransformerLayer(\n",
       "      (self_attn): MultiheadAttention(\n",
       "        (k_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (v_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (q_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (out_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (rot_emb): RotaryEmbedding()\n",
       "      )\n",
       "      (self_attn_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "      (fc1): Linear(in_features=1280, out_features=5120, bias=True)\n",
       "      (fc2): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "      (final_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (4): TransformerLayer(\n",
       "      (self_attn): MultiheadAttention(\n",
       "        (k_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (v_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (q_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (out_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (rot_emb): RotaryEmbedding()\n",
       "      )\n",
       "      (self_attn_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "      (fc1): Linear(in_features=1280, out_features=5120, bias=True)\n",
       "      (fc2): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "      (final_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (5): TransformerLayer(\n",
       "      (self_attn): MultiheadAttention(\n",
       "        (k_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (v_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (q_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (out_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (rot_emb): RotaryEmbedding()\n",
       "      )\n",
       "      (self_attn_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "      (fc1): Linear(in_features=1280, out_features=5120, bias=True)\n",
       "      (fc2): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "      (final_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (6): TransformerLayer(\n",
       "      (self_attn): MultiheadAttention(\n",
       "        (k_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (v_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (q_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (out_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (rot_emb): RotaryEmbedding()\n",
       "      )\n",
       "      (self_attn_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "      (fc1): Linear(in_features=1280, out_features=5120, bias=True)\n",
       "      (fc2): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "      (final_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (7): TransformerLayer(\n",
       "      (self_attn): MultiheadAttention(\n",
       "        (k_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (v_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (q_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (out_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (rot_emb): RotaryEmbedding()\n",
       "      )\n",
       "      (self_attn_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "      (fc1): Linear(in_features=1280, out_features=5120, bias=True)\n",
       "      (fc2): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "      (final_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (8): TransformerLayer(\n",
       "      (self_attn): MultiheadAttention(\n",
       "        (k_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (v_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (q_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (out_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (rot_emb): RotaryEmbedding()\n",
       "      )\n",
       "      (self_attn_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "      (fc1): Linear(in_features=1280, out_features=5120, bias=True)\n",
       "      (fc2): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "      (final_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (9): TransformerLayer(\n",
       "      (self_attn): MultiheadAttention(\n",
       "        (k_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (v_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (q_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (out_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (rot_emb): RotaryEmbedding()\n",
       "      )\n",
       "      (self_attn_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "      (fc1): Linear(in_features=1280, out_features=5120, bias=True)\n",
       "      (fc2): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "      (final_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (10): TransformerLayer(\n",
       "      (self_attn): MultiheadAttention(\n",
       "        (k_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (v_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (q_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (out_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (rot_emb): RotaryEmbedding()\n",
       "      )\n",
       "      (self_attn_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "      (fc1): Linear(in_features=1280, out_features=5120, bias=True)\n",
       "      (fc2): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "      (final_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (11): TransformerLayer(\n",
       "      (self_attn): MultiheadAttention(\n",
       "        (k_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (v_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (q_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (out_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (rot_emb): RotaryEmbedding()\n",
       "      )\n",
       "      (self_attn_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "      (fc1): Linear(in_features=1280, out_features=5120, bias=True)\n",
       "      (fc2): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "      (final_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (12): TransformerLayer(\n",
       "      (self_attn): MultiheadAttention(\n",
       "        (k_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (v_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (q_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (out_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (rot_emb): RotaryEmbedding()\n",
       "      )\n",
       "      (self_attn_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "      (fc1): Linear(in_features=1280, out_features=5120, bias=True)\n",
       "      (fc2): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "      (final_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (13): TransformerLayer(\n",
       "      (self_attn): MultiheadAttention(\n",
       "        (k_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (v_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (q_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (out_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (rot_emb): RotaryEmbedding()\n",
       "      )\n",
       "      (self_attn_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "      (fc1): Linear(in_features=1280, out_features=5120, bias=True)\n",
       "      (fc2): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "      (final_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (14): TransformerLayer(\n",
       "      (self_attn): MultiheadAttention(\n",
       "        (k_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (v_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (q_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (out_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (rot_emb): RotaryEmbedding()\n",
       "      )\n",
       "      (self_attn_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "      (fc1): Linear(in_features=1280, out_features=5120, bias=True)\n",
       "      (fc2): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "      (final_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (15): TransformerLayer(\n",
       "      (self_attn): MultiheadAttention(\n",
       "        (k_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (v_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (q_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (out_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (rot_emb): RotaryEmbedding()\n",
       "      )\n",
       "      (self_attn_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "      (fc1): Linear(in_features=1280, out_features=5120, bias=True)\n",
       "      (fc2): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "      (final_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (16): TransformerLayer(\n",
       "      (self_attn): MultiheadAttention(\n",
       "        (k_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (v_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (q_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (out_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (rot_emb): RotaryEmbedding()\n",
       "      )\n",
       "      (self_attn_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "      (fc1): Linear(in_features=1280, out_features=5120, bias=True)\n",
       "      (fc2): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "      (final_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (17): TransformerLayer(\n",
       "      (self_attn): MultiheadAttention(\n",
       "        (k_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (v_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (q_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (out_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (rot_emb): RotaryEmbedding()\n",
       "      )\n",
       "      (self_attn_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "      (fc1): Linear(in_features=1280, out_features=5120, bias=True)\n",
       "      (fc2): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "      (final_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (18): TransformerLayer(\n",
       "      (self_attn): MultiheadAttention(\n",
       "        (k_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (v_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (q_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (out_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (rot_emb): RotaryEmbedding()\n",
       "      )\n",
       "      (self_attn_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "      (fc1): Linear(in_features=1280, out_features=5120, bias=True)\n",
       "      (fc2): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "      (final_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (19): TransformerLayer(\n",
       "      (self_attn): MultiheadAttention(\n",
       "        (k_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (v_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (q_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (out_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (rot_emb): RotaryEmbedding()\n",
       "      )\n",
       "      (self_attn_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "      (fc1): Linear(in_features=1280, out_features=5120, bias=True)\n",
       "      (fc2): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "      (final_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (20): TransformerLayer(\n",
       "      (self_attn): MultiheadAttention(\n",
       "        (k_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (v_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (q_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (out_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (rot_emb): RotaryEmbedding()\n",
       "      )\n",
       "      (self_attn_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "      (fc1): Linear(in_features=1280, out_features=5120, bias=True)\n",
       "      (fc2): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "      (final_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (21): TransformerLayer(\n",
       "      (self_attn): MultiheadAttention(\n",
       "        (k_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (v_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (q_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (out_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (rot_emb): RotaryEmbedding()\n",
       "      )\n",
       "      (self_attn_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "      (fc1): Linear(in_features=1280, out_features=5120, bias=True)\n",
       "      (fc2): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "      (final_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (22): TransformerLayer(\n",
       "      (self_attn): MultiheadAttention(\n",
       "        (k_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (v_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (q_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (out_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (rot_emb): RotaryEmbedding()\n",
       "      )\n",
       "      (self_attn_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "      (fc1): Linear(in_features=1280, out_features=5120, bias=True)\n",
       "      (fc2): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "      (final_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (23): TransformerLayer(\n",
       "      (self_attn): MultiheadAttention(\n",
       "        (k_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (v_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (q_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (out_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (rot_emb): RotaryEmbedding()\n",
       "      )\n",
       "      (self_attn_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "      (fc1): Linear(in_features=1280, out_features=5120, bias=True)\n",
       "      (fc2): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "      (final_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (24): TransformerLayer(\n",
       "      (self_attn): MultiheadAttention(\n",
       "        (k_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (v_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (q_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (out_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (rot_emb): RotaryEmbedding()\n",
       "      )\n",
       "      (self_attn_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "      (fc1): Linear(in_features=1280, out_features=5120, bias=True)\n",
       "      (fc2): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "      (final_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (25): TransformerLayer(\n",
       "      (self_attn): MultiheadAttention(\n",
       "        (k_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (v_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (q_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (out_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (rot_emb): RotaryEmbedding()\n",
       "      )\n",
       "      (self_attn_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "      (fc1): Linear(in_features=1280, out_features=5120, bias=True)\n",
       "      (fc2): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "      (final_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (26): TransformerLayer(\n",
       "      (self_attn): MultiheadAttention(\n",
       "        (k_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (v_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (q_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (out_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (rot_emb): RotaryEmbedding()\n",
       "      )\n",
       "      (self_attn_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "      (fc1): Linear(in_features=1280, out_features=5120, bias=True)\n",
       "      (fc2): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "      (final_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (27): TransformerLayer(\n",
       "      (self_attn): MultiheadAttention(\n",
       "        (k_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (v_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (q_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (out_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (rot_emb): RotaryEmbedding()\n",
       "      )\n",
       "      (self_attn_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "      (fc1): Linear(in_features=1280, out_features=5120, bias=True)\n",
       "      (fc2): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "      (final_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (28): TransformerLayer(\n",
       "      (self_attn): MultiheadAttention(\n",
       "        (k_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (v_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (q_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (out_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (rot_emb): RotaryEmbedding()\n",
       "      )\n",
       "      (self_attn_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "      (fc1): Linear(in_features=1280, out_features=5120, bias=True)\n",
       "      (fc2): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "      (final_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (29): TransformerLayer(\n",
       "      (self_attn): MultiheadAttention(\n",
       "        (k_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (v_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (q_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (out_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (rot_emb): RotaryEmbedding()\n",
       "      )\n",
       "      (self_attn_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "      (fc1): Linear(in_features=1280, out_features=5120, bias=True)\n",
       "      (fc2): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "      (final_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (30): TransformerLayer(\n",
       "      (self_attn): MultiheadAttention(\n",
       "        (k_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (v_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (q_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (out_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (rot_emb): RotaryEmbedding()\n",
       "      )\n",
       "      (self_attn_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "      (fc1): Linear(in_features=1280, out_features=5120, bias=True)\n",
       "      (fc2): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "      (final_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (31): TransformerLayer(\n",
       "      (self_attn): MultiheadAttention(\n",
       "        (k_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (v_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (q_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (out_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (rot_emb): RotaryEmbedding()\n",
       "      )\n",
       "      (self_attn_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "      (fc1): Linear(in_features=1280, out_features=5120, bias=True)\n",
       "      (fc2): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "      (final_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (32): TransformerLayer(\n",
       "      (self_attn): MultiheadAttention(\n",
       "        (k_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (v_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (q_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (out_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (rot_emb): RotaryEmbedding()\n",
       "      )\n",
       "      (self_attn_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "      (fc1): Linear(in_features=1280, out_features=5120, bias=True)\n",
       "      (fc2): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "      (final_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "  )\n",
       "  (contact_head): ContactPredictionHead(\n",
       "    (regression): Linear(in_features=660, out_features=1, bias=True)\n",
       "    (activation): Sigmoid()\n",
       "  )\n",
       "  (emb_layer_norm_after): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "  (lm_head): RobertaLMHead(\n",
       "    (dense): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "    (layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Loads the ESM-2 model\n",
    "model, alphabet = esm.pretrained.esm2_t33_650M_UR50D()\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the sequence data already filtered\n",
    "df = pd.read_csv('data/OsmoticStress_with_binary_positions_padded_5000.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_clean = df[[\"Uniprot_ID\", \"full_sequence\", \"Qvalue(LiP)\"]]\n",
    "df_clean.reset_index(drop=True, inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fix type of full_sequence and remove NaN values\n",
    "df_clean = df_clean[~df_clean['full_sequence'].apply(lambda x: isinstance(x, float))]\n",
    "df_clean['full_sequence'] = df_clean['full_sequence'].astype(str)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Retrieve sequences for unique protein IDs to avoid generating duplicate embeddings\n",
    "df_clean_dropped = df_clean.drop_duplicates(subset='Uniprot_ID', keep='first')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Uniprot_ID</th>\n",
       "      <th>full_sequence</th>\n",
       "      <th>Qvalue(LiP)</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>P15703</td>\n",
       "      <td>MRFSTTLATAATALFFTASQVSAIGELAFNLGVKNNDGTCKSTSDY...</td>\n",
       "      <td>0.003686</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>P06169</td>\n",
       "      <td>MSEITLGKYLFERLKQVNVNTVFGLPGDFNLSLLDKIYEVEGMRWA...</td>\n",
       "      <td>0.003686</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>P38174</td>\n",
       "      <td>MTDAEIENSPASDLKELNLENEGVEQQDQAKADESDPVESKKKKNK...</td>\n",
       "      <td>0.003686</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>P00359</td>\n",
       "      <td>MVRVAINGFGRIGRLVMRIALSRPNVEVVALNDPFITNDYAAYMFK...</td>\n",
       "      <td>0.003686</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>Q04305</td>\n",
       "      <td>MSTARPRIITSKAPLLPQQTTPEQRYWRQYTSAQLVKEHNSVTHIS...</td>\n",
       "      <td>0.005126</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4957</th>\n",
       "      <td>P36000</td>\n",
       "      <td>MPPLDKRIKKFLKDSIRIAPKISGKGELSELRTGLVSQYPQTRKDA...</td>\n",
       "      <td>0.051109</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4959</th>\n",
       "      <td>P80428</td>\n",
       "      <td>MSNAALQVYGGDEVSAVVIDPGSYTTNIGYSGSDFPQSILPSVYGK...</td>\n",
       "      <td>0.051109</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4964</th>\n",
       "      <td>P33441</td>\n",
       "      <td>MPLSQKQIDQVRTKVHYSEVDTPFNKYLDILGKVTKLTGSIINGTL...</td>\n",
       "      <td>0.051150</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4987</th>\n",
       "      <td>P34240</td>\n",
       "      <td>MEKIPRWLLFSLISSVLCILGALCVPLLSVAFDSKRNSQSKLVNYG...</td>\n",
       "      <td>0.051166</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4988</th>\n",
       "      <td>P47137</td>\n",
       "      <td>MVPKFYKLSNGFKIPSIALGTYDIPRSQTAEIVYEGVKCGYRHFDT...</td>\n",
       "      <td>0.051166</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>1455 rows × 3 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "     Uniprot_ID                                      full_sequence  \\\n",
       "0        P15703  MRFSTTLATAATALFFTASQVSAIGELAFNLGVKNNDGTCKSTSDY...   \n",
       "1        P06169  MSEITLGKYLFERLKQVNVNTVFGLPGDFNLSLLDKIYEVEGMRWA...   \n",
       "2        P38174  MTDAEIENSPASDLKELNLENEGVEQQDQAKADESDPVESKKKKNK...   \n",
       "3        P00359  MVRVAINGFGRIGRLVMRIALSRPNVEVVALNDPFITNDYAAYMFK...   \n",
       "5        Q04305  MSTARPRIITSKAPLLPQQTTPEQRYWRQYTSAQLVKEHNSVTHIS...   \n",
       "...         ...                                                ...   \n",
       "4957     P36000  MPPLDKRIKKFLKDSIRIAPKISGKGELSELRTGLVSQYPQTRKDA...   \n",
       "4959     P80428  MSNAALQVYGGDEVSAVVIDPGSYTTNIGYSGSDFPQSILPSVYGK...   \n",
       "4964     P33441  MPLSQKQIDQVRTKVHYSEVDTPFNKYLDILGKVTKLTGSIINGTL...   \n",
       "4987     P34240  MEKIPRWLLFSLISSVLCILGALCVPLLSVAFDSKRNSQSKLVNYG...   \n",
       "4988     P47137  MVPKFYKLSNGFKIPSIALGTYDIPRSQTAEIVYEGVKCGYRHFDT...   \n",
       "\n",
       "      Qvalue(LiP)  \n",
       "0        0.003686  \n",
       "1        0.003686  \n",
       "2        0.003686  \n",
       "3        0.003686  \n",
       "5        0.005126  \n",
       "...           ...  \n",
       "4957     0.051109  \n",
       "4959     0.051109  \n",
       "4964     0.051150  \n",
       "4987     0.051166  \n",
       "4988     0.051166  \n",
       "\n",
       "[1455 rows x 3 columns]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_clean_dropped"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_embeddings(model, alphabet, sequences):\n",
    "    \"\"\"\n",
    "    Function to generate the embeddings.\n",
    "    \"\"\"\n",
    "    embeddings = []\n",
    "    \n",
    "    batch_converter = alphabet.get_batch_converter()\n",
    "    \n",
    "    for sequence in tqdm.tqdm(sequences, desc=\"Generating Embeddings\"):\n",
    "        data = [(0, sequence)]\n",
    "        batch_labels, batch_strs, batch_tokens = batch_converter(data)\n",
    "\n",
    "        # Generate embeddings\n",
    "        with torch.no_grad():\n",
    "            results = model(batch_tokens, repr_layers=[33]) \n",
    "            token_embeddings = results['representations'][33]\n",
    "\n",
    "        # Create full embeddings\n",
    "        full_embedding = token_embeddings.squeeze().numpy()\n",
    "        embeddings.append(full_embedding)\n",
    "    \n",
    "    return embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/data/nasif12/home_if12/l_schewinski/.conda/envs/l-schewinski-tp-4/lib/python3.7/site-packages/ipykernel_launcher.py:6: SettingWithCopyWarning: \n",
      "A value is trying to be set on a copy of a slice from a DataFrame.\n",
      "Try using .loc[row_indexer,col_indexer] = value instead\n",
      "\n",
      "See the caveats in the documentation: https://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html#returning-a-view-versus-a-copy\n",
      "  \n"
     ]
    }
   ],
   "source": [
    "def pad_sequence(sequence, target_length=1000, pad_token='<pad>'):\n",
    "    \"\"\"\n",
    "    Append the <pad> token to the sequence until it reaches a length of 1000.\n",
    "    \"\"\"\n",
    "    padding_length = target_length - len(sequence)\n",
    "    return sequence + pad_token * padding_length\n",
    "\n",
    "# Apply padding to each sequence with <pad>\n",
    "df_clean_dropped['padded_sequence'] = df_clean_dropped['full_sequence'].apply(lambda x: pad_sequence(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/data/nasif12/home_if12/l_schewinski/.conda/envs/l-schewinski-tp-4/lib/python3.7/site-packages/ipykernel_launcher.py:2: SettingWithCopyWarning: \n",
      "A value is trying to be set on a copy of a slice from a DataFrame.\n",
      "Try using .loc[row_indexer,col_indexer] = value instead\n",
      "\n",
      "See the caveats in the documentation: https://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html#returning-a-view-versus-a-copy\n",
      "  \n",
      "/data/nasif12/home_if12/l_schewinski/.conda/envs/l-schewinski-tp-4/lib/python3.7/site-packages/ipykernel_launcher.py:3: SettingWithCopyWarning: \n",
      "A value is trying to be set on a copy of a slice from a DataFrame.\n",
      "Try using .loc[row_indexer,col_indexer] = value instead\n",
      "\n",
      "See the caveats in the documentation: https://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html#returning-a-view-versus-a-copy\n",
      "  This is separate from the ipykernel package so we can avoid doing imports until\n",
      "/data/nasif12/home_if12/l_schewinski/.conda/envs/l-schewinski-tp-4/lib/python3.7/site-packages/ipykernel_launcher.py:4: SettingWithCopyWarning: \n",
      "A value is trying to be set on a copy of a slice from a DataFrame.\n",
      "Try using .loc[row_indexer,col_indexer] = value instead\n",
      "\n",
      "See the caveats in the documentation: https://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html#returning-a-view-versus-a-copy\n",
      "  after removing the cwd from sys.path.\n",
      "/data/nasif12/home_if12/l_schewinski/.conda/envs/l-schewinski-tp-4/lib/python3.7/site-packages/ipykernel_launcher.py:5: SettingWithCopyWarning: \n",
      "A value is trying to be set on a copy of a slice from a DataFrame.\n",
      "Try using .loc[row_indexer,col_indexer] = value instead\n",
      "\n",
      "See the caveats in the documentation: https://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html#returning-a-view-versus-a-copy\n",
      "  \"\"\"\n"
     ]
    }
   ],
   "source": [
    "# Check if the padding is correct\n",
    "df_clean_dropped['full_sequence_length'] = df_clean_dropped['full_sequence'].apply(len)\n",
    "df_clean_dropped['num_pad_tokens'] = df_clean_dropped['padded_sequence'].apply(lambda x: x.count('<pad>'))\n",
    "df_clean_dropped['padded_sequence_length'] = df_clean_dropped['full_sequence_length'] + df_clean_dropped['num_pad_tokens']\n",
    "df_clean_dropped['expected_padded_length'] = 1000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_clean_dropped.reset_index(drop=True, inplace=True)\n",
    "df_clean_dropped = df_clean_dropped[1000:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating Embeddings: 100%|██████████| 476/476 [1:13:50<00:00,  9.31s/it]\n"
     ]
    }
   ],
   "source": [
    "# Generate the embeddings!\n",
    "sequences = df_clean_dropped['padded_sequence'].tolist()\n",
    "embeddings = generate_embeddings(model, alphabet, sequences)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/data/nasif12/home_if12/l_schewinski/.conda/envs/l-schewinski-tp-4/lib/python3.7/site-packages/ipykernel_launcher.py:1: SettingWithCopyWarning: \n",
      "A value is trying to be set on a copy of a slice from a DataFrame.\n",
      "Try using .loc[row_indexer,col_indexer] = value instead\n",
      "\n",
      "See the caveats in the documentation: https://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html#returning-a-view-versus-a-copy\n",
      "  \"\"\"Entry point for launching an IPython kernel.\n"
     ]
    }
   ],
   "source": [
    "df_test_dropped['full_embedding'] = [e.tolist() for e in embeddings]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save embeddings\n",
    "df_test_dropped.to_pickle('data/embeddings_new_3_3.pkl', protocol=4)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:.conda-l-schewinski-tp-4]",
   "language": "python",
   "name": "conda-env-.conda-l-schewinski-tp-4-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.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}