README.md

---
license: apache-2.0
language:
- en
tags:
- OneScience
- SCNet
- DCU
- codon-optimization
- protein-design
- sequence-generation
- transformer
frameworks:
- PyTorch
---

<p align="center">
  <strong>
    <span style="font-size: 30px;">CodonTransformer</span>
  </strong>
</p>

# Model Introduction

CodonTransformer is a deep learning model for multispecies codon optimization. Given an input protein sequence and a target host organism, it generates a host-specific DNA coding sequence.

Because of the degeneracy of the genetic code, the same protein can be encoded by many different DNA sequences, while different hosts have distinct preferences for synonymous codons. CodonTransformer uses a Transformer to model the contextual relationships among proteins, codons, and host organisms, generating DNA sequences that better match the natural codon distribution of the target host while preserving the translated protein sequence as much as possible.

Paper:
> **CodonTransformer: a multispecies codon optimizer using context-aware neural networks**  
> *Nature Communications*, 2025

# Model Description

CodonTransformer is a Transformer model designed for multispecies codon optimization. Its core network is based on a BigBird masked language model and incorporates the STREAM representation proposed by the authors, which encodes target host information, amino acid information, and codon information into a unified sequence representation.

Unlike traditional optimization methods based on global codon frequencies, CodonTransformer considers not only host preferences for individual codons but also uses the contextual modeling capability of Transformers to learn local dependencies between neighboring codons. This enables it to generate DNA sequences that more closely resemble the natural coding patterns of the target host. The model can be used for multispecies codon optimization, heterologous protein expression sequence design, and further fine-tuning on custom DNA-protein-host datasets.

The official model was trained on more than one million DNA-protein paired samples spanning 164 species, including bacteria, archaea, plants, animals, and fungi.

# Use Cases

| Scenario | Description |
| --- | --- |
| Codon optimization | Redesign protein-coding DNA for a target host |
| Heterologous protein expression | Generate coding sequences that better match host codon preferences for different hosts |
| Multispecies sequence design | Switch the target organism among multiple supported hosts |
| Multiple candidate sequence generation | Generate multiple distinct DNA candidate sequences through temperature sampling |
| Batch codon optimization | Perform batch inference for multiple protein-host combinations |
| Model fine-tuning | Continue fine-tuning using custom DNA-protein-host data |
| Model pretraining | Train the model from scratch using large-scale processed training data |

# Usage

## 1. Using OneCode

You can use the OneCode online environment for intelligent one-click AI4S programming:

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

## 2. Manual Installation and Usage

### Hardware Requirements

- CodonTransformer supports inference on both CPUs and accelerator devices.
- A CPU can be used for single protein sequence inference.
- GPU/DCU devices are recommended for batch inference, long-sequence inference, fine-tuning, and pretraining.

### Environment Setup

#### DCU/SCNet 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 you encounter missing dependencies or version incompatibilities during execution, refer to the dependency versions specified in `requirements.txt` at the repository root and install or adjust the corresponding dependencies as needed.

### Model and Data Preparation

#### 1) CodonTransformer Model Weights

The official model is available on Hugging Face:

```text
https://huggingface.co/adibvafa/CodonTransformer
```

In an online environment, `from_pretrained` automatically downloads the model weights when they are not available in the local cache. In an offline environment, the weights must be downloaded in advance to the cache directory corresponding to `HF_HOME`. The recommended cache location is:

```text
/path/to/.cache/huggingface/hub/
```

Set the following environment variables when running the scripts:

```bash
export HF_HOME=/path/to/.cache/huggingface
export HF_HUB_OFFLINE=1
export TRANSFORMERS_OFFLINE=1
```

The inference and fine-tuning scripts in this repository support using the local cache through `HF_HOME` and `OFFLINE=1`.

#### 2) Training Dataset

The official training data can be obtained from Zenodo or Hugging Face Datasets:

```text
https://zenodo.org/records/12509224
https://huggingface.co/datasets/adibvafa/CodonTransformer
```

It is recommended to place the downloaded raw data in:

```text
scripts/data/raw/
```

Accordingly, the complete dataset file is recommended to be saved as:

```text
scripts/data/raw/dataset.csv
```

The processed training JSON files are recommended to be saved in:

```text
scripts/data/processed/
```
## 3. Quick Start

### Download the Model Package

```bash
hf download OneScience-Group/CodonTransformer \
  --local-dir ./CodonTransformer

cd CodonTransformer
```

# Inference Examples

## Single-Sequence Codon Optimization

Run the script:

```bash
bash scripts/slurm/run_inference_single.sh
```

The default parameters are:

```bash
PROTEIN="MFWY"
ORGANISM="Escherichia coli general"
OFFLINE=1
```

To change the input protein and host:

```bash
PROTEIN="MALWMRLLPLLALLALWGPDPAAA" \
ORGANISM="Homo sapiens" \
bash scripts/slurm/run_inference_single.sh
```

## Generate Multiple Candidate DNA Sequences

Run:

```bash
bash scripts/slurm/run_inference_multiple.sh
```

By default, multiple candidate DNA sequences are generated for the same protein sequence and saved to:

```text
outputs/multiple_predictions.csv
```

The main parameters include:

```text
deterministic=False
temperature=0.5
top_p=0.95
num_sequences=5
match_protein=True
```

where:

- `deterministic=False` enables probabilistic sampling.
- `temperature` controls sampling diversity, with a typical range of `0.2 ~ 0.8`.
- `top_p` controls nucleus sampling.
- `num_sequences` specifies the number of candidate sequences to generate.
- `match_protein=True` constrains the generated DNA to translate to the same input protein sequence.

## Change the Target Host

The target host is specified directly through the `organism` parameter, for example:

```python
organism = "Escherichia coli general"
```

It can be changed to:

```python
organism = "Homo sapiens"
```

or:

```python
organism = "Saccharomyces cerevisiae"
```

Then rerun inference to obtain a codon-optimized DNA sequence for the corresponding host.

## Batch Inference

The example data is currently located at:

```text
scripts/demo/sample_dataset.csv
```

Run directly:

```bash
bash scripts/slurm/run_inference_batch.sh
```

The default output is saved to:

```text
outputs/sample_predictions.csv
```

The input CSV must contain at least:

```text
protein_sequence
organism
```

To use a custom CSV:

```bash
INPUT_CSV=/path/to/input.csv \
OUTPUT_CSV=/path/to/output.csv \
bash scripts/slurm/run_inference_batch.sh
```

# Training

## Fine-Tuning Data Preparation

For custom fine-tuning, first prepare your own CSV file. It is recommended to place it at:

```text
scripts/data/raw/your_data.csv
```

It must contain at least:

```text
dna
protein
organism
```

Then run the data preprocessing script provided in the repository:

```bash
INPUT_CSV=$PWD/scripts/data/raw/your_data.csv \
OUTPUT_JSON=$PWD/scripts/data/processed/finetune_data.json \
bash scripts/slurm/prepare_finetune_data.sh
```

## CodonTransformer Fine-Tuning

Run the script:

```bash
bash scripts/slurm/run_finetune.sh
```

By default, it reads:

```text
scripts/data/processed/finetune_data.json
```

and saves checkpoints to:

```text
weight/checkpoints/finetune
```

To fine-tune using custom data, run:

```bash
DATASET_JSON=$PWD/scripts/data/processed/finetune_data.json \
CHECKPOINT_DIR=$PWD/weight/checkpoints/finetune \
CHECKPOINT_FILENAME=finetune.ckpt \
BATCH_SIZE=6 \
MAX_EPOCHS=15 \
NUM_WORKERS=5 \
ACCUMULATE_GRAD_BATCHES=1 \
NUM_GPUS=4 \
LEARNING_RATE=0.00005 \
WARMUP_FRACTION=0.1 \
SAVE_EVERY_N_STEPS=512 \
SEED=123 \
DEBUG=0 \
bash scripts/slurm/run_finetune.sh
```

Here, `NUM_GPUS=4`, `BATCH_SIZE=6`, and `MAX_EPOCHS=15` are the default training settings. In practice, these parameters should be adjusted according to the number of allocated GPU/DCU devices, available device memory, and dataset size.

## Export the Fine-Tuned Model and Run Inference

After fine-tuning is complete, the following script can be used to export the checkpoint into a model file suitable for inference:

```text
scripts/slurm/export_finetuned_model.sh
```

Run it as follows:

```bash
CHECKPOINT_PATH=/path/to/finetuned_checkpoint.ckpt \
OUTPUT_MODEL_PATH=/path/to/output_finetuned_model.pt \
NUM_ORGANISMS=164 \
bash scripts/slurm/export_finetuned_model.sh
```

After export, the following script can be used to load the fine-tuned model for inference:

```text
scripts/slurm/run_inference_finetuned.sh
```

Run it as follows:

```bash
PROTEIN="MFWY" \
ORGANISM="Escherichia coli general" \
MODEL_PATH=/path/to/output_finetuned_model.pt \
bash scripts/slurm/run_inference_finetuned.sh
```

## Pretraining

The current pretraining entry point is:

```text
scripts/pretrain.py
```

Pretraining is a complete model training workflow and requires large-scale processed DNA-protein-host data. Before full pretraining, the complete `dataset.csv` must first be converted into a JSONL file that can be read by the training script.

The repository provides the following preprocessing script:

```text
scripts/slurm/prepare_pretrain_data.sh
```

By default, this script reads:

```text
scripts/data/raw/dataset.csv
```

and outputs:

```text
scripts/data/processed/pretrain_data.json
```

Therefore, before full pretraining, first run:

```bash
bash scripts/slurm/prepare_pretrain_data.sh
```

If the raw data or output directory differs from the default path, modify it using environment variables:

```bash
INPUT_CSV=$PWD/scripts/data/raw/dataset.csv \
OUTPUT_JSON=$PWD/scripts/data/processed/pretrain_data.json \
bash scripts/slurm/prepare_pretrain_data.sh
```

The complete generated pretraining data is recommended to be saved as:

```text
scripts/data/processed/pretrain_data.json
```

Full pretraining can be started with:

```bash
TRAIN_DATA_PATH=$PWD/scripts/data/processed/pretrain_data.json \
CHECKPOINT_DIR=$PWD/weight/checkpoints/pretrain \
BATCH_SIZE=6 \
MAX_EPOCHS=5 \
NUM_WORKERS=5 \
ACCUMULATE_GRAD_BATCHES=1 \
NUM_GPUS=16 \
LEARNING_RATE=0.00005 \
WARMUP_FRACTION=0.1 \
SAVE_INTERVAL=5 \
SEED=123 \
DEBUG=0 \
bash scripts/slurm/run_pretrain.sh
```

Here, `NUM_GPUS=16` corresponds to the default multi-device setting in the pretraining script. In practice, adjust `NUM_GPUS` according to the number of allocated GPU/DCU devices. Pretraining checkpoints are saved by default to:

```text
weight/checkpoints/pretrain
```
# OneScience Official 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 official CodonTransformer source code repository is licensed under the **Apache License 2.0**. See the `LICENSE` file in the repository root for details.
- CodonTransformer model weights are distributed separately through Hugging Face, while the training data is available through Zenodo and Hugging Face Datasets. The model weights, training data, and related third-party resources must each be used in accordance with the licenses and terms of use specified on their respective pages.
- This repository is the **DCU-adapted version** of CodonTransformer. The use of the repository code, model weights, and related data remains subject to the licenses and terms of use of their respective original projects.