Drutai-ONNX

High-performance drug-target interaction prediction tool using ONNX models for CPU inference.

Features

Core Functionality

Multi-model Support: LSTMCNN, CNN, ConvMixer64, DSConv, MobileNetV2, ResNet18, SEResNet
ONNX Inference: Bypasses TensorFlow/Keras version dependencies, pure CPU efficient inference
Dual Input Modes: Single-table mode (recommended) and Legacy mode (compatible with original drutai)
Snap Packaging: Cross-distribution Linux support with no dependency conflicts

Performance Optimizations

Independent Feature Caching: Protein and molecule features cached separately to avoid redundant computation
Persistent Caching: SHA256 dual verification with automatic content change detection
Multithreaded Feature Extraction: Parallel processing of protein and molecule features (default max 64 threads)
Mini-batch Inference: Controls memory usage to prevent OOM (default batch_size=2000)
Explicit Memory Management: Proactive memory release at critical points

Robustness

Cache Integrity Verification: Detects and automatically cleans corrupted cache files
Friendly Error Messages: Clear error information for missing FASTA files, SMILES parsing failures, etc.
Graceful Degradation: Uses zero features with warnings when molecule parsing fails
Specific Exception Handling: Avoids catching critical errors like KeyboardInterrupt

Quick Start

1. Install

sudo snap install drutai

The snap exposes two equivalent commands:

drutai.predict: primary prediction command used in the examples below
drutai: shorter alias for the same prediction entry point

For a locally built snap:

sudo snap install --dangerous drutai_1.0.5rc3_amd64.snap

The snap uses strict confinement. Keep input, output, and FASTA files under your home directory when possible. To use files under /media or /mnt, connect the removable-media interface:

sudo snap connect drutai:removable-media

2. Prepare Inputs

Single-table mode is recommended. The input file must be a tab-separated TSV with at least these columns: sm, target, and smile.

sm	target	smile
DB00006	TREM1	CC[C@H](C)[C@H](NC(=O)...)C(O)=O
DB00014	TREM1	CC(C)C[C@H](NC(=O)...)NNC(N)=O

Each target needs one FASTA file whose filename matches the target column:

fasta/
└── TREM1.fasta

Example TREM1.fasta:

>TREM1
MALPVTALLLPLALLLHAARPSQFRVSPSGPVQGCALEVRCQEGKGHWLHSCTWLPVS

3. Run Prediction

drutai.predict -m lstmcnn -i input.tsv -t fasta/ -o output.drutai

Available model names:

lstmcnn, cnn, convmixer64, dsconv, mobilenetv2, resnet_prea18_tf2, scaresnet

The output is a TSV containing the original columns plus:

prob_inter: interaction probability from 0 to 1
pred_type: Interaction or Non-interaction, using 0.5 as the threshold

4. Common Options

Option	Meaning	Default
`-m`, `--model`	Model name without `.onnx` suffix	`lstmcnn`
`-i`, `--input`	Single-table TSV input, recommended	none
`--br` + `--smile`	Legacy two-file input	none
`-t`, `--fasta_fp`	FASTA directory, required	none
`-o`, `--output`	Output path	`predictions_onnx.drutai`
`-j`, `--threads`	Feature extraction threads	auto, up to 64
`-b`, `--batch-size`	Inference batch size; reduce if memory is limited	`2000`
`-s`, `--silence`	Hide thread-count and batch-size hints	off
`--no-cache`	Disable feature caching	off

How It Works

1. Input Processing

Single-table Mode (Recommended):

drutai.predict -m lstmcnn -i input.tsv -t fasta/ -o output.drutai

input.tsv contains three columns: sm (drug ID), target (target ID), smile (SMILES string)

Legacy Mode (Compatible with Original drutai):

drutai.predict -m lstmcnn --br relations.txt --smile smiles.txt -t fasta/ -o output.drutai

relations.txt: Drug-target relations (sm, target)
smiles.txt: Drug SMILES (sm, smile)

Dual-mode Validation:

When both modes are provided, automatically validates data consistency
Reports errors with detailed differences if inconsistent

2. Feature Extraction

Protein Features (8567 dimensions)

Extracted from FASTA files:

Composition (20 dims): Frequency of 20 amino acids
Dipeptide (400 dims): Frequency of 20×20 dipeptide combinations
Tripeptide (8000 dims): Frequency of 20×20×20 tripeptide combinations (first 8000)
CTD Features (147 dims): Composition, Transition, Distribution descriptors

Molecule Features (4293 dimensions)

Extracted from SMILES strings:

Descriptors (195 dims): RDKit molecular descriptors
Crippen Parameters (2 dims): LogP and MR
Morgan Fingerprint (2048 dims): Circular fingerprint, radius=2
Topological Torsion Fingerprint (2048 dims): Topological structure fingerprint

Feature Combination

Protein features + Molecule features = 12860 dimensions
First 11664 dimensions used as model input

3. Caching Mechanism

Cache Location

Snap Environment: ~/snap/drutai/common/.drutai_cache/
Regular Environment: ~/.cache/drutai/

Cache Key Generation

Protein: protein_{MD5(target_id)}.pkl
Molecule: molecule_{MD5(drug_id_smiles)}.pkl

Cache Data Structure

{
    'seq_hash': 'sha256_hash_of_sequence',  # Content hash
    'features': [feat1, feat2, ...]          # Feature vector
}

Dual Verification Mechanism

File Existence Check: Whether cache file exists
Content Hash Verification:
- Protein: Read FASTA sequence, compute SHA256, compare with cached seq_hash
- Molecule: Compute SMILES SHA256, compare with cached smiles_hash
Format Integrity Check: Verify required keys (seq_hash/smiles_hash and features) exist
Auto-cleanup: Automatically deletes corrupted cache files when detected

Cache Hit Conditions

File exists ✓
Format complete ✓
Content hash matches ✓

Cache Miss Scenarios

Cache file doesn't exist
Cache file corrupted (pickle load failure)
Cache format incomplete (missing required keys)
Sequence/SMILES content changed (hash mismatch)

4. Mini-batch Inference

Workflow

Build Feature Matrix: Split by batch_size, build (batch_size, 11664) matrix per batch
Data Preprocessing:
- Handle NaN/Inf: np.nan_to_num
- Numerical clipping: np.clip(-1e6, 1e6)
- Type conversion: float64 → float32
ONNX Inference: Single call processes one batch
Memory Cleanup: Immediately release float64 matrix and inference results after each batch
Progress Display: Print progress every 5 batches or at the last batch

Memory Control

Default batch_size=2000 (conservative value)
Adjustable via -b/--batch-size
Each batch processed independently to avoid loading all data at once

5. Output Format

Output is a TSV file containing original columns + prediction results:

sm	target	smile	prob_inter	pred_type
DB00006	TREM1	CC[C@H]...	0.8234	Interaction
DB00014	TREM1	CC(C)C...	0.3421	Non-interaction

prob_inter: Interaction probability (0-1)
pred_type: Interaction (>0.5) or Non-interaction (≤0.5)

Usage Examples

Basic Usage

# Single-table mode
drutai.predict -m lstmcnn -i input.tsv -t fasta/ -o output.drutai

# Legacy mode
drutai.predict -m lstmcnn --br relations.txt --smile smiles.txt -t fasta/ -o output.drutai

Performance Tuning

# Adjust thread count (default: all CPUs, max 64)
drutai.predict -m lstmcnn -i input.tsv -t fasta/ -o output.drutai -j 32

# Adjust batch size (reduce if memory insufficient)
drutai.predict -m lstmcnn -i input.tsv -t fasta/ -o output.drutai -b 1000

# Silence mode (hide thread count and batch size hints)
drutai.predict -m lstmcnn -i input.tsv -t fasta/ -o output.drutai -s

Cache Control

# Disable caching (re-extract features every time)
drutai.predict -m lstmcnn -i input.tsv -t fasta/ -o output.drutai --no-cache

# Clear cache manually
rm -rf ~/snap/drutai/common/.drutai_cache/  # Snap environment
rm -rf ~/.cache/drutai/                      # Regular environment

Multi-model Testing

for model in lstmcnn cnn dsconv mobilenetv2; do
    drutai.predict -m $model -i input.tsv -t fasta/ -o output_${model}.drutai
done

Performance Comparison

Original drutai vs drutai-onnx

Dimension	Original drutai	drutai-onnx (v1.0.5rc3)
Inference Engine	TensorFlow/Keras	ONNX Runtime
Model Format	SavedModel	ONNX
Keras Compatibility	Keras 2 only	No Keras dependency
Feature Caching	None	Persistent cache (SHA256 verification)
Memory Usage	High (full load)	Low (mini-batch)
Inference Speed	Baseline	10-15% faster
Cache Hit Speed	N/A	Extremely fast (skip feature extraction)

v1.0.5rc3 Current Capabilities

✅ Cache integrity verification (auto-cleanup corrupted cache)
✅ Friendly error messages (missing FASTA files, SMILES parsing failures)
✅ Specific exception handling (avoid catching critical errors)
✅ Code quality improvements (remove duplicate imports, update documentation)

Technical Details

Dependencies

Python ≥ 3.10
numpy ≥ 1.24.3
pandas ≥ 2.0.0
onnxruntime ≥ 1.15.1
rdkit ≥ 2023.3.2
biopython ≥ 1.81

Model List

Model Name	Architecture	Input Shape
lstmcnn	LSTM + CNN	(N, 11664)
cnn	CNN	(N, 11664)
convmixer64	ConvMixer	(N, 108, 108, 1)
dsconv	Depthwise Separable Conv	(N, 11664)
mobilenetv2	MobileNetV2	(N, 11664)
resnet_prea18_tf2	ResNet18	(N, 11664)
scaresnet	SE-ResNet	(N, 11664)

Thread Safety

Feature extraction: Multithreaded parallel (ThreadPoolExecutor)
ONNX inference: Single-threaded sequential (avoid session contention)
Cache read/write: Thread-safe (each thread handles independent files)

Troubleshooting

Common Issues

1. FASTA File Not Found

FileNotFoundError: FASTA file not found for target 'TREM1': /path/to/fasta/TREM1.fasta
Please ensure the FASTA directory contains 'TREM1.fasta'

Solution: Ensure FASTA directory contains {target_id}.fasta file.

2. SMILES Parsing Failed

UserWarning: Drug 'DB00006' with SMILES 'CC[C@H]...' failed to parse. Using zero features as fallback.

Solution: Check if SMILES string is valid, or use --no-cache to re-extract.

3. Out of Memory

MemoryError: Unable to allocate array

Solution: Reduce batch_size, e.g., -b 1000 or -b 500.

4. Corrupted Cache

RuntimeWarning: Unexpected error loading cache /path/to/cache.pkl: ...

Solution: Cache will auto-cleanup, or manually delete cache directory.

Installation

Install from Snap Store (Recommended)

sudo snap install drutai

Install from Local File

sudo snap install --dangerous drutai_1.0.5rc3_amd64.snap

Development

Build Snap Package

cd /home/husrcf/Code/drutai_snap
snapcraft pack

Local Testing

sudo snap install --dangerous drutai_1.0.5rc3_amd64.snap
drutai.predict -h
./test_snap.sh

Publish to Snap Store

snapcraft login
snapcraft upload drutai_1.0.5rc3_amd64.snap --release=beta

License

MIT License

Contributing

Issues and Pull Requests welcome: https://github.com/HUSRCF/drutai_snap

Citation

If you use this tool, please cite the original drutai project:

Sun, J. et al. (2025). Drutai: Drug-Target Interaction Prediction.
GitHub: https://github.com/2003100127/drutai

Contact

Issues: https://github.com/HUSRCF/drutai_snap/issues
Original Project: https://github.com/2003100127/drutai

Acknowledgments

Original drutai project: https://github.com/2003100127/drutai
ONNX Runtime: https://onnxruntime.ai/
RDKit: https://www.rdkit.org/

Name		Name	Last commit message	Last commit date
fasta		fasta
.gitignore		.gitignore
README.md		README.md
README_CN.md		README_CN.md
RELEASE.md		RELEASE.md
br_sm_target_test100.txt		br_sm_target_test100.txt
br_smile_test100.txt		br_smile_test100.txt
cnn.onnx		cnn.onnx
convmixer64.onnx		convmixer64.onnx
dsconv.onnx		dsconv.onnx
lstmcnn.onnx		lstmcnn.onnx
mobilenetv2.onnx		mobilenetv2.onnx
onnx_predict.py		onnx_predict.py
out_v1.drutai		out_v1.drutai
out_v2.drutai		out_v2.drutai
resnet_prea18_tf2.onnx		resnet_prea18_tf2.onnx
scaresnet.onnx		scaresnet.onnx
setup.py		setup.py
snapcraft.yaml		snapcraft.yaml
test_snap.sh		test_snap.sh

Latest commit

History

Folders and files

Repository files navigation