Switch source of OF parameters to HuggingFace

README.md

@@ -41,7 +41,8 @@ implementations, respectively.
 All Python dependencies are specified in `environment.yml`. For producing sequence 
 alignments, you'll also need `kalign`, the [HH-suite](https://github.com/soedinglab/hh-suite), 
 and one of {`jackhmmer`, [MMseqs2](https://github.com/soedinglab/mmseqs2) (nightly build)} 
-installed on on your system. Finally, some download scripts require `aria2c`.
+installed on on your system. You'll need `git-lfs` to download OpenFold parameters. 
+Finally, some download scripts require `aria2c`.
 
 For convenience, we provide a script that installs Miniconda locally, creates a 
 `conda` virtual environment, installs all Python dependencies, and downloads

notebooks/OpenFold.ipynb

@@ -31,7 +31,7 @@
         "\n",
         "OpenFold is a trainable PyTorch reimplementation of AlphaFold 2. For the purposes of inference, it is practically identical to the original (\"practically\" because ensembling is excluded from OpenFold (recycling is enabled, however)).\n",
         "\n",
-        "In this notebook, OpenFold is run with DeepMind's publicly released parameters for AlphaFold 2.\n",
+        "In this notebook, OpenFold is run with your choice of our original OpenFold parameters or DeepMind's publicly released parameters for AlphaFold 2.\n",
         "\n",
         "**Note**\n",
         "\n",
@@ -43,7 +43,7 @@
         "\n",
         "**Licenses**\n",
         "\n",
-        "This Colab uses the [AlphaFold model parameters](https://github.com/deepmind/alphafold/#model-parameters-license), made available under the Creative Commons Attribution 4.0 International ([CC BY 4.0](https://creativecommons.org/licenses/by/4.0/legalcode)) license. The Colab itself is provided under the [Apache 2.0 license](https://www.apache.org/licenses/LICENSE-2.0). See the full license statement below.\n",
+        "This Colab supports inference with the [AlphaFold model parameters](https://github.com/deepmind/alphafold/#model-parameters-license), made available under the Creative Commons Attribution 4.0 International ([CC BY 4.0](https://creativecommons.org/licenses/by/4.0/legalcode)) license. The Colab itself is provided under the [Apache 2.0 license](https://www.apache.org/licenses/LICENSE-2.0). See the full license statement below.\n",
         "\n",
         "**More information**\n",
         "\n",
@@ -132,13 +132,10 @@
         "\n",
         "GIT_REPO = 'https://github.com/aqlaboratory/openfold'\n",
         "\n",
-        "OPENFOLD_PARAM_FILE_ID = \"1OpeMrfWEUSD_KqffbPqd5p7WsJjlC3ZE\"\n",
+        "OPENFOLD_PARAM_SOURCE_URL = \"https://huggingface.co/nz/Openfold\"\n",
         "ALPHAFOLD_PARAM_SOURCE_URL = 'https://storage.googleapis.com/alphafold/alphafold_params_2022-01-19.tar'\n",
-        "OPENFOLD_PARAMS_DIR = './openfold/openfold/resources/'\n",
+        "OPENFOLD_PARAMS_DIR = './openfold/openfold/resources/openfold_params'\n",
         "ALPHAFOLD_PARAMS_DIR = './openfold/openfold/resources/params'\n",
-        "OPENFOLD_PARAMS_PATH = os.path.join(\n",
-        "  OPENFOLD_PARAMS_DIR, \"openfold_params.tar.gz\"\n",
-        ")\n",
         "ALPHAFOLD_PARAMS_PATH = os.path.join(\n",
         "  ALPHAFOLD_PARAMS_DIR, os.path.basename(ALPHAFOLD_PARAM_SOURCE_URL)\n",
         ")\n",
@@ -173,9 +170,7 @@
         "      %shell rm \"{ALPHAFOLD_PARAMS_PATH}\"\n",
         "\n",
         "      %shell mkdir --parents \"{OPENFOLD_PARAMS_DIR}\"\n",
-        "      %shell gdown --id \"{OPENFOLD_PARAM_FILE_ID}\" -O \"{OPENFOLD_PARAMS_PATH}\"\n",
-        "      %shell tar --extract --verbose --file=\"{OPENFOLD_PARAMS_PATH}\" \\\n",
-        "        --directory=\"{OPENFOLD_PARAMS_DIR}\" --preserve-permissions\n",
+        "      %shell git clone \"{OPENFOLD_PARAM_SOURCE_URL}\" \"{OPENFOLD_PARAMS_DIR}\"\n",
         "      %shell rm \"{OPENFOLD_PARAMS_PATH}\"\n",
         "      pbar.update(55)\n",
         "except subprocess.CalledProcessError:\n",
@@ -472,7 +467,6 @@
         "        of_model_name = f\"finetuning_{model_name_spl[-1]}.pt\"\n",
         "      params_name = os.path.join(\n",
         "        OPENFOLD_PARAMS_DIR,\n",
-        "        \"openfold_params\",\n",
         "        of_model_name\n",
         "      )\n",
         "      d = torch.load(params_name)\n",

scripts/download_openfold_params.sh → scripts/download_openfold_params_gdrive.sh

@@ -14,9 +14,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-# Downloads and unzips OpenFold parameters.
+# Downloads and unzips OpenFold parameters from Google Drive. Alternative to
+# the HuggingFace version.
 #
-# Usage: bash download_openfold_params.sh /path/to/download/directory
+# Usage: bash download_openfold_params_gdrive.sh /path/to/download/directory
 set -e
 
 if [[ $# -eq 0 ]]; then

scripts/download_openfold_params_huggingface.sh

+#!/bin/bash
+#
+# Copyright 2021 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Downloads and unzips OpenFold parameters.
+#
+# Usage: bash download_openfold_params_huggingface.sh /path/to/download/directory
+set -e
+
+if [[ $# -eq 0 ]]; then
+    echo "Error: download directory must be provided as an input argument."
+    exit 1
+fi
+
+URL="https://huggingface.co/nz/OpenFold"
+
+DOWNLOAD_DIR="${1}/openfold_params/"
+mkdir -p "${DOWNLOAD_DIR}"
+git clone $URL "${DOWNLOAD_DIR}"
+rm -rf "${DOWNLOAD_DIR}/.git"