DNALLM-Suite - DNA Large Language Model Toolkit

DNALLM-Suite is a comprehensive, open-source toolkit designed for fine-tuning and inference with DNA Large Language Models. It provides a unified interface for working with various DNA sequence models, supporting tasks ranging from basic sequence classification to advanced in-silico mutagenesis analysis. With built-in Model Context Protocol (MCP) support, DNALLM-Suite enables seamless communication with traditional large language models, allowing for enhanced integration and interoperability in AI-powered DNA analysis workflows.

📦 Quick Installation

pip install dnallm

See CHANGELOG.md for version history.

🚀 Key Features

• 🔄 Model Management: Load and switch between 200+ pre-trained DNA large language models from Hugging Face and ModelScope
• 🎯 Multi-Task Support: Binary/multi-class classification, regression, NER, MLM, and generation tasks
• 📊 Benchmarking: Multi-model performance comparison and evaluation metrics
• 🔧 Fine-tuning: Comprehensive training pipeline with configurable parameters
• 📱 Interactive Interfaces: Jupyter notebooks and Marimo-based interactive demos
• 🌐 MCP Support: Model Context Protocol for server/client deployment with real-time streaming
• 🧬 Advanced Analysis: In-silico mutagenesis, saturation mutation analysis, and mutation effect visualization
• 🧪 Comprehensive Testing: 1,900+ test cases covering all major functionality

🧬 Supported Models

DNALLM-Suite supports a wide range of DNA large language models including:

From-scratch baselines: any generically-loadable supported model can also be loaded with genuinely random weights via load_model_and_tokenizer(..., random_init=True). The model config and tokenizer files are fetched as usual, but no pretrained weights are ever downloaded — weights are sampled through AutoConfig.from_pretrained + Auto*.from_config with CPU-canonical seeding (random_init_seed, default 42), so identical seeds reproduce identical models. A loud "randomly initialized" banner plus a per-tensor parameter-hash table are logged at INFO level as proof that no pretrained weights leaked into the baseline. Special model families without a from-scratch path raise a ValueError; the special-family allowlist is RANDOM_INIT_SUPPORTED_FAMILIES (currently the Mamba trust-remote-code architecture, which loads generically — all other generic Auto* families are always allowed).

Masked Language Models (MLM)

• DNABERT Series: Plant DNABERT, DNABERT, DNABERT-2, DNABERT-S
• Caduceus Series: Caduceus-Ph, Caduceus-PS, PlantCaduceus
• Specialized Models: AgroNT, GENA-LM, GPN, GROVER, MutBERT, ProkBERT, PlantHelixSeek, CrossDNA

Causal Language Models (CLM)

• EVO Series: EVO-1, EVO-2
• Plant Models: Plant DNAGemma, Plant DNAGPT, Plant DNAMamba
• Other Models: GENERator, GenomeOcean, HyenaDNA, Jamba-DNA, Mistral-DNA

Model Sources

• Hugging Face Hub: Primary model repository
• ModelScope: Alternative model source with additional models
• Custom Models: Support for locally trained or custom architectures

🖥️ Supported Platforms

DNALLM-Suite has been tested on a wide range of platforms and devices:

Platforms:

• Linux
• Windows
• MacOS

Device:

• CPU
• Nvidia GPU (CUDA)
• AMD GPU (ROCm)
• Apple Silicon (MPS)
• Huawei Ascend NPU (CANN)
• Intel Arc GPU (XPU)

🛠️ Installation

Prerequisites

• Python 3.11 or higher (recommended)
• Git
• CUDA-compatible GPU (optional, for GPU acceleration)
• Environment Manager: Choose one of the following:
  • Python venv (built-in)
  • Conda/Miniconda (recommended for scientific computing)

Quick Installation with uv (Recommended)

DNALLM-Suite uses uv for dependency management and packaging.

What is uv is a fast Python package manager that is 10-100x faster than traditional tools like pip.

Method 1: Using venv + uv

# Clone repository
git clone https://github.com/zhangtaolab/DNALLM.git
cd DNALLM

# Create virtual environment
python -m venv .venv

# Activate virtual environment
source .venv/bin/activate  # Linux/MacOS
# or
.venv\Scripts\activate     # Windows

# Upgrade pip (recommended)
pip install --upgrade pip

# Install uv in virtual environment
pip install uv

# Install DNALLM with all optional dependencies
uv pip install -e '.[all]'

# Or install only base development tools
# uv pip install -e '.[base]'

# Verify installation
python -c "import dnallm; print('DNALLM installed successfully!')"

Method 2: Using conda + uv

# Clone repository
git clone https://github.com/zhangtaolab/DNALLM.git
cd DNALLM

# Create conda environment
conda create -n dnallm python=3.12 -y

# Activate conda environment
conda activate dnallm

# Install uv in conda environment
conda install uv -c conda-forge

# Install DNALLM with all optional dependencies
uv pip install -e '.[all]'

# Or install only base development tools
# uv pip install -e '.[base]'

# Verify installation
python -c "import dnallm; print('DNALLM installed successfully!')"

GPU Support

For GPU acceleration, install the appropriate CUDA version after installing the base package:

# For venv users: activate virtual environment
source .venv/bin/activate  # Linux/MacOS
# or
.venv\Scripts\activate     # Windows

# For conda users: activate conda environment
# conda activate dnallm

# CUDA 12.4 (recommended for recent GPUs)
uv pip install -e '.[all,cuda124]'

# Other supported versions: cpu, cuda121, cuda126, cuda128, cuda130
# Nvidia 5090 Please use cuda128 & torch==2.7
uv pip install -e '.[all,cuda128]'

# CUDA 13.0 (torch >= 2.9, Windows & Linux, NVIDIA driver >= 580)
uv pip install -e '.[all,cuda130]'

Warning: Hardware groups (cpu, cuda121, cuda124, cuda126, cuda128, cuda130, rocm, mamba) must not be combined — cpu/cuda* pairs are enforced as mutually exclusive by the resolver; rocm cannot share an environment with another hardware group (only one torch build installs), while mamba is designed to combine with a CUDA group (e.g. .[cuda124,mamba]). Do NOT combine multiple CUDA versions.

Dependency Groups

Group	Purpose	Includes
all	Install everything	base + dev + test + notebook + docs + ui + mcp + fla
base	Full dev environment	dev + test + notebook + mcp + extra tools
dev	Development	test + notebook + linting/typing tools
test	Testing only	pytest and plugins
notebook	Interactive notebooks	Jupyter, Marimo
docs	Build documentation	mkdocs and plugins
ui	Gradio web interface	Gradio
mcp	MCP server	(included in core)
fla	PlantHelixSeek KDA kernels	flash-linear-attention (Triton-based, uses the installed torch — combinable with any hardware group)

Hardware groups (mutually exclusive, NOT included in all):

Group	PyTorch	Use Case
cpu	2.4.0-2.11	No GPU
cuda121	2.2.0-2.11	Older NVIDIA GPUs
cuda124	2.4.0-2.11	Most modern GPUs (recommended)
cuda126	2.6.0-2.11	Ada/Hopper with Flash Attention
cuda128	2.6.0-2.11	RTX 5090 and latest hardware
cuda130	2.9.0-2.11	CUDA 13.0, Windows & Linux (driver >= 580)
rocm	2.5.0-2.11	AMD GPUs
mamba	2.6.0-2.11	Native Mamba architecture (requires CUDA)

# Examples:
uv pip install -e '.[all,cuda124]'    # Everything + CUDA 12.4
uv pip install -e '.[base,mamba]'     # Dev tools + native Mamba
uv pip install -e '.[test,cpu]'       # Testing only, no GPU (add ,mcp for the MCP example tests)

Native Mamba Support

Native Mamba architecture runs significantly faster than transformer-compatible Mamba architecture, but native Mamba depends on Nvidia GPUs.

If you need native Mamba architecture support, after installing DNALLM dependencies, use the following command:

# For venv users: activate virtual environment
source .venv/bin/activate  # Linux/MacOS

# For conda users: activate conda environment
# conda activate dnallm

# Ensure CUDA path is set correctly (nvcc version must match your PyTorch CUDA version)
export PATH=/usr/local/cuda-12/bin:$PATH
nvcc -V  # Verify CUDA compiler version

# Install Mamba support
uv pip install -e '.[mamba]' --no-cache-dir --no-build-isolation --link-mode=copy

# If encounter network issue, using the special install script for mamba (optional)
sh scripts/install_mamba.sh  # select github proxy

Note: The nvcc version must match your PyTorch CUDA version. For example, if you installed PyTorch with CUDA 12.8, you need nvcc from CUDA 12.x. Mismatched versions will cause build failures.

Please ensure your machine can connect to GitHub, otherwise Mamba dependencies may fail to download.

Note that Plant DNAMamba, Caduceus, PlantCaduceus, PlantCAD2, Jamba-DNA, JanusDNA models are all based on Mamba architecture. Therefore, the training and inference of these models can be accelerated by installing the native mamba support.

Flash-Linear-Attention (KDA) Kernels for PlantHelixSeek

PlantHelixSeek-CRE and PlantHelixSeek-Anno remote code imports fla.ops.kda.chunk.chunk_kda from flash-linear-attention for their HelixSeekDelta (KDA) layers. Install the fla extra (already included in all):

uv pip install -e '.[fla]'

Warning: without flash-linear-attention the remote code silently falls back to a pure-PyTorch path that is not KDA math (wrong gate formula, scalar decay, no delta-rule correction). Models still load and run, but produce positionally-uninformative outputs — e.g. measured p(CRE) 0.007 inside real DHS sites vs 0.009 outside (dead), versus 0.77 vs 0.22 with the kernels installed. Do not interpret PlantHelixSeek predictions without this package.

Two install notes: install the package bare (flash-linear-attention, not its [cuda]/[rocm] extras) — since v0.5 the backend extras pin their own torch and would downgrade your environment; and keep the version within 0.5.x (>=0.5.2,<0.6) — chunk_kda semantics are not guaranteed stable across minor versions.

Install Dependencies for Special Models

Several models require extra dependencies to train or inference.

These models are listed below:

Models	Model Type	Source	Dependencies
EVO-1	CausalLM	Hugging Face	GitHub
EVO2	CausalLM	Hugging Face	GitHub
PlantHelixSeek-CRE/-Anno	Token Classification / Binary	ModelScope / Hugging Face	fla extra (flash-linear-attention, see warning below)
GPN	MaskedLM	Hugging Face	GitHub
megaDNA	CausalLM	Hugging Face	GitHub
LucaOne	CausalLM	Hugging Face	GitHub
Omni-DNA	CausalLM	Hugging Face	GitHub

The installation method for the dependencies of these models can be found here.

🚀 Quick Start

1. Basic Model Loading and Inference

from dnallm import load_config, load_model_and_tokenizer, DNAInference

# Load configuration
configs = load_config("./example/notebooks/inference/inference_config.yaml")

# Load model and tokenizer
model_name = "zhangtaolab/plant-dnagpt-BPE-promoter_strength_protoplast"
model, tokenizer = load_model_and_tokenizer(
    model_name, task_config=configs["task"], source="huggingface"
)

# Initialize inference engine
inference_engine = DNAInference(config=configs, model=model, tokenizer=tokenizer)

# Make inference
sequence = "AATATATTTAATCGGTGTATAATTTCTGTGAAGATCCTCGATACTTCATATAAGAGATTTTGAGAGAGAGAGAGAACCAATTTTCGAATGGGTGAGTTGGCAAAGTATTCACTTTTCAGAACATAATTGGGAAACTAGTCACTTTACTATTCAAAATTTGCAAAGTAGTC"
inference_result = inference_engine.infer(sequence)
print(f"Inference result: {inference_result}")

2. In-silico Mutagenesis Analysis

from dnallm import Mutagenesis

# Initialize mutagenesis analyzer
mutagenesis = Mutagenesis(config=configs, model=model, tokenizer=tokenizer)

# Generate saturation mutations
mutagenesis.mutate_sequence(sequence, replace_mut=True)

# Evaluate mutation effects
predictions = mutagenesis.evaluate(strategy="mean")

# Visualize results
plot = mutagenesis.plot(predictions, save_path="mutation_effects.pdf")

3. Model Fine-tuning

from dnallm.datahandling import DNADataset
from dnallm.finetune import DNATrainer

# Prepare dataset
dataset = DNADataset.from_huggingface(
    "zhangtaolab/plant-multi-species-core-promoters",
    seq_col="sequence",
    label_col="label",
    tokenizer=tokenizer,
)

# Initialize trainer
trainer = DNATrainer(model=model, config=configs, datasets=dataset)

# Start training
trainer.train()

4. MCP Server Deployment

Start the MCP server via the CLI:

# Start with Streamable HTTP (recommended, per MCP spec 2025-11-25)
dnallm-mcp-server --transport streamable-http --host 127.0.0.1 --port 8000

The server exposes a single /mcp endpoint (POST/GET/DELETE) per MCP spec 2025-11-25.

Legacy SSE transport is still supported for backward compatibility:

# Legacy SSE transport (deprecated in MCP spec 2025-11-25, still supported)
dnallm-mcp-server --transport sse --host 127.0.0.1 --port 8000

# Start MCP server for real-time DNA sequence prediction
from dnallm.mcp import DNALLMMCPServer

# Initialize MCP server
server = DNALLMMCPServer("config/mcp_server_config.yaml")
await server.initialize()

# Start server with Streamable HTTP transport (recommended, MCP spec 2025-11-25)
server.start_server(host="0.0.0.0", port=8000, transport="streamable-http")

# Legacy SSE transport is still supported for backward compatibility
# server.start_server(host="0.0.0.0", port=8000, transport="sse")

MCP Server Features

• Real-time Streaming: Streamable HTTP for live prediction updates (MCP spec 2025-11-25)
• Multiple Transport Protocols: STDIO, Streamable HTTP (recommended), and SSE (legacy)
• Comprehensive Tools: 10+ MCP tools for DNA sequence analysis
• Model Management: Dynamic model loading and switching
• Batch Processing: Efficient handling of multiple sequences
• Health Monitoring: Built-in server diagnostics and status checks

MCP Client SDK

from dnallm.mcp.client import DNALLMMCPClient

# Connect via Streamable HTTP (recommended)
client = DNALLMMCPClient(transport="streamable-http", url="http://localhost:8000/mcp")
result = client.dna_sequence_predict("ATCGATCG", "dnabert-2")

# Connect via stdio for local CLI usage
client = DNALLMMCPClient(transport="stdio")
result = client.health_check()

# Legacy SSE transport (deprecated in MCP spec 2025-11-25, still supported)
# client = DNALLMMCPClient(transport="sse", url="http://localhost:8000/sse")

Available MCP Tools

• _dna_sequence_predict - Single sequence prediction
• _dna_batch_predict - Batch sequence processing
• _dna_multi_model_predict - Multi-model comparison
• _dna_stream_predict - Real-time streaming prediction
• _dna_stream_batch_predict - Streaming batch prediction
• _dna_stream_multi_model_predict - Streaming multi-model comparison
• _dna_mutagenesis - In-silico mutagenesis analysis
• _dna_interpret - Model interpretation (feature attribution)
• _list_loaded_models - List currently loaded models
• _get_model_info - Query a loaded model's details
• _list_models_by_task_type - List models by task type
• _get_all_available_models - List all models in the registry
• _health_check - Server monitoring

📚 Examples and Tutorials

Interactive Demos (Marimo)

# Launch Jupyter Lab
uv run --no-sync jupyter lab

# Fine-tuning demo
uv run --no-sync marimo run example/marimo/finetune/finetune_demo.py

# Inference demo
uv run --no-sync marimo run example/marimo/inference/inference_demo.py

# Benchmark demo
uv run --no-sync marimo run example/marimo/benchmark/benchmark_demo.py

Web-based UI (Gradio)

# Launch Gradio configuration generator app
uv run --no-sync python ui/run_config_app.py

# Or run the model config generator directly
uv run --no-sync python ui/model_config_generator_app.py

# For Generation, we also provide a app
uv run --no-sync python ui/generation_task_app.py

Jupyter Notebooks

# Launch Jupyter Lab
uv run --no-sync jupyter lab

# Available notebooks:
# - example/notebooks/finetune_binary/ - Binary classification fine-tuning
# - example/notebooks/finetune_multi_labels/ - Multi-label classification
# - example/notebooks/finetune_NER_task/ - Named Entity Recognition
# - example/notebooks/inference/ - Model inference
# - example/notebooks/in_silico_mutagenesis/ - Mutation analysis
# - example/notebooks/inference_for_tRNA/ - tRNA-specific analysis
# - example/notebooks/generation_evo_models/ - EVO model inference
# - example/notebooks/lora_finetune_inference/ - LoRA fine-tuning
# - example/notebooks/embedding_attention.ipynb - Embedding and attention analysis
# - example/notebooks/finetune_custom_head/ - Custom classification head
# - example/notebooks/finetune_generation/ - Sequence generation
# - example/notebooks/generation/ - Sequence generation examples
# - example/notebooks/generation_megaDNA/ - MegaDNA model inference
# - example/notebooks/interpretation/ - Model interpretation
# - example/notebooks/data_prepare/ - Data preparation examples
# - example/notebooks/benchmark/ - Model evaluation and benchmarking
# - example/notebooks/plant_helixseek_cre/ - PlantHelixSeek-CRE showcase notebook
# - example/notebooks/plant_helixseek_anno/ - PlantHelixSeek-Anno showcase notebook

🏗️ Project Structure

DNALLM/
├── dnallm/                  # Core library package
│   ├── cli/                 # Command-line interface
│   ├── configuration/       # Configuration management
│   ├── datahandling/        # Dataset processing
│   ├── finetune/            # Fine-tuning pipeline
│   ├── inference/           # Inference & analysis tools
│   ├── models/              # Model loading & registry
│   ├── tasks/               # Task definitions & metrics
│   ├── utils/               # Utility functions
│   └── mcp/                 # MCP server implementation
├── cli/                     # Legacy CLI scripts (deprecated)
├── example/                 # Examples & tutorials
│   ├── marimo/              # Interactive Marimo apps
│   └── notebooks/           # Jupyter notebooks
├── docs/                    # Documentation
├── tests/                   # Test suite
├── ui/                      # Gradio web interfaces
├── scripts/                 # Development scripts
├── .github/                 # GitHub workflows
├── pyproject.toml           # Project configuration
└── README.md                # This file

🔧 Command Line Interface

DNALLM-Suite provides convenient CLI tools:

# Main CLI with subcommands
dnallm --help

# Training
dnallm train --config path/to/config.yaml
# or
dnallm-train --config path/to/config.yaml

# Inference
dnallm inference --config path/to/config.yaml --input path/to/sequences.txt
# or
dnallm-inference --config path/to/config.yaml --input path/to/sequences.txt

# Model configuration generator
dnallm-model-config-generator

# In-silico mutagenesis
dnallm-mutagenesis --model-name <model-name> --sequence ATCGATCG

# MCP server
dnallm-mcp-server --config path/to/config.yaml

🎯 Supported Task Types

DNALLM-Suite supports the following task types:

• EMBEDDING: Extract embeddings, attention maps, and token probabilities for downstream analysis
• MASK: Masked language modeling task for pre-training
• GENERATION: Text generation task for causal language models
• BINARY: Binary classification task with two possible labels
• MULTICLASS: Multi-class classification task that specifies which class the input belongs to (more than two)
• MULTILABEL: Multi-label classification task with multiple binary labels per sample
• REGRESSION: Regression task which returns a continuous score
• NER: Token classification task which is usually for Named Entity Recognition

🧬 Zero-Shot Variant Effect Prediction (VEP)

DNALLM scores variants zero-shot from a VCF (dnallm-vep CLI, dnallm.inference.vep) under an explicit protocol:

Scoring formulas (declared here and verbatim in the module docstrings):

• Masked-LM (log-odds): delta = log P(alt_token | masked context) − log P(ref_token | masked context) at the single alignment slot.
• Causal-LM (delta-log-likelihood): delta = log P(alt_window) − log P(ref_window) over the context window with the allele substituted at the variant position.

Deltas are alt-minus-ref, so deleterious variants carry negative deltas; AUROC/AUPRC are computed over the deleteriousness score −delta (higher = more pathogenic — the evo2-clinvar/GPN field convention), placing discriminating models above the random floor.

Same-slot evaluability rule: a variant is scoreable only when the tokenized reference and alternate sequences have equal length and differ at exactly ONE token slot. Variants failing the same-slot rule are reported as structured skips with machine-readable reasons (length-changing allele, multi-slot token difference, no change) — skip-as-data with per-reason counts and a skip fraction reported as a finding, never a silent drop, and never mixed into the label/paradigm error channels. Scoring windows are uppercased (soft-masked reference input must not silently tokenize to <unk>) and clipped symmetrically around the variant; ref and alt always share one identical window.

Paradigm↔architecture guard: requesting --paradigm clm on a bidirectional model, or --paradigm mlm with a tokenizer that has no mask token, raises a ValueError — a misconfiguration must not masquerade as a near-random result.

ClinVar convention (reported alongside every result — never a bare AUROC): SNVs only (CLNVC=single_nucleotide_variant), labels Pathogenic/Likely_pathogenic (1) vs Benign/Likely_benign (0) with VUS, conflicting, and novel CLNSIG strings excluded rather than guessed, and a ≥1 review-star floor (CLNREVSTAT).

dnallm-vep --model-name <model-name> --vcf variants.vcf --reference genome.fa --paradigm mlm -o result.json

Within-convention comparables from the literature (compare only within paradigm and within convention, never as bare thresholds): Nucleotide Transformer 2.5B MLM ClinVar AUROC ≈ 0.80; Evo2-40B CLM ≈ 0.98; BPE embedding-distance baselines ≈ 0.54–0.60.

🧪 Testing

DNALLM-Suite includes a comprehensive test suite with 1,900+ test cases:

# Install test dependencies first (the mcp extra provides the langchain/pydantic_ai
# imports needed by the MCP example tests)
uv pip install -e '.[test,dev,mcp]'

# Run all tests
uv run pytest

# Run the fast lane only (as CI does on every push/PR)
uv run pytest -m "not slow"

# Run specific test categories
uv run pytest tests/inference/ -v
uv run pytest tests/mcp/ -v
uv run pytest tests/tasks/ -v

# Run with coverage
uv run pytest --cov=dnallm --cov-report=html

CI runs the fast lane (~1,900 tests) with a >=90% coverage gate on every push and pull request (Python 3.11/3.12/3.13 × numpy 2.2.0, plus a Windows leg). The slow, network-dependent suite (real-model downloads) and a nightly example-execution census that runs the example/ notebooks and demo apps end-to-end execute on a self-hosted GPU runner (see .github/workflows/ci.yml).

📖 Documentation

• Getting Started - Installation and basic usage

• API Reference - Detailed function documentation

• Concepts - Core concepts and architecture

• FAQ - Common questions and solutions

• DeepWiki - A documentation that can ask

🤝 Contributing

We welcome contributions! Please see our Contributing Guide for details on:

• Code style and standards
• Testing requirements
• Pull request process
• Development setup

📄 License

This project is licensed under the MIT License - see the LICENSE file for details.

🙏 Acknowledgments

• Hugging Face - Model hosting and transformers library
• ModelScope - Alternative model repository

📞 Support

• Documentation: docs/
• Issues: GitHub Issues
• Discussions: GitHub Discussions
• Examples: Check the example/ directory for working code

---

DNALLM - Empowering DNA sequence analysis with state-of-the-art language models.