{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Lobster Model Inference Notebook\n",
    "\n",
    "import torch\n",
    "from lobster.model import LobsterPMLM, LobsterCBMPMLM\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the test protein sequence\n",
    "test_protein = \"MGAGASAEEKHSRELEKKLKEDAEKDARTVKLLLLGAGESGKSTIVKQMKIIHQDGYSLEECLEFIAIIYGNTLQSILAIVRAMTTLNIQYGDSARQDDARKLMHMADTIEEGTMPKEMSDIIQRLWKDSGIQACFERASEYQLNDSAGYYLSDLERLVTPGYVPTEQDVLRSRVKTTGIIETQFSFKDLNFRMFDVGGQRSERKKWIHCFEGVTCIIFIAALSAYDMVLVEDDEVNRMHESLHLFNSICNHRYFATTSIVLFLNKKDVFFEKIKKAHLSICFPDYDGPNTYEDAGNYIKVQFLELNMRRDVKEIYSHMTCATDTQNVKFVFDAVTDIIIKENLKDCGLF\"\n",
    "\n",
    "# Determine the device\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the LobsterPMLM model\n",
    "lobster = LobsterPMLM(\"asalam91/lobster_24M\").to(device)\n",
    "lobster.eval()\n",
    "\n",
    "# Get MLM representation\n",
    "mlm_representation = lobster.sequences_to_latents([test_protein])[-1]\n",
    "cls_token_mlm_representation = mlm_representation[:, 0, :]\n",
    "pooled_mlm_representation = torch.mean(mlm_representation, dim=1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Load the LobsterCBMPMLM model\n",
    "cb_lobster = LobsterCBMPMLM(\"asalam91/cb_lobster_24M\").to(device)\n",
    "cb_lobster.eval()\n",
    "\n",
    "# Get CB MLM representation\n",
    "cb_mlm_representation = cb_lobster.sequences_to_latents([test_protein])[-1]\n",
    "cls_token_cb_mlm_representation = cb_mlm_representation[:, 0, :]\n",
    "pooled_cb_mlm_representation = torch.mean(cb_mlm_representation, dim=1)\n",
    "\n",
    "# Get protein concepts\n",
    "test_protein_concepts = cb_lobster.sequences_to_concepts([test_protein])[-1]\n",
    "test_protein_concepts_emb = cb_lobster.sequences_to_concepts_emb([test_protein])[-1][0]  # All of the known concepts are the same for all tokens...\n",
    "test_protein_concepts_unknown_emb"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
