OneScience-Group/LSTM_CDRs-main

LSTM_CDRs

Model Introduction

LSTM_CDRs is a CDR amino acid sequence generation model based on long short-term memory networks (Long Short-Term Memory, LSTM). Given a set of CDR sequences, the model learns the sequence distribution in the training set and generates new CDR sequences through sampling after training.

Model Description

This project uses recurrent neural networks to autoregressively model amino acid sequences. Input sequences are first padded and one-hot encoded, then passed to multilayer LSTM or GRU networks for training. After training, the script can load model weights from a specified epoch and sample new CDR sequences from the learned sequence distribution.

The official code is adapted from LSTM_peptides for VHH CDR sequence design tasks.

Use Cases

Use case	Description
CDR sequence generation	Learn the amino acid sequence distribution from a given CDR training set and sample new candidate sequences.
Local LSTM/GRU training	Train LSTM/GRU sequence generation models locally or on GPU/DCU platforms using the scripts in this repository.
Weight loading and sampling reproduction	Load checkpoint weights generated during training and generate CDR sequences with specified lengths, temperature, and sample count.

Usage

1. Using OneCode

Experience intelligent one-click AI4S programming in the OneCode online environment:

Try intelligent one-click AI4S programming

2. Manual Installation and Usage

Hardware Requirements

• A GPU or DCU is recommended for training and sampling tasks.
• A CPU can be used for small-scale connectivity checks, but full training and large-scale sampling are slower.
• DCU users need to use DTK, TensorFlow, and the OneScience environment compatible with the current cluster.

3. Quick Start

Download the Model Package

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

Set Up the Runtime Environment

DCU Environment

# Activate DTK and CONDA first
conda create -n onescience311 python=3.11 -y
conda activate onescience311
# Install with uv support
pip install onescience[bio] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai

Environment Notes

• If you encounter missing dependencies or version compatibility issues, refer to the dependency versions specified in requirements.txt and environment.yml and install or adjust the environment as needed.
• If you encounter TensorFlow-related issues during execution, uninstall TensorFlow from the current environment and resolve them as follows:

# 1. Download TensorFlow from the platform
wget --content-disposition 'https://download.sourcefind.cn:65024/file/4/tensorflow/DAS1.8/tensorflow-2.13.1+das.opt1.dtk2604-cp311-cp311-manylinux_2_28_x86_64.whl'

# 2. Install TensorFlow
pip install tensorflow* 

# 3. Load the corresponding DTK version
module load compiler/dtk/26.04 

• You can also build the runtime environment for your platform according to the dependencies declared in requirements.txt and environment.yml.

Quick Verification

python LSTM_CDRs.py --help
ls data

The training data should include:

data/Cluster1.csv
data/Cluster2.csv
data/Cluster3.csv
data/Cluster4.csv

The data files are:

Data file	Number of sequences	Sequence length
data/Cluster1.csv	2629	36
data/Cluster2.csv	4146	36
data/Cluster3.csv	2990	35
data/Cluster4.csv	11952	36

Weights and Data Preparation

The current repository already includes the training data:

data/Cluster1.csv
data/Cluster2.csv
data/Cluster3.csv
data/Cluster4.csv

The repository does not provide pretrained weight files. The weights required for sampling must first be generated through training.

After training, each experiment directory contains:

<run_name>/
  flags.txt
  <run_name>_loss_plot.pdf
  sampled_sequences_temp1.25.csv
  checkpoint/
    model.json
    model.p
    model.hdf5
    model_epoch_0.hdf5
    model_epoch_1.hdf5
    ...

For sampling or fine-tuning, --modfile should point to an existing epoch weight, for example:

Cluster1_LSTM/checkpoint/model_epoch_100.hdf5

The same checkpoint/ directory must also retain:

model.p
model.hdf5
model_epoch_*.hdf5

Training

Run a Minimal Training Check with the Script

python LSTM_CDRs.py \
  --name smoke_Cluster1 \
  --dataset data/Cluster1.csv \
  --layers 1 \
  --neurons 16 \
  --epochs 1 \
  --batch_size 64 \
  --dropout 0.1 \
  --sample 10

This command verifies that data loading, padding, one-hot encoding, model training, and sampling work end to end.

View the outputs:

ls smoke_Cluster1
ls smoke_Cluster1/checkpoint
head smoke_Cluster1/sampled_sequences_temp1.25.csv

Train the Cluster1 Model

python LSTM_CDRs.py \
  --name Cluster1_LSTM \
  --dataset data/Cluster1.csv \
  --layers 2 \
  --neurons 64 \
  --epochs 200 \
  --dropout 0.2

By default, this command samples 100 sequences after training and saves them to:

Cluster1_LSTM/sampled_sequences_temp1.25.csv

View the training logs and weights:

ls Cluster1_LSTM
ls Cluster1_LSTM/checkpoint

Train All Four Clusters

python LSTM_CDRs.py --name Cluster1_LSTM --dataset data/Cluster1.csv --layers 2 --neurons 64 --epochs 200 --dropout 0.2
python LSTM_CDRs.py --name Cluster2_LSTM --dataset data/Cluster2.csv --layers 2 --neurons 64 --epochs 200 --dropout 0.2
python LSTM_CDRs.py --name Cluster3_LSTM --dataset data/Cluster3.csv --layers 2 --neurons 64 --epochs 200 --dropout 0.2
python LSTM_CDRs.py --name Cluster4_LSTM --dataset data/Cluster4.csv --layers 2 --neurons 64 --epochs 200 --dropout 0.2

Sampling

Sampling is the generation step performed after training. The script loads existing model weights and generates new sequences from the learned CDR sequence distribution.

Cluster1 Sampling

python LSTM_CDRs.py \
  --name Cluster1_LSTM \
  --dataset data/Cluster1.csv \
  --modfile Cluster1_LSTM/checkpoint/model_epoch_100.hdf5 \
  --train False \
  --sample 10000 \
  -f 36 \
  -m 36

This command uses the following by default:

Training data: data/Cluster1.csv
Model weights: Cluster1_LSTM/checkpoint/model_epoch_100.hdf5
Sample count: 10000
Minimum length: 36
Maximum length: 36
Output file: Cluster1_LSTM/sampled_sequences_temp1.25.csv

Cluster3 Sampling

The original Cluster3 sequences have length 35, so -f 35 -m 35 is recommended for sampling:

python LSTM_CDRs.py \
  --name Cluster3_LSTM \
  --dataset data/Cluster3.csv \
  --modfile Cluster3_LSTM/checkpoint/model_epoch_100.hdf5 \
  --train False \
  --sample 10000 \
  -f 35 \
  -m 35

View Sampling Results

wc -l Cluster1_LSTM/sampled_sequences_temp1.25.csv
head Cluster1_LSTM/sampled_sequences_temp1.25.csv

Check generated sequence lengths:

awk '{print length($0)}' Cluster1_LSTM/sampled_sequences_temp1.25.csv | sort -n | uniq -c

Check the number of generated sequences duplicated in the training set:

grep -Fxf data/Cluster1.csv Cluster1_LSTM/sampled_sequences_temp1.25.csv | wc -l

Fine-Tuning

python LSTM_CDRs.py \
  --name Cluster1_to_Cluster2_finetune \
  --dataset data/Cluster2.csv \
  --modfile Cluster1_LSTM/checkpoint/model_epoch_100.hdf5 \
  --train False \
  --finetune True \
  --epochs 50 \
  --layers 2 \
  --neurons 64 \
  --dropout 0.2

Cross-Validation

python LSTM_CDRs.py \
  --name Cluster1_CV \
  --dataset data/Cluster1.csv \
  --layers 2 \
  --neurons 64 \
  --epochs 50 \
  --dropout 0.2 \
  --cv 5

Common Parameters

Training Parameters

Parameter	Description	Default/Example
--dataset	Path to the training data CSV file	data/Cluster1.csv
--name	Experiment name and output directory name	Cluster1_LSTM
--layers	Number of LSTM/GRU layers	Example: 2
--neurons	Number of neurons per layer	Example: 64
--epochs	Number of training epochs	Example: 200
--batch_size	Batch size	Default: 128
--dropout	Dropout ratio; layer n uses n * dropout	Example: 0.2
--cell	Recurrent neural network cell type	LSTM or GRU
--lr	Adam learning rate	Default: 0.01
--valsplit	Validation split ratio	Default: 0.2
--cv	Number of cross-validation folds	Disabled by default

Sampling Parameters

Parameter	Description	Default/Example
--train False	Do not train; load an existing model for sampling	Required for sampling
--modfile	Path to trained epoch weights	Cluster1_LSTM/checkpoint/model_epoch_100.hdf5
--sample	Number of sequences to generate	Example: 10000
--temp	Sampling temperature	Default: 1.25
-f, --fminlen	Minimum generated sequence length	36 for Cluster1/2/4; 35 for Cluster3
-m, --maxlen	Maximum generated sequence length	36 for Cluster1/2/4; 35 for Cluster3
--startchar	Sampling start character	Default: B

Official OneScience Information

Platform	Documentation	Main OneScience repository	Skills repository
Gitee	https://gitee.com/onescience-ai/onescience-doc	https://gitee.com/onescience-ai/onescience	https://gitee.com/onescience-ai/oneskills
GitHub	https://github.com/onescience-ai/OneScience-doc	https://github.com/onescience-ai/OneScience	https://github.com/onescience-ai/oneskills

Citation and License

• Related work: A. T. Mueller, J. A. Hiss, G. Schneider, "Recurrent Neural Network Model for Constructive Peptide Design", Journal of Chemical Information and Modeling, 2018, DOI: 10.1021/acs.jcim.7b00414.
• Application paper: P. Arras et al., "AI/ML combined with Next Generation Sequencing of VHH immune repertoires enables the rapid identification of de novo humanized and sequence-optimized single domain antibodies: a prospective case study", Frontiers in Molecular Biosciences, 2023, DOI: 10.3389/fmolb.2023.1249247.
• This project uses the MIT License; see LICENSE in the repository root. For specific terms governing the use of data and model weights, follow the information provided by the respective publishers.