README.md

---
datasets:
- OneScience/la-proteina
frameworks:
- ""
language:
- en
license: apache-2.0
tags:
- OneScience
- protein structure generation
---
<p align="center">
  <strong>
    <span style="font-size: 30px;">La-Proteina</span>
  </strong>
</p>

# Model Overview

La-Proteina is a protein structure generation model based on **Partially Latent Flow Matching**. It can directly generate all-atom protein structures together with their corresponding amino acid sequences.

Paper: _La-Proteina: Atomistic Protein Generation via Partially Latent Flow Matching_ (arXiv 2025).

- [Paper link](https://arxiv.org/abs/2507.09466)

# Model Description

La-Proteina explicitly models the protein backbone (backbone CA), while sequence and atom-level details are captured through fixed-dimensional latent variables for each residue. This effectively avoids the challenges introduced by explicit side-chain representations.

# Use Cases

| Use case | Description |
| :---: | :---: |
| Protein structure generation | Generates protein backbones (backbone CA) and local latents based on flow matching. |
| Motif-constrained generation | Supports motif position and sequence constraints for backbone design around functional motifs. |
| Diffusion model training | Trains the main La-Proteina model on PDB datasets. |
| Autoencoder training and inference | Trains the local-latent autoencoder and performs encoding, decoding, and reconstruction evaluation on PDB structures. |
| Generated result evaluation | Computes metrics such as RMSD, sequence recovery, and (co-)designability. |

# Usage

## 1. Using OneCode

You can try intelligent one-click AI4S programming through the OneCode online environment:

[Try intelligent one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)

## 2. Manual Installation and Usage

**Hardware Requirements**

- Running on a GPU or DCU is recommended.
- CPU can be used for connectivity checks, but it is relatively slow.
- DCU users need to install DTK in advance. DTK 25.04.2 or later is recommended, or the OneScience-recommended version that matches the current cluster.

**Software Requirements**

DCU users who want to learn more about adaptation details can contact liubiao@sugon.com.

**Environment Checks**

- NVIDIA GPU:

```bash
nvidia-smi
```

- Hygon DCU:

```bash
hy-smi
```

### Environment Preparation

Optional: override the default paths through environment variables:

```bash
export LAPROTEINA_ROOT=/path/to/la-proteina
export LAPROTEINA_DATASET_DIR=/path/to/dataset
export LAPROTEINA_CHECKPOINTS_DIR=/path/to/checkpoints_laproteina
export DATA_PATH=/path/to/dataset
```

## Quick Start

### 1. Install the Runtime Environment

```bash
conda create -n onescience311 python=3.11 -y
conda activate onescience311
pip install onescience[bio] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai
```

If the following code cannot find required libraries at runtime, activate CUDA as shown below.

```bash
source ${ROCM_PATH}/cuda/env.sh
export LD_LIBRARY_PATH="$CONDA_PREFIX/lib:$LD_LIBRARY_PATH"
export LD_LIBRARY_PATH="$CONDA_PREFIX/lib/python3.11/site-packages/fastpt/torch/lib:$LD_LIBRARY_PATH"
```

### 2. Download the Model Package

```bash
# By default, the package is downloaded to the model folder under the current path. To change this, adjust the path after local_dir.
hf download --model OneScience-Group/La-Proteina --local-dir ./model
cd model
```

### Training Weights and Datasets

They will be uploaded to Hugging Face soon, and command-line downloads will be supported later.

### 3. Usage Modes

It is recommended to run the scripts under the `laproteina` directory so that outputs are managed in one place.

#### 1. Train the Main La-Proteina Model (`run_train.sh`)

```bash
bash scripts/run_train.sh
```

**Common Hydra parameter overrides:**

```bash
# Single-device debugging
bash scripts/run_train.sh hardware.ngpus_per_node_=1 single=true

# Specify a run name
bash scripts/run_train.sh run_name=my_laproteina_run

# Override dataset or network configuration (for motif training scenarios)
bash scripts/run_train.sh dataset=pdb/pdb_train_motif_aa nn=local_latents_score_nn_160M_motif_idx_aa
```

Output:

- `./store/<run_name>/`: training logs, checkpoints, and Hydra configurations

#### 2. Protein Structure Generation (`run_generate.sh`)

```bash
bash scripts/run_generate.sh
```

By default, the `inference_ucond_tri` configuration is used for unconditional generation (LD2 model + AE1).

**Switching generation configurations:**

```bash
# Unconditional generation without triangle attention
bash scripts/run_generate.sh --config_name inference_ucond_notri

# Unconditional generation for long chains (300-800 residues)
bash scripts/run_generate.sh --config_name inference_ucond_notri_long

# Indexed all-atom motif scaffolding
bash scripts/run_generate.sh --config_name inference_motif_idx_aa

# Indexed tip-atom motif scaffolding
bash scripts/run_generate.sh --config_name inference_motif_idx_tip

# Unindexed all-atom motif scaffolding
bash scripts/run_generate.sh --config_name inference_motif_uidx_aa

# Unindexed tip-atom motif scaffolding
bash scripts/run_generate.sh --config_name inference_motif_uidx_tip
```

Output:

- `./inference/<config_name>/`: generated protein structure files and metadata

---

#### 3. Generated Result Evaluation (`run_evaluate.sh`)

```bash
bash scripts/run_evaluate.sh
```

By default, this evaluates the generated results corresponding to the `inference_ucond_tri` configuration.

**Preparing ProteinMPNN weights:**

ProteinMPNN weights must be downloaded before evaluation. They can be downloaded from the Hugging Face community:

```bash
hf download --model OneScience-Group/ProteinMPNN --local-dir ./weight
```

Output:

- `./inference/<config_name>/evaluation/`: evaluation result files

---

#### 4. Autoencoder Inference (`run_ae_infer.sh`)

```bash
bash scripts/run_ae_infer.sh
```

This performs encode-decode reconstruction on the PDB dataset and evaluates reconstruction metrics, such as all-atom RMSD and sequence recovery.

The script checks whether `DATA_PATH/pdb_train` and `AE1_ucond_512.ckpt` exist.

The script internally calls:

```bash
python infer_laproteina_ae.py "$@"
```

Common parameter overrides:

| Environment variable | Default value | Description |
|----------|--------|------|
| `LAPROTEINA_ROOT` | `${ONESCIENCE_DATASETS_DIR}/la-proteina` | Root directory for data and weights |
| `LAPROTEINA_CHECKPOINTS_DIR` | `${LAPROTEINA_ROOT}/checkpoints_laproteina` | Autoencoder weight directory |
| `DATA_PATH` | `${LAPROTEINA_ROOT}/dataset` | Dataset directory |

Output:

- `./inference_ae/`: reconstructed structures and evaluation metrics

### Notes

- Make sure the `ONESCIENCE_DATASETS_DIR` environment variable is correctly set before running the scripts.
- The training script checks whether `DATA_PATH/pdb_train` and `AE1_ucond_512.ckpt` exist by default. If either is missing, it will report an error and exit.
- In the current integration, `dataset=pdb` is available. `dataset=genie2` and `dataset=pdb_multimer` are not packaged in the current repository snapshot, and running them will produce an explicit error.
- The scripts automatically set ROCm/DCU-related `LD_LIBRARY_PATH` values and can run directly on Hygon DCU platforms. On CUDA platforms, these settings can be ignored or adjusted as needed.
- When generating motif scaffold structures, make sure the LD model is correctly paired with the corresponding AE model. Otherwise, generation may fail due to length or task mismatch.
- Evaluating (co-)designability requires ProteinMPNN weights. Run `script_utils/download_pmpnn_weights.sh` in advance to download them.
- It is recommended to run all scripts under `examples/biosciences/laproteina` so that output directories remain unified.
- When overriding configurations through Hydra, you can use `+CK_PATH=...` to specify the checkpoint root path. If it is not provided, the script automatically sets it to `LAPROTEINA_ROOT`.

# Official OneScience Information

| Platform | OneScience main repository | Skills repository |
| --- | --- | --- |
| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |

# Citation and License

The example code is licensed under Apache 2.0. La-Proteina model weights are licensed under the [NVIDIA Open Model License Agreement](https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-open-model-license/), and other materials are licensed under [CC-BY 4.0](https://creativecommons.org/licenses/by/4.0/legalcode).

If you use La-Proteina in your research, please cite the original paper:

```bibtex
@article{geffner2025laproteina,
  title={La-Proteina: Atomistic Protein Generation via Partially Latent Flow Matching},
  author={Geffner, Tomas and Didi, Kieran and Cao, Zhonglin and Reidenbach, Danny and Zhang, Zuobai and Dallago, Christian and Kucukbenli, Emine and Kreis, Karsten and Vahdat, Arash},
  journal={arXiv preprint arXiv:2507.09466},
  year={2025}
}
```