modelgenerator 0.1.0

ModelGenerator is an opinionated plug-and-play research framework for cross-disciplinary teams in ML & Bio

ModelGenerator

ModelGenerator is an opinionated plug-and-play research framework for cross-disciplinary teams in ML & Bio.

ModelGenerator is designed to enable rapid and reproducible prototyping with four kinds of experiments in mind:

1. Applying pre-trained foundation models to new data
2. Developing new finetuning and inference tasks for foundation models
3. Benchmarking foundation models and creating leaderboards
4. Testing new architectures for finetuning performance

while also scaling with hardware and integrating with larger data pipelines or research workflows.

ModelGenerator is built on PyTorch, HuggingFace, and Lightning, and works seamlessly with these ecosystems.

Who uses ModelGenerator?

🧬 Biologists

• Intuitive one-command CLIs for in silico experiments
• Pre-trained model zoo
• Broad data compatibility
• Pipeline-oriented workflows

🤖 ML Researchers

• Reproducible-by-design experiments
• Architecture A/B testing
• Automatic hardware scaling
• Integration with PyTorch, Lightning, HuggingFace, and WandB

☕ Software Engineers

• Extensible and modular models, tasks, and data
• Strict typing and documentation
• Fail-fast interface design
• Continuous integration and testing

🤝 Everyone benefits from

• A collaborative hub and focal point for multidisciplinary work on experiments, models, software, and data
• Community-driven development
• Permissive license for academic and non-commercial use

Projects using ModelGenerator

• Accurate and General DNA Representations Emerge from Genome Foundation Models at Scale
• A Large-Scale Foundation Model for RNA Function and Structure Prediction
• Mixture of Experts Enable Efficient and Effective Protein Understanding and Design
• Scaling Dense Representations for Single Cell with Transcriptome-Scale Context
• Balancing Locality and Reconstruction in Protein Structure Tokenizer

Quick Start

Installation

pip install modelgenerator

# To use AIDO.StructureTokenizer, also install openfold and dllogger
pip install git+https://github.com/genbio-ai/openfold.git@c4aa2fd0d920c06d3fd80b177284a22573528442
pip install git+https://github.com/NVIDIA/dllogger.git@0540a43971f4a8a16693a9de9de73c1072020769

Good for running inference, reproducing published experiments, or finetuning on new data

Developer Installation

git clone https://github.com/genbio-ai/ModelGenerator
cd modelgenerator
pip install -e .

Necessary to add new backbones, finetuning tasks, or data transformations

Quick Start

Get embeddings from a pre-trained model

mgen predict --model Embed --model.backbone aido_dna_dummy \
  --data SequencesDataModule --data.path genbio-ai/100m-random-promoters \
  --data.x_col sequence --data.id_col sequence --data.test_split_size 0.0001 \
  --config configs/examples/save_predictions.yaml

Get token probabilities from a pre-trained model

mgen predict --model Inference --model.backbone aido_dna_dummy \
  --data SequencesDataModule --data.path genbio-ai/100m-random-promoters \
  --data.x_col sequence --data.id_col sequence --data.test_split_size 0.0001 \
  --config configs/examples/save_predictions.yaml

Finetune a pre-trained model

mgen fit --model ConditionalDiffusion --model.backbone aido_dna_dummy \
  --data ConditionalDiffusionDataModule --data.path "genbio-ai/100m-random-promoters"

Evaluate a model checkpoint

mgen test --model ConditionalDiffusion --model.backbone aido_dna_dummy \
  --data ConditionalDiffusionDataModule --data.path "genbio-ai/100m-random-promoters" \
  --ckpt_path logs/lightning_logs/version_X/checkpoints/<your_model>.ckpt

Save predictions

mgen predict --model ConditionalDiffusion --model.backbone aido_dna_dummy \
  --data ConditionalDiffusionDataModule --data.path "genbio-ai/100m-random-promoters" \
  --ckpt_path logs/lightning_logs/version_X/checkpoints/<your_model>.ckpt \
  --config configs/examples/save_predictions.yaml

Configify your experiment

The command

mgen fit --model ConditionalDiffusion --model.backbone aido_dna_dummy \
  --data ConditionalDiffusionDataModule --data.path "genbio-ai/100m-random-promoters"

is equivalent to mgen fit --config my_config.yaml with

# my_config.yaml
model:
  class_path: ConditionalDiffusion
  init_args:
    backbone: aido_dna_dummy
data:
  class_path: ConditionalDiffusionDataModule
  init_args:
    path: "genbio-ai/100m-random-promoters"

Use composable configs to customize workflows

mgen fit --model SequenceRegression --data PromoterExpressionRegression \
  --config configs/defaults.yaml \
  --config configs/examples/lora_backbone.yaml \
  --config configs/examples/wandb.yaml

Configs use the LAST value for each attribute. Check the full configuration in logs/lightning_logs/your-experiment/config.yaml, or if using wandb logs/config.yaml.

Use LoRA for parameter-efficient finetuning

This also avoids saving the full model, only the LoRA weights are saved.

mgen fit --data PromoterExpressionRegression \
  --model SequenceRegression --model.backbone.use_peft true \
  --model.backbone.lora_r 16 \
  --model.backbone.lora_alpha 32 \
  --model.backbone.lora_dropout 0.1

Use continued pretraining for finetuning domain adaptation

First run pretraining objective on finetuning data

# https://arxiv.org/pdf/2310.02980
mgen fit --model MLM --model.backbone aido_dna_dummy \
  --data MLMDataModule --data.path leannmlindsey/GUE \
  --data.config_name prom_core_notata

Then finetune using the adapted model

mgen fit --model SequenceClassification --model.strict_loading false \
  --data SequenceClassificationDataModule --data.path leannmlindsey/GUE \
  --data.config_name prom_core_notata \
  --ckpt_path logs/lightning_logs/version_X/checkpoints/<your_adapted_model>.ckpt

Make sure to turn off strict_loading to replace the adapter!

Use the head/adapter that comes with the backbone

mgen fit --model SequenceClassification --data GUEClassification \
--model.use_legacy_adapter true