ems2_mutations.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Predicting Mutation Effects with ESM-2"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We would like to compute the following function, \n",
    "\n",
    "$$\n",
    "S_{\\text{masked marginal}} = \\sum_{i \\in M} \\left[ \\log p(x_i = x_{i}^{\\text{mt}}|x_{-M}) - \\log p(x_i = x_{i}^{\\text{wt}}|x_{-M}) \\right]\n",
    "$$\n",
    "\n",
    "The function `masked_marginal_scoring` calculates the change in the log probability of observing a certain mutated amino acid sequence given the rest of the protein sequence, compared to the original (wild-type) sequence. The function utilizes a language model that has been trained on protein sequences and can generate predictions for the probability of observing each amino acid at a specific position given the context (the rest of the sequence).\n",
    "\n",
    "Let's break down the function step by step:\n",
    "\n",
    "1. **Sequence Preparation:**\n",
    "The function first creates a copy of the input sequence as a list, `seq_list`. For each mutation specified in `mutations`, it checks if the wild-type (wt) amino acid at the specified position matches the amino acid in the sequence. If it does not, the function raises a ValueError. If it does, the function replaces the wild-type amino acid at the specified position with a mask token. The mask token is a special token that indicates to the language model that this position is to be predicted. This transformed sequence is then converted back into a string.\n",
    "\n",
    "2. **Tokenization:**\n",
    "The masked sequence is tokenized using the provided tokenizer, converting the sequence of amino acids into a sequence of tokens that the language model can understand. The tokenized sequence is then converted into the format expected by the language model (`inputs`).\n",
    "\n",
    "3. **Logits Calculation:**\n",
    "The function then feeds `inputs` into the language model. The output of the model is a tensor of logits. Logits are the raw, unnormalized scores output by the last layer of the model for each token in the vocabulary. In this context, the logits represent the model's belief about the likelihood of each possible amino acid at each position in the sequence, before applying any normalization to convert these scores into probabilities.\n",
    "\n",
    "4. **Mask Indices Extraction:**\n",
    "The function retrieves the indices of the mask tokens in the tokenized sequence. These indices correspond to the positions of the mutations.\n",
    "\n",
    "5. **Scoring:**\n",
    "The function initializes a score to zero. For each mutation and its corresponding mask index, it retrieves the logits for the masked position and applies the softmax function to convert these logits into probabilities. The softmax function is defined as \n",
    "\n",
    "$$\n",
    "\\text{softmax}(x)_i = \\frac{e^{x_i}}{\\sum_j e^{x_j}}\n",
    "$$\n",
    "\n",
    "where $x$ is a vector of logits, and $i$ and $j$ index over the elements of $x$. Softmax ensures that the probabilities sum to one and that they are proportional to the exponentiated logits. This is a common way to convert the unnormalized scores output by a model into probabilities.\n",
    "\n",
    "6. **Log Probabilities Calculation:**\n",
    "The function then computes the natural logarithm of these probabilities to obtain log probabilities. The reason for working with log probabilities instead of raw probabilities is to avoid numerical underflow that can occur when multiplying many small probabilities together, and to convert multiplications into additions, which are computationally easier to handle.\n",
    "\n",
    "7. **Score Update:**\n",
    "The function retrieves the log probabilities corresponding to the wild-type and mutant (mt) tokens, and subtracts the wild-type log probability from the mutant log probability. This difference in log probabilities represents the change in the log likelihood of the sequence due to the mutation, and is added to the score. The overall score returned by the function is the sum of the differences in log probabilities for all mutations. It is defined as:\n",
    "\n",
    "$$\n",
    "S_{\\text{masked marginal}} = \\sum_{i \\in M} \\left[ \\log p(x_i = x_{i}^{\\text{mt}}|x_{-M}) - \\log p(x_i = x_{i}^{\\text{wt}}|x_{-M}) \\right]\n",
    "$$\n",
    "\n",
    "where $M$ is the set of positions of the mutations, $x_i$ is the amino acid at position $i$, $x_{i}^{\\text{wt}}$ and $x_{i}^{\\text{mt}}$ are the wild-type and mutant amino acids at position $i$, respectively, and $x_{-M}$ represents the amino acids at all positions except those in $M$. This score quantifies the predicted change in the likelihood of the sequence due to the specified mutations. A positive score indicates that the mutations are predicted to make the sequence more likely (given the model's training data), while a negative score indicates that the mutations are predicted to make the sequence less likely. This approach can be used to predict the effects of multiple mutations on a protein sequence. By comparing the scores for different sets of mutations, one can rank the mutations by their predicted effects. This could be used, for example, to prioritize experimental testing of mutations, or to guide protein engineering efforts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-0.9924547672271729"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "from transformers import AutoTokenizer, EsmForMaskedLM\n",
    "from typing import List, Tuple\n",
    "\n",
    "def masked_marginal_scoring(\n",
    "    tokenizer: AutoTokenizer,\n",
    "    model: EsmForMaskedLM,\n",
    "    sequence: str,\n",
    "    mutations: List[Tuple[int, str, str]]\n",
    ") -> float:\n",
    "    \n",
    "    # Create a copy of the sequence for mutation\n",
    "    seq_list = list(sequence)\n",
    "\n",
    "    # Check and mask the positions\n",
    "    for pos, wt, mt in mutations:\n",
    "        if seq_list[pos] != wt:\n",
    "            raise ValueError(f\"The amino acid at position {pos} is {seq_list[pos]}, not {wt}.\")\n",
    "        seq_list[pos] = tokenizer.mask_token\n",
    "    \n",
    "    # Convert the mutated sequence back to string\n",
    "    masked_sequence = \"\".join(seq_list)\n",
    "    \n",
    "    # Tokenize the masked sequence\n",
    "    inputs = tokenizer(masked_sequence, return_tensors=\"pt\")\n",
    "\n",
    "    # Get the model's output\n",
    "    with torch.no_grad():\n",
    "        outputs = model(**inputs)\n",
    "\n",
    "    # Get the logits\n",
    "    logits = outputs.logits\n",
    "\n",
    "    # Get the mask token indices\n",
    "    mask_indices = torch.where(inputs.input_ids.squeeze() == tokenizer.mask_token_id)[0]\n",
    "\n",
    "    # Initialize score\n",
    "    score = 0\n",
    "\n",
    "    # Iterate over each mutation\n",
    "    for (pos, wt, mt), mask_index in zip(mutations, mask_indices):\n",
    "        # Get the logits for the masked position\n",
    "        position_logits = logits[0, mask_index]\n",
    "\n",
    "        # Apply softmax to logits to get probabilities\n",
    "        probabilities = torch.nn.functional.softmax(position_logits, dim=-1)\n",
    "\n",
    "        # Convert probabilities to log probabilities\n",
    "        log_probabilities = torch.log(probabilities)\n",
    "\n",
    "        # Get the token ids for wt and mt\n",
    "        wt_token_id = tokenizer.convert_tokens_to_ids(wt)\n",
    "        mt_token_id = tokenizer.convert_tokens_to_ids(mt)\n",
    "\n",
    "        # Retrieve the log probabilities\n",
    "        wt_log_prob = log_probabilities[wt_token_id].item()\n",
    "        mt_log_prob = log_probabilities[mt_token_id].item()\n",
    "\n",
    "        # Compute the difference and update the score\n",
    "        score += mt_log_prob - wt_log_prob\n",
    "\n",
    "    return score\n",
    "\n",
    "# Test the function\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"facebook/esm2_t6_8M_UR50D\")\n",
    "model = EsmForMaskedLM.from_pretrained(\"facebook/esm2_t6_8M_UR50D\")\n",
    "sequence = \"MAPLRKTYVLKLYVAGNTPNSVRALKTLNNILEKEFKGVYALKVIDVLKNPQLAEEDKILATPTLAKVLPPPVRRIIGDLSNREKVLIGLDLLYEEIGDQAEDDLGLE\"\n",
    "# mutations = [(67, 'V', 'R'), (82, 'R', 'D'), (83, 'E', 'A')]\n",
    "mutations = [(68, 'L', 'R'), (83, 'E', 'D'), (84, 'K', 'A')]\n",
    "\n",
    "score = masked_marginal_scoring(tokenizer, model, sequence, mutations)\n",
    "score\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Computing Sequence Alignment Scores with ESM-2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'<cls>': -16.615102767944336,\n",
       "  '<pad>': -27.01595115661621,\n",
       "  '<eos>': -16.08125114440918,\n",
       "  '<unk>': -27.027359008789062,\n",
       "  'L': -2.4245827198028564,\n",
       "  'A': -2.6821749210357666,\n",
       "  'G': -2.3036611080169678,\n",
       "  'V': -2.3579823970794678,\n",
       "  'S': -3.074674606323242,\n",
       "  'E': -2.8953990936279297,\n",
       "  'R': -3.103949785232544,\n",
       "  'T': -3.0654356479644775,\n",
       "  'I': -2.498311996459961,\n",
       "  'D': -3.350306987762451,\n",
       "  'P': -3.3211705684661865,\n",
       "  'K': -2.3286314010620117,\n",
       "  'Q': -3.354858636856079,\n",
       "  'N': -3.5923171043395996,\n",
       "  'F': -3.171297311782837,\n",
       "  'Y': -3.2451534271240234,\n",
       "  'M': -3.8390090465545654,\n",
       "  'H': -4.00257682800293,\n",
       "  'W': -4.410708904266357,\n",
       "  'C': -4.662168025970459,\n",
       "  'X': -11.923528671264648,\n",
       "  'B': -15.062100410461426,\n",
       "  'U': -15.343768119812012,\n",
       "  'Z': -15.90283489227295,\n",
       "  'O': -19.265869140625,\n",
       "  '.': -19.64056396484375,\n",
       "  '-': -19.632041931152344,\n",
       "  '<null_1>': -19.675811767578125,\n",
       "  '<mask>': -27.014863967895508},\n",
       " {'<cls>': -15.54505729675293,\n",
       "  '<pad>': -25.044166564941406,\n",
       "  '<eos>': -16.202560424804688,\n",
       "  '<unk>': -25.036680221557617,\n",
       "  'L': -2.7625975608825684,\n",
       "  'A': -3.077972650527954,\n",
       "  'G': -2.3007636070251465,\n",
       "  'V': -3.109529495239258,\n",
       "  'S': -2.5639259815216064,\n",
       "  'E': -2.3683600425720215,\n",
       "  'R': -3.1586110591888428,\n",
       "  'T': -2.9586145877838135,\n",
       "  'I': -3.3267810344696045,\n",
       "  'D': -2.1744065284729004,\n",
       "  'P': -2.4598441123962402,\n",
       "  'K': -2.66398024559021,\n",
       "  'Q': -3.402146577835083,\n",
       "  'N': -3.0627105236053467,\n",
       "  'F': -3.4811530113220215,\n",
       "  'Y': -3.684326171875,\n",
       "  'M': -4.438747406005859,\n",
       "  'H': -4.114745616912842,\n",
       "  'W': -4.799737453460693,\n",
       "  'C': -5.532626628875732,\n",
       "  'X': -10.646926879882812,\n",
       "  'B': -15.18580150604248,\n",
       "  'U': -15.58905029296875,\n",
       "  'Z': -16.164033889770508,\n",
       "  'O': -19.57425308227539,\n",
       "  '.': -19.937580108642578,\n",
       "  '-': -19.888559341430664,\n",
       "  '<null_1>': -19.944517135620117,\n",
       "  '<mask>': -25.04560661315918},\n",
       " {'<cls>': -15.621170043945312,\n",
       "  '<pad>': -24.900922775268555,\n",
       "  '<eos>': -15.876766204833984,\n",
       "  '<unk>': -24.894269943237305,\n",
       "  'L': -2.1769464015960693,\n",
       "  'A': -2.6387367248535156,\n",
       "  'G': -3.203376293182373,\n",
       "  'V': -2.6407947540283203,\n",
       "  'S': -2.9942092895507812,\n",
       "  'E': -2.6530227661132812,\n",
       "  'R': -2.7883541584014893,\n",
       "  'T': -3.3047196865081787,\n",
       "  'I': -2.6859805583953857,\n",
       "  'D': -2.7920682430267334,\n",
       "  'P': -2.5504977703094482,\n",
       "  'K': -2.5052103996276855,\n",
       "  'Q': -3.4554262161254883,\n",
       "  'N': -3.4880573749542236,\n",
       "  'F': -3.1002357006073,\n",
       "  'Y': -3.4783072471618652,\n",
       "  'M': -3.708585500717163,\n",
       "  'H': -4.19027042388916,\n",
       "  'W': -4.3441362380981445,\n",
       "  'C': -5.372040271759033,\n",
       "  'X': -10.770855903625488,\n",
       "  'B': -14.884002685546875,\n",
       "  'U': -15.256261825561523,\n",
       "  'Z': -15.815324783325195,\n",
       "  'O': -19.212984085083008,\n",
       "  '.': -19.5474796295166,\n",
       "  '-': -19.539596557617188,\n",
       "  '<null_1>': -19.567399978637695,\n",
       "  '<mask>': -24.90018653869629}]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from transformers import AutoTokenizer, EsmForMaskedLM\n",
    "import torch\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"facebook/esm2_t6_8M_UR50D\")\n",
    "model = EsmForMaskedLM.from_pretrained(\"facebook/esm2_t6_8M_UR50D\")\n",
    "\n",
    "def mask_protein_sequence(protein_sequence: str, mask_positions: list):\n",
    "    # Tokenize the protein sequence\n",
    "    inputs = tokenizer(protein_sequence, return_tensors=\"pt\")\n",
    "\n",
    "    # Mask the specified positions\n",
    "    for pos in mask_positions:\n",
    "        inputs[\"input_ids\"][0][pos] = tokenizer.mask_token_id\n",
    "\n",
    "    # Get the logits\n",
    "    with torch.no_grad():\n",
    "        logits = model(**inputs).logits\n",
    "\n",
    "    # Get the logits for the masked positions\n",
    "    mask_logits = logits[0, mask_positions]\n",
    "\n",
    "    # Apply softmax to convert logits to probabilities\n",
    "    probabilities = torch.nn.functional.softmax(mask_logits, dim=-1)\n",
    "\n",
    "    # Convert the probabilities to log probabilities\n",
    "    log_probabilities = torch.log(probabilities)\n",
    "\n",
    "    # Convert the log probabilities to a dictionary with vocabulary elements as keys\n",
    "    vocab = tokenizer.get_vocab()\n",
    "    vocab = {v: k for k, v in vocab.items()}  # reverse the key-value pairs in the vocab\n",
    "    \n",
    "    log_probabilities_dicts = []\n",
    "    for log_probs in log_probabilities:\n",
    "        log_probs_dict = {vocab[i]: log_prob.item() for i, log_prob in enumerate(log_probs)}\n",
    "        log_probabilities_dicts.append(log_probs_dict)\n",
    "\n",
    "    return log_probabilities_dicts\n",
    "\n",
    "# Test the function with a protein sequence and mask positions\n",
    "protein_sequence = \"MAPLRKTYVLKLYVAGNTPNSVRALKTLNNILEKEFKGVYALKVIDVLKNPQLAEEDKILATPTLAKVLPPPVRRIIGDLSNREKVLIGLDLLYEEIGDQAEDDLGLE\"\n",
    "mask_positions = [67, 82, 83]\n",
    "mask_protein_sequence(protein_sequence, mask_positions)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[-0.8849828243255615, -0.1657564640045166, 0.05828452110290527]\n",
      "-0.9924547672271729\n"
     ]
    }
   ],
   "source": [
    "# You can add the following function in your local environment:\n",
    "\n",
    "def compare_log_probs(protein_sequence: str, mask_positions: list, vocab_pairs: list):\n",
    "    # Get the log probabilities for the masked positions\n",
    "    log_probabilities_dicts = mask_protein_sequence(protein_sequence, mask_positions)\n",
    "    \n",
    "    # Compute the difference in log probabilities for the specified vocabulary pairs\n",
    "    log_prob_diffs = []\n",
    "    for i, (wt, mt) in enumerate(vocab_pairs):\n",
    "        log_prob_diff = log_probabilities_dicts[i][mt] - log_probabilities_dicts[i][wt]\n",
    "        log_prob_diffs.append(log_prob_diff)\n",
    "\n",
    "    return log_prob_diffs\n",
    "\n",
    "# Test the function with a protein sequence, mask positions, and vocabulary pairs\n",
    "protein_sequence = \"MAPLRKTYVLKLYVAGNTPNSVRALKTLNNILEKEFKGVYALKVIDVLKNPQLAEEDKILATPTLAKVLPPPVRRIIGDLSNREKVLIGLDLLYEEIGDQAEDDLGLE\"\n",
    "mask_positions = [69, 84, 85]\n",
    "vocab_pairs = [('L', 'R'), ('E', 'D'), ('K', 'A')]\n",
    "\n",
    "print(compare_log_probs(protein_sequence, mask_positions, vocab_pairs))\n",
    "print(sum(compare_log_probs(protein_sequence, mask_positions, vocab_pairs)))\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Computing Logits with ESM-2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'<cls>': -13.275341033935547,\n",
       "  '<pad>': -23.676189422607422,\n",
       "  '<eos>': -12.74148941040039,\n",
       "  '<unk>': -23.687597274780273,\n",
       "  'L': 0.9151788949966431,\n",
       "  'A': 0.6575865745544434,\n",
       "  'G': 1.0361005067825317,\n",
       "  'V': 0.9817789793014526,\n",
       "  'S': 0.2650868892669678,\n",
       "  'E': 0.4443625509738922,\n",
       "  'R': 0.23581159114837646,\n",
       "  'T': 0.2743259072303772,\n",
       "  'I': 0.8414493799209595,\n",
       "  'D': -0.01054537296295166,\n",
       "  'P': 0.018591076135635376,\n",
       "  'K': 1.0111300945281982,\n",
       "  'Q': -0.015097111463546753,\n",
       "  'N': -0.25255557894706726,\n",
       "  'F': 0.16846412420272827,\n",
       "  'Y': 0.0946081280708313,\n",
       "  'M': -0.4992474913597107,\n",
       "  'H': -0.6628150939941406,\n",
       "  'W': -1.070947289466858,\n",
       "  'C': -1.322406530380249,\n",
       "  'X': -8.58376693725586,\n",
       "  'B': -11.722338676452637,\n",
       "  'U': -12.004006385803223,\n",
       "  'Z': -12.56307315826416,\n",
       "  'O': -15.926107406616211,\n",
       "  '.': -16.30080223083496,\n",
       "  '-': -16.292280197143555,\n",
       "  '<null_1>': -16.336050033569336,\n",
       "  '<mask>': -23.67510223388672},\n",
       " {'<cls>': -11.88650894165039,\n",
       "  '<pad>': -21.3856201171875,\n",
       "  '<eos>': -12.544013023376465,\n",
       "  '<unk>': -21.37813377380371,\n",
       "  'L': 0.8959502577781677,\n",
       "  'A': 0.5805749893188477,\n",
       "  'G': 1.3577840328216553,\n",
       "  'V': 0.549018144607544,\n",
       "  'S': 1.0946217775344849,\n",
       "  'E': 1.2901877164840698,\n",
       "  'R': 0.49993667006492615,\n",
       "  'T': 0.6999329924583435,\n",
       "  'I': 0.33176660537719727,\n",
       "  'D': 1.4841411113739014,\n",
       "  'P': 1.1987035274505615,\n",
       "  'K': 0.9945673942565918,\n",
       "  'Q': 0.256401002407074,\n",
       "  'N': 0.5958371162414551,\n",
       "  'F': 0.17739473283290863,\n",
       "  'Y': -0.025778524577617645,\n",
       "  'M': -0.7801998853683472,\n",
       "  'H': -0.4561980366706848,\n",
       "  'W': -1.1411895751953125,\n",
       "  'C': -1.8740789890289307,\n",
       "  'X': -6.988378524780273,\n",
       "  'B': -11.527253150939941,\n",
       "  'U': -11.930501937866211,\n",
       "  'Z': -12.505486488342285,\n",
       "  'O': -15.915705680847168,\n",
       "  '.': -16.279033660888672,\n",
       "  '-': -16.230012893676758,\n",
       "  '<null_1>': -16.28597068786621,\n",
       "  '<mask>': -21.387060165405273},\n",
       " {'<cls>': -12.23246955871582,\n",
       "  '<pad>': -21.512222290039062,\n",
       "  '<eos>': -12.488065719604492,\n",
       "  '<unk>': -21.505569458007812,\n",
       "  'L': 1.2117537260055542,\n",
       "  'A': 0.7499632835388184,\n",
       "  'G': 0.18532384932041168,\n",
       "  'V': 0.7479053735733032,\n",
       "  'S': 0.3944908380508423,\n",
       "  'E': 0.735677182674408,\n",
       "  'R': 0.600346028804779,\n",
       "  'T': 0.08398042619228363,\n",
       "  'I': 0.7027195692062378,\n",
       "  'D': 0.5966317653656006,\n",
       "  'P': 0.8382022380828857,\n",
       "  'K': 0.883489727973938,\n",
       "  'Q': -0.06672608852386475,\n",
       "  'N': -0.09935729950666428,\n",
       "  'F': 0.2884642481803894,\n",
       "  'Y': -0.0896071344614029,\n",
       "  'M': -0.31988540291786194,\n",
       "  'H': -0.8015704154968262,\n",
       "  'W': -0.9554362297058105,\n",
       "  'C': -1.9833405017852783,\n",
       "  'X': -7.382155418395996,\n",
       "  'B': -11.495302200317383,\n",
       "  'U': -11.867561340332031,\n",
       "  'Z': -12.426624298095703,\n",
       "  'O': -15.824283599853516,\n",
       "  '.': -16.15877914428711,\n",
       "  '-': -16.150896072387695,\n",
       "  '<null_1>': -16.178699493408203,\n",
       "  '<mask>': -21.511486053466797}]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from transformers import AutoTokenizer, EsmForMaskedLM\n",
    "import torch\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"facebook/esm2_t6_8M_UR50D\")\n",
    "model = EsmForMaskedLM.from_pretrained(\"facebook/esm2_t6_8M_UR50D\")\n",
    "\n",
    "def mask_protein_sequence_logits(protein_sequence: str, mask_positions: list):\n",
    "    # Tokenize the protein sequence\n",
    "    inputs = tokenizer(protein_sequence, return_tensors=\"pt\")\n",
    "\n",
    "    # Mask the specified positions\n",
    "    for pos in mask_positions:\n",
    "        inputs[\"input_ids\"][0][pos] = tokenizer.mask_token_id\n",
    "\n",
    "    # Get the logits\n",
    "    with torch.no_grad():\n",
    "        logits = model(**inputs).logits\n",
    "\n",
    "    # Get the logits for the masked positions\n",
    "    mask_logits = logits[0, mask_positions]\n",
    "\n",
    "    # Convert the logits to a dictionary with vocabulary elements as keys\n",
    "    vocab = tokenizer.get_vocab()\n",
    "    vocab = {v: k for k, v in vocab.items()}  # reverse the key-value pairs in the vocab\n",
    "    \n",
    "    logits_dicts = []\n",
    "    for logit in mask_logits:\n",
    "        logits_dict = {vocab[i]: logit[i].item() for i in range(len(vocab))}\n",
    "        logits_dicts.append(logits_dict)\n",
    "\n",
    "    return logits_dicts\n",
    "\n",
    "# Test the function with a protein sequence and mask positions\n",
    "protein_sequence = \"MAPLRKTYVLKLYVAGNTPNSVRALKTLNNILEKEFKGVYALKVIDVLKNPQLAEEDKILATPTLAKVLPPPVRRIIGDLSNREKVLIGLDLLYEEIGDQAEDDLGLE\"\n",
    "mask_positions = [67, 82, 83]\n",
    "mask_protein_sequence_logits(protein_sequence, mask_positions)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Computing Probabilities with ESM-2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'<cls>': 6.083488557351302e-08,\n",
       "  '<pad>': 1.8497865999361762e-12,\n",
       "  '<eos>': 1.0375320869115967e-07,\n",
       "  '<unk>': 1.8288044689729333e-12,\n",
       "  'L': 0.08851505070924759,\n",
       "  'A': 0.06841419637203217,\n",
       "  'G': 0.09989246726036072,\n",
       "  'V': 0.09461090713739395,\n",
       "  'S': 0.04620466008782387,\n",
       "  'E': 0.055276963859796524,\n",
       "  'R': 0.04487161338329315,\n",
       "  'T': 0.0466335266828537,\n",
       "  'I': 0.08222366869449615,\n",
       "  'D': 0.03507358580827713,\n",
       "  'P': 0.036110538989305496,\n",
       "  'K': 0.09742899239063263,\n",
       "  'Q': 0.034914303570985794,\n",
       "  'N': 0.02753445692360401,\n",
       "  'F': 0.041949138045310974,\n",
       "  'Y': 0.03896258771419525,\n",
       "  'M': 0.02151491306722164,\n",
       "  'H': 0.018268506973981857,\n",
       "  'W': 0.01214656513184309,\n",
       "  'C': 0.009445960633456707,\n",
       "  'X': 6.632502390857553e-06,\n",
       "  'B': 2.874836013688764e-07,\n",
       "  'U': 2.169133779261756e-07,\n",
       "  'Z': 1.2401856963606406e-07,\n",
       "  'O': 4.294765876267093e-09,\n",
       "  '.': 2.9526474598640107e-09,\n",
       "  '-': 2.9779174681721088e-09,\n",
       "  '<null_1>': 2.850385927288812e-09,\n",
       "  '<mask>': 1.8517988791683093e-12},\n",
       " {'<cls>': 1.7736486768171744e-07,\n",
       "  '<pad>': 1.3287902896064185e-11,\n",
       "  '<eos>': 9.190035399342378e-08,\n",
       "  '<unk>': 1.338775357934141e-11,\n",
       "  'L': 0.06312758475542068,\n",
       "  'A': 0.04605252668261528,\n",
       "  'G': 0.10018230974674225,\n",
       "  'V': 0.044621944427490234,\n",
       "  'S': 0.07700184732675552,\n",
       "  'E': 0.09363416582345963,\n",
       "  'R': 0.04248471185564995,\n",
       "  'T': 0.05189075320959091,\n",
       "  'I': 0.03590850904583931,\n",
       "  'D': 0.11367559432983398,\n",
       "  'P': 0.08544827252626419,\n",
       "  'K': 0.06967036426067352,\n",
       "  'Q': 0.0333017073571682,\n",
       "  'N': 0.04676077514886856,\n",
       "  'F': 0.03077191300690174,\n",
       "  'Y': 0.025114091113209724,\n",
       "  'M': 0.011810722760856152,\n",
       "  'H': 0.016330093145370483,\n",
       "  'W': 0.008231909945607185,\n",
       "  'C': 0.003955585416406393,\n",
       "  'X': 2.3773800421622582e-05,\n",
       "  'B': 2.5403315362382273e-07,\n",
       "  'U': 1.6973119443264295e-07,\n",
       "  'Z': 9.551004609420488e-08,\n",
       "  'O': 3.1550750900777302e-09,\n",
       "  '.': 2.1939097205603275e-09,\n",
       "  '-': 2.3041364372033968e-09,\n",
       "  '<null_1>': 2.178743185865528e-09,\n",
       "  '<mask>': 1.3268781039188493e-11},\n",
       " {'<cls>': 1.6436607097602973e-07,\n",
       "  '<pad>': 1.5334400416122662e-11,\n",
       "  '<eos>': 1.272940721719351e-07,\n",
       "  '<unk>': 1.5436757774822674e-11,\n",
       "  'L': 0.11338724941015244,\n",
       "  'A': 0.07145147025585175,\n",
       "  'G': 0.04062481224536896,\n",
       "  'V': 0.07130458205938339,\n",
       "  'S': 0.050076212733983994,\n",
       "  'E': 0.0704379677772522,\n",
       "  'R': 0.06152239069342613,\n",
       "  'T': 0.0367095023393631,\n",
       "  'I': 0.06815433502197266,\n",
       "  'D': 0.0612943097949028,\n",
       "  'P': 0.07804280519485474,\n",
       "  'K': 0.08165842294692993,\n",
       "  'Q': 0.03157384693622589,\n",
       "  'N': 0.03056018427014351,\n",
       "  'F': 0.045038580894470215,\n",
       "  'Y': 0.030859608203172684,\n",
       "  'M': 0.024512173607945442,\n",
       "  'H': 0.015142186544835567,\n",
       "  'W': 0.012982714921236038,\n",
       "  'C': 0.004644644446671009,\n",
       "  'X': 2.1002782887080684e-05,\n",
       "  'B': 3.4352621014477336e-07,\n",
       "  'U': 2.3674989790833934e-07,\n",
       "  'Z': 1.353604659470875e-07,\n",
       "  'O': 4.528007746529283e-09,\n",
       "  '.': 3.240690826800119e-09,\n",
       "  '-': 3.2663383109365896e-09,\n",
       "  '<null_1>': 3.176773732960214e-09,\n",
       "  '<mask>': 1.534569346595127e-11}]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from transformers import AutoTokenizer, EsmForMaskedLM\n",
    "import torch\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"facebook/esm2_t6_8M_UR50D\")\n",
    "model = EsmForMaskedLM.from_pretrained(\"facebook/esm2_t6_8M_UR50D\")\n",
    "\n",
    "def mask_protein_sequence_probs(protein_sequence: str, mask_positions: list):\n",
    "    # Tokenize the protein sequence\n",
    "    inputs = tokenizer(protein_sequence, return_tensors=\"pt\")\n",
    "\n",
    "    # Mask the specified positions\n",
    "    for pos in mask_positions:\n",
    "        inputs[\"input_ids\"][0][pos] = tokenizer.mask_token_id\n",
    "\n",
    "    # Get the logits\n",
    "    with torch.no_grad():\n",
    "        logits = model(**inputs).logits\n",
    "\n",
    "    # Get the logits for the masked positions\n",
    "    mask_logits = logits[0, mask_positions]\n",
    "\n",
    "    # Apply softmax to convert logits to probabilities\n",
    "    probabilities = torch.nn.functional.softmax(mask_logits, dim=-1)\n",
    "\n",
    "    # Convert the probabilities to a dictionary with vocabulary elements as keys\n",
    "    vocab = tokenizer.get_vocab()\n",
    "    vocab = {v: k for k, v in vocab.items()}  # reverse the key-value pairs in the vocab\n",
    "    \n",
    "    probabilities_dicts = []\n",
    "    for probs in probabilities:\n",
    "        probs_dict = {vocab[i]: prob.item() for i, prob in enumerate(probs)}\n",
    "        probabilities_dicts.append(probs_dict)\n",
    "\n",
    "    return probabilities_dicts\n",
    "\n",
    "# Test the function with a protein sequence and mask positions\n",
    "protein_sequence = \"MAPLRKTYVLKLYVAGNTPNSVRALKTLNNILEKEFKGVYALKVIDVLKNPQLAEEDKILATPTLAKVLPPPVRRIIGDLSNREKVLIGLDLLYEEIGDQAEDDLGLE\"\n",
    "mask_positions = [67, 82, 83]\n",
    "mask_protein_sequence_probs(protein_sequence, mask_positions)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "esm2",
   "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.11.3"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}