OneScience-Group/UTRGAN

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

Model Introduction

UTRGAN is a generative and optimization framework for 5′ UTR design. It can generate candidate 5′ UTR sequences and predict and rank them according to gene expression, Mean Ribosome Load (MRL), and Translation Efficiency (TE).

The Hugging Face model package already contains the data and pretrained weights required for the basic workflow, so no additional model files need to be downloaded during inference.

Model Description

UTRGAN consists of multiple cooperating models:

• —WGAN-GP Generator: Generates candidate 5′ UTR sequences of up to 128 nt from a 40-dimensional random latent vector;
• —WGAN-GP Critic: Used during generative-model training;
• —Xpresso: Predicts gene-expression-related scores;
• —FramePool: Predicts Mean Ribosome Load (MRL);
• —MTtrans: Predicts Translation Efficiency (TE);
• —G4Boost: Performs G4-related classification and regression analyses.

WGAN, Xpresso, and FramePool use TensorFlow/Keras, while MTtrans uses PyTorch.

To maintain compatibility with the legacy H5 model files released by the original projects, Legacy Keras is used when loading the corresponding models.

Use Cases

Usage

1. OneCode

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

Try OneCode for AI4S Programming

2. Manual Installation

Hardware Requirements

• —CPU inference is supported;
• —A OneScience-supported DCU environment is recommended for batch generation, ranking, and training;
• —TensorFlow and PyTorch access the same DCU through their compatible accelerator interfaces;
• —Full training time depends on the number of candidate sequences, batch size, and device performance.

Download the Model Package

Install the Hugging Face command-line tool and download the model repository:

python -m pip install -U huggingface_hub

hf download OneScience-Group/UTRGAN --local-dir ./UTRGAN
cd UTRGAN

Install the Runtime Environment

Create and activate a Python 3.11 environment, and then install the OneScience bioscience base environment:

conda create -n utrgan python=3.11 -y
conda activate utrgan

pip install onescience[bio-dcu]

Then install the additional or replacement dependencies specified in requirements.txt:

python -m pip install --no-deps -r requirements.txt

Note: This model adaptation uses DTK 26.04. The corresponding TensorFlow runtime has also been upgraded to TensorFlow 2.18.

Check the frameworks and available devices:

python - <<'PY'
import tensorflow as tf
import torch

print("TensorFlow:", tf.__version__)
print("TensorFlow devices:", tf.config.list_physical_devices("GPU"))
print("PyTorch:", torch.__version__)
print("HIP:", torch.version.hip)
print("DCU available:", torch.cuda.is_available())
PY

Weights and Data Preparation

The Hugging Face model package already contains the resources required for basic generation, prediction, ranking, and training validation:

After applying the filtering and deduplication logic used by the official training scripts, UTRdb2 contains approximately 33,250 sequences with lengths between 65 and 128 nt.

The basic workflow does not require additional dataset downloads or precomputed features.

Optional Dependencies

requirements.txt includes XGBoost, ViennaRNA, logomaker, ruptures, and cliffs-delta for G4-related and downstream analysis scripts.

NUPACK is used only by optional MFE preprocessing scripts.

Because NUPACK 4 must be obtained and installed separately according to the licensing and installation requirements of its publisher, it is not automatically installed through standard PyPI dependencies in this model package.

NUPACK is not required for:

• —Basic inference;
• —MRL ranking;
• —TE ranking;
• —WGAN-GP training.

Quick Inference

Use the official pretrained WGAN to generate candidate sequences and rank them independently by MRL and TE.

Purpose: Generate 5′ UTR candidates in batches on DCU and output separate MRL and TE ranking results.

python scripts/predict.py \
  --device dcu \
  --device-id 0 \
  --num-candidates 1024 \
  --batch-size 128 \
  --seed 33 \
  --output-dir outputs/pretrained_batch_ranking

The results are saved to:

outputs/pretrained_batch_ranking/
├── all_candidates_scores.csv
├── ranked_by_mrl.csv
├── ranked_by_te.csv
├── generator_probabilities.npy
└── summary.json

The output files contain:

• —all_candidates_scores.csv: all candidate sequences together with their MRL and TE prediction scores;
• —ranked_by_mrl.csv: candidates ranked from highest to lowest MRL score;
• —ranked_by_te.csv: candidates ranked from highest to lowest TE score;
• —is_duplicate: indicates whether a generated sequence is duplicated;
• —MRL and TE are measured on different scales and their raw scores should not be directly added together.

MRL-Directed Optimization

Purpose: Freeze the WGAN and FramePool models and optimize the latent noise to increase the predicted MRL score of generated candidates.

python scripts/optimize_te_mrl.py \
  -gpu 0 \
  -task mrl \
  -bs 64 \
  -s 10 \
  --output-dir outputs/optimization_mrl

During this process, the WGAN and FramePool model weights remain frozen. Only the latent noise is updated.

TE-Directed Optimization

Purpose: Freeze the WGAN and MTtrans models and optimize the latent noise to increase the predicted TE score of generated candidates.

python scripts/optimize_te_mrl.py \
  -gpu 0 \
  -task te \
  -bs 64 \
  -s 10 \
  --output-dir outputs/optimization_te

During this process, the WGAN and MTtrans model weights remain frozen.

Therefore, this workflow is latent-space optimization rather than fine-tuning of the pretrained models.

Training

UTRGAN supports WGAN-GP training.

The trainable components in this workflow are:

• —Generator;
• —Critic.

Xpresso, FramePool, and MTtrans do not participate in this training entry point.

To run one epoch using the complete data-processing pipeline and the UTRdb2 dataset included in the repository:

python scripts/train.py \
  -gpu 0 \
  -bs 64 \
  -lr 5 \
  -mxl 128 \
  -dim 40 \
  --epochs 1 \
  --output-dir outputs/train_1epoch

According to the upstream README, -lr 5 corresponds to a learning rate of:

1e-5

The adapted training entry point adds the --epochs argument and corrects the learning-rate expression in the upstream script so that it is consistent with the documented definition.

To run the full upstream training configuration, use the complete UTRdb2 training dataset for 4000 epochs:

python scripts/train.py \
  -gpu 0 \
  -bs 64 \
  -lr 5 \
  -mxl 128 \
  -dim 40 \
  --epochs 4000 \
  --output-dir outputs/train_full

This retrains the WGAN-GP Generator and Critic.

Full training can require substantial runtime. Training progress should be evaluated according to available accelerator resources and training logs.

Generated checkpoints are saved in the specified output directory and do not overwrite the official pretrained weights under:

weight/

OneScience Official Resources

Citation and License

• —Upstream implementation: https://github.com/ciceklab/UTRGAN
• —Xpresso: https://github.com/vagarwal87/Xpresso
• —FramePool: https://github.com/Karollus/5UTR
• —MTtrans: https://github.com/holab-hku/MTtrans
• —The upstream UTRGAN project is distributed under the CC BY-NC-SA 2.0 license and is intended for academic and non-commercial use. Commercial use requires permission from the upstream authors.
• —Third-party models, datasets, pretrained weights, and software components remain subject to their respective original licenses and terms of use.
• —The DCU adaptation in this repository continues to use the same CC BY-NC-SA 2.0 license.