{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5476acf8",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "736f4148-e48d-40da-930b-f73560f048c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from lobster.model import LobsterCBMPMLM\n",
    "import Levenshtein"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "488c8208-7378-4846-b539-c673cbcf704a",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "\n",
    "# Load the LobsterCBMPMLM model\n",
    "cb_lobster = LobsterCBMPMLM(\"asalam91/cb_lobster_24M\").to(device)\n",
    "cb_lobster.eval()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a6bb60c-a116-43c7-8ba2-362c9d2b210d",
   "metadata": {},
   "outputs": [],
   "source": [
    "cb_lobster.list_supported_concept()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50615bfa-372e-4ad2-bf31-585d3861cb9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "concept =\"gravy\"\n",
    "test_protein =\"MGAGASAEEKHSRELEKKLKEDAEKDARTVKLLLLGAGESGKSTIVKQMKIIHQDGYSLEECLEFIAIIYGNTLQSILAIVRAMTTLNIQYGDSARQDDARKLMHMADTIEEGTMPKEMSDIIQRLWKDSGIQACFERASEYQLNDSAGYYLSDLERLVTPGYVPTEQDVLRSRVKTTGIIETQFSFKDLNFRMFDVGGQRSERKKWIHCFEGVTCIIFIAALSAYDMVLVEDDEVNRMHESLHLFNSICNHRYFATTSIVLFLNKKDVFFEKIKKAHLSICFPDYDGPNTYEDAGNYIKVQFLELNMRRDVKEIYSHMTCATDTQNVKFVFDAVTDIIIKENLKDCGLF\"\n",
    "\n",
    "[new_protien] = cb_lobster.intervene_on_sequences([test_protein],concept,edits=5,intervention_type=\"negative\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb353041-bd72-4c35-bbdd-c820ef8e7af0",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(new_protien)\n",
    "print(Levenshtein.distance(test_protein, new_protien))"
   ]
  }
 ],
 "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.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
