Fix cf-af2 notebooks to work on TPU again

sokrypton · Mar 8, 2024 · 2dca108 · 2dca108
1 parent 1abdc19
commit 2dca108
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 0 deletions.
diff --git a/AlphaFold2.ipynb b/AlphaFold2.ipynb
@@ -147,6 +147,8 @@
         "if not os.path.isfile(\"COLABFOLD_READY\"):\n",
         "  print(\"installing colabfold...\")\n",
         "  os.system(\"pip install -q --no-warn-conflicts 'colabfold[alphafold-minus-jax] @ git+https://github.com/sokrypton/ColabFold'\")\n",
+        "  if os.environ.get('TPU_NAME', False) != False:\n",
+        "    os.system(\"pip install -q --no-warn-conflicts -U dm-haiku==0.0.10 jax==0.3.25\")\n",
         "  os.system(\"ln -s /usr/local/lib/python3.*/dist-packages/colabfold colabfold\")\n",
         "  os.system(\"ln -s /usr/local/lib/python3.*/dist-packages/alphafold alphafold\")\n",
         "  os.system(\"touch COLABFOLD_READY\")\n",

diff --git a/batch/AlphaFold2_batch.ipynb b/batch/AlphaFold2_batch.ipynb
@@ -128,6 +128,9 @@
         "  # install dependencies\n",
         "  # We have to use \"--no-warn-conflicts\" because colab already has a lot preinstalled with requirements different to ours\n",
         "  pip install -q --no-warn-conflicts \"colabfold[alphafold-minus-jax] @ git+https://github.com/sokrypton/ColabFold\"\n",
+        "  if [ -n \"${TPU_NAME}\" ]; then\n",
+        "    pip install -q --no-warn-conflicts -U dm-haiku==0.0.10 jax==0.3.25\n",
+        "  fi\n",
         "  ln -s /usr/local/lib/python3.*/dist-packages/colabfold colabfold\n",
         "  ln -s /usr/local/lib/python3.*/dist-packages/alphafold alphafold\n",
         "  touch COLABFOLD_READY\n",