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

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](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.