Imranyai/CodonFM-80M-mRNA-stability

Encoder: 6 Transformer layers, hidden_size=1024, 8 attention heads, 4096 FFN
Position: Rotary Position Embeddings (RoPE, θ=10000)
Pretraining: Masked Language Modeling (MLM) on >130M coding sequences from NCBI RefSeq
Fine-tuning: Regression head (mean-pooled → Dense → Tanh → Dense → scalar)
Strategy: Freeze first 4/6 layers, unfreeze last 2 + regression head

pip install -r requirements.txt

from inference import CodonFMStabilityPredictor

# Load fine-tuned model
predictor = CodonFMStabilityPredictor.from_hub("Imranyai/CodonFM-80M-mRNA-stability")

# Predict stability
result = predictor.predict("AUGGCAGCCGAGACUCGGAACGUGGCCGGAGCAGAGGCCCCACCG...")
print(f"Stability score: {result['stability_score']:.4f}")
print(f"Sequence: {result['num_codons']} codons ({result['sequence_length_nt']} nt)")

sequences = [
    "AUGGCAGCCGAGACUCGG...",
    "AUGACAAUCGGUCAGACAAUG...",
    "AUGGGGUCUUCAUCAUCAUC...",
]
results = predictor.predict_batch(sequences, batch_size=32)
for r in results:
    print(f"Score: {r['stability_score']:.4f} ({r['num_codons']} codons)")

# Single sequence
python inference.py --sequence "AUGGCAGCCGAGACUCGG..."

# FASTA file
python inference.py --fasta input.fasta --output predictions.csv

# CSV file
python inference.py --csv data.csv --seq_column mRNA_seq --output results.csv

# Extract embeddings
python inference.py --fasta input.fasta --mode embeddings --output embeddings.npy

# Zero-shot MLM stability proxy (base model)
python inference.py --sequence "AUGGCAGCC..." --mode base_mlm

embeddings = predictor.get_embeddings(sequences, batch_size=32)
# Shape: [N, 1024] — use for clustering, classification, etc.

# Benchmark with base CodonFM model
python benchmark.py --mode base

# Benchmark with fine-tuned model
python benchmark.py --mode finetuned --model_repo Imranyai/CodonFM-80M-mRNA-stability

# Specific tasks only
python benchmark.py --mode base --tasks stability mrfp vaccine

# With GPU
python benchmark.py --mode base --device cuda --batch_size 64

# Download and audit all datasets
python data_setup.py --all

# Just download training data
python data_setup.py --training

# Preprocess, deduplicate, and export clean CSVs
python data_setup.py --preprocess --export ./processed_data

# Show the 69-token codon vocabulary
python data_setup.py --vocab

# Install dependencies
pip install -r requirements.txt

# Run training (GPU recommended)
python train_codonfm_stability.py

# Environment variables for customization:
LEARNING_RATE=5e-5 \
NUM_EPOCHS=20 \
BATCH_SIZE=16 \
FREEZE_LAYERS=4 \
MAX_LENGTH=1024 \
HUB_MODEL_ID=your-name/your-model \
python train_codonfm_stability.py

├── README.md                         # This file
├── DATASETS.md                       # Comprehensive dataset documentation
├── requirements.txt                  # Python dependencies
├── data_setup.py                     # Dataset download, preprocessing & audit
├── train_codonfm_stability.py        # Fine-tuning script
├── inference.py                      # Inference API + CLI
├── benchmark.py                      # Benchmark suite (5 tasks)
├── config.json                       # Model configuration (after training)
├── codon_vocab.json                  # Codon tokenizer vocabulary (after training)
└── pytorch_model.bin                 # Fine-tuned model weights (after training)

@article{diez2022icodon,
  title={iCodon customizes gene expression based on the codon composition},
  author={Diez, Michay and others},
  journal={Scientific Reports},
  volume={12},
  pages={12126},
  year={2022}
}

@article{li2024codonbert,
  title={CodonBERT large language model for mRNA vaccines},
  author={Li, Sizhen and others},
  journal={Genome Research},
  volume={34},
  number={7},
  pages={1027--1035},
  year={2024}
}