{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HC86rsLf-_Zt"
   },
   "source": [
    "# Generation UI\n",
    "\n",
    "This is the most flexible notebook for generating protein sequences using the ESM3 model.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "ICGSD1Jo7zAb"
   },
   "outputs": [],
   "source": [
    "# @title Input API keys, then hit `Runtime` -> `Run all`\n",
    "# @markdown Our hosted service that provides access to the full suite of ESM3 models.\n",
    "# @markdown To utilize the Forge API, users must first agree to the [Terms of Service](https://forge.evolutionaryscale.ai/termsofservice) and generate an access token via the [Forge console](https://forge.evolutionaryscale.ai/console).\n",
    "# @markdown The console also provides a comprehensive list of models available to each user.\n",
    "\n",
    "import os\n",
    "\n",
    "# @markdown ### Authentication\n",
    "# @markdown Paste your token from the [Forge console](https://forge.evolutionaryscale.ai/console)\n",
    "forge_token = \"\"  # @param {type:\"string\"}\n",
    "os.environ[\"FORGE_TOKEN\"] = forge_token\n",
    "\n",
    "# @markdown ### Model Selection\n",
    "# @markdown Enter the model name from the [Forge console page](https://forge.evolutionaryscale.ai/console) that you would like to use:\n",
    "model_name = \"esm3-medium-2024-03\"  # @param {type:\"string\"}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "03ARpZRE_N39"
   },
   "outputs": [],
   "source": [
    "# @title Install dependencies\n",
    "import os\n",
    "\n",
    "os.system(\"pip install git+https://github.com/evolutionaryscale/esm\")\n",
    "os.system(\"pip install pydssp pygtrie dna-features-viewer py3dmol\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "x1MUAuDWBAel"
   },
   "outputs": [],
   "source": [
    "# @title Create Generation UI\n",
    "# @markdown If running on Google colab, it is recommended to use the light theme and select the \"View output fullscreen\" option in the cell toolbar for the best experience\n",
    "\n",
    "from esm.widgets.utils.clients import get_forge_client\n",
    "from esm.widgets.utils.types import ClientInitContainer\n",
    "from esm.widgets.views.generation import create_generation_ui\n",
    "\n",
    "client_container = ClientInitContainer()\n",
    "create_generation_ui(get_forge_client(model_name))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "default",
   "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.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
