Skip to content

docs: how to get model embeddings (s_i, z_ij) #105

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jandom merged 4 commits into from
Jan 30, 2026
Merged

docs: how to get model embeddings (s_i, z_ij) #105

Changes from 3 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
c9e95bc
docs: how to get model embeddings (s_i, z_ij)
jandom Jan 30, 2026
eef8201
update inference.md
jandom Jan 30, 2026
7f8516b
Add runner.yml settings for outputing model embeddings
jnwei Jan 30, 2026
ad7b78b
Update docs/source/Inference.md
jandom Jan 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 21 additions & 1 deletion docs/source/Inference.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,7 +71,7 @@ Coming soon:
- Polymeric ligands such as glycans


## 2. Pre-requisites:
## 2. Pre-requisites

- OpenFold3 Conda Environment. See {ref}`OpenFold3 Installation <openfold3-installation>` for instructions on how to build this environment.
- OpenFold3 Model Parameters. See {ref}`OpenFold3 Parameters <openfold3-parameters>` for how to download these parameters.
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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. You can do so by providing `write_latent_outputs: True` 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)
Loading