From 2dca1083d6351944a32aff4f66290a961da744b0 Mon Sep 17 00:00:00 2001 From: Milot Mirdita Date: Fri, 8 Mar 2024 11:47:45 +0900 Subject: [PATCH] Fix cf-af2 notebooks to work on TPU again --- AlphaFold2.ipynb | 2 ++ batch/AlphaFold2_batch.ipynb | 3 +++ 2 files changed, 5 insertions(+) diff --git a/AlphaFold2.ipynb b/AlphaFold2.ipynb index 588f5319..9c45930d 100644 --- 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 index aefc6aae..e510e3ab 100644 --- 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",