{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Ume Multi-Modal Embeddings Tutorial\n",
    "\n",
    "This notebook show how to use the Universal Molecular Encoder (Ume) to generate embeddings for different molecular modalities: amino acids, SMILES, and nucleotides. Stay tuned for 3D coordinate embeddings and more!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/teamspace/studios/this_studio/lobster/.venv/lib/python3.12/site-packages/flash_attn/ops/triton/layer_norm.py:984: FutureWarning: `torch.cuda.amp.custom_fwd(args...)` is deprecated. Please use `torch.amp.custom_fwd(args..., device_type='cuda')` instead.\n",
      "  @custom_fwd\n",
      "/teamspace/studios/this_studio/lobster/.venv/lib/python3.12/site-packages/flash_attn/ops/triton/layer_norm.py:1043: FutureWarning: `torch.cuda.amp.custom_bwd(args...)` is deprecated. Please use `torch.amp.custom_bwd(args..., device_type='cuda')` instead.\n",
      "  @custom_bwd\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Supported modalities: ['SMILES', 'amino_acid', 'nucleotide', '3d_coordinates']\n",
      "Vocab size: 1536\n"
     ]
    }
   ],
   "source": [
    "from lobster.model import Ume\n",
    "\n",
    "ume = Ume()\n",
    "\n",
    "print(f\"Supported modalities: {ume.modalities}\")\n",
    "print(f\"Vocab size: {len(ume.get_vocab())}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load from checkpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint = \"ume-checkpoints/best.ckpt\" # Replace with the correct checkpoint path\n",
    "\n",
    "ume = Ume.load_from_checkpoint(checkpoint)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. Protein sequences\n",
    "\n",
    "Embed sample protein sequence to get full sequence embedding or per-residue embeddings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Protein embeddings shape: torch.Size([3, 768])\n",
      "Protein token-level embeddings shape: torch.Size([3, 512, 768])\n"
     ]
    }
   ],
   "source": [
    "# Example protein sequences\n",
    "protein_sequences = [\n",
    "    \"MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG\",  # Sample protein fragment\n",
    "    \"MLSRAVCGTSRQLAPVLAYLGSRQKHSLPDLPYDYGALEPHINAQIMQLHHSKHHAAYVNNLNVTEEKYQEALAKGDVTAQIALQPALKFNGGGHINHSIFWTNLSPNGGGEPKGELLEAIKRDFGSFDKFKEKLTAASVGVQGSGWGWLGFNKERGHLQIAACPNQDPLQGTTGLIPLLGIDVWEHAYYLQYKNVRPDYLKAIWNVINWENVTERYMACKK\",  # Human SOD1\n",
    "    \"MVHLTPEEKSAVTALWGKVNVDEVGGEALGRLLVVYPWTQRFFESFGDLSTPDAVMGNPKVKAHGKKVLGAFSDGLAHLDNLKGTFATLSELHCDKLHVDPENFRLLGNVLVCVLAHHFGKEFTPPVQAAYQKVVAGVANALAHKYH\"  # Hemoglobin beta chain\n",
    "]\n",
    "\n",
    "# Get embeddings for protein sequences\n",
    "protein_embeddings = ume.get_embeddings(protein_sequences, modality=\"amino_acid\")\n",
    "print(f\"Protein embeddings shape: {protein_embeddings.shape}\")\n",
    "\n",
    "# Get token-level embeddings (without aggregation)\n",
    "protein_residue_embeddings = ume.get_embeddings(protein_sequences, modality=\"amino_acid\", aggregate=False)\n",
    "print(f\"Protein token-level embeddings shape: {protein_residue_embeddings.shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. SMILES\n",
    "SMILES strings are a text-based representation of molecular structures. Here we embed common drug molecules.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SMILES embeddings shape: torch.Size([5, 768])\n"
     ]
    }
   ],
   "source": [
    "# Example SMILES strings for common molecules\n",
    "smiles_examples = [\n",
    "    \"CC(=O)OC1=CC=CC=C1C(=O)O\",  # Aspirin\n",
    "    \"CN1C=NC2=C1C(=O)N(C(=O)N2C)C\",  # Caffeine\n",
    "    \"COC1=CC=C(CCN)C=C1\",  # Dopamine\n",
    "    \"C1=CC=C(C(=C1)C(=O)O)O\",  # Salicylic acid\n",
    "    \"CCC(CC)COC(=O)[C@H](C)N[P@](=O)(OC[C@H]1O[C@@H]([C@H]([C@@H]1O)O)N1C=NC2=C1N=CN=C2N)OC1=CC=CC=C1\"  # Remdesivir\n",
    "]\n",
    "\n",
    "# Get embeddings for SMILES\n",
    "smiles_embeddings = ume.get_embeddings(smiles_examples, modality=\"SMILES\")\n",
    "print(f\"SMILES embeddings shape: {smiles_embeddings.shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. Nucleotides\n",
    "\n",
    "Embed example DNA/RNA sequences."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Nucleotide embeddings shape: torch.Size([3, 768])\n"
     ]
    }
   ],
   "source": [
    "# Example DNA/RNA sequences\n",
    "nucleotide_sequences = [\n",
    "    \"ATGCATGCATGCATGCATGCATGCATGCATGCATGCATGCATGCATGC\",  \n",
    "    \"GATTACACAGTGCTTGACCCGATCGATCGATCGATCGATCGATCGATCGA\",  \n",
    "    \"AUGCUAUGCUAGCUAGCUAGCUAGCUAUGCUAGCUAUGCUAGCUAUC\"  # RNA sequence \n",
    "]\n",
    "\n",
    "# Get embeddings for nucleotide sequences\n",
    "nucleotide_embeddings = ume.get_embeddings(nucleotide_sequences, modality=\"nucleotide\")\n",
    "print(f\"Nucleotide embeddings shape: {nucleotide_embeddings.shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using Embeddings for Downstream Tasks\n",
    "Quick example of using molecular embeddings for a classification task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "              precision    recall  f1-score   support\n",
      "\n",
      "           0       0.00      0.00      0.00       0.0\n",
      "           1       0.00      0.00      0.00       1.0\n",
      "\n",
      "    accuracy                           0.00       1.0\n",
      "   macro avg       0.00      0.00      0.00       1.0\n",
      "weighted avg       0.00      0.00      0.00       1.0\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/teamspace/studios/this_studio/lobster/.venv/lib/python3.12/site-packages/sklearn/metrics/_classification.py:1565: UndefinedMetricWarning: Precision is ill-defined and being set to 0.0 in labels with no predicted samples. Use `zero_division` parameter to control this behavior.\n",
      "  _warn_prf(average, modifier, f\"{metric.capitalize()} is\", len(result))\n",
      "/teamspace/studios/this_studio/lobster/.venv/lib/python3.12/site-packages/sklearn/metrics/_classification.py:1565: UndefinedMetricWarning: Recall is ill-defined and being set to 0.0 in labels with no true samples. Use `zero_division` parameter to control this behavior.\n",
      "  _warn_prf(average, modifier, f\"{metric.capitalize()} is\", len(result))\n",
      "/teamspace/studios/this_studio/lobster/.venv/lib/python3.12/site-packages/sklearn/metrics/_classification.py:1565: UndefinedMetricWarning: Precision is ill-defined and being set to 0.0 in labels with no predicted samples. Use `zero_division` parameter to control this behavior.\n",
      "  _warn_prf(average, modifier, f\"{metric.capitalize()} is\", len(result))\n",
      "/teamspace/studios/this_studio/lobster/.venv/lib/python3.12/site-packages/sklearn/metrics/_classification.py:1565: UndefinedMetricWarning: Recall is ill-defined and being set to 0.0 in labels with no true samples. Use `zero_division` parameter to control this behavior.\n",
      "  _warn_prf(average, modifier, f\"{metric.capitalize()} is\", len(result))\n",
      "/teamspace/studios/this_studio/lobster/.venv/lib/python3.12/site-packages/sklearn/metrics/_classification.py:1565: UndefinedMetricWarning: Precision is ill-defined and being set to 0.0 in labels with no predicted samples. Use `zero_division` parameter to control this behavior.\n",
      "  _warn_prf(average, modifier, f\"{metric.capitalize()} is\", len(result))\n",
      "/teamspace/studios/this_studio/lobster/.venv/lib/python3.12/site-packages/sklearn/metrics/_classification.py:1565: UndefinedMetricWarning: Recall is ill-defined and being set to 0.0 in labels with no true samples. Use `zero_division` parameter to control this behavior.\n",
      "  _warn_prf(average, modifier, f\"{metric.capitalize()} is\", len(result))\n"
     ]
    }
   ],
   "source": [
    "# Dummy classification setup \n",
    "import numpy as np\n",
    "from sklearn.ensemble import RandomForestClassifier\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.metrics import classification_report\n",
    "\n",
    "# SMILES with some property labels\n",
    "inputs = [\"CC(=O)OC1=CC=CC=C1C(=O)O\", \"CN1C=NC2=C1C(=O)N(C(=O)N2C)C\", \"COC1=CC=C(CCN)C=C1\", \n",
    "              \"C1=CC=C(C(=C1)C(=O)O)O\", \"CC12CCC(CC1)CC(C2)C(C)CN\"]\n",
    "labels = [0, 1, 0, 1, 0]  # Binary classification example\n",
    "\n",
    "# Get embeddings\n",
    "X = ume.get_embeddings(inputs, modality=\"SMILES\").cpu().numpy()\n",
    "y = np.array(labels)\n",
    "\n",
    "# Train a simple classifier\n",
    "X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)\n",
    "clf = RandomForestClassifier(n_estimators=100, random_state=42)\n",
    "clf.fit(X_train, y_train)\n",
    "\n",
    "# Evaluate\n",
    "y_pred = clf.predict(X_test)\n",
    "print(classification_report(y_test, y_pred))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluate on Molecular Property Prediction Tasks\n",
    "\n",
    "Here is how to evaluate Ume on tasks defined as callbacks. Note that training on and evaluating these tasks will take a few minutes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "MoleculeACELinearProbeCallback: 100%|██████████| 30/30 [07:58<00:00, 15.94s/it]\n"
     ]
    }
   ],
   "source": [
    "import pandas as pd\n",
    "\n",
    "from lobster.callbacks import CalmLinearProbeCallback, MoleculeACELinearProbeCallback\n",
    "\n",
    "molecule_ace_probe = MoleculeACELinearProbeCallback(\n",
    "    max_length=ume.embedding_dim\n",
    ")\n",
    "molecule_ace_scores = molecule_ace_probe.evaluate(ume)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "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>CHEMBL238_Ki</th>\n",
       "      <th>CHEMBL1862_Ki</th>\n",
       "      <th>CHEMBL237_EC50</th>\n",
       "      <th>CHEMBL231_Ki</th>\n",
       "      <th>CHEMBL214_Ki</th>\n",
       "      <th>CHEMBL4616_EC50</th>\n",
       "      <th>CHEMBL234_Ki</th>\n",
       "      <th>CHEMBL228_Ki</th>\n",
       "      <th>CHEMBL2034_Ki</th>\n",
       "      <th>CHEMBL233_Ki</th>\n",
       "      <th>...</th>\n",
       "      <th>CHEMBL2147_Ki</th>\n",
       "      <th>CHEMBL218_EC50</th>\n",
       "      <th>CHEMBL236_Ki</th>\n",
       "      <th>CHEMBL244_Ki</th>\n",
       "      <th>CHEMBL2047_EC50</th>\n",
       "      <th>CHEMBL4203_Ki</th>\n",
       "      <th>CHEMBL4005_Ki</th>\n",
       "      <th>CHEMBL204_Ki</th>\n",
       "      <th>CHEMBL235_EC50</th>\n",
       "      <th>mean</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>mse</th>\n",
       "      <td>2.095025</td>\n",
       "      <td>5.312742</td>\n",
       "      <td>3.668455</td>\n",
       "      <td>2.874007</td>\n",
       "      <td>0.977828</td>\n",
       "      <td>1.559722</td>\n",
       "      <td>0.929982</td>\n",
       "      <td>1.373721</td>\n",
       "      <td>2.358630</td>\n",
       "      <td>1.170707</td>\n",
       "      <td>...</td>\n",
       "      <td>2.118004</td>\n",
       "      <td>3.143562</td>\n",
       "      <td>1.167833</td>\n",
       "      <td>1.422116</td>\n",
       "      <td>1.564671</td>\n",
       "      <td>3.172571</td>\n",
       "      <td>3.062797</td>\n",
       "      <td>1.224741</td>\n",
       "      <td>0.990469</td>\n",
       "      <td>2.021026</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>r2</th>\n",
       "      <td>-0.544668</td>\n",
       "      <td>-1.430170</td>\n",
       "      <td>-0.756787</td>\n",
       "      <td>-0.672591</td>\n",
       "      <td>0.243814</td>\n",
       "      <td>-0.861025</td>\n",
       "      <td>0.278650</td>\n",
       "      <td>0.044004</td>\n",
       "      <td>-1.337632</td>\n",
       "      <td>0.312313</td>\n",
       "      <td>...</td>\n",
       "      <td>0.460752</td>\n",
       "      <td>-1.877037</td>\n",
       "      <td>0.349264</td>\n",
       "      <td>0.454258</td>\n",
       "      <td>-0.626796</td>\n",
       "      <td>-1.758412</td>\n",
       "      <td>-1.626563</td>\n",
       "      <td>0.476720</td>\n",
       "      <td>0.145517</td>\n",
       "      <td>-0.457946</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>spearman</th>\n",
       "      <td>0.473802</td>\n",
       "      <td>0.539632</td>\n",
       "      <td>0.380362</td>\n",
       "      <td>0.486178</td>\n",
       "      <td>0.546139</td>\n",
       "      <td>0.273161</td>\n",
       "      <td>0.608297</td>\n",
       "      <td>0.501315</td>\n",
       "      <td>0.387129</td>\n",
       "      <td>0.602608</td>\n",
       "      <td>...</td>\n",
       "      <td>0.760325</td>\n",
       "      <td>0.273797</td>\n",
       "      <td>0.642653</td>\n",
       "      <td>0.727715</td>\n",
       "      <td>0.476775</td>\n",
       "      <td>0.262011</td>\n",
       "      <td>0.314889</td>\n",
       "      <td>0.709204</td>\n",
       "      <td>0.543911</td>\n",
       "      <td>0.489678</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>3 rows × 31 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "          CHEMBL238_Ki  CHEMBL1862_Ki  CHEMBL237_EC50  CHEMBL231_Ki  \\\n",
       "mse           2.095025       5.312742        3.668455      2.874007   \n",
       "r2           -0.544668      -1.430170       -0.756787     -0.672591   \n",
       "spearman      0.473802       0.539632        0.380362      0.486178   \n",
       "\n",
       "          CHEMBL214_Ki  CHEMBL4616_EC50  CHEMBL234_Ki  CHEMBL228_Ki  \\\n",
       "mse           0.977828         1.559722      0.929982      1.373721   \n",
       "r2            0.243814        -0.861025      0.278650      0.044004   \n",
       "spearman      0.546139         0.273161      0.608297      0.501315   \n",
       "\n",
       "          CHEMBL2034_Ki  CHEMBL233_Ki  ...  CHEMBL2147_Ki  CHEMBL218_EC50  \\\n",
       "mse            2.358630      1.170707  ...       2.118004        3.143562   \n",
       "r2            -1.337632      0.312313  ...       0.460752       -1.877037   \n",
       "spearman       0.387129      0.602608  ...       0.760325        0.273797   \n",
       "\n",
       "          CHEMBL236_Ki  CHEMBL244_Ki  CHEMBL2047_EC50  CHEMBL4203_Ki  \\\n",
       "mse           1.167833      1.422116         1.564671       3.172571   \n",
       "r2            0.349264      0.454258        -0.626796      -1.758412   \n",
       "spearman      0.642653      0.727715         0.476775       0.262011   \n",
       "\n",
       "          CHEMBL4005_Ki  CHEMBL204_Ki  CHEMBL235_EC50      mean  \n",
       "mse            3.062797      1.224741        0.990469  2.021026  \n",
       "r2            -1.626563      0.476720        0.145517 -0.457946  \n",
       "spearman       0.314889      0.709204        0.543911  0.489678  \n",
       "\n",
       "[3 rows x 31 columns]"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pd.DataFrame(molecule_ace_scores).head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating train split: 5222 examples [00:00, 74980.33 examples/s]\n",
      "Generating train split: 562 examples [00:00, 37857.53 examples/s]s/it]\n",
      "Generating train split: 5696 examples [00:00, 151651.71 examples/s]it]\n",
      "Generating train split: 452 examples [00:00, 21029.68 examples/s]s/it]\n",
      "Generating train split: 14772 examples [00:00, 125365.49 examples/s]t]\n",
      "Generating train split: 369 examples [00:00, 23669.81 examples/s]s/it]\n",
      "CalmLinearProbeCallback: 100%|██████████| 8/8 [02:38<00:00, 19.83s/it]\n"
     ]
    }
   ],
   "source": [
    "calm_probe = CalmLinearProbeCallback(\n",
    "    max_length=ume.embedding_dim\n",
    ")\n",
    "calm_scores = calm_probe.evaluate(ume)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "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>solubility</th>\n",
       "      <th>localization</th>\n",
       "      <th>meltome</th>\n",
       "      <th>mean</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>mse</th>\n",
       "      <td>3.649214</td>\n",
       "      <td>NaN</td>\n",
       "      <td>56.701900</td>\n",
       "      <td>30.175557</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>r2</th>\n",
       "      <td>-0.273509</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.464096</td>\n",
       "      <td>0.095294</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>spearman</th>\n",
       "      <td>0.038007</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.463611</td>\n",
       "      <td>0.250809</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>accuracy</th>\n",
       "      <td>NaN</td>\n",
       "      <td>0.890833</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.890833</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>f1</th>\n",
       "      <td>NaN</td>\n",
       "      <td>0.398531</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.398531</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "          solubility  localization    meltome       mean\n",
       "mse         3.649214           NaN  56.701900  30.175557\n",
       "r2         -0.273509           NaN   0.464096   0.095294\n",
       "spearman    0.038007           NaN   0.463611   0.250809\n",
       "accuracy         NaN      0.890833        NaN   0.890833\n",
       "f1               NaN      0.398531        NaN   0.398531"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pd.DataFrame(calm_scores).head()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
