OneScience-Group/CodonTransformer

CodonTransformer

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

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

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:

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:

/path/to/.cache/huggingface/hub/

Set the following environment variables when running the scripts:

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:

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

It is recommended to place the downloaded raw data in:

scripts/data/raw/

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

scripts/data/raw/dataset.csv

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

scripts/data/processed/

3. Quick Start

Download the Model Package

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

cd CodonTransformer

Inference Examples

Single-Sequence Codon Optimization

Run the script:

bash scripts/slurm/run_inference_single.sh

The default parameters are:

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

To change the input protein and host:

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

Generate Multiple Candidate DNA Sequences

Run:

bash scripts/slurm/run_inference_multiple.sh

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

outputs/multiple_predictions.csv

The main parameters include:

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:

organism = "Escherichia coli general"

It can be changed to:

organism = "Homo sapiens"

or:

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:

scripts/demo/sample_dataset.csv

Run directly:

bash scripts/slurm/run_inference_batch.sh

The default output is saved to:

outputs/sample_predictions.csv

The input CSV must contain at least:

protein_sequence
organism

To use a custom CSV:

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:

scripts/data/raw/your_data.csv

It must contain at least:

dna
protein
organism

Then run the data preprocessing script provided in the repository:

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 scripts/slurm/run_finetune.sh

By default, it reads:

scripts/data/processed/finetune_data.json

and saves checkpoints to:

weight/checkpoints/finetune

To fine-tune using custom data, run:

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:

scripts/slurm/export_finetuned_model.sh

Run it as follows:

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:

scripts/slurm/run_inference_finetuned.sh

Run it as follows:

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:

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:

scripts/slurm/prepare_pretrain_data.sh

By default, this script reads:

scripts/data/raw/dataset.csv

and outputs:

scripts/data/processed/pretrain_data.json

Therefore, before full pretraining, first run:

bash scripts/slurm/prepare_pretrain_data.sh

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

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:

scripts/data/processed/pretrain_data.json

Full pretraining can be started with:

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:

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.