Merge branch 'main' into main

.github/workflows/docker-image.yml

@@ -11,5 +11,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@v4
+    - name: Cleanup  # https://github.com/actions/virtual-environments/issues/2840
+      run: sudo rm -rf /usr/share/dotnet && sudo rm -rf /opt/ghc && sudo rm -rf "/usr/local/share/boost" && sudo rm -rf "$AGENT_TOOLSDIRECTORY"
     - name: Build the Docker image
-      run: docker build . --file Dockerfile --tag openfold:$(date +%s)
\ No newline at end of file
+      run: docker build . --file Dockerfile --tag openfold:$(date +%s)

.readthedocs.yaml

+version: 2
+
+# Set the OS, Python version and other tools you might need
+build:
+  os: ubuntu-22.04
+  tools:
+    python: "mambaforge-4.10" 
+
+# Build documentation in the "docs/" directory with Sphinx
+sphinx:
+  configuration: docs/source/conf.py
+
+conda:
+  environment: docs/environment.yml

LICENSE

@@ -187,7 +187,7 @@
       same "printed page" as the copyright notice for easier
       identification within third-party archives.
 
-   Copyright [yyyy] [name of copyright owner]
+   Copyright 2024 AlQuraishi Laboratory 
 
    Licensed under the Apache License, Version 2.0 (the "License");
    you may not use this file except in compliance with the License.

docs/Makefile

+# Minimal makefile for Sphinx documentation
+#
+
+# You can set these variables from the command line, and also
+# from the environment for the first two.
+SPHINXOPTS    ?=
+SPHINXBUILD   ?= sphinx-build
+SOURCEDIR     = source
+BUILDDIR      = build
+
+# Put it first so that "make" without argument is like "make help".
+help:
+	@$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
+
+.PHONY: help Makefile
+
+# Catch-all target: route all unknown targets to Sphinx using the new
+# "make mode" option.  $(O) is meant as a shortcut for $(SPHINXOPTS).
+%: Makefile
+	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)

docs/environment.yml

+name: openfold-docs
+channels:
+  - conda-forge
+dependencies:
+  - sphinx=7
+  - myst-parser=3
+  - furo
\ No newline at end of file

docs/source/Aux_seq_files.md

+# Auxiliary Sequence Files for OpenFold Training
+
+The training dataset of OpenFold is very large. The `pdb` directory alone contains 185,000 mmcifs;  each chain for has multiple sequence alignment files and mmcif files. 
+
+OpenFold introduces a few new file structures for faster access to alignments and mmcif data.
+
+This documentation will explain the benefits of having the condensed file structure, and explain the contents of each of the files.
+
+## Default alignment file structure 
+
+One way to store mmcifs and alignments files would be to have a directory for each mmcif chain. 
+
+For example, consider two protein as a case study
+```
+- OpenProteinSet
+  └── mmcifs 
+	 ├── 3lrm.cif
+	 └── 6kwc.cif
+	 ...
+```
+
+In the `alignments` directory, [PDB:6KWC](https://www.rcsb.org/structure/6KWC) is a monomer with one chain, and thus would have one alignment direcotry. [PDB:3LRM](https://www.rcsb.org/structure/3lrm), a homotetramer, would have one alignment directory for each of its four chains.
+```
+- OpenProteinSet 
+  └── alignments 
+	  └── 3lrm_A 
+		  ├── bfd_uniclust_hits.a3m
+		  ├── mgnify_hits.a3m
+		  ├── pdb70_hits.hhr
+		  └── uniref90_hits.a3m 
+	  └── 3lrm_B 
+		  ├── bfd_uniclust_hits.a3m
+		  ├── mgnify_hits.a3m
+		  ├── pdb70_hits.hhr
+		  └── uniref90_hits.a3m 
+	  └── 3lrm_C 
+		  ├── bfd_uniclust_hits.a3m
+		  ├── mgnify_hits.a3m
+		  ├── pdb70_hits.hhr
+		  └── uniref90_hits.a3m 
+	  └── 3lrm_D 
+		  ├── bfd_uniclust_hits.a3m
+		  ├── mgnify_hits.a3m
+		  ├── pdb70_hits.hhr
+		  └── uniref90_hits.a3m 
+	  └── 6kwc_A 
+		  ├── bfd_uniclust_hits.a3m
+		  ├── mgnify_hits.a3m
+		  ├── pdb70_hits.hhr
+		  └── uniref90_hits.a3m 
+		...
+```
+
+In practice, the IO overhead of having one directory per protein chain makes accessing the alignments slow. 
+
+## OpenFold DB file structure 
+
+ Here we describe a new filesystem that can be used by OpenFold for more efficient access of alignment file and index file contents
+
+All together, the file directory would look like:
+```
+- OpenProteinSet 
+  ├── duplicate_pdb_chains.txt
+  └── pdb
+	  ├── mmcif_cache.json 
+	  └── mmcifs 
+		  ├── 3lrm.cif
+		  └── 6kwc.cif
+	  └── alignment_db
+		  ├── alignment_db_0.db 
+                  ├── alignment_db_1.db
+                   ...
+                  ├── alignment_db_9.db
+		  └── alignment_db.index 
+```
+
+We will describe each of the file types here.
+
+### Alignments db files and index files
+
+To speed up access of MSAs, OpenFold has an alternate alignments storage procedure. Instead of storing dedicated files for each single alignment, we consolidate large sets of alignments to single files referred to as _alignments_db's_. This can reduce I/O overhead and in practice we recommend using around 10 `alignments_db_x.db` files to store the total training set of OpenFold. During training, OpenFold can access each alignment using byte index pointers that are stored in a separate index file (`alignments_db.index`). The alignments for the `3LRM` and `6KWC` examples would be recorded in the index file as follows:
+
+```alignments_db.index
+{
+	...
+	"3lrm_A": {
+        "db": "alignment_db_0.db",
+        "files": [
+            ["bfd_uniclust_hits.a3m", 212896478938, 1680200],
+            ["mgnify_hits.a3m", 212893696883, 2782055],
+            ["pdb70_hits.hhr", 212898159138, 614978],
+            ["uniref90_hits.a3m", 212898774116, 6165789]
+        ]
+    },
+    "6kwc_A": {
+        "db": "alignment_db_1.db",
+        "files": [
+            ["bfd_uniclust_hits.a3m", 415618723280, 380289],
+            ["mgnify_hits.a3m", 415618556077, 167203],
+            ["pdb70_hits.hhr", 415619103569, 148672],
+            ["uniref90_hits.a3m", 415617547852, 1008225]
+        ]
+    }
+	...
+}
+```
+
+For each entry, the corresponding `alignment_db` file and the byte start location and number of bytes to read the respective alignments are given. For example, the alignment information in `bfd_uniclust_hits.a3m` for chain `3lrm_A` can be found in the database file `alignment_db_0.db`, starting at byte location `212896478938` and reading in the next `1680200` bytes.
+
+### Chain cache files and mmCIF cache files
+
+Information from the mmcif files can be parsed in advance to create a `chain_cache.json` or a `mmcif_cache.json`. For OpenFold, the `chain_cache.json` is used to sample chains for training, and the `mmcif_cache.json` is used to prefilter templates. 
+
+Here's what the chain_cache.json entry looks like for our examples:
+
+```chain_cache.json
+{
+	...
+	    "3lrm_A": {
+        "release_date": "2010-06-30",
+        "seq": "MFAFYFLTACISLKGVFGVSPSYNGLGLTPQMGWDNWNTFACDVSEQLLLDTADRISDLGLKDMGYKYIILDDCWSSGRDSDGFLVADEQKFPNGMGHVADHLHNNSFLFGMYSSAGEYTCAGYPGSLGREEEDAQFFANNRVDYLKYANCYNKGQFGTPEISYHRYKAMSDALNKTGRPVFYSLCNWGQDLTFYWGSGIANSWRMSGDVTAEFTRPDSRCPCDGDEYDCKYAGFHCSIMNILNKAAPMGQNAGVGGWNDLDNLEVGVGNLTDDEEKAHFSMWAMVKSPLIIGANVNNLKASSYSIYSQASVIAINQDSNGIPATRVWRYYVSDTDEYGQGEIQMWSGPLDNGDQVVALLNGGSVSRPMNTTLEEIFFDSNLGSKKLTSTWDIYDLWANRVDNSTASAILGRNKTATGILYNATEQSYKDGLSKNDTRLFGQKIGSLSPNAILNTTVPAHGIAFYRLRPSSDYKDDDDK",
+        "resolution": 2.7,
+        "cluster_size": 6
+    },
+    "3lrm_B": {
+        "release_date": "2010-06-30",
+        "seq": "MFAFYFLTACISLKGVFGVSPSYNGLGLTPQMGWDNWNTFACDVSEQLLLDTADRISDLGLKDMGYKYIILDDCWSSGRDSDGFLVADEQKFPNGMGHVADHLHNNSFLFGMYSSAGEYTCAGYPGSLGREEEDAQFFANNRVDYLKYANCYNKGQFGTPEISYHRYKAMSDALNKTGRPVFYSLCNWGQDLTFYWGSGIANSWRMSGDVTAEFTRPDSRCPCDGDEYDCKYAGFHCSIMNILNKAAPMGQNAGVGGWNDLDNLEVGVGNLTDDEEKAHFSMWAMVKSPLIIGANVNNLKASSYSIYSQASVIAINQDSNGIPATRVWRYYVSDTDEYGQGEIQMWSGPLDNGDQVVALLNGGSVSRPMNTTLEEIFFDSNLGSKKLTSTWDIYDLWANRVDNSTASAILGRNKTATGILYNATEQSYKDGLSKNDTRLFGQKIGSLSPNAILNTTVPAHGIAFYRLRPSSDYKDDDDK",
+        "resolution": 2.7,
+        "cluster_size": 6
+    },
+    "3lrm_C": {
+        "release_date": "2010-06-30",
+        "seq": "MFAFYFLTACISLKGVFGVSPSYNGLGLTPQMGWDNWNTFACDVSEQLLLDTADRISDLGLKDMGYKYIILDDCWSSGRDSDGFLVADEQKFPNGMGHVADHLHNNSFLFGMYSSAGEYTCAGYPGSLGREEEDAQFFANNRVDYLKYANCYNKGQFGTPEISYHRYKAMSDALNKTGRPVFYSLCNWGQDLTFYWGSGIANSWRMSGDVTAEFTRPDSRCPCDGDEYDCKYAGFHCSIMNILNKAAPMGQNAGVGGWNDLDNLEVGVGNLTDDEEKAHFSMWAMVKSPLIIGANVNNLKASSYSIYSQASVIAINQDSNGIPATRVWRYYVSDTDEYGQGEIQMWSGPLDNGDQVVALLNGGSVSRPMNTTLEEIFFDSNLGSKKLTSTWDIYDLWANRVDNSTASAILGRNKTATGILYNATEQSYKDGLSKNDTRLFGQKIGSLSPNAILNTTVPAHGIAFYRLRPSSDYKDDDDK",
+        "resolution": 2.7,
+        "cluster_size": 6
+    },
+    "3lrm_D": {
+        "release_date": "2010-06-30",
+        "seq": "MFAFYFLTACISLKGVFGVSPSYNGLGLTPQMGWDNWNTFACDVSEQLLLDTADRISDLGLKDMGYKYIILDDCWSSGRDSDGFLVADEQKFPNGMGHVADHLHNNSFLFGMYSSAGEYTCAGYPGSLGREEEDAQFFANNRVDYLKYANCYNKGQFGTPEISYHRYKAMSDALNKTGRPVFYSLCNWGQDLTFYWGSGIANSWRMSGDVTAEFTRPDSRCPCDGDEYDCKYAGFHCSIMNILNKAAPMGQNAGVGGWNDLDNLEVGVGNLTDDEEKAHFSMWAMVKSPLIIGANVNNLKASSYSIYSQASVIAINQDSNGIPATRVWRYYVSDTDEYGQGEIQMWSGPLDNGDQVVALLNGGSVSRPMNTTLEEIFFDSNLGSKKLTSTWDIYDLWANRVDNSTASAILGRNKTATGILYNATEQSYKDGLSKNDTRLFGQKIGSLSPNAILNTTVPAHGIAFYRLRPSSDYKDDDDK",
+        "resolution": 2.7,
+        "cluster_size": 6
+    },
+	"6kwc_A": {
+        "release_date": "2021-01-27",
+        "seq": "GSTIQPGTGYNNGYFYSYWNDGHGGVTYTNGPGGQFSVNWSNSGEFVGGKGWQPGTKNKVINFSGSYNPNGNSYLSVYGWSRNPLIEYYIVENFGTYNPSTGATKLGEVTSDGSVYDIYRTQRVNQPSIIGTATFYQYWSVRRNHRSSGSVNTANHFNAWAQQGLTLGTMDYQIVAVQGYFSSGSASITVS",
+        "resolution": 1.297,
+        "cluster_size": 195
+    },
+	...
+}
+```
+
+The mmcif_cache.json file would contain similar information, but condensed by mmcif id, e.g.
+
+```mmcif_cache.json
+{
+    "3lrm": {
+        "release_date": "2010-06-30",
+        "chain_ids": [
+            "A",
+            "B",
+            "C",
+            "D"
+        ],
+        "seqs": [
+            "MFAFYFLTACISLKGVFGVSPSYNGLGLTPQMGWDNWNTFACDVSEQLLLDTADRISDLGLKDMGYKYIILDDCWSSGRDSDGFLVADEQKFPNGMGHVADHLHNNSFLFGMYSSAGEYTCAGYPGSLGREEEDAQFFANNRVDYLKYANCYNKGQFGTPEISYHRYKAMSDALNKTGRPVFYSLCNWGQDLTFYWGSGIANSWRMSGDVTAEFTRPDSRCPCDGDEYDCKYAGFHCSIMNILNKAAPMGQNAGVGGWNDLDNLEVGVGNLTDDEEKAHFSMWAMVKSPLIIGANVNNLKASSYSIYSQASVIAINQDSNGIPATRVWRYYVSDTDEYGQGEIQMWSGPLDNGDQVVALLNGGSVSRPMNTTLEEIFFDSNLGSKKLTSTWDIYDLWANRVDNSTASAILGRNKTATGILYNATEQSYKDGLSKNDTRLFGQKIGSLSPNAILNTTVPAHGIAFYRLRPSSDYKDDDDK",
+            "MFAFYFLTACISLKGVFGVSPSYNGLGLTPQMGWDNWNTFACDVSEQLLLDTADRISDLGLKDMGYKYIILDDCWSSGRDSDGFLVADEQKFPNGMGHVADHLHNNSFLFGMYSSAGEYTCAGYPGSLGREEEDAQFFANNRVDYLKYANCYNKGQFGTPEISYHRYKAMSDALNKTGRPVFYSLCNWGQDLTFYWGSGIANSWRMSGDVTAEFTRPDSRCPCDGDEYDCKYAGFHCSIMNILNKAAPMGQNAGVGGWNDLDNLEVGVGNLTDDEEKAHFSMWAMVKSPLIIGANVNNLKASSYSIYSQASVIAINQDSNGIPATRVWRYYVSDTDEYGQGEIQMWSGPLDNGDQVVALLNGGSVSRPMNTTLEEIFFDSNLGSKKLTSTWDIYDLWANRVDNSTASAILGRNKTATGILYNATEQSYKDGLSKNDTRLFGQKIGSLSPNAILNTTVPAHGIAFYRLRPSSDYKDDDDK",
+            "MFAFYFLTACISLKGVFGVSPSYNGLGLTPQMGWDNWNTFACDVSEQLLLDTADRISDLGLKDMGYKYIILDDCWSSGRDSDGFLVADEQKFPNGMGHVADHLHNNSFLFGMYSSAGEYTCAGYPGSLGREEEDAQFFANNRVDYLKYANCYNKGQFGTPEISYHRYKAMSDALNKTGRPVFYSLCNWGQDLTFYWGSGIANSWRMSGDVTAEFTRPDSRCPCDGDEYDCKYAGFHCSIMNILNKAAPMGQNAGVGGWNDLDNLEVGVGNLTDDEEKAHFSMWAMVKSPLIIGANVNNLKASSYSIYSQASVIAINQDSNGIPATRVWRYYVSDTDEYGQGEIQMWSGPLDNGDQVVALLNGGSVSRPMNTTLEEIFFDSNLGSKKLTSTWDIYDLWANRVDNSTASAILGRNKTATGILYNATEQSYKDGLSKNDTRLFGQKIGSLSPNAILNTTVPAHGIAFYRLRPSSDYKDDDDK",
+            "MFAFYFLTACISLKGVFGVSPSYNGLGLTPQMGWDNWNTFACDVSEQLLLDTADRISDLGLKDMGYKYIILDDCWSSGRDSDGFLVADEQKFPNGMGHVADHLHNNSFLFGMYSSAGEYTCAGYPGSLGREEEDAQFFANNRVDYLKYANCYNKGQFGTPEISYHRYKAMSDALNKTGRPVFYSLCNWGQDLTFYWGSGIANSWRMSGDVTAEFTRPDSRCPCDGDEYDCKYAGFHCSIMNILNKAAPMGQNAGVGGWNDLDNLEVGVGNLTDDEEKAHFSMWAMVKSPLIIGANVNNLKASSYSIYSQASVIAINQDSNGIPATRVWRYYVSDTDEYGQGEIQMWSGPLDNGDQVVALLNGGSVSRPMNTTLEEIFFDSNLGSKKLTSTWDIYDLWANRVDNSTASAILGRNKTATGILYNATEQSYKDGLSKNDTRLFGQKIGSLSPNAILNTTVPAHGIAFYRLRPSSDYKDDDDK"
+        ],
+        "no_chains": 4,
+        "resolution": 2.7
+    },
+	    "6kwc": {
+        "release_date": "2021-01-27",
+        "chain_ids": [
+            "A"
+        ],
+        "seqs": [
+            "GSTIQPGTGYNNGYFYSYWNDGHGGVTYTNGPGGQFSVNWSNSGEFVGGKGWQPGTKNKVINFSGSYNPNGNSYLSVYGWSRNPLIEYYIVENFGTYNPSTGATKLGEVTSDGSVYDIYRTQRVNQPSIIGTATFYQYWSVRRNHRSSGSVNTANHFNAWAQQGLTLGTMDYQIVAVQGYFSSGSASITVS"
+        ],
+        "no_chains": 1,
+        "resolution": 1.297
+    },
+    ...
+}
+```
+
+
+### Duplicate pdb chain files 
+
+Duplicate chains occur across pdb entries. Some of these chains are the homomeric units of a multimer, others are subunits that are shared across different protein.
+
+To reduce storage overhead of creating / storing identical data for duplicate entries, we have a duplicate chain file. Each line stores the all chains that are identical. Our `6kwc` and `3lrm` examples would be stored as follows.
+
+```duplicate_pdb_chains.txt
+...
+6kwc_A
+3lrm_A 3lrm_B 3lrm_C 3lrm_D
+...
+```
+
+

docs/source/FAQ.md

+# FAQ
+
+Frequently asked questions or encountered issues when running OpenFold. 
+
+## Setup
+
+- When running unit tests (e.g. [`./scripts/run_unit_tests.sh`](https://github.com/aqlaboratory/openfold/blob/main/scripts/run_unit_tests.sh)), I see an error such as  
+	```
+	ImportError: version GLIBCXX_3.4.30 not found
+	```
+
+	> Solution: Make sure that the `$LD_LIBRARY_PATH` environment has been set to include the conda path, e.g. `export $LD_LIBRARY_PATH=$CONDA_PREFIX/lib:$LD_LIBRARY_PATH`
+
+- I see a CUDA mismatch error, eg. 
+```
+The detected CUDA version (11.8) mismatches the version that was used to compile
+PyTorch (12.1). Please make sure to use the same CUDA versions.
+```
+ 
+ > 	Solution: Ensure that your system's CUDA driver and toolkit match your intended OpenFold installation (CUDA 11 by default).  You can check the CUDA driver version with a command such as `nvidia-smi`
+
+- I get some error involving `fatal error: cuda_runtime.h: No such file or directory` and or `ninja: build stopped: subcommand failed.`. 
+
+> Solution: Something went wrong with setting up some of the custom kernels. Try running `install_third_party_dependencies.sh` again or try `python3 setup.py install` from inside the OpenFold folder. Make sure to prepend the conda environment as described above before running this.
+
+## Training
+
+- My model training is hanging on the data loading step:
+	 > Solution: While each system is different, a few general suggestions:
+		 - Check your `$KMP_AFFINITY` environment setting and see if it is suitable for your system.
+  		 - Adjust the number of data workers used to prepare data with the `--num_workers` setting. Increasing the number could help with dataset processing speed. However, to many workers could cause an OOM issue. 
+
+- When I reload my pretrained model weights or checkpoints, I get `RuntimeError: Error(s) in loading state_dict for OpenFoldWrapper: Unexpected key(s) in state_dict:`
+	> Solution: This suggests that your checkpoint / model weights are in OpenFold v1 format with outdated model layer names. Convert your weights/checkpoints following [this guide](convert_of_v1_weights.md).
\ No newline at end of file

docs/source/Inference.md

@@ -94,10 +94,10 @@ where `$PRECOMPUTED_ALIGNMENTS` is a directory that contains alignments. A sampl
 ```
 alignments
 └── 6KWC_1 
-    ├── bfd_uniclust_hits.a3m
-    ├── hhsearch_output.hhr
-    ├── mgnify_hits.sto
-    └── uniref90_hits.sto
+    ├── bfd_uniclust_hits.a3m
+    ├── hhsearch_output.hhr
+    ├── mgnify_hits.sto
+    └── uniref90_hits.sto
 ```
 
 `bfd_uniclust_hits.a3m`, `mgnify_hits.sto`, and `uniref90_hits.sto` are all alignments of the query structure against the BFD, Mgnify, and Uniref90 datasets respsectively. `hhsearch_output.hhr` contains hits against the PDB70 database used for template matching. The example directory `examples/monomer/alignments` shows examples of expected directories.
@@ -145,17 +145,17 @@ Some commonly used command line flags are here. A full list of flags can be view
 
 #### Speeding up inference 
 
-The **DeepSpeed DS4Sci_EvoformerAttention kernel** is a memory-efficient attention kernel developed as part of a collaboration between OpenFold and the DeepSpeed4Science initiative. 
+The **DeepSpeed DS4Sci_EvoformerAttention kernel** is a memory-efficient attention kernel developed as part of a collaboration between OpenFold and the DeepSpeed4Science initiative. 
 
 If your system supports deepseed, using deepspeed generally leads an inference speedup of 2 - 3x without significant additional memory use. You may specify this option by selecting the `--use_deepspeed_inference` argument. 
 
-If DeepSpeed is unavailable for your system, you may also try using [FlashAttention](https://github.com/HazyResearch/flash-attention) by adding `globals.use_flash = True` to the `--experiment_config_json`. Note that FlashAttention appears to work best for sequences with < 1000 residues.
+If DeepSpeed is unavailable for your system, you may also try using [FlashAttention](https://github.com/HazyResearch/flash-attention) by adding `globals.use_flash = True` to the `--experiment_config_json`. Note that FlashAttention appears to work best for sequences with < 1000 residues.
 
 #### Large-scale batch inference 
-For large-scale batch inference, we offer an optional tracing mode, which massively improves runtimes at the cost of a lengthy model compilation process. To enable it, add `--trace_model` to the inference command.
+For large-scale batch inference, we offer an optional tracing mode, which massively improves runtimes at the cost of a lengthy model compilation process. To enable it, add `--trace_model` to the inference command.
 
 #### Configuring the chunk size for sequence alignments
-Note that chunking (as defined in section 1.11.8 of the AlphaFold 2 supplement) is enabled by default in inference mode. To disable it, set `globals.chunk_size` to `None` in the config. If a value is specified, OpenFold will attempt to dynamically tune it, considering the chunk size specified in the config as a minimum. This tuning process automatically ensures consistently fast runtimes regardless of input sequence length, but it also introduces some runtime variability, which may be undesirable for certain users. It is also recommended to disable this feature for very long chains (see below). To do so, set the `tune_chunk_size` option in the config to `False`.
+Note that chunking (as defined in section 1.11.8 of the AlphaFold 2 supplement) is enabled by default in inference mode. To disable it, set `globals.chunk_size` to `None` in the config. If a value is specified, OpenFold will attempt to dynamically tune it, considering the chunk size specified in the config as a minimum. This tuning process automatically ensures consistently fast runtimes regardless of input sequence length, but it also introduces some runtime variability, which may be undesirable for certain users. It is also recommended to disable this feature for very long chains (see below). To do so, set the `tune_chunk_size` option in the config to `False`.
 
 #### Long sequence inference 
 To minimize memory usage during inference on long sequences, consider the following changes:

docs/source/Installation.md

+# Setting Up OpenFold 
+
+In this guide, we will OpenFold and its dependencies.
+
+**Pre-requisites**
+
+This package is currently supported for CUDA 11 and Pytorch 1.12. All dependencies are listed in the [`environment.yml`](https://github.com/aqlaboratory/openfold/blob/main/environment.yml)
+
+At this time, only Linux systems are supported.
+
+## Instructions
+:::
+
+### Installation:
+1. Clone the repository, e.g. `git clone https://github.com/aqlaboratory/openfold.git`
+1. From the `openfold` repo:
+    - Create a [Mamba]("https://github.com/conda-forge/miniforge/releases/latest/download/) environment, e.g.
+        `mamba env create -n openfold_env -f environment.yml`
+      Mamba is recommended as the dependencies required by OpenFold are quite large and mamba can speed up the process.
+    - Activate the environment, e.g `conda activate openfold_env`
+1. Run the setup script to configure kernels and folding resources.
+	> scripts/install_third_party_dependencies.sh`
+1. Prepend the conda environment to the `$LD_LIBRARY_PATH`., e.g. 
+		`export $LD_LIBRARY_PATH=$CONDA_PREFIX/lib:$LD_LIBRARY_PATH`. You may optionally set this as a conda environment variable according to the [conda docs](https://conda.io/projects/conda/en/latest/user-guide/tasks/manage-environments.html#saving-environment-variables) to activate each time the environment is used.
+1. Download parameters. We recommend using a destination as `openfold/resources` as our unittests will look for the weights there.
+	-  For AlphaFold2 weights, use 
+		> ./scripts/download_alphafold_params.sh <dest>
+	 - For OpenFold weights, use : 
+		>  ./scripts/download_openfold_params.sh <dest>
+	 - For OpenFold SoloSeq weights, use: 
+		> ./scripts/download_openfold_soloseq_params.sh <dest>
+
+### Checking your build with unit tests: 
+
+To test your installation, you can run OpenFold unit tests. Make sure that the OpenFold and AlphaFold parameters have been downloaded, and that they are located (or symlinked) in the directory `openfold/resources` 
+
+Run with the following script:
+> scripts/run_unit_tests.sh
+
+The script is a thin wrapper around Python's `unittest` suite, and recognizes `unittest` arguments. E.g., to run a specific test verbosely:
+
+> scripts/run_unit_tests.sh -v tests.test_model
+
+**Alphafold Comparison tests:**
+Certain tests perform equivalence comparisons with the AlphaFold implementation. Instructions to run this level of tests requires an environment with both AlphaFold 2.0.1 and OpenFold installed, and is not covered in this guide. These tests are skipped by default if no installation of AlphaFold is found. 
+
+## Environment specific modifications 
+
+### CUDA 12
+To use OpenFold on CUDA 12 environment rather than a CUDA 11 environment.
+	In step 1, use the branch [`pl_upgrades`](https://github.com/aqlaboratory/openfold/tree/pl_upgrades) rather than the main branch, i.e. replace the URL in step 1 with https://github.com/aqlaboratory/openfold/tree/pl_upgrades
+	Follow the rest of the steps of [Installation Guide](#Installation)
+
+
+### MPI
+To use OpenFold with MPI support, you will need to add the package [`mpi4py`](https://pypi.org/project/mpi4py/). This can be done with pip in your OpenFold environment, e.g. `$ pip install mpi4py`. 
+
+
+### Install OpenFold parameters without aws
+If you don't have access to `aws` on your system, you can use a different download source:
+
+- HuggingFace (requires `git-lts`):	`scripts/download_openfold_params_huggingface.sh`
+- Google Drive: `scripts/download_openfold_params_gdrive.sh`
+
+### Docker setup
+
+A [`Dockerfile`] is provided to build an OpenFold Docker image. Additional notes for setting up a docker container for OpenFold and running inference can be found [here](original_readme.md#building-and-using-the-docker-container).

docs/source/Multimer_Inference.md

+# Multimer Inference
+
+To run inference on a complex or multiple complexes using a set of DeepMind's pretrained parameters, you will need:
+
+- AlphaFold Multimer v2.3 parameters
+- Updated sequence databases, with UniRef and PDB Seqres databases. 
+
+
+## Upgrade from a previous OpenFold Installation 
+
+If you had previously downloaded OpenFold parameters and or AlphaFold databases, you will need to download updated versions. Here are some instructions for upgrading from an existing openfold installations. 
+
+### Download AlphaFold-Multimer v2.3 Model Parameters 
+1. Re-download the alphafold parameters to get the latest
+AlphaFold-Multimer v2.3 weights:
+    
+   ```bash
+    bash scripts/download_alphafold_params.sh openfold/resources
+   ```
+
+### Download AlphaFold Databases for Multimer 
+
+1. Download the [UniProt](https://www.uniprot.org/uniprotkb/) 
+and [PDB SeqRes](https://www.rcsb.org/) databases: 
+    
+   ```bash
+    bash scripts/download_uniprot.sh data/
+   ```
+    
+    The PDB SeqRes and PDB databases must be from the same date to avoid potential 
+    errors during template searching. Remove the existing `data/pdb_mmcif` directory 
+    and download both databases:
+    
+   ```bash
+    bash scripts/download_pdb_mmcif.sh data/
+    bash scripts/download_pdb_seqres.sh data/
+   ```
+
+1. Additionally, AlphaFold-Multimer uses upgraded versions of the [MGnify](https://www.ebi.ac.uk/metagenomics) 
+and [UniRef30](https://uniclust.mmseqs.com/) (previously UniClust30) databases. To download the upgraded databases, run:
+    
+   ```bash
+    bash scripts/download_uniref30.sh data/
+    bash scripts/download_mgnify.sh data/
+   ```
+
+```{note}
+Multimer inference can also run with the older database versions if desired. 
+```
+
+## Running Multimer Inference 
+
+The [`run_pretrained_openfold.py`](https://github.com/aqlaboratory/openfold/blob/main/run_pretrained_openfold.py) script can be used to run multimer inference with the follwoing command.
+
+```bash
+python3 run_pretrained_openfold.py \
+    fasta_dir \
+    data/pdb_mmcif/mmcif_files/ \
+    --uniref90_database_path data/uniref90/uniref90.fasta \
+    --mgnify_database_path data/mgnify/mgy_clusters_2022_05.fa \
+    --pdb_seqres_database_path data/pdb_seqres/pdb_seqres.txt \
+    --uniref30_database_path data/uniref30/UniRef30_2021_03 \
+    --uniprot_database_path data/uniprot/uniprot.fasta \
+    --bfd_database_path data/bfd/bfd_metaclust_clu_complete_id30_c90_final_seq.sorted_opt \
+    --jackhmmer_binary_path lib/conda/envs/openfold_venv/bin/jackhmmer \
+    --hhblits_binary_path lib/conda/envs/openfold_venv/bin/hhblits \
+    --hmmsearch_binary_path lib/conda/envs/openfold_venv/bin/hmmsearch \
+    --hmmbuild_binary_path lib/conda/envs/openfold_venv/bin/hmmbuild \
+    --kalign_binary_path lib/conda/envs/openfold_venv/bin/kalign \
+    --config_preset "model_1_multimer_v3" \
+    --model_device "cuda:0" \
+    --output_dir ./ 
+```
+
+**Notes:**
+- Template searching in the multimer pipeline uses HMMSearch with the PDB SeqRes database, replacing HHSearch and PDB70 used in the monomer pipeline.
+- As with monomer inference, if you've already computed alignments for the query, you can use the `--use_precomputed_alignments` option.
+- At this time, only AlphaFold parameter weights are available for multimer mode. 
\ No newline at end of file

docs/source/OpenFold_Parameters.md

+# Notes on OpenFold Training and Parameters
+
+For OpenFold model parameters, v. 06_22.
+
+## Training details
+
+OpenFold was trained using OpenFold on 44 A100s using the training schedule from Table 4 in
+the AlphaFold supplement. AlphaFold was used as the pre-distillation model.
+Training data is hosted publicly in the "OpenFold Training Data" RODA repository.
+
+To improve model diversity, we forked training after the initial training phase
+and finetuned an additonal branch without templates.
+
+## Parameter files
+
+Parameter files fall into the following categories:
+
+    initial_training.pt:
+        OpenFold at the end of the initial training phase.
+    finetuning_x.pt:
+        Checkpoints in chronological order corresponding to peaks in the
+        validation LDDT-Ca during the finetuning phase. Roughly evenly spaced
+        across the 45 finetuning epochs.
+
+        NOTE: finetuning_1.pt, which was included in a previous release, has
+        been deprecated.
+    finetuning_no_templ_x.pt
+        Checkpoints in chronological order corresponding to peaks during an
+        additional finetuning phase also starting from the 'initial_training.pt'
+        checkpoint but with templates disabled.
+    finetuning_no_templ_ptm_x.pt
+        Checkpoints in chronological order corresponding to peaks during the
+        pTM training phase of the `no_templ` branch. Models in this category
+        include the pTM module and comprise the most recent of the checkpoints
+        in said branch.
+    finetuning_ptm_x.pt:
+        Checkpoints in chronological order corresponding to peaks in the pTM
+        training phase of the mainline branch. Models in this category include
+        the pTM module and comprise the most recent of the checkpoints in said
+        branch.
+
+Average validation LDDT-Ca scores for each of the checkpoints are listed below.
+The validation set contains approximately 180 chains drawn from CAMEO over a
+three-month period at the end of 2021.
+
+    initial_training: 0.9088
+    finetuning_2: 0.9061
+    finetuning_3: 0.9075
+    finetuning_4: 0.9059
+    finetuning_5: 0.9054
+    finetuning_no_templ_1: 0.9014
+    finetuning_no_templ_2: 0.9032
+    finetuning_no_templ_ptm_1: 0.9025
+    finetuning_ptm_1: 0.9075
+    finetuning_ptm_2: 0.9097
\ No newline at end of file

docs/source/OpenFold_Training_Setup.md

+# Setting up the OpenFold PDB training set from RODA
+
+The multiple sequence alignments of OpenProteinSet and mmCIF structure files required to train OpenFold are freely available at the [Registry of Open Data on AWS (RODA)](https://registry.opendata.aws/openfold/). Additionally, OpenFold requires some postprocessing and [auxiliary files](Aux_seq_files.md) for training that need to be generated from the AWS data manually. This documentation is intended to give a full overview of those steps starting from the data download.
+
+### Pre-Requisites:
+- OpenFold conda environment. See [OpenFold Installation](Installation.md) for instructions on how to build this environment. 
+- In particular, the [AWS CLI](https://aws.amazon.com/cli/) is used to download data from RODA.
+- For this guide, we assume that the OpenFold codebase is located at `$OF_DIR`.
+
+## 1. Downloading alignments and structure files
+To fetch all the alignments corresponding to the original PDB training set of OpenFold alongside their mmCIF 3D structures, you can run the following commands:
+
+```bash
+mkdir -p alignment_data/alignment_dir_roda
+aws s3 cp s3://openfold/pdb/ alignment_data/alignment_dir_roda/ --recursive --no-sign-request
+
+mkdir pdb_data
+aws s3 cp s3://openfold/pdb_mmcif.zip pdb_data/ --no-sign-request
+aws s3 cp s3://openfold/duplicate_pdb_chains.txt . --no-sign-request
+unzip pdb_mmcif.zip -d pdb_data
+```
+
+The nested alignment directory structure is not yet exactly what OpenFold expects, so you can run the `flatten_roda.sh` script to convert them to the correct format:
+
+```bash
+bash $OF_DIR/scripts/flatten_roda.sh alignment_data/alignment_dir_roda alignment_data/
+```
+
+Afterwards, the old directory can be safely removed:
+
+```bash
+rm -r alignment_data/alignment_dir_roda
+```
+
+## 2. Creating alignment DBs (optional)
+As further explained in [Auxiliary Sequence Files in OpenFold](Aux_seq_files.md), OpenFold supports an alternate format for storing alignments that can increase training performance in I/O bottlenecked systems. These so-called `alignment_db` files can be generated with the following script:
+
+```bash
+python $OF_DIR/scripts/alignment_db_scripts/create_alignment_db_sharded.py \
+    alignment_data/alignments \
+    alignment_data/alignment_dbs \
+    alignment_db \
+    --n_shards 10 \
+    --duplicate_chains_file pdb_data/duplicate_pdb_chains.txt
+```
+
+We recommend creating 10 total `alignment_db` files (= "shards") for better
+filesystem health and fast preprocessing, but note that this script will only run
+optimally if the number of CPUs on your machine is at least as big as the number
+of shards you are creating.
+
+As an optional check, you can run the following command which should return $634,434$:
+
+```bash
+grep "files" alignment_data/alignment_dbs/alignment_db.index | wc -l
+```
+
+## 3. Adding duplicate chains to alignments (skip if step 2 was used)
+To save space, the OpenProteinSet alignment database is stored without duplicates, meaning that only one representative alignment is stored for all chains with identical sequences in the PDB and duplicate instances are tracked with a [`duplicate_chains.txt`](Aux_seq_files.md#duplicate-pdb-chain-files) file. As OpenFold will select chains during training based on the chains in the alignment directory (or `alignment_db`), we therefore need to add those duplicate chains back in in order to train on the full conformational diversity of chains in the PDB.
+
+If you've followed the optional Step 2, the `.index` file of your `alignment_db` files will have already been adjusted for duplicates and you can proceed to the next step. Otherwise, the standard alignment directory can be expanded to accommodate duplicates by inserting symlinked directories for the duplicate chains that point to their representative alignments:
+
+```bash
+python $OF_DIR/scripts/expand_alignment_duplicates.py \
+    alignment_data/alignments \
+    pdb_data/duplicate_pdb_chains.txt
+```
+
+As an optional check, the following command should return $634,434$:
+
+```bash
+ls alignment_data/alignments/ | wc -l
+```
+
+## 4. Generating cluster-files
+The AlphaFold dataloader adjusts the sampling probability of chains by their inverse cluster size, so we need to generate these sequence clusters for our training set.
+
+As a first step, we'll need a `.fasta` file of all sequences in the training set. This can be generated with the following scripts, depending on how you set up your alignment data in the previous steps:
+
+**Use this if you set up the duplicate-expanded alignment directory (faster):**
+```bash
+python $OF_DIR/scripts/alignment_data_to_fasta.py \
+    alignment_data/all-seqs.fasta \
+    --alignment_dir alignment_data/alignments
+```
+
+**Use this if you set up the `alignment_db` files:**
+```bash
+python $OF_DIR/scripts/alignment_data_to_fasta.py \
+    alignment_data/all-seqs.fasta \
+    --alignment_db_index alignment_data/alignment_dbs/alignment_db.index
+```
+
+Next, we need to generate a cluster file at 40% sequence identity, which will contain all chains in a particular cluster on the same line. You'll need [MMSeqs2](https://github.com/soedinglab/MMseqs2?tab=readme-ov-file#installation) for this as well, which can be set up either in a conda environment or as a binary.
+
+```bash
+python $OF_DIR/scripts/fasta_to_clusterfile.py \
+    alignment_data/all-seqs.fasta \
+    alignment_data/all-seqs_clusters-40.txt \
+    /path/to/mmseqs \
+    --seq-id 0.4
+```
+
+## 5. Generating cluster-files
+As a last step, OpenFold requires ["cache" files](Aux_seq_files.md#chain-cache-files-and-mmcif-cache-files) with metadata information for each chain that are used for choosing templates and samples during training.
+
+The data caches for OpenProteinSet can be downloaded from RODA with the following:
+
+```bash
+aws s3 cp s3://openfold/data_caches/ pdb_data/ --recursive --no-sign-request
+```
+If you wish to create data caches for your own datasets, the steps to generate the cache are as follows:
+
+```bash
+mkdir pdb_data/data_caches
+
+python $OF_DIR/scripts/generate_mmcif_cache.py \
+    pdb_data/mmcif_files \
+    pdb_data/data_caches/mmcif_cache.json \
+    --no_workers 16
+```
+
+The chain-data-cache is used for filtering training samples and adjusting per-chain sampling probabilities and can be generated with the following script:
+
+```bash
+python $OF_DIR/scripts/generate_chain_data_cache.py \
+    pdb_data/mmcif_files \
+    pdb_data/data_caches/chain_data_cache.json \
+    --cluster_file alignment_data/all-seqs_clusters-40.txt \
+    --no_workers 16
+```

docs/source/Single_Sequence_Inference.md

+### Soloseq inference
+
+MSA-free sequence to structure prediction using the [ESM-1b model](https://github.com/facebookresearch/esm) embeddings.
+
+To run inference for a sequence using the SoloSeq single-sequence model, you can either precompute ESM-1b embeddings in bulk, or you can generate them during inference.
+
+For generating ESM-1b embeddings in bulk, use the provided script: [`scripts/precompute_embeddings.py`](https://github.com/aqlaboratory/openfold/blob/main/scripts/precompute_embeddings.py). The script takes a directory of FASTA files (one sequence per file) and generates ESM-1b embeddings in the same format and directory structure as required by SoloSeq. Following is an example command to use the script:
+
+```shell
+python scripts/precompute_embeddings.py fasta_dir/ embeddings_output_dir/
+```
+
+In the same per-label subdirectories inside `embeddings_output_dir`, you can also place `*.hhr` files (outputs from HHSearch), which can contain the details about the structures that you want to use as templates. If you do not place any such file, templates will not be used and only the ESM-1b embeddings will be used to predict the structure. If you want to use templates, you need to pass the PDB MMCIF dataset to the command.
+
+Then download the SoloSeq model weights, e.g.:
+
+```shell
+bash scripts/download_openfold_soloseq_params.sh openfold/resources
+```
+
+Now, you are ready to run inference:
+
+```shell
+python run_pretrained_openfold.py \
+    fasta_dir \
+    data/pdb_mmcif/mmcif_files/ \
+    --use_precomputed_alignments embeddings_output_dir \
+    --output_dir ./ \
+    --model_device "cuda:0" \
+    --config_preset "seq_model_esm1b_ptm" \
+    --openfold_checkpoint_path openfold/resources/openfold_soloseq_params/seq_model_esm1b_ptm.pt
+```
+
+For generating the embeddings during inference, skip the `--use_precomputed_alignments` argument. The `*.hhr` files will be generated as well if you pass the paths to the relevant databases and tools, as specified in the command below. If you skip the database and tool arguments, HHSearch will not be used to find templates and only generated ESM-1b embeddings will be used to predict the structure.
+
+```shell
+python3 run_pretrained_openfold.py \
+    fasta_dir \
+    data/pdb_mmcif/mmcif_files/ \
+    --output_dir ./ \
+    --model_device "cuda:0" \
+    --config_preset "seq_model_esm1b_ptm" \
+    --openfold_checkpoint_path openfold/resources/openfold_soloseq_params/seq_model_esm1b_ptm.pt \
+    --uniref90_database_path data/uniref90/uniref90.fasta \
+    --pdb70_database_path data/pdb70/pdb70 \
+    --jackhmmer_binary_path lib/conda/envs/openfold_venv/bin/jackhmmer \
+    --hhsearch_binary_path lib/conda/envs/openfold_venv/bin/hhsearch \
+    --kalign_binary_path lib/conda/envs/openfold_venv/bin/kalign \
+```
+
+For generating template information, you will need the UniRef90 and PDB70 databases and the JackHmmer and HHSearch binaries.
+
+SoloSeq allows you to use the same flags and optimizations as the MSA-based OpenFold. For example, you can skip relaxation using `--skip_relaxation`, save all model outputs using `--save_outputs`, and generate output files in MMCIF format using `--cif_output`.
+
+```{note}
+Due to the nature of the ESM-1b embeddings, the sequence length for inference using the SoloSeq model is limited to 1022 residues. Sequences longer than that will be truncated.
+```
\ No newline at end of file

docs/source/Training_OpenFold.md

+# Training OpenFold
+## Background
+
+This guide covers how to train an OpenFold model for monomers. Some additional instructions are provided at the end for fine-tuning your model.
+
+### Pre-requisites: 
+
+This guide requires the following:
+- [Installation of OpenFold and dependencies](Installation.md) (Including jackhmmer and hhblits depedencies)
+- A preprocessed dataset:
+	- For this guide, we will use the original OpenFold dataset which is available on RODA, processed with [these instructions](OpenFold_Training_Setup.md).
+- GPUs configured with CUDA. Training OpenFold with CPUs only is not supported. 
+
+## Training a new OpenFold model 
+
+#### Basic command
+
+For a dataset that has the default alignment file structure, e.g. 
+
+```
+-$DATA_DIR
+  └── pdb_data 
+ 	   ├──  mmcifs 
+ 	   	    ├── 3lrm.cif
+ 	   	    └── 6kwc.cif
+			...
+	   ├── obsolete.dat 	
+	   ├── duplicate_pdb_chains.txt 
+	   └── data_caches 
+	   		├── duplicate_pdb_chains.txt 
+	   		└── data_caches 
+  └── alignment_data
+	  └── alignments
+   	  		├── 3lrm_A/ 
+   	  		├── 3lrm_B/ 
+	  	    └── 6kwc_A/
+	  		...
+```
+
+The basic command to train a new OpenFold model is: 
+
+```
+python3 train_openfold.py $DATA_DIR/pdb/mmcifs $DATA_DIR/alignment_data/alignments $TEMPLATE_MMCIF_DIR $OUTPUT_DIR \
+    --max_template_date 2021-10-10 \ 
+    --train_chain_data_cache_path $DATA_DIR/pdb_data/data_caches/chain_data_cache.json \
+    --template_release_dates_cache_path $DATA_DIR/pdb_data/data_caches/mmcif_cache.json \ 
+	--config_preset initial_training \
+    --seed 42 \
+    --obsolete_pdbs_file_path $DATA_DIR/pdb_data/obsolete.dat \
+    --num_nodes 1 \
+    --gpus 4 \
+    --num_workers 4
+```
+
+The required arguments are:
+- `mmcif_dir` : Mmcif files for the training set.
+- `alignments_dir`: Alignments for the sequences in `mmcif_dir`, see expected directory structure 
+- `template_mmcif_dir`:  Template mmcif files with structures, which can be the same directory as mmcif_dir. The `max_template_date` and `template_release_dates_cache_path` will specify which templates will be allowed based on a date cutoff
+- `output_dir` : Where model checkpoint files and other outputs will be saved. 
+
+Commonly used flags include:
+- `config_preset`: Specifies which selection of hyperparameters should be used for initial model training. Commonly used configs are defined in [`openfold/config.py`](https://github.com/aqlaboratory/openfold)
+- `num_nodes` and `gpus`:  Specifies number of nodes and GPUs available to train OpenFold.
+- `seed` - Specifies random seed
+- `num_workers`: Number of CPU workers to assign for creating dataset examples
+- `obsolete_pdbs_file_path`: Specifies obsolete pdb IDs that should be excluded from training.
+- `val_data_dir` and `val_alignment_dir`: Specifies data directory and alignments for validation dataset. 
+
+```{note}
+Note that `--seed` must be specified to correctly configure training examples on multi-GPU training runs
+```
+
+#### Train with OpenFold Dataset Configuration
+
+If the [OpenFold alignment database](OpenFold_Training_Setup.md#2-creating-alignment-dbs-optional) setup is used, resulting in a data directory such as:
+```
+- $DATA_DIR 
+  ├── duplicate_pdb_chains.txt
+  ├── pdb_data
+	  └── mmcifs 
+		  ├── 3lrm.cif
+		  └── 6kwc.cif
+  └── alignment_data 
+	  └── alignment_db
+		  ├── alignment_db_0.db 
+          ├── alignment_db_1.db
+          ...
+          ├── alignment_db_9.db
+		  └── alignment_db.index 
+```
+
+The training command will use the `alignment_index_path` argument to specify `db.index` files, e.g.: 
+
+```
+python3 train_openfold.py $DATA_DIR/pdb_data/mmcifs $DATA_DIR/alignment_data/alignment_db $TEMPLATE_MMCIF_DIR $OUTPUT_DIR \
+    --max_template_date 2021-10-10 \ 
+    --train_chain_data_cache_path $DATA_DIR/pdb_data/data_caches/chain_data_cache.json \
+    --template_release_dates_cache_path $DATA_DIR/pdb_data/data_caches/mmcif_cache.json \ 
+	--alignment_index_path $DATA_DIR/pdb/alignment_db.index 
+	--config_preset initial_training \
+    --seed 42 \
+    --obsolete_pdbs_file_path $DATA_DIR/pdb/obsolete.dat \
+    --num_nodes 1 \
+    --gpus 4 \
+    --num_workers 4
+```
+
+#### Additional command line flag options:
+
+Here we provide brief descriptions for customizing your training run of OpenFold. A full description of all flags can be accessed by using the `--help` option in the script 
+
+- **Use Deepspeed acceleration strategy:** `--deepspeed_config` This option configures OpenFold to use custom Deepspeed kernels. This option requires a deepspeed_config.json, you can create your own, or use the one in the OpenFold directory 
+
+- **Use a validation dataset:** Specify validation database paths with `--val_data_dir` + `--val_alignment_dir`. Validation metrics will be evaluated on these datasets.
+
+- **Use a self-distillation dataset:**  Specify paths with `--distillation_data_dir` and `--distillation_alignment_dir` flags
+
+- **Change specific parameters in the model or data setup:**  `--experiment_config_json`. These parameters must be defined in the [`openfold/config.py`](https://github.com/aqlaboratory/openfold/blob/main/openfold/config.py). For example to change the crop size for training a model, you can write the following json:
+	```cropsize.json
+	{
+			"data.train.crop_size": 128
+	}
+	```
+
+- **Configure training settings with PyTorch Lightning** 
+	
+	Some flags e.g. `--precision`, `--max_epochs` configure training behavior. See the Pytorch Lightning Trainer args section in the `--help`  menu for more information and consult [Pytorch lightning documentation](https://lightning.ai/docs/pytorch/stable/)
+	
+	- Precision: On A100s, OpenFold training works best with bfloat 16 precision (e.g. `--precision bf16-mixed`) 
+	
+- **Restart training from an existing checkpoint:** Use the `--resume_from_ckpt` to restart training from an existing checkpoint.
+
+## Advanced Training Configurations 
+:::
+
+### Fine tuning from existing model weights 
+
+If you have existing model weights, you can fine tune the model by specifying a checkpoint path with `--resume_from_ckpt` and `--resume_model_weights_only` arguments, e.g. 
+
+```
+python3 train_openfold.py $DATA_DIR/mmcifs $DATA_DIR/alignment.db $TEMPLATE_MMCIF_DIR $OUTPUT_DIR \
+    --max_template_date 2021-10-10 \ 
+    --train_chain_data_cache_path chain_data_cache.json \
+    --template_release_dates_cache_path mmcif_cache.json \ 
+	--config_preset finetuning \
+	--alignment_index_path $DATA_DIR/pdb/alignment_db.index \ 
+    --seed 4242022 \
+    --obsolete_pdbs_file_path obsolete.dat \
+    --num_nodes 1 \
+    --gpus 4 \
+    --num_workers 4 \
+	--resume_from_ckpt $CHECKPOINT_PATH \
+	--resume_model_weights_only
+```
+
+If you have model parameters from OpenFold v1.x, you may need to convert your checkpoint file or parameter. See [Converting OpenFold v1 Weights](convert_of_v1_weights.md) for more details. 
+
+### Using MPI
+
+If MPI is configured on your system, and you would like to use MPI to train OpenFold models, you may do so with the following step:
+
+ 1. Add the `mpi4py` package, which are available through pip and conda. Please see [mpi4py documentation](https://pypi.org/project/mpi4py/) for more instructions on installation.
+2. Add the `--mpi_plugin` flag to your training command.
+
+
+### Training Multimer  models
+
+```{note}
+Coming soon.
+```
\ No newline at end of file

docs/source/conf.py

+# Configuration file for the Sphinx documentation builder.
+#
+# For the full list of built-in configuration values, see the documentation:
+# https://www.sphinx-doc.org/en/master/usage/configuration.html
+
+# -- Project information -----------------------------------------------------
+# https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information
+
+project = 'OpenFold'
+copyright = '2024, OpenFold Team'
+author = 'OpenFold Team'
+
+# -- General configuration ---------------------------------------------------
+# https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration
+
+extensions = [
+	'myst_parser',
+]
+
+templates_path = ['_templates']
+exclude_patterns = []
+
+
+
+# -- Options for HTML output -------------------------------------------------
+# https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output
+
+html_theme = 'furo'
+html_static_path = ['_static']
+myst_enable_extensions = ["colon_fence", "dollarmath", "amsmath"]
+

docs/source/convert_of_v1_weights.md

+## Weights Renaming 
+
+As part of the [OpenFold v2 update](https://github.com/aqlaboratory/openfold/releases/tag/v2.0.0) with the integration of multimer prediction, certain model layers of the AlphaFold model were renamed. For example.
+
+`module.model.template_angle_embedder.*` is now referred to as 
+`module.model.template_embedder.template_single_embedder.*`
+
+If you have some checkpoints that were trained using OpenFold v1 or older, and now want to resume training on OpenFold v2, you may need to convert your checkpoints.
+
+## FAQ
+
+### Do I need to convert my checkpoints / model weights?
+
+If you want to run inference or resume training from a checkpoint that was trained with OpenFold V1, you will need to convert your checkpoint.
+
+If you want load model weights only, without starting from a specific time step, then you should not need to convert your checkpoints. The training of the model will begin from `step=0` in this case. To do so, you'll need both the `--resume_from_ckpt` and `--resume_model_weights_only`  flags. This example allows you train starting from the pre-trained openfold weights:
+
+```bash
+$ python3 $OPENFOLD_DIR/train_openfold.py test_data_epoch/mmcifs test_data_epoch/alignments test_data_epoch/template_mmcifs $OUTPUT_DIR 2021-09-30 \
+	...
+	--resume_from_ckpt openfold/resources/openfold_params/finetuning_2.pt \
+	--resume_model_weights_only
+
+```
+
+### How do I convert my checkpoints? 
+
+Use [`scripts/convert_v1_to_v2_weights.py`](https://github.com/aqlaboratory/openfold/blob/main/scripts/convert_v1_to_v2_weights.py) e.g.
+
+	`python scripts/convert_v1_to_v2_weights.py checkpoints/6-209.ckpt checkpoints/6-209.ckpt.converted`
+
+Then, to resume training, set the following flags:
+
+```bash
+$ python3 $OPENFOLD_DIR/train_openfold.py test_data_epoch/mmcifs test_data_epoch/alignments test_data_epoch/template_mmcifs $OUTPUT_DIR 2021-09-30 \
+	...
+	--resume_from_ckpt checkpoints/6-209.ckpt.converted
+```
\ No newline at end of file

docs/source/index.md

+# OpenFold
+
+```{figure} ../imgs/of_banner.png
+:width: 900px
+:align: center
+:alt: Comparison of OpenFold and AlphaFold2 predictions to the experimental structure of PDB 7KDX, chain B._
+```
+Welcome to the Documentation for [OpenFold](https://github.com/aqlaboratory/openfold), the fully open source, trainable, PyTorch-based reproduction of DeepMind's 
+[AlphaFold 2](https://github.com/deepmind/alphafold).
+
+Here, you will find guides and documentation for:
+- [Getting started with OpenFold](installation.md)!
+- Learn how to [run inference with OpenFold](Inference.md)
+- [Train your own OpenFold models](Training_OpenFold.md)
+- Find guidance for setup and running OpenFold in the [FAQ](FAQ.md).
+
+We also have a [Colab notebook](https://colab.research.google.com/github/aqlaboratory/openfold/blob/main/notebooks/OpenFold.ipynb) that can be used for single structure / multimer prediction.
+
+Some portions of the documentation are still under migration from the original README, which can be found [here](original_readme.md).
+
+# Features
+
+OpenFold carefully reproduces (almost) all of the features of the original open
+source monomer (v2.0.1) and multimer (v2.3.2) inference code. The sole exception is 
+model ensembling, which fared poorly in DeepMind's own ablation testing and is being 
+phased out in future DeepMind experiments. It is omitted here for the sake of reducing 
+clutter. In cases where the *Nature* paper differs from the source, we always defer to the 
+latter.
+
+OpenFold is trainable in full precision, half precision, or `bfloat16` with or without DeepSpeed, 
+and we've trained it from scratch, matching the performance of the original. 
+We've publicly released model weights and our training data — some 400,000 
+MSAs and PDB70 template hit files — under a permissive license. Model weights 
+are available via scripts in this repository while the MSAs are hosted by the 
+[Registry of Open Data on AWS (RODA)](https://registry.opendata.aws/openfold). 
+Try out running inference for yourself with our [Colab notebook](https://colab.research.google.com/github/aqlaboratory/openfold/blob/main/notebooks/OpenFold.ipynb).
+
+OpenFold also supports inference using AlphaFold's official parameters, and 
+vice versa (see `scripts/convert_of_weights_to_jax.py`).
+
+OpenFold has the following advantages over the reference implementation:
+
+- **Faster inference** on GPU, sometimes by as much as 2x. The greatest speedups are achieved on Ampere or higher architecture GPUs.
+- **Inference on extremely long chains**, made possible by our implementation of low-memory attention 
+([Rabe & Staats 2021](https://arxiv.org/pdf/2112.05682.pdf)). OpenFold can predict the structures of
+  sequences with more than 4000 residues on a single A100, and even longer ones with CPU offloading.
+- **Custom CUDA attention kernels** modified from [FastFold](https://github.com/hpcaitech/FastFold)'s 
+kernels support in-place attention during inference and training. They use 
+4x and 5x less GPU memory than equivalent FastFold and stock PyTorch 
+implementations, respectively.
+- **Efficient alignment scripts** using the original AlphaFold HHblits/JackHMMER pipeline or [ColabFold](https://github.com/sokrypton/ColabFold)'s, which uses the faster MMseqs2 instead. We've used them to generate millions of alignments.
+- **FlashAttention** support greatly speeds up MSA attention.
+- **DeepSpeed DS4Sci_EvoformerAttention kernel** is a memory-efficient attention kernel developed as part of a collaboration between OpenFold and the DeepSpeed4Science initiative. The kernel provides substantial speedups for training and inference, and significantly reduces the model's peak device memory requirement by 13X. The model is 15% faster during the initial training and finetuning stages, and up to 4x faster during inference.
+
+# Copyright Notice
+
+While AlphaFold's and, by extension, OpenFold's source code is licensed under
+the permissive Apache Licence, Version 2.0, DeepMind's pretrained parameters 
+fall under the CC BY 4.0 license, a copy of which is downloaded to 
+`openfold/resources/params` by the installation script. Note that the latter
+replaces the original, more restrictive CC BY-NC 4.0 license as of January 2022.
+
+## Contributing
+
+If you encounter problems using OpenFold, feel free to create an issue! We also
+welcome pull requests from the community.
+
+## Citing this Work
+
+Please cite our paper:
+
+```bibtex
+@article {Ahdritz2022.11.20.517210,
+	author = {Ahdritz, Gustaf and Bouatta, Nazim and Floristean, Christina and Kadyan, Sachin and Xia, Qinghui and Gerecke, William and O{\textquoteright}Donnell, Timothy J and Berenberg, Daniel and Fisk, Ian and Zanichelli, Niccolò and Zhang, Bo and Nowaczynski, Arkadiusz and Wang, Bei and Stepniewska-Dziubinska, Marta M and Zhang, Shang and Ojewole, Adegoke and Guney, Murat Efe and Biderman, Stella and Watkins, Andrew M and Ra, Stephen and Lorenzo, Pablo Ribalta and Nivon, Lucas and Weitzner, Brian and Ban, Yih-En Andrew and Sorger, Peter K and Mostaque, Emad and Zhang, Zhao and Bonneau, Richard and AlQuraishi, Mohammed},
+	title = {{O}pen{F}old: {R}etraining {A}lpha{F}old2 yields new insights into its learning mechanisms and capacity for generalization},
+	elocation-id = {2022.11.20.517210},
+	year = {2022},
+	doi = {10.1101/2022.11.20.517210},
+	publisher = {Cold Spring Harbor Laboratory},
+	URL = {https://www.biorxiv.org/content/10.1101/2022.11.20.517210},
+	eprint = {https://www.biorxiv.org/content/early/2022/11/22/2022.11.20.517210.full.pdf},
+	journal = {bioRxiv}
+}
+```
+If you use OpenProteinSet, please also cite:
+
+```bibtex
+@misc{ahdritz2023openproteinset,
+      title={{O}pen{P}rotein{S}et: {T}raining data for structural biology at scale}, 
+      author={Gustaf Ahdritz and Nazim Bouatta and Sachin Kadyan and Lukas Jarosch and Daniel Berenberg and Ian Fisk and Andrew M. Watkins and Stephen Ra and Richard Bonneau and Mohammed AlQuraishi},
+      year={2023},
+      eprint={2308.05326},
+      archivePrefix={arXiv},
+      primaryClass={q-bio.BM}
+}
+```
+Any work that cites OpenFold should also cite [AlphaFold](https://www.nature.com/articles/s41586-021-03819-2) and [AlphaFold-Multimer](https://www.biorxiv.org/content/10.1101/2021.10.04.463034v1) if applicable.
+
+
+```{toctree}
+:hidden: 
+:caption: Guides
+Installation.md
+Inference.md
+Single_Sequence_Inference.md
+Multimer_Inference.md
+OpenFold_Training_Setup.md
+Training_OpenFold.md
+```
+
+```{toctree}
+:hidden: 
+:caption: Reference 
+Aux_seq_files.md
+OpenFold_Parameters.md
+FAQ.md
+original_readme.md
+```

docs/source/original_readme.md

+# Original OpenFold README 
+
+A faithful but trainable PyTorch reproduction of DeepMind's 
+[AlphaFold 2](https://github.com/deepmind/alphafold).
+
+## Contents
+
+- [OpenFold]
+  - [Contents]
+  - [Features](#features)
+  - [Installation (Linux)](#installation-linux)
+  - [Download Alignment Databases](#download-alignment-databases)
+  - [Inference](#inference)
+    - [Monomer inference](#monomer-inference)
+    - [Multimer Inference](#multimer-inference)
+    - [Soloseq Inference](#soloseq-inference)
+  - [Training](#training)
+  - [Testing](#testing)
+  - [Building and Using the Docker Container](#building-and-using-the-docker-container)
+  - [Copyright Notice](#copyright-notice)
+  - [Contributing](#contributing)
+  - [Citing this Work](#citing-this-work)
+
+## Features
+
+OpenFold carefully reproduces (almost) all of the features of the original open
+source monomer (v2.0.1) and multimer (v2.3.2) inference code. The sole exception is 
+model ensembling, which fared poorly in DeepMind's own ablation testing and is being 
+phased out in future DeepMind experiments. It is omitted here for the sake of reducing 
+clutter. In cases where the *Nature* paper differs from the source, we always defer to the 
+latter.
+
+OpenFold is trainable in full precision, half precision, or `bfloat16` with or without DeepSpeed, 
+and we've trained it from scratch, matching the performance of the original. 
+We've publicly released model weights and our training data — some 400,000 
+MSAs and PDB70 template hit files — under a permissive license. Model weights 
+are available via scripts in this repository while the MSAs are hosted by the 
+[Registry of Open Data on AWS (RODA)](https://registry.opendata.aws/openfold). 
+Try out running inference for yourself with our [Colab notebook](https://colab.research.google.com/github/aqlaboratory/openfold/blob/main/notebooks/OpenFold.ipynb).
+
+OpenFold also supports inference using AlphaFold's official parameters, and 
+vice versa (see `scripts/convert_of_weights_to_jax.py`).
+
+OpenFold has the following advantages over the reference implementation:
+
+- **Faster inference** on GPU, sometimes by as much as 2x. The greatest speedups are achieved on Ampere or higher architecture GPUs.
+- **Inference on extremely long chains**, made possible by our implementation of low-memory attention 
+([Rabe & Staats 2021](https://arxiv.org/pdf/2112.05682.pdf)). OpenFold can predict the structures of
+  sequences with more than 4000 residues on a single A100, and even longer ones with CPU offloading.
+- **Custom CUDA attention kernels** modified from [FastFold](https://github.com/hpcaitech/FastFold)'s 
+kernels support in-place attention during inference and training. They use 
+4x and 5x less GPU memory than equivalent FastFold and stock PyTorch 
+implementations, respectively.
+- **Efficient alignment scripts** using the original AlphaFold HHblits/JackHMMER pipeline or [ColabFold](https://github.com/sokrypton/ColabFold)'s, which uses the faster MMseqs2 instead. We've used them to generate millions of alignments.
+- **FlashAttention** support greatly speeds up MSA attention.
+- **DeepSpeed DS4Sci_EvoformerAttention kernel** is a memory-efficient attention kernel developed as part of a collaboration between OpenFold and the DeepSpeed4Science initiative. The kernel provides substantial speedups for training and inference, and significantly reduces the model's peak device memory requirement by 13X. The model is 15% faster during the initial training and finetuning stages, and up to 4x faster during inference. To use this feature, simply set the `use_deepspeed_evo_attention` option in `openfold/config.py`.
+
+## Installation (Linux)
+
+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 your system. You'll need `git-lfs` to download OpenFold parameters. 
+Finally, some download scripts require `aria2c` and `aws`.
+
+This package is currently supported for CUDA 11 and Pytorch 1.12
+
+To install:
+1. Clone the repository, e.g. `git clone https://github.com/aqlaboratory/openfold.git`
+1. From the `openfold` repo: 
+    - Create a [Mamba]("https://github.com/conda-forge/miniforge/releases/latest/download/) environment, e.g. 
+        `mamba env create -n openfold_env -f environment.yml`
+      Mamba is recommended as the dependencies required by OpenFold are quite large and mamba can speed up the process.
+    - Activate the environment, e.g `conda activate openfold_env` 
+1. Run `scripts/install_third_party_dependencies.sh` to configure kernels and folding resources.
+
+For some systems, it may help to append the Conda environment library path to `$LD_LIBRARY_PATH`. The `install_third_party_dependencies.sh` script does this once, but you may need this for each bash instance.
+
+
+## Download Alignment Databases
+
+If you intend to generate your own alignments, e.g. for inference, you have two 
+choices for downloading protein databases, depending on whether you want to use 
+DeepMind's MSA generation pipeline (w/ HMMR & HHblits) or 
+[ColabFold](https://github.com/sokrypton/ColabFold)'s, which uses the faster
+MMseqs2 instead. For the former, run:
+
+```bash
+bash scripts/download_alphafold_dbs.sh data/
+```
+
+For the latter, run:
+
+```bash
+bash scripts/download_mmseqs_dbs.sh data/    # downloads .tar files
+bash scripts/prep_mmseqs_dbs.sh data/        # unpacks and preps the databases
+```
+
+Make sure to run the latter command on the machine that will be used for MSA
+generation (the script estimates how the precomputed database index used by
+MMseqs2 should be split according to the memory available on the system).
+
+If you're using your own precomputed MSAs or MSAs from the RODA repository, 
+there's no need to download these alignment databases. Simply make sure that
+the `alignment_dir` contains one directory per chain and that each of these
+contains alignments (.sto, .a3m, and .hhr) corresponding to that chain. You
+can use `scripts/flatten_roda.sh` to reformat RODA downloads in this way.
+Note that the RODA alignments are NOT compatible with the recent .cif ground
+truth files downloaded by `scripts/download_alphafold_dbs.sh`. To fetch .cif 
+files that match the RODA MSAs, once the alignments are flattened, use 
+`scripts/download_roda_pdbs.sh`. That script outputs a list of alignment dirs 
+for which matching .cif files could not be found. These should be removed from 
+the alignment directory.
+
+Alternatively, you can use raw MSAs from 
+[ProteinNet](https://github.com/aqlaboratory/proteinnet). After downloading
+that database, use `scripts/prep_proteinnet_msas.py` to convert the data 
+into a format recognized by the OpenFold parser. The resulting directory 
+becomes the `alignment_dir` used in subsequent steps. Use 
+`scripts/unpack_proteinnet.py` to extract `.core` files from ProteinNet text 
+files.
+
+For both inference and training, the model's hyperparameters can be tuned from
+`openfold/config.py`. Of course, if you plan to perform inference using 
+DeepMind's pretrained parameters, you will only be able to make changes that
+do not affect the shapes of model parameters. For an example of initializing
+the model, consult `run_pretrained_openfold.py`.
+
+## Inference
+
+OpenFold now supports three inference modes: 
+- [Monomer Inference](#monomer-inference): OpenFold reproduction of AlphaFold2. Inference available with either DeepMind's pretrained parameters or OpenFold trained parameters. 
+- [Multimer Inference](#multimer-inference): OpenFold reproduction of AlphaFold-Multimer. Inference available with DeepMind's pre-trained parameters. 
+- [Single Sequence Inference (SoloSeq)](#soloseq-inference): Language Model based structure prediction, using [ESM-1b](https://github.com/facebookresearch/esm) embeddings.  
+
+More instructions for each inference mode are provided below:
+
+### Monomer inference
+
+To run inference on a sequence or multiple sequences using a set of DeepMind's 
+pretrained parameters, first download the OpenFold weights e.g.:
+
+```bash 
+bash scripts/download_openfold_params.sh openfold/resources
+```
+
+then run e.g.:
+
+```bash
+python3 run_pretrained_openfold.py \
+    fasta_dir \
+    data/pdb_mmcif/mmcif_files/ \
+    --uniref90_database_path data/uniref90/uniref90.fasta \
+    --mgnify_database_path data/mgnify/mgy_clusters_2018_12.fa \
+    --pdb70_database_path data/pdb70/pdb70 \
+    --uniclust30_database_path data/uniclust30/uniclust30_2018_08/uniclust30_2018_08 \
+    --bfd_database_path data/bfd/bfd_metaclust_clu_complete_id30_c90_final_seq.sorted_opt \
+    --jackhmmer_binary_path lib/conda/envs/openfold_venv/bin/jackhmmer \
+    --hhblits_binary_path lib/conda/envs/openfold_venv/bin/hhblits \
+    --hhsearch_binary_path lib/conda/envs/openfold_venv/bin/hhsearch \
+    --kalign_binary_path lib/conda/envs/openfold_venv/bin/kalign \
+    --config_preset "model_1_ptm" \
+    --model_device "cuda:0" \
+    --output_dir ./ \
+    --openfold_checkpoint_path openfold/resources/openfold_params/finetuning_ptm_2.pt
+```
+
+where `data` is the same directory as in the previous step. If `jackhmmer`, 
+`hhblits`, `hhsearch` and `kalign` are available at the default path of 
+`/usr/bin`, their `binary_path` command-line arguments can be dropped.
+If you've already computed alignments for the query, you have the option to 
+skip the expensive alignment computation here with 
+`--use_precomputed_alignments`.
+
+`--openfold_checkpoint_path` or `--jax_param_path` accept comma-delineated lists
+of .pt/DeepSpeed OpenFold checkpoints and AlphaFold's .npz JAX parameter files, 
+respectively. For a breakdown of the differences between the different parameter 
+files, see the README downloaded to `openfold/resources/openfold_params/`. Since 
+OpenFold was trained under a newer training schedule than the one from which the 
+`model_n` config presets are derived, there is no clean correspondence between 
+`config_preset` settings and OpenFold checkpoints; the only restraints are that 
+`*_ptm` checkpoints must be run with `*_ptm` config presets and that `_no_templ_`
+checkpoints are only compatible with template-less presets (`model_3` and above).
+
+Note that chunking (as defined in section 1.11.8 of the AlphaFold 2 supplement)
+is enabled by default in inference mode. To disable it, set `globals.chunk_size`
+to `None` in the config. If a value is specified, OpenFold will attempt to 
+dynamically tune it, considering the chunk size specified in the config as a 
+minimum. This tuning process automatically ensures consistently fast runtimes 
+regardless of input sequence length, but it also introduces some runtime 
+variability, which may be undesirable for certain users. It is also recommended
+to disable this feature for very long chains (see below). To do so, set the 
+`tune_chunk_size` option in the config to `False`.
+
+For large-scale batch inference, we offer an optional tracing mode, which
+massively improves runtimes at the cost of a lengthy model compilation process.
+To enable it, add `--trace_model` to the inference command.
+
+To get a speedup during inference, enable [FlashAttention](https://github.com/HazyResearch/flash-attention)
+in the config. Note that it appears to work best for sequences with < 1000 residues.
+
+To minimize memory usage during inference on long sequences, consider the
+following changes:
+
+- As noted in the AlphaFold-Multimer paper, the AlphaFold/OpenFold template
+stack is a major memory bottleneck for inference on long sequences. OpenFold
+supports two mutually exclusive inference modes to address this issue. One,
+`average_templates` in the `template` section of the config, is similar to the
+solution offered by AlphaFold-Multimer, which is simply to average individual
+template representations. Our version is modified slightly to accommodate 
+weights trained using the standard template algorithm. Using said weights, we
+notice no significant difference in performance between our averaged template 
+embeddings and the standard ones. The second, `offload_templates`, temporarily 
+offloads individual template embeddings into CPU memory. The former is an 
+approximation while the latter is slightly slower; both are memory-efficient 
+and allow the model to utilize arbitrarily many templates across sequence 
+lengths. Both are disabled by default, and it is up to the user to determine 
+which best suits their needs, if either.
+- Inference-time low-memory attention (LMA) can be enabled in the model config.
+This setting trades off speed for vastly improved memory usage. By default,
+LMA is run with query and key chunk sizes of 1024 and 4096, respectively.
+These represent a favorable tradeoff in most memory-constrained cases.
+Powerusers can choose to tweak these settings in 
+`openfold/model/primitives.py`. For more information on the LMA algorithm,
+see the aforementioned Staats & Rabe preprint.
+- Disable `tune_chunk_size` for long sequences. Past a certain point, it only
+wastes time.
+- As a last resort, consider enabling `offload_inference`. This enables more
+extensive CPU offloading at various bottlenecks throughout the model.
+- Disable FlashAttention, which seems unstable on long sequences.
+
+Using the most conservative settings, we were able to run inference on a 
+4600-residue complex with a single A100. Compared to AlphaFold's own memory 
+offloading mode, ours is considerably faster; the same complex takes the more 
+efficent AlphaFold-Multimer more than double the time. Use the
+`long_sequence_inference` config option to enable all of these interventions
+at once. The `run_pretrained_openfold.py` script can enable this config option with the 
+`--long_sequence_inference` command line option
+
+Input FASTA files containing multiple sequences are treated as complexes. In
+this case, the inference script runs AlphaFold-Gap, a hack proposed
+[here](https://twitter.com/minkbaek/status/1417538291709071362?lang=en), using
+the specified stock AlphaFold/OpenFold parameters (NOT AlphaFold-Multimer).
+
+### Multimer Inference
+
+To run inference on a complex or multiple complexes using a set of DeepMind's pretrained parameters, run e.g.:
+
+```bash
+python3 run_pretrained_openfold.py \
+    fasta_dir \
+    data/pdb_mmcif/mmcif_files/ \
+    --uniref90_database_path data/uniref90/uniref90.fasta \
+    --mgnify_database_path data/mgnify/mgy_clusters_2022_05.fa \
+    --pdb_seqres_database_path data/pdb_seqres/pdb_seqres.txt \
+    --uniref30_database_path data/uniref30/UniRef30_2021_03 \
+    --uniprot_database_path data/uniprot/uniprot.fasta \
+    --bfd_database_path data/bfd/bfd_metaclust_clu_complete_id30_c90_final_seq.sorted_opt \
+    --jackhmmer_binary_path lib/conda/envs/openfold_venv/bin/jackhmmer \
+    --hhblits_binary_path lib/conda/envs/openfold_venv/bin/hhblits \
+    --hmmsearch_binary_path lib/conda/envs/openfold_venv/bin/hmmsearch \
+    --hmmbuild_binary_path lib/conda/envs/openfold_venv/bin/hmmbuild \
+    --kalign_binary_path lib/conda/envs/openfold_venv/bin/kalign \
+    --config_preset "model_1_multimer_v3" \
+    --model_device "cuda:0" \
+    --output_dir ./ 
+```
+
+As with monomer inference, if you've already computed alignments for the query, you can use 
+the `--use_precomputed_alignments` option. Note that template searching in the multimer pipeline 
+uses HMMSearch with the PDB SeqRes database, replacing HHSearch and PDB70 used in the monomer pipeline.
+
+**Upgrade from an existing OpenFold installation** 
+
+The above command requires several upgrades to existing openfold installations. 
+
+1. Re-download the alphafold parameters to get the latest
+AlphaFold-Multimer v3 weights:
+    
+   ```bash
+    bash scripts/download_alphafold_params.sh openfold/resources
+   ```
+
+2. Download the [UniProt](https://www.uniprot.org/uniprotkb/) 
+and [PDB SeqRes](https://www.rcsb.org/) databases: 
+    
+   ```bash
+    bash scripts/download_uniprot.sh data/
+   ```
+    
+    The PDB SeqRes and PDB databases must be from the same date to avoid potential 
+    errors during template searching. Remove the existing `data/pdb_mmcif` directory 
+    and download both databases:
+    
+   ```bash
+    bash scripts/download_pdb_mmcif.sh data/
+    bash scripts/download_pdb_seqres.sh data/
+   ```
+
+3. Additionally, AlphaFold-Multimer uses upgraded versions of the [MGnify](https://www.ebi.ac.uk/metagenomics) 
+and [UniRef30](https://uniclust.mmseqs.com/) (previously UniClust30) databases. To download the upgraded databases, run:
+    
+   ```bash
+    bash scripts/download_uniref30.sh data/
+    bash scripts/download_mgnify.sh data/
+   ```
+   Multimer inference can also run with the older database versions if desired. 
+
+
+### Soloseq Inference
+
+To run inference for a sequence using the SoloSeq single-sequence model, you can either precompute ESM-1b embeddings in bulk, or you can generate them during inference.
+
+For generating ESM-1b embeddings in bulk, use the provided script: `scripts/precompute_embeddings.py`. The script takes a directory of FASTA files (one sequence per file) and generates ESM-1b embeddings in the same format and directory structure as required by SoloSeq. Following is an example command to use the script:
+
+```bash
+python scripts/precompute_embeddings.py fasta_dir/ embeddings_output_dir/
+```
+
+In the same per-label subdirectories inside `embeddings_output_dir`, you can also place `*.hhr` files (outputs from HHSearch), which can contain the details about the structures that you want to use as templates. If you do not place any such file, templates will not be used and only the ESM-1b embeddings will be used to predict the structure. If you want to use templates, you need to pass the PDB MMCIF dataset to the command.
+
+Then download the SoloSeq model weights, e.g.:
+
+
+```bash 
+bash scripts/download_openfold_soloseq_params.sh openfold/resources
+```
+
+
+Now, you are ready to run inference:
+```bash
+python run_pretrained_openfold.py \
+    fasta_dir \
+    data/pdb_mmcif/mmcif_files/ \
+    --use_precomputed_alignments embeddings_output_dir \
+    --output_dir ./ \
+    --model_device "cuda:0" \
+    --config_preset "seq_model_esm1b_ptm" \
+    --openfold_checkpoint_path openfold/resources/openfold_soloseq_params/seq_model_esm1b_ptm.pt
+```
+
+For generating the embeddings during inference, skip the `--use_precomputed_alignments` argument. The `*.hhr` files will be generated as well if you pass the paths to the relevant databases and tools, as specified in the command below. If you skip the database and tool arguments, HHSearch will not be used to find templates and only generated ESM-1b embeddings will be used to predict the structure.
+```bash
+python3 run_pretrained_openfold.py \
+    fasta_dir \
+    data/pdb_mmcif/mmcif_files/ \
+    --output_dir ./ \
+    --model_device "cuda:0" \
+    --config_preset "seq_model_esm1b_ptm" \
+    --openfold_checkpoint_path openfold/resources/openfold_soloseq_params/seq_model_esm1b_ptm.pt \
+    --uniref90_database_path data/uniref90/uniref90.fasta \
+    --pdb70_database_path data/pdb70/pdb70 \
+    --jackhmmer_binary_path lib/conda/envs/openfold_venv/bin/jackhmmer \
+    --hhsearch_binary_path lib/conda/envs/openfold_venv/bin/hhsearch \
+    --kalign_binary_path lib/conda/envs/openfold_venv/bin/kalign \
+```
+
+For generating template information, you will need the UniRef90 and PDB70 databases and the JackHmmer and HHSearch binaries. 
+
+SoloSeq allows you to use the same flags and optimizations as the MSA-based OpenFold. For example, you can skip relaxation using `--skip_relaxation`, save all model outputs using `--save_outputs`, and generate output files in MMCIF format using `--cif_output`.
+
+**NOTE:** Due to the nature of the ESM-1b embeddings, the sequence length for inference using the SoloSeq model is limited to 1022 residues. Sequences longer than that will be truncated.
+
+## Training
+
+To train the model, you will first need to precompute protein alignments. 
+
+You have two options. You can use the same procedure DeepMind used by running
+the following:
+
+```bash
+python3 scripts/precompute_alignments.py mmcif_dir/ alignment_dir/ \
+    --uniref90_database_path data/uniref90/uniref90.fasta \
+    --mgnify_database_path data/mgnify/mgy_clusters_2018_12.fa \
+    --pdb70_database_path data/pdb70/pdb70 \
+    --uniclust30_database_path data/uniclust30/uniclust30_2018_08/uniclust30_2018_08 \
+    --bfd_database_path data/bfd/bfd_metaclust_clu_complete_id30_c90_final_seq.sorted_opt \
+    --cpus_per_task 16 \
+    --jackhmmer_binary_path lib/conda/envs/openfold_venv/bin/jackhmmer \
+    --hhblits_binary_path lib/conda/envs/openfold_venv/bin/hhblits \
+    --hhsearch_binary_path lib/conda/envs/openfold_venv/bin/hhsearch \
+    --kalign_binary_path lib/conda/envs/openfold_venv/bin/kalign
+```
+
+As noted before, you can skip the `binary_path` arguments if these binaries are 
+at `/usr/bin`. Expect this step to take a very long time, even for small 
+numbers of proteins.
+
+Alternatively, you can generate MSAs with the ColabFold pipeline (and templates
+with HHsearch) with:
+
+```bash
+python3 scripts/precompute_alignments_mmseqs.py input.fasta \
+    data/mmseqs_dbs \
+    uniref30_2103_db \
+    alignment_dir \
+    ~/MMseqs2/build/bin/mmseqs \
+    /usr/bin/hhsearch \
+    --env_db colabfold_envdb_202108_db
+    --pdb70 data/pdb70/pdb70
+```
+
+where `input.fasta` is a FASTA file containing one or more query sequences. To 
+generate an input FASTA from a directory of mmCIF and/or ProteinNet .core 
+files, we provide `scripts/data_dir_to_fasta.py`.
+
+Next, generate a cache of certain datapoints in the template mmCIF files:
+
+```bash
+python3 scripts/generate_mmcif_cache.py \
+    mmcif_dir/ \
+    mmcif_cache.json \
+    --no_workers 16
+```
+
+This cache is used to pre-filter templates. 
+
+Next, generate a separate chain-level cache with data used for training-time 
+data filtering:
+
+```bash
+python3 scripts/generate_chain_data_cache.py \
+    mmcif_dir/ \
+    chain_data_cache.json \
+    --cluster_file clusters-by-entity-40.txt \
+    --no_workers 16
+```
+
+where the `cluster_file` argument is a file of chain clusters, one cluster
+per line.
+
+Optionally, download an AlphaFold-style validation set from 
+[CAMEO](https://cameo3d.org) using `scripts/download_cameo.py`. Use the 
+resulting FASTA files to generate validation alignments and then specify 
+the validation set's location using the `--val_...` family of training script 
+flags.
+
+Finally, call the training script:
+
+```bash
+python3 train_openfold.py mmcif_dir/ alignment_dir/ template_mmcif_dir/ output_dir/ \
+    2021-10-10 \ 
+    --template_release_dates_cache_path mmcif_cache.json \ 
+    --precision bf16 \
+    --gpus 8 --replace_sampler_ddp=True \
+    --seed 4242022 \ # in multi-gpu settings, the seed must be specified
+    --deepspeed_config_path deepspeed_config.json \
+    --checkpoint_every_epoch \
+    --resume_from_ckpt ckpt_dir/ \
+    --train_chain_data_cache_path chain_data_cache.json \
+    --obsolete_pdbs_file_path obsolete.dat
+```
+
+where `--template_release_dates_cache_path` is a path to the mmCIF cache. 
+Note that `template_mmcif_dir` can be the same as `mmcif_dir` which contains
+training targets. A suitable DeepSpeed configuration file can be generated with 
+`scripts/build_deepspeed_config.py`. The training script is 
+written with [PyTorch Lightning](https://github.com/PyTorchLightning/pytorch-lightning) 
+and supports the full range of training options that entails, including 
+multi-node distributed training, validation, and so on. For more information, 
+consult PyTorch Lightning documentation and the `--help` flag of the training 
+script.
+
+Note that, despite its variable name, `mmcif_dir` can also contain PDB files 
+or even ProteinNet .core files. 
+
+To emulate the AlphaFold training procedure, which uses a self-distillation set 
+subject to special preprocessing steps, use the family of `--distillation` flags.
+
+In cases where it may be burdensome to create separate files for each chain's
+alignments, alignment directories can be consolidated using the scripts in 
+`scripts/alignment_db_scripts/`. First, run `create_alignment_db.py` to
+consolidate an alignment directory into a pair of database and index files.
+Once all alignment directories (or shards of a single alignment directory)
+have been compiled, unify the indices with `unify_alignment_db_indices.py`. The
+resulting index, `super.index`, can be passed to the training script flags
+containing the phrase `alignment_index`. In this scenario, the `alignment_dir`
+flags instead represent the directory containing the compiled alignment
+databases. Both the training and distillation datasets can be compiled in this
+way. Anecdotally, this can speed up training in I/O-bottlenecked environments.
+
+## Testing
+
+To run unit tests, use
+
+```bash
+scripts/run_unit_tests.sh
+```
+
+The script is a thin wrapper around Python's `unittest` suite, and recognizes
+`unittest` arguments. E.g., to run a specific test verbosely:
+
+```bash
+scripts/run_unit_tests.sh -v tests.test_model
+```
+
+Certain tests require that AlphaFold (v2.0.1) be installed in the same Python
+environment. These run components of AlphaFold and OpenFold side by side and
+ensure that output activations are adequately similar. For most modules, we
+target a maximum pointwise difference of `1e-4`.
+
+## Building and Using the Docker Container
+
+**Building the Docker Image**
+
+Openfold can be built as a docker container using the included dockerfile. To build it, run the following command from the root of this repository:
+
+```bash
+docker build -t openfold .
+```
+
+**Running the Docker Container** 
+
+The built container contains both `run_pretrained_openfold.py` and `train_openfold.py` as well as all necessary software dependencies. It does not contain the model parameters, sequence, or structural databases. These should be downloaded to the host machine following the instructions in the Usage section above. 
+
+The docker container installs all conda components to the base conda environment in `/opt/conda`, and installs openfold itself in `/opt/openfold`,
+
+Before running the docker container, you can verify that your docker installation is able to properly communicate with your GPU by running the following command:
+
+
+```bash
+docker run --rm --gpus all nvidia/cuda:11.0-base nvidia-smi
+```
+
+Note the `--gpus all` option passed to `docker run`. This option is necessary in order for the container to use the GPUs on the host machine.
+
+To run the inference code under docker, you can use a command like the one below.  In this example, parameters and sequences from the alphafold dataset are being used and are located at `/mnt/alphafold_database` on the host machine, and the input files are located in the current working directory. You can adjust the volume mount locations as needed to reflect the locations of your data. 
+
+```bash
+docker run \
+--gpus all \
+-v $PWD/:/data \
+-v /mnt/alphafold_database/:/database \
+-ti openfold:latest \
+python3 /opt/openfold/run_pretrained_openfold.py \
+/data/fasta_dir \
+/database/pdb_mmcif/mmcif_files/ \
+--uniref90_database_path /database/uniref90/uniref90.fasta \
+--mgnify_database_path /database/mgnify/mgy_clusters_2018_12.fa \
+--pdb70_database_path /database/pdb70/pdb70 \
+--uniclust30_database_path /database/uniclust30/uniclust30_2018_08/uniclust30_2018_08 \
+--output_dir /data \
+--bfd_database_path /database/bfd/bfd_metaclust_clu_complete_id30_c90_final_seq.sorted_opt \
+--model_device cuda:0 \
+--jackhmmer_binary_path /opt/conda/bin/jackhmmer \
+--hhblits_binary_path /opt/conda/bin/hhblits \
+--hhsearch_binary_path /opt/conda/bin/hhsearch \
+--kalign_binary_path /opt/conda/bin/kalign \
+--openfold_checkpoint_path /database/openfold_params/finetuning_ptm_2.pt
+```
+
+## Copyright Notice
+
+While AlphaFold's and, by extension, OpenFold's source code is licensed under
+the permissive Apache Licence, Version 2.0, DeepMind's pretrained parameters 
+fall under the CC BY 4.0 license, a copy of which is downloaded to 
+`openfold/resources/params` by the installation script. Note that the latter
+replaces the original, more restrictive CC BY-NC 4.0 license as of January 2022.
+
+## Contributing
+
+If you encounter problems using OpenFold, feel free to create an issue! We also
+welcome pull requests from the community.
+
+## Citing this Work
+
+Please cite our paper:
+
+```bibtex
+@article {Ahdritz2022.11.20.517210,
+	author = {Ahdritz, Gustaf and Bouatta, Nazim and Floristean, Christina and Kadyan, Sachin and Xia, Qinghui and Gerecke, William and O{\textquoteright}Donnell, Timothy J and Berenberg, Daniel and Fisk, Ian and Zanichelli, Niccolò and Zhang, Bo and Nowaczynski, Arkadiusz and Wang, Bei and Stepniewska-Dziubinska, Marta M and Zhang, Shang and Ojewole, Adegoke and Guney, Murat Efe and Biderman, Stella and Watkins, Andrew M and Ra, Stephen and Lorenzo, Pablo Ribalta and Nivon, Lucas and Weitzner, Brian and Ban, Yih-En Andrew and Sorger, Peter K and Mostaque, Emad and Zhang, Zhao and Bonneau, Richard and AlQuraishi, Mohammed},
+	title = {{O}pen{F}old: {R}etraining {A}lpha{F}old2 yields new insights into its learning mechanisms and capacity for generalization},
+	elocation-id = {2022.11.20.517210},
+	year = {2022},
+	doi = {10.1101/2022.11.20.517210},
+	publisher = {Cold Spring Harbor Laboratory},
+	URL = {https://www.biorxiv.org/content/10.1101/2022.11.20.517210},
+	eprint = {https://www.biorxiv.org/content/early/2022/11/22/2022.11.20.517210.full.pdf},
+	journal = {bioRxiv}
+}
+```
+If you use OpenProteinSet, please also cite:
+
+```bibtex
+@misc{ahdritz2023openproteinset,
+      title={{O}pen{P}rotein{S}et: {T}raining data for structural biology at scale}, 
+      author={Gustaf Ahdritz and Nazim Bouatta and Sachin Kadyan and Lukas Jarosch and Daniel Berenberg and Ian Fisk and Andrew M. Watkins and Stephen Ra and Richard Bonneau and Mohammed AlQuraishi},
+      year={2023},
+      eprint={2308.05326},
+      archivePrefix={arXiv},
+      primaryClass={q-bio.BM}
+}
+```
+Any work that cites OpenFold should also cite [AlphaFold](https://www.nature.com/articles/s41586-021-03819-2) and [AlphaFold-Multimer](https://www.biorxiv.org/content/10.1101/2021.10.04.463034v1) if applicable.

environment.yml

-name: openfold-venv
+name: openfold-env 
 channels:
   - conda-forge
   - bioconda
@@ -10,25 +10,26 @@ dependencies:
   - pip
   - openmm=7.7
   - pdbfixer
-  - cudatoolkit==11.3.*
-  - pytorch-lightning==1.5.10
-  - biopython==1.79
-  - numpy==1.21
-  - pandas==2.0
+  - pytorch-lightning
+  - biopython
+  - numpy
+  - pandas
   - PyYAML==5.4.1
   - requests
   - scipy==1.7
   - tqdm==4.62.2
-  - typing-extensions==3.10
-  - wandb==0.12.21
+  - typing-extensions==4.0
+  - wandb
   - modelcif==0.7
   - awscli
   - ml-collections
   - aria2
+  - mkl=2024.0
   - git
   - bioconda::hmmer==3.3.2
   - bioconda::hhsuite==3.3.0
   - bioconda::kalign2==2.04
+  - bioconda::mmseqs2
   - pytorch::pytorch=1.12.*
   - pip:
       - deepspeed==0.12.4