fixing esmfold api

sokrypton · Feb 25, 2024 · d36504f · d36504f
1 parent f1f6a57
commit d36504f
Showing 1 changed file with 52 additions and 15 deletions.
diff --git a/beta/ESMFold_api.ipynb b/beta/ESMFold_api.ipynb
@@ -4,7 +4,7 @@
   "metadata": {
     "colab": {
       "provenance": [],
-      "authorship_tag": "ABX9TyP2LmoAi6MCx6sKIv8Ecumw",
+      "authorship_tag": "ABX9TyOTvNvYm0TYMGMVWT4fM7rz",
       "include_colab_link": true
     },
     "kernelspec": {
@@ -38,9 +38,7 @@
         "#### **Limitations**\n",
         "- max total length=400\n",
         "- for lengths > 400, try our [colab notebook](https://colab.research.google.com/github/sokrypton/ColabFold/blob/main/ESMFold.ipynb)\n",
-        "---\n",
-        "\n",
-        "*credit: [@drchristhorpe](https://twitter.com/drchristhorpe)*"
+        "---"
       ],
       "metadata": {
         "id": "v2jU2RLyacXQ"
@@ -58,28 +56,67 @@
         "#@title run **ESMFold**\n",
         "import os, time, re, requests\n",
         "if not os.path.isfile(\"colabfold.py\"):\n",
+        "  os.system(\"pip -q install pyppeteer nest_asyncio\")\n",
+        "  import asyncio\n",
+        "  import nest_asyncio\n",
+        "  from pyppeteer import launch\n",
+        "  import base64\n",
+        "\n",
+        "  # Apply nest_asyncio to enable nested event loops\n",
+        "  nest_asyncio.apply()\n",
+        "\n",
+        "  async def fetch_blob_content(page, blob_url):\n",
+        "    blob_to_base64 = \"\"\"\n",
+        "    async (blobUrl) => {\n",
+        "        const blob = await fetch(blobUrl).then(r => r.blob());\n",
+        "        return new Promise((resolve) => {\n",
+        "            const reader = new FileReader();\n",
+        "            reader.onloadend = () => resolve(reader.result);\n",
+        "            reader.readAsDataURL(blob);\n",
+        "        });\n",
+        "    }\n",
+        "    \"\"\"\n",
+        "    base64_data = await page.evaluate(blob_to_base64, blob_url)\n",
+        "    _, encoded = base64_data.split(',', 1)\n",
+        "    return base64.b64decode(encoded)\n",
+        "\n",
+        "  async def extract_pdb_file_download_link_and_content(url):\n",
+        "    browser = await launch(headless=True, args=['--no-sandbox', '--disable-setuid-sandbox'])\n",
+        "    page = await browser.newPage()\n",
+        "    await page.goto(url, {'waitUntil': 'networkidle0'})\n",
+        "    elements = await page.querySelectorAll('a.btn.bg-purple')\n",
+        "    for element in elements:\n",
+        "        href = await page.evaluate('(element) => element.getAttribute(\"href\")', element)\n",
+        "        if 'blob:https://esmatlas.com/' in href:\n",
+        "            content = await fetch_blob_content(page, href)\n",
+        "            await browser.close()\n",
+        "            return href, content\n",
+        "    await browser.close()\n",
+        "    return \"No PDB file link found.\", None\n",
+        "\n",
+        "  def esmfold_api(sequence):\n",
+        "    url = f'https://esmatlas.com/resources/fold/result?fasta_header=%3Eunnamed&sequence={sequence}'\n",
+        "    result = asyncio.get_event_loop().run_until_complete(extract_pdb_file_download_link_and_content(url))\n",
+        "    if result[1]:\n",
+        "      pdb_str = result[1].decode('utf-8')\n",
+        "      return pdb_str\n",
+        "    else:\n",
+        "      return \"Failed to retrieve PDB content.\"\n",
+        "\n",
         "  os.system(\"pip -q install py3Dmol\")\n",
         "  os.system(\"wget -qnc https://raw.githubusercontent.com/sokrypton/ColabFold/main/beta/colabfold.py\")\n",
         "  os.system(\"mkdir -p tmp\")\n",
         "\n",
         "import colabfold as cf\n",
         "\n",
-        "def predict(sequence):\n",
-        "  esmfold_api_url = 'https://api.esmatlas.com/foldSequence/v1/pdb/'\n",
-        "  r = requests.post(esmfold_api_url, data=sequence)\n",
-        "  while r.status_code != 200:\n",
-        "    time.sleep(5)\n",
-        "    r = requests.post(esmfold_api_url, data=sequence)\n",
-        "  structure = r.text\n",
-        "  return structure\n",
         "\n",
         "sequence = \"GWSTELEKHREELKEFLKKEGITNVEIRIDNGRLEVRVEGGTERLKRFLEELRQKLEKKGYTVDIKIE\" #@param {type:\"string\"}\n",
         "sequence = re.sub(\"[^A-Z]\", \"\", sequence.upper())\n",
         "#assert len(sequence) <= 400, \"error: max length supported is 400\"\n",
         "hash = cf.get_hash(sequence)[:5]\n",
         "pdb_filename = f\"tmp/prediction_{hash}.pdb\"\n",
         "if not os.path.isfile(pdb_filename):\n",
-        "  pdb_str = predict(sequence)\n",
+        "  pdb_str = esmfold_api(sequence)\n",
         "  with open(pdb_filename,\"w\") as out:\n",
         "    out.write(pdb_str)\n",
         "\n",
@@ -92,11 +129,11 @@
         "                show_mainchains, color,\n",
         "                color_HP=True,\n",
         "                size=(800,480),\n",
-        "                vmin=0.5, vmax=0.9)       \n",
+        "                vmin=0.5, vmax=0.9)\n",
         "v.setHoverable({}, True,\n",
         "               '''function(atom,viewer,event,container){if(!atom.label){atom.label=viewer.addLabel(\"      \"+atom.resn+\":\"+atom.resi,{position:atom,backgroundColor:'mintcream',fontColor:'black'});}}''',\n",
         "               '''function(atom,viewer){if(atom.label){viewer.removeLabel(atom.label);delete atom.label;}}''')\n",
-        "v.show()           \n",
+        "v.show()\n",
         "if color == \"lDDT\":\n",
         "  cf.plot_plddt_legend().show()"
       ]