{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# [Tutorial 2](https://github.com/evolutionaryscale/esm/tree/main/cookbook/tutorials): Embedding with ESM C\n",
    "\n",
    "In this notebook we will see how to embed a batch of sequences using ESM C, as well as explore its different layers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'default (Python 3.11.11)' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: '/home/wxi/wt4/.pixi/envs/default/bin/python -m pip install ipykernel -U --force-reinstall'"
     ]
    }
   ],
   "source": [
    "# Install esm and other dependencies\n",
    "! pip install esm\n",
    "! pip install matplotlib"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Set up Forge client for ESM C"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Grab a token from [the Forge console](https://forge.evolutionaryscale.ai/console) and add it below. Note that your token is like a password for your account and you should take care to protect it. For this reason it is recommended to frequently create a new token and delete old, unused ones. It is also recommended to paste the token directly into an environment variable or use a utility like `getpass` as shown below so tokens are not accidentally shared or checked into code repositories."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from getpass import getpass\n",
    "\n",
    "token = getpass(\"Token from Forge console: \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from esm.sdk import client\n",
    "\n",
    "model = client(\n",
    "    model=\"esmc-300m-2024-12\", url=\"https://forge.evolutionaryscale.ai\", token=token\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Set up utilities for embedding sequences\n",
    "\n",
    "Since we're embedding more than a few sequences, we're going to use a threaded async call to Forge and let Forge take care of batching and parallelization on the backend."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from concurrent.futures import ThreadPoolExecutor\n",
    "from typing import Sequence\n",
    "\n",
    "from esm.sdk.api import (\n",
    "    ESM3InferenceClient,\n",
    "    ESMProtein,\n",
    "    ESMProteinError,\n",
    "    LogitsConfig,\n",
    "    LogitsOutput,\n",
    "    ProteinType,\n",
    ")\n",
    "\n",
    "EMBEDDING_CONFIG = LogitsConfig(\n",
    "    sequence=True, return_embeddings=True, return_hidden_states=True\n",
    ")\n",
    "\n",
    "\n",
    "def embed_sequence(model: ESM3InferenceClient, sequence: str) -> LogitsOutput:\n",
    "    protein = ESMProtein(sequence=sequence)\n",
    "    protein_tensor = model.encode(protein)\n",
    "    output = model.logits(protein_tensor, EMBEDDING_CONFIG)\n",
    "    return output\n",
    "\n",
    "\n",
    "def batch_embed(\n",
    "    model: ESM3InferenceClient, inputs: Sequence[ProteinType]\n",
    ") -> Sequence[LogitsOutput]:\n",
    "    \"\"\"Forge supports auto-batching. So batch_embed() is as simple as running a collection\n",
    "    of embed calls in parallel using asyncio.\n",
    "    \"\"\"\n",
    "    with ThreadPoolExecutor() as executor:\n",
    "        futures = [\n",
    "            executor.submit(embed_sequence, model, protein) for protein in inputs\n",
    "        ]\n",
    "        results = []\n",
    "        for future in futures:\n",
    "            try:\n",
    "                results.append(future.result())\n",
    "            except Exception as e:\n",
    "                results.append(ESMProteinError(500, str(e)))\n",
    "    return results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Requesting a specific hidden layer\n",
    "\n",
    "ESM C 6B's hidden states are really large, so we only allow one specific layer to be requested per API call. This also works for other ESM C models, but it is required for ESM C 6B. \n",
    "Refer to https://forge.evolutionaryscale.ai/console to find the number of hidden layers for each model. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ESMC_6B_EMBEDDING_CONFIG = LogitsConfig(return_hidden_states=True, ith_hidden_layer=55)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load dataset\n",
    "\n",
    "This dataset is taken from Muir, et al. 2024 [\"Evolutionary-Scale Enzymology Enables Biochemical Constant Prediction Across a Multi-Peaked Catalytic Landscape\"](https://doi.org/10.1101/2024.10.23.619915) which explores a model enzyme called Adenylate Kinase (ADK). Adenylate Kinase appears in many different organisms with different structural classes (referred to as its \"lid type\"). We'll embed this set of ADK sequences and see if we can recover known structural classes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2025-01-16 00:34:10--  https://docs.google.com/uc?export=download&id=1SpOkL11MJxIgy99dqufvUNJuCiuhxuyg\n",
      "Resolving docs.google.com (docs.google.com)... 142.251.46.238, 2607:f8b0:4005:811::200e\n",
      "Connecting to docs.google.com (docs.google.com)|142.251.46.238|:443... connected.\n",
      "HTTP request sent, awaiting response... 303 See Other\n",
      "Location: https://drive.usercontent.google.com/download?id=1SpOkL11MJxIgy99dqufvUNJuCiuhxuyg&export=download [following]\n",
      "--2025-01-16 00:34:10--  https://drive.usercontent.google.com/download?id=1SpOkL11MJxIgy99dqufvUNJuCiuhxuyg&export=download\n",
      "Resolving drive.usercontent.google.com (drive.usercontent.google.com)... 142.250.191.33, 2607:f8b0:4005:80f::2001\n",
      "Connecting to drive.usercontent.google.com (drive.usercontent.google.com)|142.250.191.33|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 43132 (42K) [application/octet-stream]\n",
      "Saving to: ‘adk.csv’\n",
      "\n",
      "adk.csv             100%[===================>]  42.12K  --.-KB/s    in 0.02s   \n",
      "\n",
      "2025-01-16 00:34:13 (1.93 MB/s) - ‘adk.csv’ saved [43132/43132]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "!wget --no-check-certificate \"https://docs.google.com/uc?export=download&id=1SpOkL11MJxIgy99dqufvUNJuCiuhxuyg\" -O adk.csv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "\n",
    "adk_path = \"adk.csv\"\n",
    "df = pd.read_csv(adk_path)\n",
    "df = df[[\"org_name\", \"sequence\", \"lid_type\", \"temperature\"]]\n",
    "df = df[df[\"lid_type\"] != \"other\"]  # drop one structural class for simplicity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Retrying... Attempt 1 after 1.0s due to: (502, 'Failure in logits: {\"status\":\"error\",\"message\":\"Model unavailable - please retry\"}')\n",
      "Retrying... Attempt 1 after 1.0s due to: (502, 'Failure in logits: {\"status\":\"error\",\"message\":\"Model unavailable - please retry\"}')\n",
      "Retrying... Attempt 1 after 1.0s due to: (502, 'Failure in encode: {\"status\":\"error\",\"message\":\"Model unavailable - please retry\"}')\n",
      "Retrying... Attempt 1 after 1.0s due to: (429, 'Failure in logits: {\"status\":\"error\",\"message\":\"You have exceeded your usage cap of 50 requests during the last minute with 50 requests for esmc-300m-2024-12 and logits.\"}')\n",
      "Retrying... Attempt 1 after 1.0s due to: (429, 'Failure in logits: {\"status\":\"error\",\"message\":\"You have exceeded your usage cap of 50 requests during the last minute with 50 requests for esmc-300m-2024-12 and logits.\"}')\n",
      "Retrying... Attempt 1 after 1.0s due to: (429, 'Failure in logits: {\"status\":\"error\",\"message\":\"You have exceeded your usage cap of 50 requests during the last minute with 50 requests for esmc-300m-2024-12 and logits.\"}')\n",
      "Retrying... Attempt 1 after 1.0s due to: (429, 'Failure in logits: {\"status\":\"error\",\"message\":\"You have exceeded your usage cap of 50 requests during the last minute with 50 requests for esmc-300m-2024-12 and logits.\"}')\n",
      "Retrying... Attempt 1 after 1.0s due to: (429, 'Failure in logits: {\"status\":\"error\",\"message\":\"You have exceeded your usage cap of 50 requests during the last minute with 50 requests for esmc-300m-2024-12 and logits.\"}')\n",
      "Retrying... Attempt 2 after 2.0s due to: (429, 'Failure in logits: {\"status\":\"error\",\"message\":\"You have exceeded your usage cap of 50 requests during the last minute with 50 requests for esmc-300m-2024-12 and logits.\"}')\n",
      "Retrying... Attempt 2 after 2.0s due to: (429, 'Failure in logits: {\"status\":\"error\",\"message\":\"You have exceeded your usage cap of 50 requests during the last minute with 50 requests for esmc-300m-2024-12 and logits.\"}')\n",
      "Retrying... Attempt 2 after 2.0s due to: (429, 'Failure in logits: {\"status\":\"error\",\"message\":\"You have exceeded your usage cap of 50 requests during the last minute with 50 requests for esmc-300m-2024-12 and logits.\"}')\n",
      "Retrying... Attempt 2 after 2.0s due to: (429, 'Failure in logits: {\"status\":\"error\",\"message\":\"You have exceeded your usage cap of 50 requests during the last minute with 50 requests for esmc-300m-2024-12 and logits.\"}')\n",
      "Retrying... Attempt 2 after 2.0s due to: (429, 'Failure in logits: {\"status\":\"error\",\"message\":\"You have exceeded your usage cap of 50 requests during the last minute with 50 requests for esmc-300m-2024-12 and logits.\"}')\n",
      "Retrying... Attempt 3 after 4.0s due to: (429, 'Failure in logits: {\"status\":\"error\",\"message\":\"You have exceeded your usage cap of 50 requests during the last minute with 50 requests for esmc-300m-2024-12 and logits.\"}')\n",
      "Retrying... Attempt 3 after 4.0s due to: (429, 'Failure in logits: {\"status\":\"error\",\"message\":\"You have exceeded your usage cap of 50 requests during the last minute with 50 requests for esmc-300m-2024-12 and logits.\"}')\n",
      "Retrying... Attempt 3 after 4.0s due to: (429, 'Failure in logits: {\"status\":\"error\",\"message\":\"You have exceeded your usage cap of 50 requests during the last minute with 50 requests for esmc-300m-2024-12 and logits.\"}')\n",
      "Retrying... Attempt 3 after 4.0s due to: (429, 'Failure in logits: {\"status\":\"error\",\"message\":\"You have exceeded your usage cap of 50 requests during the last minute with 50 requests for esmc-300m-2024-12 and logits.\"}')\n",
      "Retrying... Attempt 1 after 1.0s due to: (429, 'Failure in encode: {\"status\":\"error\",\"message\":\"You have exceeded your usage cap of 50 requests during the last minute with 52 requests for esmc-300m-2024-12 and encode.\"}')\n",
      "Retrying... Attempt 1 after 1.0s due to: (429, 'Failure in encode: {\"status\":\"error\",\"message\":\"You have exceeded your usage cap of 50 requests during the last minute with 52 requests for esmc-300m-2024-12 and encode.\"}')\n",
      "Retrying... Attempt 1 after 1.0s due to: (429, 'Failure in encode: {\"status\":\"error\",\"message\":\"You have exceeded your usage cap of 50 requests during the last minute with 54 requests for esmc-300m-2024-12 and encode.\"}')\n",
      "Retrying... Attempt 1 after 1.0s due to: (429, 'Failure in encode: {\"status\":\"error\",\"message\":\"You have exceeded your usage cap of 50 requests during the last minute with 54 requests for esmc-300m-2024-12 and encode.\"}')\n",
      "Retrying... Attempt 1 after 1.0s due to: (429, 'Failure in logits: {\"status\":\"error\",\"message\":\"You have exceeded your usage cap of 50 requests during the last minute with 52 requests for esmc-300m-2024-12 and logits.\"}')\n",
      "Retrying... Attempt 1 after 1.0s due to: (429, 'Failure in logits: {\"status\":\"error\",\"message\":\"You have exceeded your usage cap of 50 requests during the last minute with 52 requests for esmc-300m-2024-12 and logits.\"}')\n",
      "Retrying... Attempt 2 after 2.0s due to: (429, 'Failure in encode: {\"status\":\"error\",\"message\":\"You have exceeded your usage cap of 50 requests during the last minute with 54 requests for esmc-300m-2024-12 and encode.\"}')\n",
      "Retrying... Attempt 2 after 2.0s due to: (429, 'Failure in encode: {\"status\":\"error\",\"message\":\"You have exceeded your usage cap of 50 requests during the last minute with 54 requests for esmc-300m-2024-12 and encode.\"}')\n",
      "Retrying... Attempt 1 after 1.0s due to: (429, 'Failure in logits: {\"status\":\"error\",\"message\":\"You have exceeded your usage cap of 50 requests during the last minute with 52 requests for esmc-300m-2024-12 and logits.\"}')\n",
      "Retrying... Attempt 1 after 1.0s due to: (429, 'Failure in encode: {\"status\":\"error\",\"message\":\"You have exceeded your usage cap of 50 requests during the last minute with 54 requests for esmc-300m-2024-12 and encode.\"}')\n",
      "Retrying... Attempt 2 after 2.0s due to: (429, 'Failure in encode: {\"status\":\"error\",\"message\":\"You have exceeded your usage cap of 50 requests during the last minute with 54 requests for esmc-300m-2024-12 and encode.\"}')\n",
      "Retrying... Attempt 1 after 1.0s due to: (429, 'Failure in encode: {\"status\":\"error\",\"message\":\"You have exceeded your usage cap of 50 requests during the last minute with 51 requests for esmc-300m-2024-12 and encode.\"}')\n",
      "Retrying... Attempt 1 after 1.0s due to: (429, 'Failure in encode: {\"status\":\"error\",\"message\":\"You have exceeded your usage cap of 50 requests during the last minute with 51 requests for esmc-300m-2024-12 and encode.\"}')\n",
      "Retrying... Attempt 1 after 1.0s due to: (429, 'Failure in encode: {\"status\":\"error\",\"message\":\"You have exceeded your usage cap of 50 requests during the last minute with 51 requests for esmc-300m-2024-12 and encode.\"}')\n",
      "Retrying... Attempt 1 after 1.0s due to: (429, 'Failure in encode: {\"status\":\"error\",\"message\":\"You have exceeded your usage cap of 50 requests during the last minute with 51 requests for esmc-300m-2024-12 and encode.\"}')\n",
      "Retrying... Attempt 2 after 2.0s due to: (429, 'Failure in encode: {\"status\":\"error\",\"message\":\"You have exceeded your usage cap of 50 requests during the last minute with 51 requests for esmc-300m-2024-12 and encode.\"}')\n",
      "Retrying... Attempt 2 after 2.0s due to: (502, 'Failure in logits: {\"status\":\"error\",\"message\":\"Model unavailable - please retry\"}')\n",
      "Retrying... Attempt 1 after 1.0s due to: (429, 'Failure in encode: {\"status\":\"error\",\"message\":\"You have exceeded your usage cap of 50 requests during the last minute with 50 requests for esmc-300m-2024-12 and encode.\"}')\n",
      "Retrying... Attempt 1 after 1.0s due to: (429, 'Failure in encode: {\"status\":\"error\",\"message\":\"You have exceeded your usage cap of 50 requests during the last minute with 50 requests for esmc-300m-2024-12 and encode.\"}')\n",
      "Retrying... Attempt 1 after 1.0s due to: (429, 'Failure in encode: {\"status\":\"error\",\"message\":\"You have exceeded your usage cap of 50 requests during the last minute with 50 requests for esmc-300m-2024-12 and encode.\"}')\n",
      "Retrying... Attempt 2 after 2.0s due to: (429, 'Failure in encode: {\"status\":\"error\",\"message\":\"You have exceeded your usage cap of 50 requests during the last minute with 50 requests for esmc-300m-2024-12 and encode.\"}')\n",
      "Retrying... Attempt 1 after 1.0s due to: (429, 'Failure in encode: {\"status\":\"error\",\"message\":\"You have exceeded your usage cap of 50 requests during the last minute with 50 requests for esmc-300m-2024-12 and encode.\"}')\n",
      "Retrying... Attempt 2 after 2.0s due to: (429, 'Failure in encode: {\"status\":\"error\",\"message\":\"You have exceeded your usage cap of 50 requests during the last minute with 50 requests for esmc-300m-2024-12 and encode.\"}')\n",
      "Retrying... Attempt 1 after 1.0s due to: (429, 'Failure in encode: {\"status\":\"error\",\"message\":\"You have exceeded your usage cap of 50 requests during the last minute with 50 requests for esmc-300m-2024-12 and encode.\"}')\n",
      "Retrying... Attempt 1 after 1.0s due to: (502, 'Failure in logits: {\"status\":\"error\",\"message\":\"Model unavailable - please retry\"}')\n",
      "Retrying... Attempt 1 after 1.0s due to: (502, 'Failure in logits: {\"status\":\"error\",\"message\":\"Model unavailable - please retry\"}')Retrying... Attempt 1 after 1.0s due to: (502, 'Failure in encode: {\"status\":\"error\",\"message\":\"Model unavailable - please retry\"}')\n",
      "\n",
      "Retrying... Attempt 2 after 2.0s due to: (502, 'Failure in encode: {\"status\":\"error\",\"message\":\"Model unavailable - please retry\"}')\n"
     ]
    }
   ],
   "source": [
    "# You may see some error messages due to rate limits on each Forge account,\n",
    "# but this will retry until the embedding job is complete\n",
    "# This may take a few minutes to run\n",
    "outputs = batch_embed(model, df[\"sequence\"].tolist())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "embedding shape [num_layers, hidden_size]: torch.Size([31, 960])\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "# we'll summarize the embeddings using their mean across the sequence dimension\n",
    "# which allows us to compare embeddings for sequences of different lengths\n",
    "all_mean_embeddings = [\n",
    "    torch.mean(output.hidden_states, dim=-2).squeeze() for output in outputs\n",
    "]\n",
    "\n",
    "# now we have a list of tensors of [num_layers, hidden_size]\n",
    "print(\"embedding shape [num_layers, hidden_size]:\", all_mean_embeddings[0].shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Examine the performance of different layer embeddings\n",
    "\n",
    "For this example, we're going to use PCA to visualize whether the embeddings separate our proteins by their structural class. To assess the quality of our PCA, we fit a K means classifier with three clusters, corresponding to the three structural classes of our enzyme, and compute the [rand index](https://en.wikipedia.org/wiki/Rand_index), a measure of the quality of the clustering."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.cluster import KMeans\n",
    "from sklearn.decomposition import PCA\n",
    "from sklearn.metrics import adjusted_rand_score\n",
    "\n",
    "N_KMEANS_CLUSTERS = 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_embeddings_at_layer(all_mean_embeddings: torch.Tensor, layer_idx: int):\n",
    "    stacked_mean_embeddings = torch.stack(\n",
    "        [embedding[layer_idx, :] for embedding in all_mean_embeddings]\n",
    "    ).numpy()\n",
    "\n",
    "    # project all the embeddings to 2D using PCA\n",
    "    pca = PCA(n_components=2)\n",
    "    pca.fit(stacked_mean_embeddings)\n",
    "    projected_mean_embeddings = pca.transform(stacked_mean_embeddings)\n",
    "\n",
    "    # compute kmeans purity as a measure of how good the clustering is\n",
    "    kmeans = KMeans(n_clusters=N_KMEANS_CLUSTERS, random_state=0).fit(\n",
    "        projected_mean_embeddings\n",
    "    )\n",
    "    rand_index = adjusted_rand_score(df[\"lid_type\"], kmeans.labels_)\n",
    "\n",
    "    # plot the clusters\n",
    "    plt.figure(figsize=(4, 4))\n",
    "    sns.scatterplot(\n",
    "        x=projected_mean_embeddings[:, 0],\n",
    "        y=projected_mean_embeddings[:, 1],\n",
    "        hue=df[\"lid_type\"],\n",
    "    )\n",
    "    plt.title(\n",
    "        f\"PCA of mean embeddings at layer {layer_idx}.\\nRand index: {rand_index:.2f}\"\n",
    "    )\n",
    "    plt.xlabel(\"PC 1\")\n",
    "    plt.ylabel(\"PC 2\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 61008 characters]",
      "text/plain": [
       "<Figure size 400x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 50636 characters]",
      "text/plain": [
       "<Figure size 400x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_embeddings_at_layer(all_mean_embeddings, layer_idx=30)\n",
    "plot_embeddings_at_layer(all_mean_embeddings, layer_idx=12)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see that the top principal components of layer 12 separate structural classes better than that of layer 30. Embed away! And keep in mind that different layers may be better or worse for your particular use-case."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "default",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
