CaLM.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "54a16c13-4807-46c9-adec-fd3bc6889de7",
   "metadata": {},
   "source": [
    "# CaLM\n",
    "## The Codon adaptation Language Model\n",
    "\n",
    "Implementation of the [Codon adaptation Language Model (CaLM)](https://www.biorxiv.org/content/10.1101/2022.12.15.519894v1.abstract), a large protein language model has been trained on coding DNA, instead of amino acids. Embeddings from CaLM can be used to build automatic features for downstream machine learning tasks — like predicting a protein's function, the species it belongs to, or others. This Colab Notebook enables automatic embedding of an arbitrary number of proteins using the GPUs available at Google Colaboratory.  \n",
    "\n",
    "CaLM is a research tool developed by Carlos Outeiral, at the Oxford Protein Informatics Group."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d2aa439-93ce-4ec9-a8d6-726ebd9730c7",
   "metadata": {},
   "source": [
    "## Set up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d6d9b83f-60a2-4b8b-9d3e-a2f133610ab0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Looking in indexes: https://pypi.org/simple, https://download.pytorch.org/whl/cpu\n",
      "Requirement already satisfied: torch in /opt/conda/lib/python3.9/site-packages (1.13.1+cpu)\n",
      "Requirement already satisfied: torchvision in /opt/conda/lib/python3.9/site-packages (0.14.1+cpu)\n",
      "Requirement already satisfied: torchaudio in /opt/conda/lib/python3.9/site-packages (0.13.1+cpu)\n",
      "Requirement already satisfied: typing-extensions in /opt/conda/lib/python3.9/site-packages (from torch) (4.4.0)\n",
      "Requirement already satisfied: requests in /opt/conda/lib/python3.9/site-packages (from torchvision) (2.27.1)\n",
      "Requirement already satisfied: pillow!=8.3.*,>=5.3.0 in /opt/conda/lib/python3.9/site-packages (from torchvision) (9.4.0)\n",
      "Requirement already satisfied: numpy in /opt/conda/lib/python3.9/site-packages (from torchvision) (1.24.1)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /opt/conda/lib/python3.9/site-packages (from requests->torchvision) (2022.12.7)\n",
      "Requirement already satisfied: charset-normalizer~=2.0.0 in /opt/conda/lib/python3.9/site-packages (from requests->torchvision) (2.0.4)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /opt/conda/lib/python3.9/site-packages (from requests->torchvision) (1.26.8)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /opt/conda/lib/python3.9/site-packages (from requests->torchvision) (3.3)\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\n",
      "Requirement already satisfied: einops in /opt/conda/lib/python3.9/site-packages (0.6.0)\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\n",
      "Requirement already satisfied: rotary_embedding_torch in /opt/conda/lib/python3.9/site-packages (0.2.1)\n",
      "Requirement already satisfied: einops>=0.3 in /opt/conda/lib/python3.9/site-packages (from rotary_embedding_torch) (0.6.0)\n",
      "Requirement already satisfied: torch>=1.6 in /opt/conda/lib/python3.9/site-packages (from rotary_embedding_torch) (1.13.1+cpu)\n",
      "Requirement already satisfied: typing-extensions in /opt/conda/lib/python3.9/site-packages (from torch>=1.6->rotary_embedding_torch) (4.4.0)\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\n",
      "Requirement already satisfied: biopython in /opt/conda/lib/python3.9/site-packages/biopython-1.80-py3.9-linux-x86_64.egg (1.80)\n",
      "Requirement already satisfied: numpy in /opt/conda/lib/python3.9/site-packages (from biopython) (1.24.1)\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\n",
      "running install\n",
      "/opt/conda/lib/python3.9/site-packages/setuptools/command/install.py:34: SetuptoolsDeprecationWarning: setup.py install is deprecated. Use build and pip and other standards-based tools.\n",
      "  warnings.warn(\n",
      "/opt/conda/lib/python3.9/site-packages/setuptools/command/easy_install.py:144: EasyInstallDeprecationWarning: easy_install command is deprecated. Use build and pip and other standards-based tools.\n",
      "  warnings.warn(\n",
      "running bdist_egg\n",
      "running egg_info\n",
      "writing CaLM.egg-info/PKG-INFO\n",
      "writing dependency_links to CaLM.egg-info/dependency_links.txt\n",
      "writing requirements to CaLM.egg-info/requires.txt\n",
      "writing top-level names to CaLM.egg-info/top_level.txt\n",
      "reading manifest file 'CaLM.egg-info/SOURCES.txt'\n",
      "writing manifest file 'CaLM.egg-info/SOURCES.txt'\n",
      "installing library code to build/bdist.linux-x86_64/egg\n",
      "running install_lib\n",
      "running build_py\n",
      "copying calm/pretrained.py -> build/lib/calm\n",
      "creating build/bdist.linux-x86_64/egg\n",
      "creating build/bdist.linux-x86_64/egg/calm\n",
      "copying build/lib/calm/alphabet.py -> build/bdist.linux-x86_64/egg/calm\n",
      "copying build/lib/calm/multihead_attention.py -> build/bdist.linux-x86_64/egg/calm\n",
      "copying build/lib/calm/pipeline.py -> build/bdist.linux-x86_64/egg/calm\n",
      "copying build/lib/calm/dataset.py -> build/bdist.linux-x86_64/egg/calm\n",
      "copying build/lib/calm/__init__.py -> build/bdist.linux-x86_64/egg/calm\n",
      "copying build/lib/calm/model.py -> build/bdist.linux-x86_64/egg/calm\n",
      "copying build/lib/calm/pretrained.py -> build/bdist.linux-x86_64/egg/calm\n",
      "copying build/lib/calm/modules.py -> build/bdist.linux-x86_64/egg/calm\n",
      "copying build/lib/calm/sequence.py -> build/bdist.linux-x86_64/egg/calm\n",
      "byte-compiling build/bdist.linux-x86_64/egg/calm/alphabet.py to alphabet.cpython-39.pyc\n",
      "byte-compiling build/bdist.linux-x86_64/egg/calm/multihead_attention.py to multihead_attention.cpython-39.pyc\n",
      "byte-compiling build/bdist.linux-x86_64/egg/calm/pipeline.py to pipeline.cpython-39.pyc\n",
      "byte-compiling build/bdist.linux-x86_64/egg/calm/dataset.py to dataset.cpython-39.pyc\n",
      "byte-compiling build/bdist.linux-x86_64/egg/calm/__init__.py to __init__.cpython-39.pyc\n",
      "byte-compiling build/bdist.linux-x86_64/egg/calm/model.py to model.cpython-39.pyc\n",
      "byte-compiling build/bdist.linux-x86_64/egg/calm/pretrained.py to pretrained.cpython-39.pyc\n",
      "byte-compiling build/bdist.linux-x86_64/egg/calm/modules.py to modules.cpython-39.pyc\n",
      "byte-compiling build/bdist.linux-x86_64/egg/calm/sequence.py to sequence.cpython-39.pyc\n",
      "creating build/bdist.linux-x86_64/egg/EGG-INFO\n",
      "copying CaLM.egg-info/PKG-INFO -> build/bdist.linux-x86_64/egg/EGG-INFO\n",
      "copying CaLM.egg-info/SOURCES.txt -> build/bdist.linux-x86_64/egg/EGG-INFO\n",
      "copying CaLM.egg-info/dependency_links.txt -> build/bdist.linux-x86_64/egg/EGG-INFO\n",
      "copying CaLM.egg-info/requires.txt -> build/bdist.linux-x86_64/egg/EGG-INFO\n",
      "copying CaLM.egg-info/top_level.txt -> build/bdist.linux-x86_64/egg/EGG-INFO\n",
      "zip_safe flag not set; analyzing archive contents...\n",
      "calm.__pycache__.pretrained.cpython-39: module references __file__\n",
      "creating 'dist/CaLM-0.1.1-py3.9.egg' and adding 'build/bdist.linux-x86_64/egg' to it\n",
      "removing 'build/bdist.linux-x86_64/egg' (and everything under it)\n",
      "Processing CaLM-0.1.1-py3.9.egg\n",
      "removing '/opt/conda/lib/python3.9/site-packages/CaLM-0.1.1-py3.9.egg' (and everything under it)\n",
      "creating /opt/conda/lib/python3.9/site-packages/CaLM-0.1.1-py3.9.egg\n",
      "Extracting CaLM-0.1.1-py3.9.egg to /opt/conda/lib/python3.9/site-packages\n",
      "CaLM 0.1.1 is already the active version in easy-install.pth\n",
      "\n",
      "Installed /opt/conda/lib/python3.9/site-packages/CaLM-0.1.1-py3.9.egg\n",
      "Processing dependencies for CaLM==0.1.1\n",
      "Searching for rotary-embedding-torch==0.2.1\n",
      "Best match: rotary-embedding-torch 0.2.1\n",
      "Adding rotary-embedding-torch 0.2.1 to easy-install.pth file\n",
      "\n",
      "Using /opt/conda/lib/python3.9/site-packages\n",
      "Searching for einops==0.6.0\n",
      "Best match: einops 0.6.0\n",
      "Adding einops 0.6.0 to easy-install.pth file\n",
      "\n",
      "Using /opt/conda/lib/python3.9/site-packages\n",
      "Searching for biopython==1.80\n",
      "Best match: biopython 1.80\n",
      "Processing biopython-1.80-py3.9-linux-x86_64.egg\n",
      "biopython 1.80 is already the active version in easy-install.pth\n",
      "\n",
      "Using /opt/conda/lib/python3.9/site-packages/biopython-1.80-py3.9-linux-x86_64.egg\n",
      "Searching for torch==1.13.1+cpu\n",
      "Best match: torch 1.13.1+cpu\n",
      "Adding torch 1.13.1+cpu to easy-install.pth file\n",
      "Installing convert-caffe2-to-onnx script to /opt/conda/bin\n",
      "Installing convert-onnx-to-caffe2 script to /opt/conda/bin\n",
      "Installing torchrun script to /opt/conda/bin\n",
      "\n",
      "Using /opt/conda/lib/python3.9/site-packages\n",
      "Searching for requests==2.27.1\n",
      "Best match: requests 2.27.1\n",
      "Adding requests 2.27.1 to easy-install.pth file\n",
      "\n",
      "Using /opt/conda/lib/python3.9/site-packages\n",
      "Searching for numpy==1.24.1\n",
      "Best match: numpy 1.24.1\n",
      "Adding numpy 1.24.1 to easy-install.pth file\n",
      "Installing f2py script to /opt/conda/bin\n",
      "Installing f2py3 script to /opt/conda/bin\n",
      "Installing f2py3.9 script to /opt/conda/bin\n",
      "\n",
      "Using /opt/conda/lib/python3.9/site-packages\n",
      "Searching for typing-extensions==4.4.0\n",
      "Best match: typing-extensions 4.4.0\n",
      "Adding typing-extensions 4.4.0 to easy-install.pth file\n",
      "\n",
      "Using /opt/conda/lib/python3.9/site-packages\n",
      "Searching for idna==3.3\n",
      "Best match: idna 3.3\n",
      "Adding idna 3.3 to easy-install.pth file\n",
      "\n",
      "Using /opt/conda/lib/python3.9/site-packages\n",
      "Searching for charset-normalizer==2.0.4\n",
      "Best match: charset-normalizer 2.0.4\n",
      "Adding charset-normalizer 2.0.4 to easy-install.pth file\n",
      "Installing normalizer script to /opt/conda/bin\n",
      "\n",
      "Using /opt/conda/lib/python3.9/site-packages\n",
      "Searching for certifi==2022.12.7\n",
      "Best match: certifi 2022.12.7\n",
      "Adding certifi 2022.12.7 to easy-install.pth file\n",
      "\n",
      "Using /opt/conda/lib/python3.9/site-packages\n",
      "Searching for urllib3==1.26.8\n",
      "Best match: urllib3 1.26.8\n",
      "Adding urllib3 1.26.8 to easy-install.pth file\n",
      "\n",
      "Using /opt/conda/lib/python3.9/site-packages\n",
      "Finished processing dependencies for CaLM==0.1.1\n"
     ]
    }
   ],
   "source": [
    "! pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cpu\n",
    "! pip3 install einops\n",
    "! pip3 install rotary_embedding_torch\n",
    "! pip3 install biopython\n",
    "! python setup.py install"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e88d3928-6f8f-4329-8def-88f41f385abf",
   "metadata": {},
   "source": [
    "## Useful functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9e8f5459-fecb-4493-9872-83889cb06a87",
   "metadata": {},
   "outputs": [],
   "source": [
    "import string\n",
    "\n",
    "\n",
    "def _split_into_codons(seq):\n",
    "  for i in range(0, len(seq), 3):\n",
    "    yield seq[i:i + 3]\n",
    "        \n",
    "def split_into_codons(seq):\n",
    "  \"\"\"Yield successive 3-letter chunks of a string/sequence.\"\"\"\n",
    "  return list(_split_into_codons(seq))\n",
    "\n",
    "def randstr(N=20):\n",
    "  \"\"\"Returns a random alphanumerical string.\"\"\"\n",
    "  return ''.join(random.choice(string.ascii_uppercase + string.digits)\n",
    "    for _ in range(N))\n",
    "  \n",
    "def is_dna(seq):\n",
    "  \"\"\"Checks if a sequence matches the DNA alphabet.\"\"\"\n",
    "  return set(seq.upper()).issubset({'A', 'T', 'C', 'G'})\n",
    "\n",
    "def is_rna(seq):\n",
    "  \"\"\"Checks if a sequence matches the RNA alphabet.\"\"\"\n",
    "  return set(seq.upper()).issubset({'A', 'U', 'C', 'G'})\n",
    "\n",
    "def check_valid_codons(seq):\n",
    "  \"\"\"Checks if a sequence has a number of symbols\n",
    "  that is a multiple of three.\"\"\"\n",
    "  return len(seq) % 3 == 0\n",
    "\n",
    "def check_start_codon(seq):\n",
    "  \"\"\"Checks that the start codon is correct.\"\"\"\n",
    "  if is_dna(seq):\n",
    "    return seq[:3] == 'ATG'\n",
    "  elif is_rna(seq):\n",
    "    return seq[:3] == 'AUG'\n",
    "  else:\n",
    "    return False\n",
    "\n",
    "def check_stop_codon(seq):\n",
    "  \"\"\"Checks that the stop codon is correct.\"\"\"\n",
    "  if is_dna(seq):\n",
    "    return seq[-3:] in ['TGA', 'TAG', 'TAA']\n",
    "  elif is_rna(seq):\n",
    "    return seq[-3:] in ['UGA', 'UAG', 'UAA']\n",
    "  else:\n",
    "    return False\n",
    "\n",
    "def check_no_interstitial_stop_codons(seq):\n",
    "  codons = split_into_codons(seq)\n",
    "  coding_positions = set(codons[:-1])\n",
    "  if is_dna(seq):\n",
    "    return {'TGA', 'TAG', 'TAA'} not in coding_positions\n",
    "  elif is_rna(seq):\n",
    "    return {'UGA', 'UAG', 'UAA'} not in coding_positions\n",
    "  else:\n",
    "    return False"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61da5df9-17cc-44a2-b524-33416db6d9e6",
   "metadata": {},
   "source": [
    "## Embed a sequence\n",
    "\n",
    "The following code exemplifies how to embed a sequence using CaLM. In order to embed a specific cDNA sequence, just modify the value of the `cDNA_sequence` variable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "3a2bdaee-ba79-4206-a8da-987ff482f538",
   "metadata": {},
   "outputs": [],
   "source": [
    "cDNA_sequence = 'ATGGTATAGAGGCATTGA'\n",
    "alphabet = 'DNA'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "201038f0-7702-41f6-a369-96082095398c",
   "metadata": {},
   "outputs": [],
   "source": [
    "if cDNA_sequence == '':\n",
    "  raise ValueError('No cDNA sequence was provided.')\n",
    "if alphabet == 'DNA':\n",
    "  if not is_dna(cDNA_sequence):\n",
    "    raise ValueError('Provided DNA sequence did not pass sanity check.')\n",
    "elif alphabet == 'RNA':\n",
    "  if not is_rna(cDNA_sequence):\n",
    "    raise ValueError('Provided RNA sequence did not pass sanity check.')\n",
    "if not check_valid_codons(cDNA_sequence):\n",
    "  raise ValueError('The number of nucleotides is not a multiple of three.')\n",
    "if not check_start_codon(cDNA_sequence):\n",
    "  raise ValueError('Invalid start codon.')\n",
    "if not check_stop_codon(cDNA_sequence):\n",
    "  raise ValueError('Invalid stop codon.')\n",
    "if not check_no_interstitial_stop_codons(cDNA_sequence):\n",
    "  raise ValueError('Provided sequence contains interstitial stop codons.')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08f96b9f-671a-4227-9611-18e1768203d8",
   "metadata": {},
   "source": [
    "In the paper we have used sequence representations, where the representations of all codons in a protein are averaged. This is the default output of CaLM:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a96b6b42-f7f7-4ff2-a150-8a81d645a30d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Embeddings shape:  torch.Size([1, 768])\n",
      "tensor([[-6.3635e-02,  9.2567e-02,  3.2795e-01, -9.2734e-02,  1.4253e-01,\n",
      "          4.5020e-01,  3.4477e-02,  2.9611e-01, -2.6117e-01, -2.3144e-01,\n",
      "          8.8645e-02, -9.5536e-03, -1.6554e-01, -5.2163e-03, -1.2918e-01,\n",
      "          5.2430e-01,  1.1867e-01,  2.6617e-02,  3.1034e-02, -3.2623e-01,\n",
      "          3.2302e-01,  7.4128e-01,  2.0493e-01,  3.0914e-01,  5.4850e-02,\n",
      "         -5.8367e-02,  2.8206e-01,  9.0826e-02, -3.7607e-01, -5.7403e-02,\n",
      "         -2.5932e-01, -3.9007e-02,  7.0167e-02, -4.0534e-02,  2.1532e-02,\n",
      "          2.3047e-01,  8.3513e-02, -3.9460e-02,  2.4399e-02, -1.3421e-02,\n",
      "         -2.1616e-02,  1.2546e-03, -5.6222e-02, -1.9381e-01, -5.3109e-02,\n",
      "          5.6924e-02, -3.8164e-02,  9.5740e-02, -2.2428e-01,  6.0630e-02,\n",
      "         -4.6556e-01, -1.8238e-01, -3.3394e-02,  1.6678e-02,  1.0147e-01,\n",
      "          4.5168e-01, -4.5880e-02,  1.6808e-01, -8.8289e-02,  3.2918e-02,\n",
      "         -4.5759e-02,  1.2530e-01, -1.2781e-01,  1.8653e-01,  1.9146e-01,\n",
      "          3.4138e-02, -1.6460e-01, -2.4801e-02,  4.7263e-02,  2.9361e-01,\n",
      "          2.6679e-02, -1.5638e-01,  4.6845e-03, -2.7963e-01, -3.2994e-02,\n",
      "         -1.2314e-01,  1.1020e-01, -5.2079e-02,  1.4723e-01,  6.2587e-01,\n",
      "         -2.5697e-03, -1.8606e-01, -1.0766e-01,  6.9701e-02,  2.4899e-01,\n",
      "         -2.9258e-01,  7.9931e-02, -7.5113e-02, -8.8572e-02,  1.2220e-01,\n",
      "         -2.7352e-01, -1.1069e-01, -2.4840e-01,  1.9511e-02, -1.1779e-01,\n",
      "          1.9397e-01, -4.8832e-01,  1.5015e-01,  1.5246e-03,  6.4072e-02,\n",
      "         -4.3058e-02, -7.7309e-02, -4.0519e-02,  4.4803e-02,  7.3168e-02,\n",
      "          2.9387e-01,  7.8205e-02, -4.1183e-02,  1.5855e-02, -1.6530e-01,\n",
      "          1.4072e-01,  3.5963e-03, -2.1548e-01,  1.0451e-01, -1.8126e-01,\n",
      "         -1.2519e-01, -4.6533e-02,  5.6209e-02,  2.0633e-01,  9.0930e-03,\n",
      "          4.0429e-02,  3.4890e-02,  2.8618e-01,  1.6479e-01, -7.1263e-02,\n",
      "         -1.3708e-02,  9.2621e-02, -1.0473e-01, -2.1588e-01, -4.5165e-02,\n",
      "         -1.0330e-01, -1.7761e-01, -1.4028e-01, -2.5426e-01,  5.2739e-02,\n",
      "         -2.4323e-01, -1.9355e-01, -3.4573e-03,  2.5063e-01, -1.2548e-01,\n",
      "          1.0942e-01, -1.7734e-01,  9.1511e-02, -2.7072e-02,  3.9006e-01,\n",
      "          1.0756e-01, -8.0799e-02, -5.8764e-02,  2.5559e-01, -2.0900e-01,\n",
      "         -2.8431e-01,  8.6169e-02,  1.3630e-01,  5.5191e-02,  1.5407e-01,\n",
      "         -6.2874e-02,  1.6068e-01,  1.5486e-01,  6.0477e-02,  1.5401e-01,\n",
      "         -7.3754e-02,  1.5853e-01, -1.0406e-01,  3.8215e-01, -2.4619e-01,\n",
      "          1.6592e-01,  2.3102e-02,  2.6929e-01,  1.7070e-01, -7.2734e-02,\n",
      "          2.0474e-01, -2.1306e-01, -1.8840e-01, -1.6485e-01,  3.6863e-02,\n",
      "         -1.1650e-01, -5.7621e-02,  3.8492e-02,  4.2533e-02, -6.6603e-03,\n",
      "          8.4657e-03, -1.8841e-02, -1.8426e-01,  2.1079e-01,  4.4668e-01,\n",
      "         -1.1725e-01,  2.3318e-01,  2.7303e-02,  1.3215e-01, -1.6148e-01,\n",
      "          1.0169e-01,  2.2254e-02, -7.3922e-02, -3.6615e-01, -4.8721e-02,\n",
      "         -4.1029e-02, -5.2290e-02,  2.0962e-01, -1.5366e-01,  9.4728e-02,\n",
      "          7.2255e-02,  2.0900e-01,  3.5956e-01, -2.0349e-01,  2.9403e-01,\n",
      "          9.1223e-02,  1.4258e-01,  2.2456e-01,  1.3469e-01,  1.5547e-01,\n",
      "         -5.7662e-02,  1.7157e-01, -5.0694e-01, -2.4274e-01,  4.3141e-02,\n",
      "          3.9609e-02, -4.6076e-02, -3.3312e-01, -2.7940e-01,  1.7395e-02,\n",
      "         -1.5324e-01, -4.7143e-02, -5.2218e-02,  2.9709e-01, -1.8466e-01,\n",
      "          1.5983e-02,  2.3126e-01, -1.0341e-01,  1.0249e-01, -1.2571e-01,\n",
      "          3.6654e-03,  8.1768e-02, -2.2741e-01, -7.1914e-02, -3.0221e-01,\n",
      "          2.0988e-02, -5.2731e-02, -8.7555e-03,  4.9642e-02,  1.9268e-01,\n",
      "          1.7470e-01,  5.6471e-03, -6.3470e-02, -4.3939e-01, -1.0744e-01,\n",
      "         -1.0433e-01,  1.3393e-01,  1.3556e-02, -1.9502e-01, -1.7937e-01,\n",
      "         -1.5695e-01,  1.0315e-01, -4.3637e-01, -4.5748e-01,  3.1487e-02,\n",
      "         -6.5000e-02,  1.2256e-01, -1.5674e-01, -2.8490e-01,  8.7936e-02,\n",
      "          6.4411e-02,  1.6045e-01,  2.3017e-02,  2.2646e-01, -1.7998e-01,\n",
      "          1.0508e-01,  2.9600e-01, -5.4279e-01,  3.7999e-01,  3.3001e-02,\n",
      "         -5.0419e-02,  2.4464e-01,  5.7338e-04,  1.4268e-02, -9.0515e-02,\n",
      "         -1.1967e-01, -2.7860e-02, -2.4143e-01,  1.7259e-01,  1.7723e-01,\n",
      "         -1.8084e-01, -5.6170e-02, -5.1971e-02,  6.8723e-02, -1.8080e-01,\n",
      "         -6.0325e-02,  2.1605e-02, -3.3520e-02,  2.4578e-02, -2.3962e-01,\n",
      "          4.1545e-01, -1.2746e-01, -8.0299e-02,  2.5067e-02,  2.0960e-02,\n",
      "         -8.3884e-02,  8.0515e-02, -1.9178e-01, -3.1370e-01,  9.1367e-02,\n",
      "         -1.0046e-01, -3.5596e-02,  1.0703e-01,  1.3427e-01, -1.0074e-02,\n",
      "         -1.1704e-01, -5.1560e-02, -1.9614e-01, -2.0788e-01,  4.7909e-02,\n",
      "          2.9853e-01, -1.9393e-01, -4.8335e-01,  1.4031e-01, -1.8514e-01,\n",
      "          3.2592e-01, -2.4868e-02, -5.7686e-02, -2.6939e-01, -1.2551e-01,\n",
      "          8.6345e-02,  4.8496e-01,  1.3209e-01, -3.3972e-01, -1.4162e-01,\n",
      "         -4.7442e-02, -6.4960e-02, -5.0003e-02,  9.3553e-03,  1.5132e-01,\n",
      "         -4.1151e-02,  2.1611e-01,  1.2316e-01, -2.0018e-01,  5.6060e-02,\n",
      "          7.7024e-02, -9.0570e-02,  2.8947e-01, -2.8218e-02,  1.6329e-02,\n",
      "          2.3487e-01, -3.5280e-02,  4.0298e-01, -3.1017e-01, -7.5724e-02,\n",
      "         -2.7377e-01, -1.7645e-01, -1.9716e-01,  1.4999e-01,  5.0614e-02,\n",
      "         -2.4139e-01, -8.4491e-02, -7.7502e-01,  1.6107e-01, -1.3077e-01,\n",
      "         -5.2211e-02,  1.3107e-01,  1.7110e-01, -2.0435e-01, -5.7051e-02,\n",
      "         -7.0301e-02, -1.3122e-01, -8.7522e-02, -1.3065e-01,  6.1465e-02,\n",
      "         -1.2973e-01, -1.6553e-01, -7.3124e-02, -1.5800e-01,  9.5326e-02,\n",
      "         -1.0746e-01,  6.9956e-02, -1.2881e-01, -2.3991e-02,  2.9143e-01,\n",
      "         -2.8760e-01,  6.3024e-03,  5.1014e-02, -5.8826e-02, -6.4520e-02,\n",
      "         -1.8639e-01, -5.8983e-02, -1.1459e-01,  7.0554e-02,  9.3255e-02,\n",
      "         -7.9684e-03, -8.6842e-02, -9.3317e-02, -2.6913e-01,  1.0956e-01,\n",
      "          2.0494e-01,  1.0202e-01, -1.6011e-01, -4.5220e-02,  2.2970e-01,\n",
      "          5.6257e-02, -1.4488e-01, -7.4258e-03,  2.1213e-01, -7.4599e-02,\n",
      "          5.7717e-02,  1.1219e-01, -2.3243e-02,  3.6814e-02,  1.8650e-01,\n",
      "         -2.7773e-01, -2.0738e-01,  1.7161e-01, -1.1415e-01, -1.1874e-01,\n",
      "         -3.7919e-01,  9.2678e-03,  1.5927e-02,  2.2591e-02, -1.5486e-01,\n",
      "          3.7391e-02, -2.1637e-02,  2.0140e-01, -4.9309e-02,  2.2354e-01,\n",
      "          1.3154e-02, -1.8001e-02, -8.7475e-02, -1.2430e-01, -1.3644e-01,\n",
      "          1.2649e-01, -4.5096e-01,  4.7869e-02, -8.1404e-02,  1.5475e-02,\n",
      "         -3.0876e-02,  2.9531e-02,  8.2456e-02,  4.3768e-02,  1.2841e-01,\n",
      "         -5.1760e-02,  2.4799e-01, -9.8049e-02, -7.3253e-02, -4.7375e-02,\n",
      "          3.2339e-02, -3.7066e-01, -4.0623e-01,  3.9738e-01, -2.7329e-02,\n",
      "         -1.6930e-01, -5.8657e-02, -5.3689e-02, -6.8112e-02, -3.2449e-02,\n",
      "          8.8816e-02, -2.9662e-01, -1.0410e-01, -1.3745e-01, -2.7309e-01,\n",
      "          1.7027e-01,  6.2757e-02, -1.4127e-01, -2.8371e-02,  3.0332e-01,\n",
      "         -1.8054e-01, -6.2166e-04, -8.6959e-04,  3.8476e-01,  6.3244e-02,\n",
      "         -2.0215e-01,  1.2555e-01, -1.5984e-02,  2.0604e-02,  8.7541e-03,\n",
      "          6.3996e-02,  1.2963e-01,  1.7385e-01,  1.1763e-01, -7.1408e-02,\n",
      "         -4.6543e-02, -2.9189e-01, -6.1168e-02,  1.5007e-01,  3.8530e-01,\n",
      "         -4.9687e-01, -2.9291e-01, -3.5480e-02, -1.5042e-02, -1.4120e-01,\n",
      "          1.4511e-01, -3.7045e-02,  8.9234e-02, -1.1243e-01, -3.9437e-01,\n",
      "          3.3098e-01,  1.9299e-01, -2.1868e-01, -2.5297e-01, -5.4342e-02,\n",
      "          3.5754e-02,  7.6895e-02, -2.6860e-02, -2.0628e-01, -3.0720e-01,\n",
      "         -8.6704e-02,  1.3900e-01, -2.2204e-01, -6.1683e-02, -6.1634e-02,\n",
      "         -8.4303e-02, -8.9016e-02, -2.0297e-01, -1.1271e-01, -1.4278e-01,\n",
      "         -1.2793e-01,  4.3103e-01,  2.2654e-01, -2.7308e-01, -9.9331e-02,\n",
      "          2.3322e-02,  6.3023e-03,  1.7135e-01, -2.1307e-01,  1.9977e-01,\n",
      "         -4.6958e-02, -1.4373e-01, -1.1334e-01, -3.4244e-01,  5.4103e-02,\n",
      "         -3.8872e-02,  2.9599e-01,  1.7934e-01, -5.4940e-02,  1.6384e-01,\n",
      "         -7.8572e-02, -5.0933e-02, -2.0482e-01, -2.6196e-01,  1.4794e-01,\n",
      "          1.8799e-01, -5.0779e-02, -1.0097e-01,  1.5996e-01, -1.6789e-01,\n",
      "          6.8034e-02,  1.8981e-01,  1.0673e-01,  1.4535e-01, -1.1403e-01,\n",
      "         -3.3646e-01, -1.0874e-01,  1.6886e-01,  3.5199e-02, -2.0116e-01,\n",
      "         -1.5037e-01, -1.7399e-01,  2.0305e-01,  1.4353e-01, -1.1139e-01,\n",
      "          3.5664e-01, -1.8341e-01,  3.5331e-02,  2.1268e-02,  9.0671e-02,\n",
      "          9.0504e-02, -9.3841e-02,  1.3239e-01,  1.9952e-01,  5.8175e-02,\n",
      "         -4.2319e-02,  4.6191e-02, -1.2276e-01, -1.5645e-01, -1.0735e-01,\n",
      "          4.2778e-01, -1.2920e-01,  4.9049e-01, -2.5486e-01,  1.3907e-02,\n",
      "         -1.0983e-01, -1.5628e-01,  3.6488e-02,  3.3667e-01, -8.9549e-02,\n",
      "          2.1153e-01, -5.5962e-04,  6.1066e-02,  2.6326e-02,  6.2022e-02,\n",
      "          3.0317e-02, -6.5001e-02,  2.6180e-01, -4.7237e-02, -1.1445e-01,\n",
      "          4.9714e-02,  1.0454e-01,  1.4672e-01,  1.5318e-01,  1.6015e-01,\n",
      "          1.9983e-01, -2.0120e-01, -3.9972e-02, -8.0200e-04,  6.5315e-02,\n",
      "          1.4778e-01, -9.6505e-02, -1.3803e-01, -8.6644e-02,  2.8043e-01,\n",
      "          1.1591e-01,  2.7147e-01,  3.4964e-02,  2.6816e-01,  1.6876e-01,\n",
      "         -9.3193e-02, -3.0501e-02, -1.0417e-01,  2.2172e-01, -2.7514e-01,\n",
      "         -3.0198e-02, -6.4659e-02,  1.8636e-01,  3.8278e-02, -3.3058e-01,\n",
      "          7.2328e-02, -4.2641e-02, -7.7800e-02, -1.5568e-01, -9.8057e-02,\n",
      "          4.5618e-02, -2.1135e-02,  1.2437e-01,  1.1939e+00,  3.3564e-01,\n",
      "          3.5382e-01, -1.3239e-01,  6.7778e-02,  2.9517e-01, -7.3989e-02,\n",
      "         -1.3045e-01,  5.3311e-02,  1.4405e-01,  2.0137e-01,  4.4489e-02,\n",
      "         -1.3991e-01, -2.5243e-01, -1.6363e-01, -4.0235e-02, -2.5422e-01,\n",
      "          1.9907e-01, -6.2998e-02, -1.9029e-01, -1.5215e-01, -9.1656e-02,\n",
      "          1.6257e-01,  1.2240e-01, -1.1206e-02, -1.8399e-01,  1.9794e-01,\n",
      "         -1.3729e-01,  8.1769e-02,  2.3035e-02, -7.4267e-02,  2.0242e-01,\n",
      "          2.2295e-01,  2.5712e-01, -1.0401e-01, -4.7103e-02,  7.4051e-02,\n",
      "         -3.4324e-01, -4.4846e-02, -1.3105e-01, -6.0563e-02, -2.0920e-01,\n",
      "          3.7044e-02, -1.0496e-01, -2.0170e-01,  4.7160e-02,  1.1022e-01,\n",
      "         -1.7099e-01,  6.9416e-02,  1.9117e-01, -2.1948e-01,  2.8040e-01,\n",
      "         -1.9076e-01, -1.7727e-01,  1.2369e-01,  2.3997e-01,  3.1148e-02,\n",
      "          1.0161e-01, -5.6211e-01, -1.9064e-01,  1.4695e-01,  6.7425e-02,\n",
      "          5.9223e-02, -2.1138e-01,  2.8456e-01,  5.0065e-02, -9.6161e-02,\n",
      "          1.7948e-03,  1.2727e-01, -9.2363e-02,  4.1763e-01, -8.4012e-02,\n",
      "         -1.1330e-01,  7.9156e-02, -7.7612e-02, -9.8303e-03,  9.9261e-02,\n",
      "         -7.3926e-02, -1.6396e-01,  1.2717e-01, -8.7491e-02,  2.9173e-01,\n",
      "         -1.9873e-01,  6.9414e-02,  1.4317e-02,  6.8591e-02, -4.1381e-02,\n",
      "          1.2949e-01, -1.5732e-01, -3.7030e-01,  1.3178e-01, -2.0526e-01,\n",
      "         -4.0536e-04,  1.7382e-01,  2.6689e-02, -5.3026e-02,  9.7106e-02,\n",
      "          1.2278e-01, -9.4680e-02, -2.8574e-01,  1.3593e-01, -2.1333e-01,\n",
      "         -8.1438e-04, -1.4076e-01,  2.5413e-02,  2.1724e-01, -9.0343e-02,\n",
      "          9.9731e-02,  9.6838e-04,  1.0976e-01,  5.2334e-02,  5.3187e-01,\n",
      "          2.9523e-01, -6.6753e-02, -4.4593e-02,  2.5245e-03,  3.1611e-02,\n",
      "          6.0745e-02,  1.7864e-01, -2.4272e-01,  1.4914e-01,  7.6441e-02,\n",
      "          2.3497e-01,  1.2486e-01,  4.2835e-02,  5.8753e-02, -1.8538e-01,\n",
      "         -3.7354e-02, -1.1340e-01,  2.0788e-02, -1.4170e-01,  1.6536e-01,\n",
      "          1.2090e-01, -1.5087e-02, -1.8951e-01, -7.8547e-02,  2.3427e-01,\n",
      "          3.6591e-01, -1.3912e-01,  1.3829e-01]], grad_fn=<MeanBackward1>)\n"
     ]
    }
   ],
   "source": [
    "from calm import CaLM\n",
    "\n",
    "model = CaLM()\n",
    "vector = model.embed_sequence(cDNA_sequence)\n",
    "print('Embeddings shape: ', vector.shape)\n",
    "print(vector)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4db20c8-7ade-47b1-a972-b467291d37f6",
   "metadata": {},
   "source": [
    "It is also possible to obtain the raw, unaveraged embeddings:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "5e75f241-fd8e-4dee-8f0d-37eab8f2b549",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Embeddings shape:  torch.Size([1, 8, 768])\n",
      "tensor([[[-0.0320, -0.0483,  0.5161,  ...,  0.5614, -0.1594,  0.4159],\n",
      "         [ 0.0960,  0.2022,  0.1533,  ...,  0.4696,  0.0534, -0.1295],\n",
      "         [-0.0489,  0.1945,  0.2206,  ...,  0.4299, -0.3417,  0.0037],\n",
      "         ...,\n",
      "         [-0.1136, -0.1073,  0.2894,  ...,  0.4902, -0.0640,  0.0426],\n",
      "         [-0.2687,  0.1268,  0.3846,  ...,  0.2381, -0.1316,  0.3572],\n",
      "         [ 0.1666, -0.0149,  0.4737,  ...,  0.2082, -0.1321,  0.4308]]],\n",
      "       grad_fn=<TransposeBackward0>)\n"
     ]
    }
   ],
   "source": [
    "vector = model.embed_sequence(cDNA_sequence, average=False)\n",
    "print('Embeddings shape: ', vector.shape)\n",
    "print(vector)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}