prism-antibody 1.2.0
PRISM - Partitioning Residue Identity in Somatic Maturation for antibody language modeling
License: MIT
Requires Python: >=3.8

# PRISM

**P**artitioning **R**esidue **I**dentity in **S**omatic **M**aturation

This is the official repository for the paper:

> **Explicit representation of germline and non-germline residues improves antibody language modeling**

PRISM is a PyTorch Lightning-based framework for supervised fine-tuning of ESM2 protein language models on antibody sequences. It features a multi-head architecture that jointly learns amino acid identity prediction and germline/non-germline (GL/NGL) position classification.

## Installation

### Option 1: pip install (Recommended)

```bash
# Clone the repository
git clone https://github.com/your-org/prism.git
cd prism

# Create a virtual environment (recommended)
python -m venv venv
source venv/bin/activate  # On Windows: venv\Scripts\activate

# Install in editable mode
pip install -e .

# For development tools (testing, linting)
pip install -e ".[dev]"

# For analysis scripts (plotting, ML)
pip install -e ".[analysis]"

# For everything
pip install -e ".[dev,analysis]"
```

### Option 2: Using requirements.txt

```bash
# Install from requirements file
pip install -r requirements.txt

# Then install the package
pip install -e .
```

### Verify Installation

```python
from prism import SFT_ESM2, GeneVocabulary
print("PRISM installed successfully!")
```

## Project Structure

```
prism/
├── src/prism/                    # Core Python package
│   ├── __init__.py               # Package exports
│   ├── model.py                  # SFT_ESM2 PyTorch Lightning module
│   ├── io_utils.py               # Dataset & DataModule classes
│   ├── multimodal_io.py          # Gene vocabulary & antibody dataset
│   └── utils.py                  # Utility functions
│
├── script/                       # Executable scripts
│   ├── train_esm.py              # Main training script
│   ├── inference_esm.py          # Basic inference
│   ├── inference_esm_with_logprobs.py  # Inference with per-position log probabilities
│   │
│   ├── configs/                  # Training configuration files
│   │   ├── config_esm2_v34_unpaired_balanced.yaml           # Stage 1: Pretraining
│   │   ├── config_esm2_v34.1b_paired_cdr_ngl_focus.yaml     # Stage 2: Finetuning
│   │   └── config_esm2_v34.1b_ablation*.yaml                # Ablation study configs
│   │
│   ├── data/                     # Data processing pipeline (1-8)
│   │   ├── 1.processing_and_filtering_*.py   # Filter OAS, identify NGL
│   │   ├── 2.visualize_unpaired_statistics.py
│   │   ├── 3.filter_by_p90_and_save.py
│   │   ├── 4-6.cluster_*.py                  # MMseqs2 clustering
│   │   ├── 7.extract_gene_information.py
│   │   └── 8.extract_data_for_probe.py
│   │
│   └── analyze/                  # Analysis & evaluation scripts
│       ├── 1.pppl_calculation/   # Perplexity calculations
│       │   └── plot_ppl_comparison.py
│       │
│       ├── 2.gl-ngl_calculation/ # GL/NGL embeddings & linear probes
│       │   ├── extract_residue_embeddings.py
│       │   ├── train_probe.py
│       │   ├── inference_linear_probes.py
│       │   └── visualize_residue_umap_centered.py
│       │
│       ├── 3.zero-shot/          # Zero-shot prediction benchmarks
│       │   ├── benchmark_evo_ab_fixed.py        # PRISM affinity scoring
│       │   ├── benchmark_developability_evo_ab.py  # PRISM PPL scoring
│       │   ├── benchmark_baselines_fixed.py     # ESM2/AbLang2 baselines
│       │   ├── plot_correlation_comparison.py   # Binding affinity plots
│       │   ├── plot_developability_correlations.py
│       │   └── plot_immunogenicity_correlation.py
│       │
│       └── 4.thera-sabdab/       # Controllable generation experiments
│           ├── extract_random_masking_prism.py
│           ├── extract_random_masking_baselines.py
│           └── visualize_generation_results.py
│
├── setup.py                      # Package setup (pip install -e .)
├── pyproject.toml                # Modern Python packaging config
├── requirements.txt              # Explicit dependencies
└── README.md                     # This file
```

## Core Module Overview

### `src/prism/model.py` - SFT_ESM2

The main PyTorch Lightning module for ESM2 fine-tuning:

```python
from prism import SFT_ESM2

model = SFT_ESM2(
    model_identifier="esm2_t12_35M_UR50D",
    use_multihead_architecture=True,
    use_germline_genes=True,
    use_region_embedding=True,
    ngl_loss_alpha=1.0,  # Weight for NGL positions
)
```

### `src/prism/io_utils.py` - Data Loading

Dataset and DataModule classes for training:

```python
from prism import SFTDataModule, LazyShardedDataModule

# Standard dataset
data_module = SFTDataModule(
    data_path="path/to/data.pkl",
    tokenizer=tokenizer,
    batch_size=32,
)

# Memory-efficient sharded loading
data_module = LazyShardedDataModule(
    shard_dir="path/to/shards/",
    tokenizer=tokenizer,
    batch_size=256,
)
```

### `src/prism/multimodal_io.py` - Gene & Region Handling

```python
from prism import GeneVocabulary, AntibodyDataset

# Build gene vocabulary from data
gene_vocab = GeneVocabulary.from_dataframe(
    df,
    v_gene_col='v_gene_heavy',
    j_gene_col='j_gene_heavy'
)

# Encode genes
gene_ids = gene_vocab.encode(['IGHV1-2*01', 'IGHJ4*02'])
```

## Training

### Quick Start

```bash
# Navigate to script directory
cd script

# Run training with config file
python train_esm.py --config configs/config_esm2_v34_unpaired_balanced.yaml
```

### Training Stages

**Stage 1: Pretraining** (on large unpaired antibody dataset)
```bash
python train_esm.py --config configs/config_esm2_v34_unpaired_balanced.yaml
```

**Stage 2: Finetuning** (on paired antibody sequences)
```bash
# Edit config to point to pretrained checkpoint:
# pretrained_checkpoint_path: "outputs/.../best.ckpt"

python train_esm.py --config configs/config_esm2_v34.1b_paired_cdr_ngl_focus.yaml
```

### Ablation Study Configs

The repository includes configs for a 2×2 factorial ablation study:

|                     | Multihead (Full) | Simple LM Head |
|---------------------|------------------|----------------|
| **With Pretraining**| v34.1b (Best)    | Ablation 2     |
| **No Pretraining**  | Ablation 1       | Ablation 3     |

**Ablation 1** (`config_esm2_v34.1b_ablation1_no_pretrain.yaml`):
- Full multihead architecture but NO pretraining on 66M sequences
- Tests the contribution of pretraining

**Ablation 2** (two-stage):
- Stage 2a (`config_esm2_v34_ablation2_simple_unpaired.yaml`): Pretrain simple LM head
- Stage 2b (`config_esm2_v34.1b_ablation2_simple_paired.yaml`): Finetune simple LM head
- Tests the contribution of multihead architecture (with pretraining)

**Ablation 3** (`config_esm2_v34.1b_ablation3_simple_no_pretrain.yaml`):
- Simple LM head + NO pretraining
- Worst-case baseline (neither component)

### Configuration Options

Key parameters in config YAML:

```yaml
# Data
data:
  data_path: "path/to/data"
  batch_size: 256
  mask_prob: 0.15

# Model
model:
  model_identifier: "esm2_t12_35M_UR50D"
  use_multihead_architecture: true
  ngl_loss_alpha: 1.0      # NGL loss weight (1.0 = balanced)
  aa_loss_ngl_weight: 1.0  # Include NGL in AA loss (1.0 = yes)

# Training
training:
  max_epochs: 100
  learning_rate: 1e-4
  gradient_accumulation_steps: 8
```

### Multi-GPU Training

Training automatically uses PyTorch Lightning's DDP strategy:

```bash
# Uses all available GPUs
python train_esm.py --config configs/config.yaml

# Specify GPUs
CUDA_VISIBLE_DEVICES=0,1,2,3 python train_esm.py --config configs/config.yaml
```

## Inference

### Basic Inference

```bash
python inference_esm.py \
    --config configs/config.yaml \
    --checkpoint path/to/checkpoint.ckpt \
    --data_path path/to/test_data.pkl
```

### Inference with Log Probabilities

For detailed per-position analysis:

```bash
python inference_esm_with_logprobs.py \
    --checkpoint_path path/to/checkpoint.ckpt \
    --config_path configs/config.yaml \
    --input_csv data/sequences.csv \
    --output_csv results/predictions.csv
```

### Programmatic Inference

```python
from prism import SFT_ESM2
from transformers import AutoTokenizer
import torch

# Load model
model = SFT_ESM2.load_from_checkpoint("path/to/checkpoint.ckpt")
model.eval()

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t12_35M_UR50D")

# Prepare input (heavy + light concatenated)
sequence = "EVQLVESGGGLVQPGGSLRLSCAASGFTFS<cls>DIQMTQSPSSLSASVGDRVTITC"
inputs = tokenizer(sequence, return_tensors="pt", padding=True)

# Run inference
with torch.no_grad():
    outputs = model(inputs["input_ids"], inputs["attention_mask"])

# outputs contains:
# - aa_logits: [batch, seq_len, vocab_size] - AA predictions
# - origin_logits: [batch, seq_len, 2] - GL/NGL predictions
```

## Analysis Scripts

### 1. Perplexity Calculation (`script/analyze/1.pppl_calculation/`)

Evaluate model perplexity on test sequences:

```bash
cd script/analyze/1.pppl_calculation
python plot_ppl_comparison.py --data_path path/to/ppl_data.pkl --output_dir ./output
```

### 2. GL/NGL Analysis (`script/analyze/2.gl-ngl_calculation/`)

Analyze germline vs non-germline position predictions:

```bash
cd script/analyze/2.gl-ngl_calculation

# Extract embeddings from models
python extract_residue_embeddings.py --checkpoint path/to/model.ckpt

# Train linear probes for GL/NGL classification
python train_probe.py --data_dir data/linear_probe_data --embedding_dim 480

# Run inference with trained probes
python inference_linear_probes.py --probe_dir runs/linear_probe

# Visualize embeddings with UMAP
python visualize_residue_umap_centered.py --data_path data/embeddings.pkl
```

### 3. Zero-shot Benchmarks (`script/analyze/3.zero-shot/`)

Evaluate zero-shot prediction capabilities:

```bash
cd script/analyze/3.zero-shot

# PRISM scoring for binding affinity
python benchmark_evo_ab_fixed.py --mode affinity \
    --config configs/config.yaml --checkpoint model.ckpt \
    --data_path data/binding_data.csv

# PRISM PPL scoring for developability
python benchmark_developability_evo_ab.py \
    --config configs/config.yaml --checkpoint model.ckpt \
    --data_path data/developability_data.csv

# Plot correlation results
python plot_correlation_comparison.py --csv_path data/results.csv
python plot_developability_correlations.py --csv_path data/developability_with_ppl.csv
python plot_immunogenicity_correlation.py --csv_path data/immunogenicity.csv
```

### 4. Controllable Generation (`script/analyze/4.thera-sabdab/`)

Experiments for controllable antibody sequence recovery:

```bash
cd script/analyze/4.thera-sabdab

# PRISM predictions at random masked positions
python extract_random_masking_prism.py \
    --config configs/config.yaml --checkpoint model.ckpt \
    --data_path data/therasabdab_germline.csv

# Baseline model predictions for comparison
python extract_random_masking_baselines.py \
    --data_path data/therasabdab_germline.csv

# Visualize results
python visualize_generation_results.py --prism_csv results/prism.csv
```

## Data Processing Pipeline

The `script/data/` folder contains a numbered pipeline for processing OAS (Observed Antibody Space) data:

| Script | Purpose |
|--------|---------|
| `1.processing_and_filtering_*` | Filter OAS data, identify NGL positions |
| `2.visualize_unpaired_statistics.py` | Visualize data distribution |
| `3.filter_by_p90_and_save.py` | Quality filtering by sequence coverage |
| `4.cluster_paired_and_unpaired_*` | MMseqs2-based sequence clustering |
| `5.cluster_using_fasta.py` | Alternative FASTA-based clustering |
| `6.deduplicate_clusters.py` | Remove redundant clusters |
| `7.extract_gene_information.py` | Parse V/J gene annotations |
| `8.extract_data_for_probe.py` | Prepare data for linear probe training |

## License

MIT License - see [LICENSE](LICENSE) for details.