diff --git a/docs/source/Inference.md b/docs/source/Inference.md index 0aa607ce..2782836e 100644 --- a/docs/source/Inference.md +++ b/docs/source/Inference.md @@ -71,7 +71,7 @@ Coming soon: - Polymeric ligands such as glycans -## 2. Pre-requisites: +## 2. Pre-requisites - OpenFold3 Conda Environment. See {ref}`OpenFold3 Installation ` for instructions on how to build this environment. - OpenFold3 Model Parameters. See {ref}`OpenFold3 Parameters ` for how to download these parameters. @@ -258,6 +258,16 @@ output_writer_settings: --- +(output-model-embeddings)= +#### 📦 Output single and pair model embeddings +To save model embeddings, add the field `write_latent_outputs` to the `output_writer_settings` i.e.: +```yaml +output_writer_settings: + write_latent_outputs: True +``` + +--- + (inference-low-memory-mode)= #### 🧠 Low Memory Mode To run inference on larger queries to run on limited memory, add the following to apply the [model presets](https://github.com/aqlaboratory/openfold-3/blob/main/openfold3/projects/of3_all_atom/config/model_setting_presets.yml) to run in low memory mode. @@ -512,6 +522,16 @@ When processing multimer inputs (e.g., hemoglobin α + β chains), OpenFold3 aut - Stores raw alignments in `raw/paired/` temporarily - Converts them into per-chain `.npz` alignments in [`paired/`](https://huggingface.co/OpenFold/OpenFold3/tree/main/examples/output_multimer_with_colabfold_msas/colabfold_msas/paired) +### 4.5 Embeddings +At inference you can instruct the model to produce the pair-rep and single-rep embeddings by {ref}`adjusting the output_writer_settings in your runner.yaml `. +This will cause the model to produce a `*_latent_output.pt`, which can be loaded like so and has the following shape. + +```python +output = torch.load("*_latent_output.pt") +print(output.keys()) +["si_trunk", "zij_trunk", "atom_positions_predicted"] +``` +If you want to change these outputs, the code is in the [model.py](../../openfold3/projects/of3_all_atom/model.py#L384-L388)