Add desolubilization, multipass, and olig tasks

.gitignore

@@ -13,12 +13,12 @@
 /src/lm/evoflow_playground.ipynb
 /src/utils/ubuntu_font
 
-/src/sampling/old_guidance.py
-
 /MeMDLM_v2.egg-info/
 *.pth
 *.ckpt
 *.err
 *.out
 *.csv
+*.ipynb
+*.png
 __pycache__/

README.md

@@ -2,25 +2,373 @@
 
 ![MemDLM diagram](./memdlm_schematic.png)
 
+**arXiv preprint:** [link TBD]
 
-arXiv preprint: ...
+Reparameterized diffusion models (RDMs) have recently matched autoregressive methods in protein generation, motivating their use for challenging tasks such as designing membrane proteins, which possess interleaved soluble and transmembrane (TM) regions.
 
-Reparameterized diffusion models (RDMs) have recently matched autoregressive methods in protein generation, motivating their use for challenging tasks such as designing membrane proteins, which possess interleaved soluble and transmembrane (TM) regions.  
+We introduce ***Membrane Diffusion Language Model (MemDLM)***, a fine-tuned RDM-based protein language model that enables controllable membrane protein sequence design. MemDLM-generated sequences recapitulate the TM residue density and structural features of natural membrane proteins, achieving comparable biological plausibility and outperforming state-of-the-art diffusion baselines in motif scaffolding tasks by producing:
 
-We introduce ***Membrane Diffusion Language Model (MemDLM)***, a fine-tuned RDM-based protein language model that enables controllable membrane protein sequence design. MemDLM-generated sequences recapitulate the TM residue density and structural features of natural membrane proteins, achieving comparable biological plausibility and outperforming state-of-the-art diffusion baselines in motif scaffolding tasks by producing:  
+- Lower perplexity
+- Higher BLOSUM-62 scores
+- Improved pLDDT confidence
 
-- Lower perplexity  
-- Higher BLOSUM-62 scores  
-- Improved pLDDT confidence  
+To enhance controllability, we develop ***Per-Token Guidance (PET)***, a novel classifier-guided sampling strategy that selectively solubilizes residues while preserving conserved TM domains. This yields sequences with reduced TM density but intact functional cores.
 
-To enhance controllability, we develop ***Per-Token Guidance (PET)***, a novel classifier-guided sampling strategy that selectively solubilizes residues while preserving conserved TM domains. This yields sequences with reduced TM density but intact functional cores.  
+Importantly, MemDLM designs validated in TOXCAT β-lactamase growth assays demonstrate successful TM insertion, distinguishing high-quality generated sequences from poor ones.
 
-Importantly, MemDLM designs validated in TOXCAT β-lactamase growth assays demonstrate successful TM insertion, distinguishing high-quality generated sequences from poor ones.  
+Together, our framework establishes the first experimentally validated diffusion-based model for rational membrane protein generation, integrating *de novo* design, motif scaffolding, and targeted property optimization.
 
-Together, our framework establishes the first experimentally validated diffusion-based model for rational membrane protein generation, integrating *de novo* design, motif scaffolding, and targeted property optimization.  
+---
 
+## Repository Authors
 
+- <u>[Shrey Goel](https://shreygoel09.github.io/)</u> – undergraduate student at Duke University
+- <u>[Pranam Chatterjee](mailto:pranam@seas.upenn.edu)</u> – Assistant Professor at University of Pennsylvania
 
-## **Repository Authors**
-- <u>[Shrey Goel](https://shreygoel09.github.io/)</u> – undergraduate student at Duke University  
-- <u>[Pranam Chatterjee](mailto:pranam@seas.upenn.edu)</u> – Assistant Professor at University of Pennsylvania  
+---
+
+## Table of Contents
+
+1. [Installation](#installation)
+2. [Configuration](#configuration)
+3. [Data Preparation](#data-preparation)
+4. [Training](#training)
+   - [MeMDLM diffusion model](#1-memdlm-diffusion-model)
+   - [Solubility classifier](#2-solubility-classifier)
+   - [Oligomerization classifier](#3-oligomerization-classifier)
+   - [Multipass classifier](#4-multipass-classifier)
+5. [Sampling](#sampling)
+   - [Unconditional generation](#1-unconditional-generation)
+   - [Solubilization (PET)](#2-solubilization-pet)
+   - [Desolubilization (PET)](#3-desolubilization-pet)
+   - [Multipass generation](#4-multipass-generation)
+   - [Oligomerization](#5-oligomerization)
+6. [Citation](#citation)
+
+---
+
+## Installation
+
+Clone the repository and install it in editable mode from the repository root:
+
+```bash
+git clone https://github.com/<your-org>/MeMDLM_v2.git
+cd MeMDLM_v2
+pip install -e .
+pip install -r requirements.txt
+```
+
+Log in to Weights & Biases (used for experiment logging during training):
+
+```bash
+wandb login
+# or: export WANDB_API_KEY=<your-key>
+```
+
+**Requirements:** CUDA-capable GPU(s), Python 3.10+, and sufficient disk space for pretrained checkpoints (EvoFlow-650M, ESM2-650M) and saved model weights.
+
+---
+
+## Configuration
+
+All hyperparameters live in YAML files under `src/configs/`. Before running anything, set `base_dir` in each config to your local clone path:
+
+```yaml
+base_dir: /path/to/MeMDLM_v2
+```
+
+| Config file | Purpose |
+|---|---|
+| `src/configs/lm.yaml` | Fine-tune the MeMDLM diffusion language model |
+| `src/configs/solubility.yaml` | Train/evaluate the per-residue solubility classifier (used for PET solubilization) |
+| `src/configs/desolubilize.yaml` | PET desolubilization sampling settings (shares the solubility classifier checkpoint) |
+| `src/configs/oligo.yaml` | Train/evaluate the oligomerization classifier |
+| `src/configs/multipass.yaml` | Train/evaluate the multipass TM-segment classifier |
+
+**Training modes** are controlled by `training.mode` in each config:
+
+| Value | Behavior |
+|---|---|
+| `train` | Run training from scratch |
+| `test` | Load `best_model.ckpt` and evaluate on the test set |
+| `resume_from_checkpoint` | Resume MeMDLM training from `checkpointing.resume_ckpt_path` (`lm.yaml` only) |
+
+Checkpoints are saved to `checkpoints/<wandb.name>/best_model.ckpt`. After training the diffusion model, ensure `lm.ft_evoflow` in the classifier and sampling configs matches the `wandb.name` from `lm.yaml`.
+
+---
+
+## Data Preparation
+
+Place your datasets under `<base_dir>/data/`. Expected CSV formats:
+
+### MeMDLM & solubility classifier (`data/train.csv`, `data/test.csv`, `data/val.csv`)
+
+| Column | Description |
+|---|---|
+| `Sequence` | Amino acid sequence. For the solubility classifier, use **uppercase** for soluble residues and **lowercase** for TM residues. For MeMDLM pretraining, sequences are uppercased automatically. |
+
+### Oligomerization classifier (`data/olig_clf/`)
+
+| Column | Description |
+|---|---|
+| `Sequence` | Input sequence |
+| `Binary Label` | `0` or `1` oligomerization label |
+
+### Multipass classifier (`data/multipass/`)
+
+| Column | Description |
+|---|---|
+| `Sequence` | Input sequence |
+| `TM_segments` | Integer count of predicted TM segments |
+
+---
+
+## Training
+
+All training jobs are launched from the **repository root**. We recommend running long jobs with `nohup` so they persist after logout. Create a `logs/` directory first:
+
+```bash
+mkdir -p logs
+```
+
+### 1. MeMDLM diffusion model
+
+**Config:** `src/configs/lm.yaml`  
+**Entry point:** `src/lm/memdlm/main.py`
+
+Set `training.mode: train` and adjust `training.devices` to the number of GPUs available.
+
+```bash
+nohup python src/lm/memdlm/main.py > logs/memdlm_train.out 2>&1 &
+```
+
+Key settings in `lm.yaml` (do not change unless reproducing a new experiment):
+
+| Parameter | Default |
+|---|---|
+| Base model | `fredzzp/EvoFlow-650M-context-3070` |
+| `training.max_steps` | 5000 |
+| `training.n_layers` | 3 (unfrozen transformer layers) |
+| `optim.lr` | 4e-5 |
+| `data.batch_size` | 32 |
+| `lm.num_diffusion_timesteps` | 500 |
+
+Checkpoint saved to: `checkpoints/<wandb.name>/best_model.ckpt`
+
+To evaluate a trained checkpoint, set `training.mode: test`.
+
+```bash
+nohup python src/lm/memdlm/main.py > logs/memdlm_test.out 2>&1 &
+```
+
+---
+
+### 2. Solubility classifier
+
+**Config:** `src/configs/solubility.yaml`  
+**Entry point:** `src/guidance/solubility/main.py`
+
+Trains a per-residue ESM-based classifier that predicts soluble vs. TM residues. This checkpoint is required for PET solubilization and desolubilization sampling.
+
+Set `training.mode: train`.
+
+```bash
+nohup python src/guidance/solubility/main.py > logs/solubility_train.out 2>&1 &
+```
+
+Key settings:
+
+| Parameter | Default |
+|---|---|
+| `training.max_steps` | 3000 |
+| `model.num_layers` | 4 |
+| `optim.lr` | 3e-5 |
+| `data.batch_size` | 32 |
+
+Checkpoint saved to: `checkpoints/<wandb.name>/best_model.ckpt`
+
+---
+
+### 3. Oligomerization classifier
+
+**Config:** `src/configs/oligo.yaml`  
+**Entry point:** `src/guidance/oligo/main.py`
+
+Set `training.mode: train`.
+
+```bash
+nohup python src/guidance/oligo/main.py > logs/oligo_train.out 2>&1 &
+```
+
+Key settings:
+
+| Parameter | Default |
+|---|---|
+| `training.max_steps` | 3000 |
+| `model.num_layers` | 1 |
+| `optim.lr` | 1e-4 |
+| `data.batch_size` | 64 |
+| `data.max_seq_len` | 54 |
+
+Checkpoint saved to: `checkpoints/<wandb.name>/best_model.ckpt`
+
+---
+
+### 4. Multipass classifier
+
+**Config:** `src/configs/multipass.yaml`  
+**Entry point:** `src/guidance/multipass/main.py`
+
+Predicts the number of TM segments in a sequence. Required for multipass-guided generation.
+
+Set `training.mode: train`.
+
+```bash
+nohup python src/guidance/multipass/main.py > logs/multipass_train.out 2>&1 &
+```
+
+Key settings:
+
+| Parameter | Default |
+|---|---|
+| `training.max_steps` | 3000 |
+| `model.num_layers` | 4 |
+| `optim.lr` | 3e-5 |
+| `data.batch_size` | 32 |
+
+Checkpoint saved to: `checkpoints/<wandb.name>/best_model.ckpt`
+
+---
+
+### Recommended training order
+
+```
+1. MeMDLM (lm.yaml)
+       ↓
+2. Classifiers in any order:
+   • solubility.yaml
+   • oligo.yaml
+   • multipass.yaml
+```
+
+Update `lm.ft_evoflow` in all downstream configs to match the `wandb.name` produced by step 1.
+
+---
+
+## Sampling
+
+Sampling scripts generate sequences and write CSVs under `<base_dir>/results/`. Each script loads the fine-tuned MeMDLM checkpoint (`checkpoints/<lm.ft_evoflow>/best_model.ckpt`) and, where applicable, the corresponding classifier checkpoint (`checkpoints/<wandb.name>/best_model.ckpt`).
+
+Run all sampling jobs from the **repository root**:
+
+```bash
+mkdir -p logs
+```
+
+### 1. Unconditional generation
+
+**Config:** `src/configs/lm.yaml`  
+**Script:** `src/sampling/unconditional_generator.py`
+
+Generates *de novo* membrane protein sequences from a fully masked prior.
+
+```bash
+nohup python src/sampling/unconditional_generator.py > logs/unconditional_sample.out 2>&1 &
+```
+
+**Output:** `results/denovo/<wandb.name>/<date>_multipass/seqs_with_ppl.csv`
+
+Columns: `Generated Sequence`, `ESM PPL`, `MeMDLM PPL`
+
+---
+
+### 2. Solubilization (PET)
+
+**Config:** `src/configs/solubility.yaml`  
+**Script:** `src/sampling/pet_generator.py`
+
+Uses Per-Token Guidance to redesign **uppercase (soluble)** positions in a scaffold while preserving **lowercase (TM)** residues. Update the input CSV path in `pet_generator.py` (default: `results/heme/cybtx.csv`) to point to your scaffold sequences.
+
+```bash
+nohup python src/sampling/pet_generator.py > logs/solubilize_sample.out 2>&1 &
+```
+
+**Output:** `results/heme/<lm.ft_evoflow>/solubilize/<date>/<prior-params>/infilled_seqs.csv`
+
+---
+
+### 3. Desolubilization (PET)
+
+**Config:** `src/configs/desolubilize.yaml`  
+**Script:** `src/sampling/desolubilize_generator.py`
+
+The inverse of solubilization: redesigns **lowercase (soluble)** positions while preserving **uppercase (TM)** residues. Uses the same solubility classifier checkpoint. Update the input CSV path in `desolubilize_generator.py` (default: `results/heme/4d2.csv`).
+
+```bash
+nohup python src/sampling/desolubilize_generator.py > logs/desolubilize_sample.out 2>&1 &
+```
+
+**Output:** `results/heme/<lm.ft_evoflow>/desolubilize/<date>/<prior-params>/infilled_seqs.csv`
+
+---
+
+### 4. Multipass generation
+
+**Config:** `src/configs/multipass.yaml`  
+**Script:** `src/sampling/multipass_generator.py`
+
+Generates sequences guided toward a target multipass TM topology using the multipass classifier.
+
+```bash
+nohup python src/sampling/multipass_generator.py > logs/multipass_sample.out 2>&1 &
+```
+
+**Output:** `results/multipass/<wandb.name>/<date>/lamb=<reg_strength>_tau=<sampling_temperature>/seqs_with_ppl.csv`
+
+Columns: `Generated Sequence`, `ESM PPL`, `MeMDLM PPL`, `Pred TM Segments`
+
+---
+
+### 5. Oligomerization
+
+**Config:** `src/configs/oligo.yaml`  
+**Script:** `src/sampling/olig_generator.py`
+
+Redesigns sequences to increase predicted oligomerization propensity using the oligomerization classifier (NOS guidance).
+
+```bash
+nohup python src/sampling/olig_generator.py > logs/oligo_sample.out 2>&1 &
+```
+
+**Output:** `results/oligo/<wandb.name>/<date>/seqs_with_ppl.csv`
+
+Columns: `Original Sequence`, `Generated Sequence`, `OG Olig Value`, `New Olig Value`, `Olig Increase`, `ESM PPL`, `MeMDLM PPL`, `MemDLM Blosum`
+
+---
+
+### Sampling quick-reference
+
+| Task | Config | Script | Classifier required |
+|---|---|---|---|
+| Unconditional | `lm.yaml` | `unconditional_generator.py` | No |
+| Solubilization | `solubility.yaml` | `pet_generator.py` | Solubility |
+| Desolubilization | `desolubilize.yaml` | `desolubilize_generator.py` | Solubility |
+| Multipass | `multipass.yaml` | `multipass_generator.py` | Multipass |
+| Oligomerization | `oligo.yaml` | `olig_generator.py` | Oligomerization |
+
+---
+
+## Citation
+
+If you use this repository in your research, please cite:
+
+```bibtex
+@article{goel2026memdlm,
+  title   = {Token-Level Guided Discrete Diffusion for Membrane Protein Design},
+  author  = {Goel, Shrey and Chatterjee, Pranam},
+  journal = {arXiv preprint arXiv:XXXX.XXXXX},
+  year    = {2026}
+}
+```
+
+If you enjoyed this repo, please cite it and star the repository. We appreciate your support!

requirements.txt

@@ -0,0 +1,176 @@
+aiohappyeyeballs==2.6.1
+aiohttp==3.12.15
+aiosignal==1.4.0
+annotated-types==0.7.0
+antlr4-python3-runtime==4.9.3
+anyio==4.10.0
+argon2-cffi==25.1.0
+argon2-cffi-bindings==25.1.0
+arrow==1.3.0
+asttokens==3.0.0
+async-lru==2.0.5
+attrs==25.3.0
+babel==2.17.0
+beautifulsoup4==4.13.5
+bio==1.8.0
+biopython==1.85
+biothings_client==0.4.1
+bleach==6.2.0
+blosum==2.2.0
+certifi==2025.8.3
+cffi==2.0.0
+charset-normalizer==3.4.3
+click==8.3.0
+comm==0.2.3
+contourpy==1.3.3
+cycler==0.12.1
+datasets==4.4.2
+debugpy==1.8.16
+decorator==5.2.1
+defusedxml==0.7.1
+dill==0.4.0
+executing==2.2.1
+fastjsonschema==2.21.2
+filelock==3.19.1
+fonttools==4.61.1
+fqdn==1.5.1
+frozenlist==1.7.0
+fsspec==2025.9.0
+gitdb==4.0.12
+GitPython==3.1.45
+gprofiler-official==1.0.0
+h11==0.16.0
+hf-xet==1.1.10
+httpcore==1.0.9
+httpx==0.28.1
+huggingface-hub==0.34.4
+idna==3.10
+ipykernel==6.30.1
+ipython==9.5.0
+ipython_pygments_lexers==1.1.1
+isoduration==20.11.0
+jedi==0.19.2
+Jinja2==3.1.6
+joblib==1.5.2
+json5==0.12.1
+jsonpointer==3.0.0
+jsonschema==4.25.1
+jsonschema-specifications==2025.9.1
+jupyter_client==8.6.3
+jupyter_core==5.8.1
+jupyter-events==0.12.0
+jupyter-lsp==2.3.0
+jupyter_server==2.17.0
+jupyter_server_terminals==0.5.3
+jupyterlab==4.4.7
+jupyterlab_pygments==0.3.0
+jupyterlab_server==2.27.3
+kiwisolver==1.4.9
+lark==1.2.2
+lightning==2.5.5
+lightning-utilities==0.15.2
+MarkupSafe==3.0.2
+matplotlib==3.10.8
+matplotlib-inline==0.1.7
+MeMDLM_v2==1.0
+mistune==3.1.4
+mpmath==1.3.0
+multidict==6.6.4
+multiprocess==0.70.18
+mygene==3.2.2
+nbclient==0.10.2
+nbconvert==7.16.6
+nbformat==5.10.4
+nest-asyncio==1.6.0
+networkx==3.5
+notebook==7.4.5
+notebook_shim==0.2.4
+numpy==2.3.3
+nvidia-cublas-cu12==12.8.4.1
+nvidia-cuda-cupti-cu12==12.8.90
+nvidia-cuda-nvrtc-cu12==12.8.93
+nvidia-cuda-runtime-cu12==12.8.90
+nvidia-cudnn-cu12==9.10.2.21
+nvidia-cufft-cu12==11.3.3.83
+nvidia-cufile-cu12==1.13.1.3
+nvidia-curand-cu12==10.3.9.90
+nvidia-cusolver-cu12==11.7.3.90
+nvidia-cusparse-cu12==12.5.8.93
+nvidia-cusparselt-cu12==0.7.1
+nvidia-nccl-cu12==2.27.3
+nvidia-nvjitlink-cu12==12.8.93
+nvidia-nvtx-cu12==12.8.90
+omegaconf==2.3.0
+packaging==25.0
+pandas==2.3.2
+pandocfilters==1.5.1
+parso==0.8.5
+pexpect==4.9.0
+pillow==12.1.0
+pip==25.2
+platformdirs==4.4.0
+pooch==1.8.2
+prometheus_client==0.22.1
+prompt_toolkit==3.0.52
+propcache==0.3.2
+protobuf==6.32.1
+psutil==7.0.0
+ptyprocess==0.7.0
+pure_eval==0.2.3
+pyarrow==22.0.0
+pycparser==2.23
+pydantic==2.11.9
+pydantic_core==2.33.2
+Pygments==2.19.2
+pyparsing==3.3.2
+python-dateutil==2.9.0.post0
+python-json-logger==3.3.0
+pytorch-lightning==2.5.5
+pytz==2025.2
+PyYAML==6.0.2
+pyzmq==27.1.0
+referencing==0.36.2
+regex==2025.9.1
+requests==2.32.5
+rfc3339-validator==0.1.4
+rfc3986-validator==0.1.1
+rfc3987-syntax==1.1.0
+rpds-py==0.27.1
+safetensors==0.6.2
+scikit-learn==1.7.2
+scipy==1.16.2
+seaborn==0.13.2
+Send2Trash==1.8.3
+sentry-sdk==2.38.0
+setuptools==78.1.1
+six==1.17.0
+smmap==5.0.2
+sniffio==1.3.1
+soupsieve==2.8
+stack-data==0.6.3
+sympy==1.14.0
+terminado==0.18.1
+threadpoolctl==3.6.0
+tinycss2==1.4.0
+tokenizers==0.20.3
+torch==2.8.0
+torchmetrics==1.8.2
+tornado==6.5.2
+tqdm==4.67.1
+traitlets==5.14.3
+transformers==4.46.0
+triton==3.4.0
+types-python-dateutil==2.9.0.20250822
+typing_extensions==4.15.0
+typing-inspection==0.4.1
+tzdata==2025.2
+uri-template==1.3.0
+urllib3==2.5.0
+wandb==0.22.0
+wcwidth==0.2.13
+webcolors==24.11.1
+webencodings==0.5.1
+websocket-client==1.8.0
+wheel==0.45.1
+xxhash==3.6.0
+yarl==1.20.1

src/configs/{guidance.yaml → desolubilize.yaml}

@@ -1,14 +1,14 @@
 
 
 seed: 42
-base_dir: /scratch/sgoel/MeMDLM_v2
+base_dir: /scratch/pranamlab/sgoel/MeMDLM_v2
 
 
 lm:
   pretrained_esm: facebook/esm2_t33_650M_UR50D
   pretrained_evoflow: fredzzp/EvoFlow-650M-context-3070
   pretrained_dplm: airkingbd/dplm_650m
-  ft_evoflow: ft_eflow-3070-650M_steps=50k_layers=3_lr=0.00004_wd=.01_polynom_pwr=1_betas=.9-.98_bsz=8_gclip=1.0
+  ft_evoflow: ft_eflow-3070-650M_steps=5k_layers=3_lr=1e-5_lr-end=2e-6_wd=.01_polynom_pwr=1_betas=.9-.98_bsz=8_gclip=1.0_ml=1024
   ft_dplm: ft_dplm-650M_steps=5k_layers=3_lr=0.00004_wd=.01_polynom_pwr=1_betas=.9-.98_bsz=32_gclip=1.0
 
 model:
@@ -38,32 +38,33 @@ training:
   val_check_interval: 250
   enable_progress_bar: true
   grad_clip_val: 1.0
-  devices: [0]  # list of GPU IDs from 0-7
+  devices: 2
 
 guidance:
+  task: desolubilize
   n_steps: 128
-  alpha: 3
+  alpha: 5
   gamma: 0.3
   saliency_eps: 1e-4
   saliency_t: 2.0
-  sampling_t: 0.7
-  boltzmann_t: 0.3
-  top_p: 0.2
+  sampling_t: 0.8
+  boltzmann_t: 0.2
+  top_p: 0.3
   steps: 128
-  prior: lm_probs  # lm_probs / boltzmann
+  prior: boltzmann
 
 data:
   batch_size: 32
   max_seq_len: 1024
-  train: ${base_dir}/data/classifier/train.csv
-  test: ${base_dir}/data/classifier/test.csv
-  val: ${base_dir}/data/classifier/val.csv
+  train: ${base_dir}/data/train.csv
+  test: ${base_dir}/data/test.csv
+  val: ${base_dir}/data/val.csv
 
 
 wandb:
-  project: memdlm_guidance
+  project: memdlm_again
   group: programmablebio
-  name: new_data_cleaned_steps3k_lr3e-5_bsz32_heads2_drpt0.5_layers4
+  name: clf_steps3k_lr3e-5_bsz32_heads2_drpt0.5_layers4
   id: ${.name}_${seed}
 
 

src/configs/lm.yaml

@@ -17,7 +17,7 @@ optim:
   type: adamw
   scheduler: polynomial
   lr: 0.00004
-  lr_end: 1e-5
+  lr_end: 2e-6
   warmup_init_lr: 1e-07
   weight_decay: 0.01
   beta1: 0.9
@@ -35,24 +35,24 @@ training:
   val_check_interval: 250
   enable_progress_bar: true
   grad_clip_val: 1.0
-  devices: [0,1,2]  # list of GPU IDs
+  devices: 3 # number of GPUs
 
 sampling:
   n_steps: 128
 
 
 data:
-  batch_size: 8
+  batch_size: 32
   max_seq_len: 1024
-  train: ${base_dir}/data/new/train.csv
-  test: ${base_dir}/data/new/test.csv
-  val: ${base_dir}/data/new/val.csv
+  train: ${base_dir}/data/train.csv
+  test: ${base_dir}/data/test.csv
+  val: ${base_dir}/data/val.csv
 
 
 wandb:
-  project: memdlm
+  project: memdlm_again
   group: programmablebio
-  name: ft_eflow-3070-650M_steps=5k_layers=3_lr=0.00004_wd=.01_polynom_pwr=1_betas=.9-.98_bsz=8_gclip=1.0_ml=1024
+  name: ft_eflow-3070-650M_steps=5k_layers=3_lr=1e-5_lr-end=2e-6_wd=.01_polynom_pwr=1_betas=.9-.98_bsz=8_gclip=1.0_ml=1024
   # name: ft_progen-base-764M_steps=50k_layers=2_lr=0.00004_wd=.1_cosine-to-frac_betas=.9-.999_bsz=8_gclip=0.8
   # name: ft_dplm-650M_steps=5k_layers=3_lr=0.00004_wd=.01_polynom_pwr=1_betas=.9-.98_bsz=32_gclip=1.0
   # name: ft_esm-650M_steps=3k_layers=3_lr=0.00004_wd=.01_polynom_pwr=1_betas=.9-.98_bsz=32_gclip=1.0

src/configs/multipass.yaml

@@ -0,0 +1,70 @@
+
+
+seed: 42
+base_dir: /scratch/pranamlab/sgoel/MeMDLM_v2
+
+
+lm:
+  pretrained_esm: facebook/esm2_t33_650M_UR50D
+  pretrained_evoflow: fredzzp/EvoFlow-650M-context-3070
+  pretrained_dplm: airkingbd/dplm_650m
+  ft_evoflow: ft_eflow-3070-650M_steps=5k_layers=3_lr=1e-5_lr-end=2e-6_wd=.01_polynom_pwr=1_betas=.9-.98_bsz=8_gclip=1.0_ml=1024
+  ft_dplm: ft_dplm-650M_steps=5k_layers=3_lr=0.00004_wd=.01_polynom_pwr=1_betas=.9-.98_bsz=32_gclip=1.0
+  num_diffusion_timesteps: 500
+
+model:
+  d_model: 1280 
+  num_heads: 2
+  dropout: 0.5
+  num_layers: 4
+  label_pad_value: -100
+
+optim:
+  type: adamw
+  lr: 3e-5
+  lr_end: 1e-5
+  weight_decay: 0.01
+  beta1: 0.9
+  beta2: 0.98
+  power: 1
+
+
+training:
+  mode: test  # train / test
+  n_layers: 4
+  max_steps: 3000
+  warmup_steps: 150
+  log_every_n_steps: 10
+  num_sanity_val_steps: 2
+  val_check_interval: 250
+  enable_progress_bar: true
+  grad_clip_val: 1.0
+  devices: 2
+
+guidance:
+  guide_steps: 12
+  diffusion_steps: 128
+  reg_strength: 3 # 3 vs 0.1
+  step_size: 0.1
+  sampling_temperature: 1.15
+
+data:
+  batch_size: 32
+  max_seq_len: 1024
+  train: ${base_dir}/data/multipass/train_tm_segment_counts.csv
+  test: ${base_dir}/data/multipass/test_tm_segment_counts.csv
+  val: ${base_dir}/data/multipass/val_tm_segment_counts.csv
+
+
+wandb:
+  project: memdlm_again
+  group: programmablebio
+  name: multipass_clf_steps3k_lr3e-5_bsz32_heads2_drpt0.5_layers4
+  id: ${.name}_${seed}
+
+
+checkpointing:
+  save_every_n_steps: 250
+  save_dir: ${base_dir}/checkpoints/${wandb.name}
+  resume_ckpt_path: ${checkpointing.save_dir}/last.ckpt
+  best_ckpt_path: ${checkpointing.save_dir}/best_model.ckpt

src/configs/oligo.yaml

@@ -0,0 +1,86 @@
+
+
+
+seed: 42
+base_dir: /scratch/pranamlab/sgoel/MeMDLM_v2
+
+
+lm:
+  pretrained_esm: facebook/esm2_t33_650M_UR50D
+  pretrained_evoflow: fredzzp/EvoFlow-650M-context-3070
+  pretrained_dplm: airkingbd/dplm_650m
+  ft_evoflow: ft_eflow-3070-650M_steps=5k_layers=3_lr=1e-5_lr-end=2e-6_wd=.01_polynom_pwr=1_betas=.9-.98_bsz=8_gclip=1.0_ml=1024
+  ft_dplm: ft_dplm-650M_steps=5k_layers=3_lr=0.00004_wd=.01_polynom_pwr=1_betas=.9-.98_bsz=32_gclip=1.0
+
+model:
+  esm_dim: 1280
+  d_model: 128
+  num_heads: 2
+  dropout: 0.3
+  num_layers: 1
+  label_pad_value: -100
+
+optim:
+  type: adamw
+  lr: 1e-4
+  lr_end: 1e-5
+  weight_decay: 0.02
+  beta1: 0.9
+  beta2: 0.98
+  power: 1
+
+
+training:
+  mode: test  # train / test
+  n_layers: 4
+  max_steps: 3000
+  warmup_steps: 200
+  log_every_n_steps: 10
+  num_sanity_val_steps: 2
+  val_check_interval: 100
+  enable_progress_bar: true
+  grad_clip_val: 1.0
+  devices: 2
+
+guidance:
+  n_steps: 128
+  alpha: 3
+  gamma: 0.3
+  saliency_eps: 1e-4
+  saliency_t: 2.0
+  sampling_t: 0.7
+  boltzmann_t: 0.3
+  top_p: 0.2
+  steps: 128
+  prior: lm_probs  # lm_probs / boltzmann
+
+olig_guidance:
+  guide_steps: 12
+  diffusion_steps: 64
+  reg_strength: 1
+  step_size: 0.1
+  temperature: 1.0
+  eps: 1e-8
+  residue_thresh: 0.5
+  topk_frac: 0.2
+
+data:
+  batch_size: 64
+  max_seq_len: 54
+  train: ${base_dir}/data/olig_clf/train.csv
+  test: ${base_dir}/data/olig_clf/test.csv
+  val: ${base_dir}/data/olig_clf/val.csv
+
+
+wandb:
+  project: memdlm_oligo
+  group: programmablebio
+  name: oligo-clf_steps3k_lr1e-4_bsz64_drpt0.3_wd0.02_dmodel128_simple-mlp_mean-pool
+  id: ${.name}_${seed}
+
+
+checkpointing:
+  save_every_n_steps: 250
+  save_dir: ${base_dir}/checkpoints/${wandb.name}
+  resume_ckpt_path: ${checkpointing.save_dir}/last.ckpt
+  best_ckpt_path: ${checkpointing.save_dir}/best_model.ckpt

src/configs/solubility.yaml

@@ -0,0 +1,75 @@
+
+
+seed: 42
+base_dir: /scratch/pranamlab/sgoel/MeMDLM_v2
+
+
+lm:
+  pretrained_esm: facebook/esm2_t33_650M_UR50D
+  pretrained_evoflow: fredzzp/EvoFlow-650M-context-3070
+  pretrained_dplm: airkingbd/dplm_650m
+  ft_evoflow: ft_eflow-3070-650M_steps=5k_layers=3_lr=1e-5_lr-end=2e-6_wd=.01_polynom_pwr=1_betas=.9-.98_bsz=8_gclip=1.0_ml=1024
+  ft_dplm: ft_dplm-650M_steps=5k_layers=3_lr=0.00004_wd=.01_polynom_pwr=1_betas=.9-.98_bsz=32_gclip=1.0
+
+model:
+  d_model: 1280 
+  num_heads: 2
+  dropout: 0.5
+  num_layers: 4
+  label_pad_value: -100
+
+optim:
+  type: adamw
+  lr: 3e-5
+  lr_end: 1e-5
+  weight_decay: 0.01
+  beta1: 0.9
+  beta2: 0.98
+  power: 1
+
+
+training:
+  mode: test  # train / test
+  n_layers: 4
+  max_steps: 3000
+  warmup_steps: 150
+  log_every_n_steps: 10
+  num_sanity_val_steps: 2
+  val_check_interval: 250
+  enable_progress_bar: true
+  grad_clip_val: 1.0
+  devices: 2
+
+guidance:
+  task: solubilize
+  n_steps: 128
+  alpha: 5 # 3
+  gamma: 0.3
+  saliency_eps: 1e-4
+  saliency_t: 2.0
+  sampling_t: 1.0  # 0.7
+  boltzmann_t: 0.15 # 0.3
+  top_p: 0.3 # 0.2
+  steps: 128
+  prior: boltzmann
+
+data:
+  batch_size: 32
+  max_seq_len: 1024
+  train: ${base_dir}/data/train.csv
+  test: ${base_dir}/data/test.csv
+  val: ${base_dir}/data/val.csv
+
+
+wandb:
+  project: memdlm_again
+  group: programmablebio
+  name: clf_steps3k_lr3e-5_bsz32_heads2_drpt0.5_layers4
+  id: ${.name}_${seed}
+
+
+checkpointing:
+  save_every_n_steps: 250
+  save_dir: ${base_dir}/checkpoints/${wandb.name}
+  resume_ckpt_path: ${checkpointing.save_dir}/last.ckpt
+  best_ckpt_path: ${checkpointing.save_dir}/best_model.ckpt

src/guidance/multipass/dataloader.py

@@ -0,0 +1,92 @@
+import torch
+import pandas as pd
+import lightning.pytorch as pl
+
+from transformers import AutoModel, AutoTokenizer
+from torch.utils.data import Dataset, DataLoader
+
+
+class MembraneDataset(Dataset):
+    def __init__(self, config, data_path):
+        self.config = config
+        self.data = pd.read_csv(data_path)
+        self.tokenizer = AutoTokenizer.from_pretrained(self.config.lm.pretrained_esm)
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        sequence = self.data.iloc[idx]["Sequence"]
+        tm_segs = self.data.iloc[idx]["TM_segments"]
+
+        tokens = self.tokenizer(
+            sequence.upper(),
+            return_tensors='pt',
+            padding='max_length',
+            truncation=True,
+            max_length=self.config.data.max_seq_len,
+        )
+
+        labels = torch.tensor(1 if tm_segs > 0 else 0, dtype=torch.float)
+
+        return {
+            "input_ids": tokens['input_ids'],
+            "attention_mask": tokens['attention_mask'],
+            "labels": labels
+        }
+
+
+def collate_fn(batch):
+    input_ids = torch.stack([item['input_ids'].squeeze(0) for item in batch])
+    masks = torch.stack([item['attention_mask'].squeeze(0) for item in batch])
+    labels = torch.stack([item['labels'] for item in batch])
+
+    return {
+        'input_ids': input_ids,
+        'attention_mask': masks,
+        'labels': labels
+    }
+
+
+class MembraneDataModule(pl.LightningDataModule):
+    def __init__(self, config, train_dataset, val_dataset, test_dataset, collate_fn=collate_fn):
+        super().__init__()
+        self.train_dataset = train_dataset
+        self.val_dataset = val_dataset
+        self.test_dataset = test_dataset
+        self.collate_fn = collate_fn
+        self.batch_size = config.data.batch_size
+
+    def train_dataloader(self):
+        return DataLoader(self.train_dataset,
+                          batch_size=self.batch_size,
+                          collate_fn=self.collate_fn,
+                          num_workers=8,
+                          pin_memory=True)
+    
+    def val_dataloader(self):
+        return DataLoader(self.val_dataset,
+                          batch_size=self.batch_size,
+                          collate_fn=self.collate_fn,
+                          num_workers=8,
+                          pin_memory=True)
+    
+    def test_dataloader(self):
+        return DataLoader(self.test_dataset,
+                          batch_size=self.batch_size,
+                          collate_fn=self.collate_fn,
+                          num_workers=8,
+                          pin_memory=True)
+    
+
+def get_datasets(config):
+    """Helper method to grab datasets to quickly init data module in main.py"""
+    train_dataset = MembraneDataset(config, config.data.train)
+    val_dataset = MembraneDataset(config, config.data.val)
+    test_dataset = MembraneDataset(config, config.data.test)
+    
+    return  {
+        "train": train_dataset,
+        "val": val_dataset,
+        "test": test_dataset
+    }

src/guidance/multipass/main.py

@@ -0,0 +1,74 @@
+#!/usr/bin/env python3
+
+import os
+import wandb
+import lightning.pytorch as pl
+
+from omegaconf import OmegaConf
+from lightning.pytorch.loggers import WandbLogger
+from lightning.pytorch.callbacks import ModelCheckpoint, LearningRateMonitor
+
+from src.utils.model_utils import _print
+from src.utils.config_utils import load_config
+from src.guidance.multipass.multipass_module import MultipassClassifier
+from src.guidance.multipass.dataloader import MembraneDataModule, get_datasets
+from src.lm.memdlm.diffusion_module import MembraneDiffusion
+
+
+config = load_config("multipass.yaml")
+wandb.login()
+
+# data
+datasets = get_datasets(config)
+data_module = MembraneDataModule(
+    config=config,
+    train_dataset=datasets['train'],
+    val_dataset=datasets['val'],
+    test_dataset=datasets['test'],
+)
+
+# wandb logging
+#wandb.init(project=config.wandb.project, name=config.wandb.name)
+wandb_logger = WandbLogger(**config.wandb)
+
+# lightning checkpoints
+lr_monitor = LearningRateMonitor(logging_interval="step")
+checkpoint_callback = ModelCheckpoint(
+    monitor="val/loss",
+    save_top_k=1,
+    mode="min",
+    dirpath=config.checkpointing.save_dir,
+    filename="best_model",
+)
+
+# lightning trainer
+trainer = pl.Trainer(
+    max_steps=config.training.max_steps,
+    accelerator="cuda",
+    devices=1, #config.training.devices if config.training.mode=='train' else [0],
+    #strategy=DDPStrategy(find_unused_parameters=True),
+    callbacks=[checkpoint_callback, lr_monitor],
+    logger=wandb_logger,
+    log_every_n_steps=config.training.log_every_n_steps
+)
+
+# Folder to save checkpoints
+ckpt_dir = config.checkpointing.save_dir
+os.makedirs(ckpt_dir, exist_ok=True)
+
+# instantiate model
+model = MultipassClassifier(config, MembraneDiffusion(config))
+
+# train or evalute the model
+if config.training.mode == "train":
+    trainer.fit(model, datamodule=data_module)
+
+elif config.training.mode == "test":
+    ckpt_path = os.path.join(ckpt_dir, "best_model.ckpt")
+    state_dict = model.get_state_dict(ckpt_path)
+    model.load_state_dict(state_dict)
+    trainer.test(model, datamodule=data_module, ckpt_path=ckpt_path)
+else:
+    raise ValueError(f"{config.training.mode} is invalid. Must be 'train' or 'test'")
+
+wandb.finish()

src/guidance/multipass/multipass_module.py

@@ -0,0 +1,167 @@
+import gc
+import torch
+import torch.nn as nn
+import lightning.pytorch as pl
+import torch.nn.functional as F
+
+from omegaconf import OmegaConf
+
+from src.utils.model_utils import _print
+from src.guidance.solubility.utils import CosineWarmup
+
+from sklearn.metrics import roc_auc_score, accuracy_score
+
+
+config = OmegaConf.load("/scratch/pranamlab/sgoel/MeMDLM_v2/src/configs/multipass.yaml")
+
+
+class MultipassClassifier(pl.LightningModule):
+    def __init__(self, config, diffusion_model):
+        super().__init__()
+        self.config = config
+        self.loss_fn = nn.BCEWithLogitsLoss(reduction='none')
+        self.all_preds = []
+        self.all_labels = []
+
+        self.diffusion_model = diffusion_model
+        for p in self.diffusion_model.model.parameters():
+            p.requires_grad = False
+        self.diffusion_model.eval()
+
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=config.model.d_model,
+            nhead=config.model.num_heads,
+            dropout=config.model.dropout,
+            batch_first=True
+        )
+        self.encoder = nn.TransformerEncoder(encoder_layer, config.model.num_layers)
+        self.layer_norm = nn.LayerNorm(config.model.d_model)
+        self.dropout = nn.Dropout(config.model.dropout)
+        self.mlp = nn.Sequential(
+            nn.Linear(config.model.d_model, config.model.d_model // 2),
+            nn.ReLU(),
+            nn.Dropout(config.model.dropout),
+            nn.Linear(config.model.d_model // 2, 1),
+        )
+
+
+    # -------# Classifier step #-------- #
+    def forward(self, x_t, attn_mask, embeds=None, with_hidden=None):
+        if embeds is None:
+            with torch.no_grad():
+                embeds = self.diffusion_model.forward(
+                    input_ids=x_t,
+                    attention_mask=attn_mask,
+                    return_hidden=True
+                )
+        else:
+            assert with_hidden is not None
+
+        encodings = self.encoder(embeds, src_key_padding_mask=(attn_mask== 0))
+        encodings = self.dropout(self.layer_norm(encodings))
+        mask = attn_mask.unsqueeze(-1)
+        pooled = (encodings * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1)
+        logits = self.mlp(pooled).squeeze(-1)
+        return logits
+
+    def step(self, batch):
+        input_ids = batch['input_ids']
+        attention_mask = batch['attention_mask']
+        labels = batch['labels']
+
+        t1 = self.diffusion_model.sample_t(input_ids)
+        maskable = self.diffusion_model.is_maskable(input_ids)
+
+        x_t, _ = self.diffusion_model.noise_x0(input_ids, t1, maskable_mask=maskable)
+        
+        logits = self.forward(x_t, attention_mask)
+        loss = self.compute_loss(logits, labels)
+        
+        return loss, logits
+
+    # -------# Training / Evaluation #-------- #
+    def training_step(self, batch, batch_idx):
+        train_loss, _ = self.step(batch)
+        self.log(name="train/loss", value=train_loss.item(), on_step=True, on_epoch=False, logger=True, sync_dist=True)
+        self.save_ckpt()
+        return train_loss
+
+    def validation_step(self, batch, batch_idx):
+        val_loss, _ = self.step(batch)
+        self.log(name="val/loss", value=val_loss.item(), on_step=False, on_epoch=True, logger=True, sync_dist=True)
+        return val_loss
+
+    def test_step(self, batch):
+        test_loss, logits = self.step(batch)
+        preds = F.sigmoid(logits)
+        self.all_preds.append(preds.detach().cpu())
+        self.all_labels.append(batch['labels'].detach().cpu())
+        self.log(name="test/loss", value=test_loss.item(), on_step=False, on_epoch=True, logger=True, sync_dist=True)
+        return test_loss
+
+    def on_test_epoch_start(self):
+        self.all_preds = []
+        self.all_labels = []
+
+    def on_test_epoch_end(self):
+        all_preds = torch.cat(self.all_preds).numpy()
+        all_labels = torch.cat(self.all_labels).numpy()
+        
+        auroc = roc_auc_score(all_labels, all_preds)
+        binaries = (all_preds > 0.5).astype(int)
+        accuracy = accuracy_score(all_labels, binaries)
+
+        self.log(name="test/AUROC", value=auroc, on_step=False, on_epoch=True, logger=True, sync_dist=True)
+        self.log(name="test/accuracy", value=accuracy, on_step=False, on_epoch=True, logger=True, sync_dist=True)
+    
+    def optimizer_step(self, *args, **kwargs):
+        super().optimizer_step(*args, **kwargs)
+        gc.collect()
+        torch.cuda.empty_cache()
+
+    def configure_optimizers(self):
+        path = self.config.training
+        optimizer = torch.optim.AdamW(self.parameters(), lr=self.config.optim.lr)
+        lr_scheduler = CosineWarmup(
+            optimizer,
+            warmup_steps=path.warmup_steps,
+            total_steps=path.max_steps,
+        )
+        scheduler_dict = {
+            "scheduler": lr_scheduler,
+            "interval": 'step',
+            'frequency': 1,
+            'monitor': 'val/loss',
+            'name': 'learning_rate'
+        }
+        return [optimizer], [scheduler_dict]
+    
+    def save_ckpt(self):
+        curr_step = self.global_step
+        save_every = self.config.training.val_check_interval
+        if curr_step % save_every == 0 and curr_step > 0:  # Save every 250 steps
+            ckpt_path = f"{self.config.checkpointing.save_dir}/step={curr_step}.ckpt"
+            self.trainer.save_checkpoint(ckpt_path)
+    
+    # -------# Loss and Test Set Metrics #-------- #
+    def compute_loss(self, logits, labels):
+        """Helper method to handle loss calculation"""
+        loss = self.loss_fn(logits, labels.float()).mean()
+        return loss
+
+    # -------# Helper Functions #-------- #
+    def get_state_dict(self, ckpt_path):
+        """Helper method to load and process a trained model's state dict from saved checkpoint"""
+        def remove_model_prefix(state_dict):
+            for k in state_dict.keys():
+                if "model." in k:
+                    k.replace('model.', '')
+            return state_dict  
+
+        checkpoint = torch.load(ckpt_path, weights_only=False)#, map_location='cuda' if torch.cuda.is_available() else 'cpu')
+        state_dict = checkpoint.get("state_dict", checkpoint)
+
+        if any(k.startswith("model.") for k in state_dict.keys()):
+            state_dict = remove_model_prefix(state_dict)
+        
+        return state_dict

src/guidance/oligo/dataloader.py

@@ -0,0 +1,91 @@
+import torch
+import pandas as pd
+import lightning.pytorch as pl
+
+from transformers import AutoTokenizer
+from torch.utils.data import Dataset, DataLoader
+
+class OligomerDataset(Dataset):
+    def __init__(self, config, data_path):
+        self.config = config
+        self.data = pd.read_csv(data_path)
+        self.tokenizer = AutoTokenizer.from_pretrained(self.config.lm.pretrained_esm)
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        sequence = self.data.iloc[idx]["Sequence"]
+
+        tokens = self.tokenizer(
+            sequence,
+            return_tensors='pt',
+            padding='max_length',
+            truncation=True,
+            max_length=self.config.data.max_seq_len,
+        )
+
+        labels = torch.tensor(self.data.iloc[idx]['Binary Label'], dtype=torch.float)
+
+        return {
+            "input_ids": tokens['input_ids'],
+            "attention_mask": tokens['attention_mask'],
+            "labels": labels
+        }
+
+
+def collate_fn(batch):
+    input_ids = torch.stack([item['input_ids'].squeeze(0) for item in batch])
+    masks = torch.stack([item['attention_mask'].squeeze(0) for item in batch])
+    labels = torch.stack([item['labels'] for item in batch])
+
+    return {
+        'input_ids': input_ids,
+        'attention_mask': masks,
+        'labels': labels
+    }
+
+
+class OligomerDataModule(pl.LightningDataModule):
+    def __init__(self, config, train_dataset, val_dataset, test_dataset, collate_fn=collate_fn):
+        super().__init__()
+        self.train_dataset = train_dataset
+        self.val_dataset = val_dataset
+        self.test_dataset = test_dataset
+        self.collate_fn = collate_fn
+        self.batch_size = config.data.batch_size
+
+    def train_dataloader(self):
+        return DataLoader(self.train_dataset,
+                          batch_size=self.batch_size,
+                          collate_fn=self.collate_fn,
+                          num_workers=8,
+                          pin_memory=True,
+                          shuffle=True)
+    
+    def val_dataloader(self):
+        return DataLoader(self.val_dataset,
+                          batch_size=self.batch_size,
+                          collate_fn=self.collate_fn,
+                          num_workers=8,
+                          pin_memory=True)
+    
+    def test_dataloader(self):
+        return DataLoader(self.test_dataset,
+                          batch_size=self.batch_size,
+                          collate_fn=self.collate_fn,
+                          num_workers=8,
+                          pin_memory=True)
+    
+
+def get_datasets(config):
+    """Helper method to grab datasets to quickly init data module in main.py"""
+    train_dataset = OligomerDataset(config, config.data.train)
+    val_dataset = OligomerDataset(config, config.data.val)
+    test_dataset = OligomerDataset(config, config.data.test)
+    
+    return  {
+        "train": train_dataset,
+        "val": val_dataset,
+        "test": test_dataset
+    }

src/guidance/oligo/main.py

@@ -0,0 +1,74 @@
+#!/usr/bin/env python3
+
+import os
+import wandb
+import lightning.pytorch as pl
+
+from omegaconf import OmegaConf
+from lightning.pytorch.strategies import DDPStrategy
+from lightning.pytorch.loggers import WandbLogger
+from lightning.pytorch.callbacks import ModelCheckpoint, LearningRateMonitor
+
+from src.utils.model_utils import _print
+from src.utils.config_utils import load_config
+from src.guidance.oligo.oligo_module import OligomerClassifier
+from src.guidance.oligo.dataloader import get_datasets, OligomerDataModule
+
+
+config = load_config("oligo.yaml")
+wandb.login()
+
+# data
+datasets = get_datasets(config)
+data_module = OligomerDataModule(
+    config=config,
+    train_dataset=datasets['train'],
+    val_dataset=datasets['val'],
+    test_dataset=datasets['test'],
+)
+
+# wandb logging
+#wandb.init(project=config.wandb.project, name=config.wandb.name)
+wandb_logger = WandbLogger(**config.wandb)
+
+# lightning checkpoints
+lr_monitor = LearningRateMonitor(logging_interval="step")
+checkpoint_callback = ModelCheckpoint(
+    monitor="val/loss",
+    save_top_k=1,
+    mode="min",
+    dirpath=config.checkpointing.save_dir,
+    filename="best_model",
+)
+
+# lightning trainer
+trainer = pl.Trainer(
+    max_steps=config.training.max_steps,
+    accelerator="cuda",
+    devices=1, #config.training.devices if config.training.mode=='train' else [0],
+    #strategy=DDPStrategy(find_unused_parameters=True),
+    callbacks=[checkpoint_callback, lr_monitor],
+    logger=wandb_logger,
+    log_every_n_steps=config.training.log_every_n_steps
+)
+
+# Folder to save checkpoints
+ckpt_dir = config.checkpointing.save_dir
+os.makedirs(ckpt_dir, exist_ok=True)
+
+# instantiate model
+model = OligomerClassifier(config)
+
+# train or evalute the model
+if config.training.mode == "train":
+    trainer.fit(model, datamodule=data_module)
+
+elif config.training.mode == "test":
+    ckpt_path = os.path.join(ckpt_dir, "best_model.ckpt")
+    state_dict = model.get_state_dict(ckpt_path)
+    model.load_state_dict(state_dict)
+    trainer.test(model, datamodule=data_module, ckpt_path=ckpt_path)
+else:
+    raise ValueError(f"{config.training.mode} is invalid. Must be 'train' or 'test'")
+
+wandb.finish()

src/guidance/oligo/oligo_module.py

@@ -0,0 +1,155 @@
+import gc
+import torch
+import torch.nn as nn
+import lightning.pytorch as pl
+
+from omegaconf import OmegaConf
+from transformers import AutoModel
+from torchmetrics.classification import BinaryAUROC, BinaryAccuracy
+
+from src.utils.model_utils import _print
+from src.guidance.oligo.utils import CosineWarmup, HelixRoPE
+
+
+config = OmegaConf.load("/scratch/pranamlab/sgoel/MeMDLM_v2/src/configs/oligo.yaml")
+
+class OligomerClassifier(pl.LightningModule):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.loss_fn = nn.BCEWithLogitsLoss()
+        self.auroc = BinaryAUROC()
+        self.accuracy = BinaryAccuracy()
+
+        self.esm_model = AutoModel.from_pretrained(config.lm.pretrained_esm)
+        for p in self.esm_model.parameters():
+            p.requires_grad = False
+
+        self.down_proj = nn.Linear(config.model.esm_dim, config.model.d_model)
+        self.layer_norm = nn.LayerNorm(config.model.d_model)
+        self.dropout = nn.Dropout(config.model.dropout)
+
+        self.mlp = nn.Sequential(
+            nn.Linear(config.model.d_model, config.model.d_model // 2),
+            nn.ReLU(),
+            nn.Dropout(config.model.dropout),
+            nn.Linear(config.model.d_model // 2, 1),
+        )
+
+    def forward(self, batch):
+        attention_mask = batch["attention_mask"]
+
+        if "input_ids" in batch:
+            esm_embeds = self.get_esm_embeddings(batch["input_ids"], attention_mask)
+        else:
+            esm_embeds = batch["embeds"]
+
+        if esm_embeds.ndim == 2:
+            esm_embeds = esm_embeds.unsqueeze(0)
+        if attention_mask.ndim == 1:
+            attention_mask = attention_mask.unsqueeze(0)
+
+        x = self.down_proj(esm_embeds)
+        x = self.layer_norm(x)
+        x = self.dropout(x)
+
+        mask = attention_mask.unsqueeze(-1).float()
+        x = x * mask
+        
+        pooled = x.sum(dim=1) / mask.sum(dim=1).clamp(min=1.0)
+        logits = self.mlp(pooled).squeeze(-1)
+
+        return logits
+    
+    # -------# Training / Evaluation #-------- #
+    def training_step(self, batch, batch_idx):
+        train_loss, _ = self.compute_loss(batch)
+        self.log(name="train/loss", value=train_loss.item(), on_step=True, on_epoch=False, logger=True, sync_dist=True)
+        self.save_ckpt()
+        return train_loss
+
+    def validation_step(self, batch, batch_idx):
+        val_loss, _ = self.compute_loss(batch)
+        self.log(name="val/loss", value=val_loss.item(), on_step=False, on_epoch=True, logger=True, sync_dist=True)
+        return val_loss
+
+    def test_step(self, batch):
+        test_loss, preds = self.compute_loss(batch)
+        auroc, accuracy = self.get_metrics(batch, preds)
+        self.log(name="test/loss", value=test_loss.item(), on_step=False, on_epoch=True, logger=True, sync_dist=True)
+        self.log(name="test/AUROC", value=auroc.item(), on_step=False, on_epoch=True, logger=True, sync_dist=True)
+        self.log(name="test/accuracy", value=accuracy.item(), on_step=False, on_epoch=True, logger=True, sync_dist=True)
+        return test_loss
+
+    def on_test_epoch_end(self):
+        self.auroc.reset()
+        self.accuracy.reset()
+    
+    def optimizer_step(self, *args, **kwargs):
+        super().optimizer_step(*args, **kwargs)
+        gc.collect()
+        torch.cuda.empty_cache()
+
+    def configure_optimizers(self):
+        path = self.config.training
+        optimizer = torch.optim.AdamW(self.parameters(), lr=self.config.optim.lr)
+        lr_scheduler = CosineWarmup(
+            optimizer,
+            warmup_steps=path.warmup_steps,
+            total_steps=path.max_steps,
+        )
+        scheduler_dict = {
+            "scheduler": lr_scheduler,
+            "interval": 'step',
+            'frequency': 1,
+            'monitor': 'val/loss',
+            'name': 'learning_rate'
+        }
+        return [optimizer], [scheduler_dict]
+    
+    def save_ckpt(self):
+        curr_step = self.global_step
+        save_every = self.config.training.val_check_interval
+        if curr_step % save_every == 0 and curr_step > 0:  # Save every 250 steps
+            ckpt_path = f"{self.config.checkpointing.save_dir}/step={curr_step}.ckpt"
+            self.trainer.save_checkpoint(ckpt_path)
+    
+    # -------# Loss and Test Set Metrics #-------- #
+    @torch.no_grad
+    def get_esm_embeddings(self, input_ids, attention_mask):
+        outputs = self.esm_model(input_ids=input_ids, attention_mask=attention_mask)
+        embeddings = outputs.last_hidden_state
+        return embeddings
+
+    def compute_loss(self, batch):
+        """Helper method to handle loss calculation"""
+        labels = batch['labels']
+        preds = self.forward(batch)
+        loss = self.loss_fn(preds, labels)
+        return loss, preds
+
+    def get_metrics(self, batch, preds):
+        """Helper method to compute metrics"""
+        labels = batch['labels']
+        _print(f"labels {labels.shape}")
+        _print(f"preds {preds.shape}")
+        auroc = self.auroc.forward(preds, labels)
+        accuracy = self.accuracy.forward(preds, labels)
+        return auroc, accuracy
+
+    # -------# Helper Functions #-------- #
+    def get_state_dict(self, ckpt_path):
+        """Helper method to load and process a trained model's state dict from saved checkpoint"""
+        def remove_model_prefix(state_dict):
+            for k in state_dict.keys():
+                if "model." in k:
+                    k.replace('model.', '')
+            return state_dict  
+
+        checkpoint = torch.load(ckpt_path, weights_only=False)#, map_location='cuda' if torch.cuda.is_available() else 'cpu')
+        state_dict = checkpoint.get("state_dict", checkpoint)
+
+        if any(k.startswith("model.") for k in state_dict.keys()):
+            state_dict = remove_model_prefix(state_dict)
+        
+        return state_dict

src/guidance/oligo/utils.py

@@ -0,0 +1,50 @@
+
+import torch
+import torch.nn as nn
+import numpy as np
+from torch.optim.lr_scheduler import _LRScheduler
+
+
+
+class HelixRoPE(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        pos = torch.arange(config.data.max_seq_len).float()
+        thetas = (2 * torch.pi / 3.6) * pos 
+        self.register_buffer("thetas", thetas)
+
+    def forward(self, x, mask):
+        B, L, D = x.shape
+        assert D % 2 == 0
+
+        thetas = self.thetas[:L]
+        cos = torch.cos(thetas).unsqueeze(0).unsqueeze(-1)
+        sin = torch.sin(thetas).unsqueeze(0).unsqueeze(-1)
+
+        x_double = x.view(B, L, D//2, 2)
+        x1 = x_double[..., 0]
+        x2 = x_double[..., 1]
+
+        r1 = cos * x1 - sin * x2
+        r2 = sin * x1 + cos * x2
+
+        ropes = torch.stack([r1, r2], dim=-1).view(B, L, D)
+        return ropes * mask  # attention mask to ignore pad tokens
+
+
+class CosineWarmup(_LRScheduler):
+    def __init__(self, optimizer, warmup_steps, total_steps, eta_ratio=0.1, last_epoch=-1):
+        self.warmup_steps = warmup_steps
+        self.total_steps = total_steps
+        self.eta_ratio = eta_ratio  # The ratio of minimum to maximum learning rate
+        super(CosineWarmup, self).__init__(optimizer, last_epoch)
+
+    def get_lr(self):
+        if self.last_epoch < self.warmup_steps:
+            return [base_lr * self.last_epoch / self.warmup_steps for base_lr in self.base_lrs]
+
+        progress = (self.last_epoch - self.warmup_steps) / (self.total_steps - self.warmup_steps)
+        cosine_decay = 0.5 * (1 + np.cos(np.pi * progress))
+        decayed_lr = (1 - self.eta_ratio) * cosine_decay + self.eta_ratio
+
+        return [decayed_lr * base_lr for base_lr in self.base_lrs]

src/guidance/{main.py → solubility/main.py}

@@ -10,12 +10,13 @@ from lightning.pytorch.loggers import WandbLogger
 from lightning.pytorch.callbacks import ModelCheckpoint, LearningRateMonitor
 
 from src.utils.model_utils import _print
-from src.guidance.solubility_module import SolubilityClassifier
-from src.guidance.dataloader import MembraneDataModule, get_datasets
+from src.utils.config_utils import load_config
+from src.guidance.solubility.solubility_module import SolubilityClassifier
+from src.guidance.solubility.dataloader import MembraneDataModule, get_datasets
 
 
-config = OmegaConf.load("/scratch/sgoel/MeMDLM_v2/src/configs/guidance.yaml")
-wandb.login(key='2b76a2fa2c1cdfddc5f443602c17b011fefb0a8f')
+config = load_config("solubility.yaml")
+wandb.login()
 
 # data
 datasets = get_datasets(config)

src/guidance/{solubility_module.py → solubility/solubility_module.py}

@@ -8,10 +8,10 @@ from transformers import AutoModel
 from torchmetrics.classification import BinaryAUROC, BinaryAccuracy
 
 from src.utils.model_utils import _print
-from src.guidance.utils import CosineWarmup
+from src.guidance.solubility.utils import CosineWarmup
 
 
-config = OmegaConf.load("/scratch/sgoel/MeMDLM_v2/src/configs/guidance.yaml")
+config = OmegaConf.load("/scratch/pranamlab/sgoel/MeMDLM_v2/src/configs/solubility.yaml")
 
 class SolubilityClassifier(pl.LightningModule):
     def __init__(self, config):
@@ -146,7 +146,7 @@ class SolubilityClassifier(pl.LightningModule):
                     k.replace('model.', '')
             return state_dict  
 
-        checkpoint = torch.load(ckpt_path, map_location='cuda' if torch.cuda.is_available() else 'cpu')
+        checkpoint = torch.load(ckpt_path, weights_only=False)#, map_location='cuda' if torch.cuda.is_available() else 'cpu')
         state_dict = checkpoint.get("state_dict", checkpoint)
 
         if any(k.startswith("model.") for k in state_dict.keys()):

src/guidance/solubility/utils.py

@@ -0,0 +1,21 @@
+
+
+import numpy as np
+from torch.optim.lr_scheduler import _LRScheduler
+
+class CosineWarmup(_LRScheduler):
+    def __init__(self, optimizer, warmup_steps, total_steps, eta_ratio=0.1, last_epoch=-1):
+        self.warmup_steps = warmup_steps
+        self.total_steps = total_steps
+        self.eta_ratio = eta_ratio  # The ratio of minimum to maximum learning rate
+        super(CosineWarmup, self).__init__(optimizer, last_epoch)
+
+    def get_lr(self):
+        if self.last_epoch < self.warmup_steps:
+            return [base_lr * self.last_epoch / self.warmup_steps for base_lr in self.base_lrs]
+
+        progress = (self.last_epoch - self.warmup_steps) / (self.total_steps - self.warmup_steps)
+        cosine_decay = 0.5 * (1 + np.cos(np.pi * progress))
+        decayed_lr = (1 - self.eta_ratio) * cosine_decay + self.eta_ratio
+
+        return [decayed_lr * base_lr for base_lr in self.base_lrs]

src/lm/memdlm/diffusion_module.py

@@ -28,7 +28,7 @@ class MembraneDiffusion(pl.LightningModule):
         self.mask_id = self.tokenizer.mask_token_id
         self.pad_id = self.tokenizer.pad_token_id
 
-    def forward(self, input_ids, attention_mask, guidance: Optional[bool] = False):
+    def forward(self, input_ids, attention_mask, return_hidden=None):
         """
         Forward pass through language model.
 
@@ -38,7 +38,12 @@ class MembraneDiffusion(pl.LightningModule):
         Returns:
             - logits (torch.Tensor): [B, L, V], unnormalized model outputs
         """
-        return self.model(input_ids=input_ids, attention_mask=attention_mask).logits
+        if return_hidden:
+            out = self.model(input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True)
+            return out.hidden_states[-1]
+        
+        else:
+            return self.model(input_ids=input_ids, attention_mask=attention_mask).logits
 
     # -------# Diffusion #-------- #
     def step(self, batch):
@@ -78,7 +83,6 @@ class MembraneDiffusion(pl.LightningModule):
         u = torch.rand_like(x0, dtype=torch.float) 
         t1_mask = (u < (t1 / self.config.lm.num_diffusion_timesteps)[:, None]) & maskable_mask
         x_t1 = x0.masked_fill(t1_mask, self.mask_id)
-        x_t1 = x_t1.masked_fill(t1_mask, self.mask_id)
         return x_t1, t1_mask
 
     def get_weight(self, t, weight_type):
@@ -100,7 +104,7 @@ class MembraneDiffusion(pl.LightningModule):
         
         Args:
             logits (torch.Tensor): [B, L, vocab_size], unnormalized model outputs
-            labels (torch.Tensor): [B, L], target labels (with padding tokens as -100)
+            labels (torch.Tensor): [B, L], target labels with pad tokens
             weight (torch.Tensor): [B, 1], per-sample weight for loss calculation
         Returns:
             loss (torch.Tensor): Averaged loss over the batch
@@ -127,22 +131,22 @@ class MembraneDiffusion(pl.LightningModule):
     # -------# Training / Evaluation #-------- #
     def training_step(self, batch):
         loss, ppl = self.step(batch)
-        self.log("train/loss", loss.item(), on_step=True, on_epoch=False, prog_bar=True)
-        self.log("train/ppl", ppl.item(), on_step=True, on_epoch=False, prog_bar=False)
+        self.log("train/loss", loss, on_step=True, on_epoch=False, prog_bar=True)
+        self.log("train/ppl", ppl, on_step=True, on_epoch=False, prog_bar=False)
         return loss
     
     def validation_step(self, batch):
         loss, ppl = self.step(batch)
         self.cleanup()
-        self.log("val/loss", loss.item(), on_step=False, on_epoch=True, prog_bar=True, sync_dist=True)
-        self.log("val/ppl", ppl.item(), on_step=False, on_epoch=True, prog_bar=False, sync_dist=True)
+        self.log("val/loss", loss, on_step=False, on_epoch=True, prog_bar=True, sync_dist=True)
+        self.log("val/ppl", ppl, on_step=False, on_epoch=True, prog_bar=False, sync_dist=True)
         return loss
 
     def test_step(self, batch):
         loss, ppl = self.step(batch)
         self.cleanup()
-        self.log('test/loss', loss.item(), on_step=False, on_epoch=True, prog_bar=True, sync_dist=True)
-        self.log("test/ppl", ppl.item(), on_step=False, on_epoch=True, prog_bar=False, sync_dist=True)
+        self.log('test/loss', loss, on_step=False, on_epoch=True, prog_bar=True, sync_dist=True)
+        self.log("test/ppl", ppl, on_step=False, on_epoch=True, prog_bar=False, sync_dist=True)
         return loss
 
 
@@ -176,7 +180,7 @@ class MembraneDiffusion(pl.LightningModule):
                     k.replace('model.', '')
             return state_dict  
 
-        checkpoint = torch.load(ckpt_path, map_location='cuda' if torch.cuda.is_available() else 'cpu')
+        checkpoint = torch.load(ckpt_path, weights_only=False)#, map_location='cuda' if torch.cuda.is_available() else 'cpu')
         state_dict = checkpoint.get("state_dict", checkpoint)
 
         if any(k.startswith("model.") for k in state_dict.keys()):

src/lm/memdlm/main.py

@@ -14,12 +14,13 @@ from lightning.pytorch.callbacks import ModelCheckpoint, LearningRateMonitor
 from src.lm.memdlm.diffusion_module import MembraneDiffusion
 from src.lm.memdlm.dataloader import MembraneDataModule, get_datasets
 from src.utils.model_utils import apply_rdm_freezing
+from src.utils.config_utils import load_config
 
-wandb.login(key='2b76a2fa2c1cdfddc5f443602c17b011fefb0a8f')
+wandb.login()
 
 
 # Load yaml config
-config = OmegaConf.load("/scratch/pranamlab/sgoel/MeMDLM_v2/src/configs/lm.yaml")
+config = load_config("lm.yaml")
 
 # Get datasets
 datasets = get_datasets(config)

src/sampling/{guided_generator.py → desolubilize_generator.py}

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 
-import sys
+
 import os
 import torch
 import pandas as pd
@@ -9,20 +9,30 @@ from datetime import datetime
 from omegaconf import OmegaConf
 from transformers import AutoTokenizer, AutoModelForMaskedLM
 
-from src.lm.memdlm.diffusion_module import MembraneFlow
-from src.utils.model_utils import _print
-from src.sampling.guided_sampler import GuidedSampler
+from src.lm.memdlm.diffusion_module import MembraneDiffusion
+from src.sampling.pet_sampler import PETSampler
 from src.utils.generate_utils import (
     mask_for_scaffold,
     calc_blosum_score,
-    calc_ppl
+    calc_ppl,
+    calc_tm_enrich
 )
 
-config = OmegaConf.load("/home/a03-sgoel/MeMDLM_v2/src/configs/guidance.yaml")
+from src.utils.model_utils import _print
+from src.utils.config_utils import load_config, repo_path
+
+
+config = load_config("desolubilize.yaml")
+task = config.guidance.task
 
-os.chdir(f'/home/a03-sgoel/MeMDLM_v2/results/infilling/guided/{config.lm.ft_evoflow}/test_set/')
+results_root = repo_path("results", "heme", config.lm.ft_evoflow)
 todays_date = datetime.today().strftime('%Y-%m-%d')
-csv_save_path = f'./{todays_date}_boltzmann-soft_new_clf_data_cleaned/'
+
+if config.guidance.prior == 'boltzmann':
+    csv_save_path = results_root / task / todays_date / f"{config.guidance.prior}-t={config.guidance.boltzmann_t}_p={config.guidance.top_p}_tau={config.guidance.sampling_t}"
+elif config.guidance.prior == 'lm_probs':
+    csv_save_path = results_root / task / todays_date / f"{config.guidance.prior}_p={config.guidance.top_p}_tau={config.guidance.sampling_t}"
+
 try: os.makedirs(csv_save_path, exist_ok=False)
 except FileExistsError: pass
 
@@ -33,25 +43,28 @@ def main():
     tokenizer = AutoTokenizer.from_pretrained(config.lm.pretrained_esm)
     esm_model = AutoModelForMaskedLM.from_pretrained(config.lm.pretrained_esm).eval().to(device)
 
-    diffusion = MembraneFlow(config).to(device)
-    state_dict = diffusion.get_state_dict(f"/home/a03-sgoel/MeMDLM_v2/checkpoints/{config.lm.ft_evoflow}/best_model.ckpt")
+    diffusion = MembraneDiffusion(config).to(device)
+    state_dict = diffusion.get_state_dict(str(repo_path("checkpoints", config.lm.ft_evoflow, "best_model.ckpt")))
     diffusion.load_state_dict(state_dict)
     diffusion.eval().to(device)
 
-    sampler = GuidedSampler(config, esm_model, tokenizer, diffusion, device)
+    sampler = PETSampler(config, esm_model, tokenizer, diffusion, device)
 
-    df = pd.read_csv('/home/a03-sgoel/MeMDLM_v2/data/classifier/test.csv')
+    # Update this path to your input CSV of scaffold sequences (uppercase = TM, lowercase = soluble).
+    df = pd.read_csv(str(repo_path("results", "heme", "4d2.csv")))
     sequences = df['Sequence'].tolist()
 
-    gen_seqs, ppls, blosums = [], [], []
+    gen_seqs, ppls, blosums, og_tms, gen_tms, delta_tms = [], [], [], [], [], []
 
 
-    for seq in tqdm(sequences, desc='Infilling Sequences'):
+    for seq in tqdm(sequences, desc='Desolubilizing Sequences'):
         masked_seq = mask_for_scaffold(seq, generate_type='uppercase', mask_token='<mask>')
         tokens = tokenizer(masked_seq, return_tensors='pt')
         input_ids, attn_masks = tokens['input_ids'].to(device), tokens['attention_mask'].to(device)
-        
-        soluble_idxs = [i for i in range(len(seq)) if seq[i].isupper()]
+
+        tm_idxs = [i for i in range(len(seq)) if seq[i].isupper()]
+        soluble_idxs = [i + 1 for i in range(len(seq)) if seq[i].islower()]
+
         infilled_tokens = sampler.optimize_sequence(
             input_ids=input_ids,
             attn_masks=attn_masks,
@@ -59,32 +72,49 @@ def main():
         )
         infilled_seq = tokenizer.decode(infilled_tokens).replace(" ", "")[5:-5]
         
-        bl = calc_blosum_score(seq.upper(), infilled_seq, soluble_idxs)
+        try:
+            bl = calc_blosum_score(seq.upper(), infilled_seq, tm_idxs)
+        except:
+            bl = float('inf')
+        
         try:
             ppl = calc_ppl(esm_model, tokenizer, infilled_seq, [i for i in range(len(seq))], model_type='esm')
         except:
             ppl = float('inf')
 
+        try:
+            og_tm, gen_tm, delta_tm = calc_tm_enrich(seq.upper(), infilled_seq, tm_idxs)
+        except:
+            og_tm, gen_tm, delta_tm = float('inf'), float('inf'), float('inf')
+
         gen_seqs.append(infilled_seq)
         ppls.append(ppl)
         blosums.append(bl)
+        og_tms.append(og_tm)
+        gen_tms.append(gen_tm)
+        delta_tms.append(delta_tm)
 
         _print(seq)
         _print(infilled_seq)
         _print(ppl)
         _print(bl)
+        _print(og_tm)   
+        _print(gen_tm)
+        _print(delta_tm)
         _print('\n')
 
 
     df['MeMDLM Sequence'] = gen_seqs
     df['MeMDLM PPL'] = ppls
     df['MeMDLM BLOSUM'] = blosums
+    df['OG TM Enrichment'] = og_tms
+    df['MeMDLM TM Enrichment'] = gen_tms
+    df['Delta TM Enrichment'] = delta_tms
 
     _print(df)
-    df.to_csv(f'./{csv_save_path}/t=0.7_new-data-cleaned_infilled_seqs.csv', index=False)
+    df.to_csv(csv_save_path / "infilled_seqs.csv", index=False)
 
 
     
 if __name__ == "__main__":
     main()
-    

src/sampling/multipass_generator.py

@@ -0,0 +1,100 @@
+#!/usr/bin/env python3
+
+import sys
+import os
+
+import random
+import torch
+import pandas as pd
+import numpy as np
+
+from tqdm import tqdm
+from collections import Counter
+from omegaconf import OmegaConf
+from datetime import datetime
+from src.utils.generate_utils import mask_for_de_novo
+from transformers import AutoTokenizer, AutoModelForMaskedLM
+
+from src.lm.memdlm.diffusion_module import MembraneDiffusion
+from src.sampling.multipass_sampler import MultipassSampler
+
+from src.utils.generate_utils import calc_ppl
+from src.utils.model_utils import _print
+from src.utils.config_utils import load_config, repo_path
+
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+config = load_config("multipass.yaml")
+
+date = datetime.now().strftime("%Y-%m-%d")
+
+
+
+
+def main():
+    csv_save_path = repo_path(
+        'results', 'multipass', config.wandb.name, date,
+        f"lamb={config.guidance.reg_strength}_tau={config.guidance.sampling_temperature}"
+    )
+    
+    try: os.makedirs(csv_save_path, exist_ok=False)
+    except FileExistsError: pass
+
+    tokenizer = AutoTokenizer.from_pretrained(config.lm.pretrained_evoflow)
+    
+    memdlm = MembraneDiffusion(config).to(device)
+    state_dict = memdlm.get_state_dict(str(repo_path("checkpoints", config.lm.ft_evoflow, "best_model.ckpt")))
+    memdlm.load_state_dict(state_dict)
+    memdlm.eval()
+
+    esm_pth = config.lm.pretrained_esm
+    esm_model = AutoModelForMaskedLM.from_pretrained(esm_pth).to(device)
+    esm_model.eval()
+
+    generator = MultipassSampler(config, device, memdlm, tokenizer)
+
+    seq_lengths = [x for x in range(60, 161) for _ in range(3)]
+    #seq_lengths = random.sample([x for x in range(120, 161)], 15)
+
+    generation_results = []
+    for seq_len in tqdm(seq_lengths, desc=f"Generating sequences: "):
+        seq_res = []
+
+        masked_seq = mask_for_de_novo(seq_len) # Sequence of all <mask> tokens
+        tokens = tokenizer(masked_seq, return_tensors='pt')
+
+        gen_seq = ""
+        attempts = 0
+
+        while len(gen_seq) != seq_len and attempts < 3:
+            gen_seq, pred_tm_segs = generator.sample_guidance(
+                tokens,
+                config.guidance.guide_steps,
+                config.guidance.diffusion_steps
+            )
+            attempts += 1
+
+        if len(gen_seq) != seq_len:
+            esm_ppl, memdlm_ppl = None, None
+        else:
+            esm_ppl = calc_ppl(esm_model, tokenizer, gen_seq, [i for i in range(len(gen_seq))], model_type='esm')
+            memdlm_ppl = calc_ppl(memdlm, tokenizer, gen_seq, [i for i in range(len(gen_seq))], model_type='diffusion')
+
+        _print(f'seq: {gen_seq}')
+        _print(f'pred_tm_segs: {pred_tm_segs}')
+        _print(f"ESM PPL: {esm_ppl}")
+        _print(f"MeMDLM PPL: {memdlm_ppl}")
+        _print('\n')
+
+        seq_res.append(gen_seq)
+        seq_res.append(esm_ppl)
+        seq_res.append(memdlm_ppl)
+        seq_res.append(pred_tm_segs)
+        generation_results.append(seq_res)
+
+    df = pd.DataFrame(generation_results, columns=['Generated Sequence', 'ESM PPL', 'MeMDLM PPL', 'Pred TM Segments'])
+    df.to_csv(str(csv_save_path / "seqs_with_ppl.csv"), index=False)
+
+
+if __name__ == "__main__":
+    main()

src/sampling/multipass_sampler.py

@@ -0,0 +1,166 @@
+import os
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from src.guidance.multipass.multipass_module import MultipassClassifier
+from src.utils.model_utils import _print
+from src.utils.config_utils import repo_path
+
+
+
+class MultipassSampler:
+    def __init__(self, config, device, mdlm, tokenizer):
+        self.config = config
+        self.device = device
+
+        self.diffusion = mdlm
+        self.memdlm_lm = self.diffusion.model.lm_head
+        self.tokenizer = self.diffusion.tokenizer
+
+        ckpt_path = str(repo_path("checkpoints", config.wandb.name, "best_model.ckpt"))
+        self.classifier_model = MultipassClassifier(config, self.diffusion).eval().to(self.device)
+        state_dict = self.classifier_model.get_state_dict(ckpt_path)
+        self.classifier_model.load_state_dict(state_dict)
+
+        self.SPECIAL_TOKEN_IDS = {0, 1, 2, 3, 29}
+
+
+    def stochastic_sample_from_categorical(self, logits, temperature, noise_scale=1.0):
+        """
+        Sample from a categorical distribution with optional temperature scaling and Gumbel noise.
+        Returns the sampled tokens and their log-probabilities (used as confidence scores).
+        """
+        logits = logits.double()
+        if temperature != 0:
+            gumbel_noise = -torch.log(-torch.log(torch.rand_like(logits) + 1e-8) + 1e-8)
+            logits = (logits / temperature) + (noise_scale * gumbel_noise)
+        scores, tokens = logits.log_softmax(dim=-1).max(dim=-1)
+        return tokens, scores
+
+
+    def topk_lowest_masking(self, scores, cutoff_len):
+        """
+        scores: [b, n]; cutoff_len: [b, 1]
+        Returns a [b, n] bool mask, True at the cutoff_len lowest-scoring positions.
+        """
+        sorted_index = scores.sort(-1)[0]
+        cutoff = sorted_index.gather(dim=-1, index=cutoff_len)
+        return scores < cutoff
+    
+
+    def classifier_score(self, hidden_states, attention_mask):
+        return self.classifier_model(x_t=None, attn_mask=attention_mask, with_hidden=True, embeds=hidden_states)
+
+
+    def guidance_loss(self, og_hidden, og_logits, attention_mask, delta):
+        """
+        Implementation of explore-exploi  guidance as in LaMBO-2 (https://arxiv.org/pdf/2305.20009).
+        Disregarding use of saliency map as we have one score per sequence.
+        """
+        lamb = self.config.guidance.reg_strength
+
+        h_current = og_hidden + delta
+        new_logits = self.memdlm_lm(h_current)
+        score = self.classifier_score(h_current, attention_mask)
+
+        kl = F.kl_div(
+            F.log_softmax(new_logits, dim=-1),
+            F.softmax(og_logits, dim=-1),
+            reduction='sum'
+        )
+
+        loss = lamb * kl - score.sum()
+        return loss
+
+
+    def optimized_sampling(self, og_logits, og_hidden, attention_mask, n_steps):
+        """
+        At each diffusion timestep, take n_steps of gradient-based guidance on the hidden states
+        """
+        eta = self.config.guidance.step_size
+        og_logits = og_logits.detach()
+
+        delta = nn.Parameter(torch.zeros_like(og_hidden), requires_grad=True)
+        optimizer = torch.optim.Adagrad([delta], lr=eta)
+
+        with torch.enable_grad():
+            for _ in range(n_steps):
+                optimizer.zero_grad()
+                loss = self.guidance_loss(og_hidden, og_logits, attention_mask, delta)
+                loss.backward()
+                optimizer.step()
+
+        h_new = og_hidden + delta.data
+        new_logits = self.memdlm_lm(h_new)
+        return new_logits, h_new
+    
+
+    def guided_logits(self, xt, attention_mask, guide_steps):
+        """
+        Run the diffusion model on the current sequence and apply classifier guidance to the
+        hidden states, returning the guided LM-head logits used for scoring/sampling this step.
+        """
+        with torch.no_grad():
+            hidden_states = self.diffusion(xt, attention_mask, return_hidden=True)
+            hidden_states = hidden_states.unsqueeze(0) if hidden_states.ndim != 3 else hidden_states
+            base_logits = self.memdlm_lm(hidden_states)
+
+        logits, _ = self.optimized_sampling(base_logits, hidden_states, attention_mask, guide_steps)
+        return logits
+
+
+    def sample_guidance(self, tokens, guide_steps, diffusion_steps,
+                        kappa_fn=lambda t: t, eta=1.0, alpha=1.0):
+        """
+        Confidence-based progressive-unmasking denoising (as in the unconditional sampler),
+        but each step uses classifier-guided logits instead of the raw model logits.
+
+        Args:
+            kappa_fn: unmasking schedule, kappa(t) in [0,1]; fraction of positions committed by step t.
+            eta:      re-masking scale applied to already-unmasked candidate positions.
+            alpha:    blends token log-prob (alpha=1) and negative entropy (alpha=0) in the score.
+        """
+        tau = self.config.guidance.sampling_temperature
+
+        xt = tokens['input_ids'].to(self.device)
+        attention_mask = torch.ones_like(xt).to(self.device)
+
+        dt = 1 / diffusion_steps
+        fix_mask = (xt != self.tokenizer.mask_token_id) # would be none for de novo setting
+
+        x0 = xt
+        for i in range(1, diffusion_steps + 1):
+            kappa_t = kappa_fn(i * dt)
+
+            logits = self.guided_logits(xt, attention_mask, guide_steps)
+
+            with torch.no_grad():
+                last_mask = xt == self.tokenizer.mask_token_id          
+                unmask_t = ~last_mask & ~fix_mask                       
+
+                x0, logp = self.stochastic_sample_from_categorical(logits, tau)
+
+                entropy = torch.distributions.Categorical(logits=logits).entropy()
+                score = alpha * logp + (1 - alpha) * -entropy
+                score = score.masked_fill(fix_mask, float('inf'))  # never remask fixed tokens
+                score[unmask_t] = score[unmask_t] * eta
+
+                num_to_mask = ((~fix_mask).sum(1, keepdim=True).float() * (1 - kappa_t)).long()
+                lowest_k_mask = self.topk_lowest_masking(score, num_to_mask)
+
+                xt[lowest_k_mask] = self.tokenizer.mask_token_id  # remask lowest-confidence
+                mask_2_x0 = last_mask & ~lowest_k_mask  # commit newly-confident positions
+                xt[mask_2_x0] = x0[mask_2_x0]
+
+        xt[xt == self.tokenizer.mask_token_id] = x0[xt == self.tokenizer.mask_token_id]  # commit remaining
+
+        seq = xt.squeeze()[1:-1]
+        optim_tokens = self.tokenizer.decode(seq, skip_special_tokens=True).replace(" ", "")
+
+        with torch.no_grad():
+            final_hidden = self.diffusion(xt, attention_mask, return_hidden=True)
+            final_hidden = final_hidden.unsqueeze(0) if final_hidden.ndim != 3 else final_hidden
+            final_preds = self.classifier_score(final_hidden, attention_mask)
+
+        return optim_tokens, torch.sigmoid(final_preds).item()

src/sampling/olig_generator.py

@@ -0,0 +1,95 @@
+#!/usr/bin/env python3
+
+import sys
+import os
+
+import random
+import torch
+import pandas as pd
+import numpy as np
+
+from tqdm import tqdm
+from datetime import datetime
+from transformers import AutoTokenizer, AutoModelForMaskedLM
+
+from src.lm.memdlm.diffusion_module import MembraneDiffusion
+from src.sampling.olig_sampler import NOSSampler
+
+from src.utils.generate_utils import calc_blosum_score, calc_ppl
+from src.utils.model_utils import _print
+from src.utils.config_utils import load_config, repo_path
+
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+config = load_config("oligo.yaml")
+
+date = datetime.now().strftime("%Y-%m-%d")
+
+
+
+
+def main():
+    csv_save_path = repo_path('results', 'oligo', config.wandb.name, date)
+    
+    try: os.makedirs(csv_save_path, exist_ok=False)
+    except FileExistsError: pass
+
+    tokenizer = AutoTokenizer.from_pretrained(config.lm.pretrained_evoflow)
+    
+    memdlm = MembraneDiffusion(config).to(device)
+    state_dict = memdlm.get_state_dict(str(repo_path("checkpoints", config.lm.ft_evoflow, "best_model.ckpt")))
+    memdlm.load_state_dict(state_dict)
+    memdlm.eval()
+
+    esm_pth = config.lm.pretrained_esm
+    esm_model = AutoModelForMaskedLM.from_pretrained(esm_pth).to(device)
+    esm_model.eval()
+
+    generator = NOSSampler(config, device, memdlm, esm_model, tokenizer)
+
+    # Determine length from positive controls
+    df = pd.read_csv(str(repo_path('data', 'olig_clf', 'test.csv')))
+    seqs = df['Sequence'].tolist()
+
+
+    generation_results = []
+    for seq in tqdm(seqs, desc=f"Generating sequences: "):
+        seq_res = []
+
+        seq_len = len(seq)
+        tokens = tokenizer(seq, return_tensors='pt')
+
+        gen_seq = ""
+        attempts = 0
+
+        while len(gen_seq) != seq_len and attempts < 3:
+            gen_seq, og_pred, final_pred = generator.sample_guidance(
+                tokens,
+                config.olig_guidance.guide_steps,
+                config.olig_guidance.diffusion_steps
+            )
+            attempts += 1
+
+        if len(gen_seq) != seq_len:
+            esm_ppl, memdlm_ppl = None, None
+        else:
+            esm_ppl = calc_ppl(esm_model, tokenizer, gen_seq, [i for i in range(len(gen_seq))], model_type='esm')
+            memdlm_ppl = calc_ppl(memdlm, tokenizer, gen_seq, [i for i in range(len(gen_seq))], model_type='diffusion')
+            blosum = calc_blosum_score(seq, gen_seq, indices=[i for i in range(len(gen_seq))])
+
+        seq_res.append(seq)
+        seq_res.append(gen_seq)
+        seq_res.append(og_pred)
+        seq_res.append(final_pred)
+        seq_res.append(final_pred - og_pred)
+        seq_res.append(esm_ppl)
+        seq_res.append(memdlm_ppl)
+        seq_res.append(blosum)
+        generation_results.append(seq_res)
+
+    df = pd.DataFrame(generation_results, columns=['Original Sequence', 'Generated Sequence', 'OG Olig Value', 'New Olig Value', 'Olig Increase', 'ESM PPL', 'MeMDLM PPL', 'MemDLM Blosum'])
+    df.to_csv(str(csv_save_path / "seqs_with_ppl.csv"), index=False)
+
+
+if __name__ == "__main__":
+    main()

src/sampling/olig_sampler.py

@@ -0,0 +1,187 @@
+import os
+import torch
+from torch import nn
+import torch.nn.functional as F
+from transformers import AutoModelForMaskedLM, AutoTokenizer
+
+from src.guidance.oligo.oligo_module import OligomerClassifier
+from src.sampling.unconditional_sampler import UnconditionalSampler
+from src.lm.memdlm.diffusion_module import MembraneDiffusion
+from src.utils.model_utils import _print
+from src.utils.config_utils import repo_path
+
+
+
+class NOSSampler:
+    def __init__(self, config, device, mdlm, esm, tokenizer):
+        self.config = config
+        self.device = device
+
+        self.diffusion = mdlm
+        self.memdlm_lm = self.diffusion.model.lm_head
+        self.tokenizer = self.diffusion.tokenizer
+
+        ckpt_path = str(repo_path("checkpoints", config.wandb.name, "best_model.ckpt"))
+        self.classifier_model = OligomerClassifier(config).eval().to(self.device)
+        state_dict = self.classifier_model.get_state_dict(ckpt_path)
+        self.classifier_model.load_state_dict(state_dict)
+
+        self.SPECIAL_TOKEN_IDS = {0, 1, 2, 3, 29}
+
+
+    def sample_from_categorical(self, logits):
+        gumbel_noise = -torch.log(-torch.log(torch.rand_like(logits) + 1e-8) + 1e-8)
+        logits += gumbel_noise
+        log_probs = F.log_softmax(logits, dim=-1)
+        _, tokens = log_probs.max(dim=-1)
+        return tokens, log_probs
+    
+
+    def get_clf_preds(self, hidden_states, attention_mask):
+        """Obtain diffusion model logits and classifier predictions from hidden states"""
+        batch = {"embeds": hidden_states.squeeze(), "attention_mask": attention_mask.squeeze()}
+        preds = self.classifier_model(batch)
+        return self.memdlm_lm(hidden_states), preds
+
+
+    def embed_and_run_clf(self, input_ids, attention_masks):
+        """Get sequence embeddings and classifier model predictions"""
+        outputs = self.esm(input_ids=input_ids, attention_mask=attention_masks)
+        sequence_embeddings = outputs.last_hidden_state.squeeze(0)
+
+        batch = {"embeds": sequence_embeddings, "attention_mask": attention_masks}
+        preds = self.classifier_model(batch)
+        
+        return {
+            "clf_preds": preds.requires_grad_(True), # Enable gradients for backprop
+            "embeds": sequence_embeddings 
+        }
+
+
+    def compute_saliency(self, embeddings, attention_masks):
+        """
+        Compute a saliency map using gradients as defined in LaMBO-2 (https://arxiv.org/pdf/2305.20009)
+        """
+        embeddings = embeddings.detach().requires_grad_(True)
+        batch = {
+            "embeds": embeddings.squeeze(),
+            "attention_mask": attention_masks.squeeze()
+        }
+        out = self.classifier_model(batch).sum() 
+        out.backward(retain_graph=True)
+
+        # Creating the saliency map (Eq.5 in LaMBO-2 paper)
+        grads = embeddings.grad.abs().sum(dim=-1)  # Summation across hidden dim. Abs value for mangitude only
+        saliency = grads.pow(1.0 / self.config.olig_guidance.temperature).clamp(min=self.config.olig_guidance.eps)
+        return saliency.squeeze()
+    
+
+    def determine_edit_positions(self, saliency_map, preds):
+        """
+        Create a one-hot mask that indicates the top-k low-value residue positions.
+        We defind low-value positions as those with high saliency scores and
+        thus a high edit probability.
+        """
+        probabilities = saliency_map.masked_fill(preds >= self.config.olig_guidance.residue_thresh, 0.0) # exclude high-value tokens
+        probabilities = probabilities / probabilities.sum()
+        
+        nonzero = probabilities.count_nonzero().item()
+        seq_len = probabilities.shape[0]
+        topk = max(1, int(seq_len * self.config.olig_guidance.topk_frac))
+        _, topk_edit_pos = torch.topk(probabilities, min(topk, nonzero))
+        
+        mask = torch.zeros_like(probabilities).scatter(0, topk_edit_pos, torch.ones_like(probabilities))
+        return mask.unsqueeze(-1)
+
+
+    def update_logits(self, og_hidden, og_logits, hidden_state_mask, attention_mask, optimizer, delta):
+        """
+        Shift logits distribution towards only high-quality residues by applying the explore-exploit loss.
+        """
+        lamb = self.config.olig_guidance.reg_strength
+        
+        h_current = og_hidden + hidden_state_mask * delta
+        new_logits, v_ht_prime = self.get_clf_preds(h_current, attention_mask)
+
+        kl = F.kl_div(
+            F.log_softmax(new_logits, dim=-1),
+            F.softmax(og_logits, dim=-1),
+            reduction='sum'
+        )
+        
+        loss = lamb * kl - v_ht_prime.sum()
+        loss.backward(retain_graph=True)
+        optimizer.step()
+        optimizer.zero_grad()
+        
+        return delta
+    
+
+    def optimized_sampling(self, og_logits, og_hidden, attention_mask, n_steps):
+        """Main entry point to optimize a generated sequence."""
+        eta = self.config.olig_guidance.step_size
+
+        # Calculate initial clf predictions
+        batch = {"embeds": og_hidden, "attention_mask": attention_mask}
+        preds = self.classifier_model(batch)
+        
+        delta = nn.Parameter(torch.zeros_like(og_hidden), requires_grad=True)
+        optimizer = torch.optim.Adagrad([delta], lr=eta)
+        optimizer.zero_grad()
+        
+        with torch.enable_grad():
+            for n in range(n_steps):
+
+                # Compute saliency map and edit positions using updated hidden states
+                saliency_map = self.compute_saliency(og_hidden + delta.data, attention_mask)
+                
+                # One-hot mask that encodes the saliency-selected edit positions
+                mask = self.determine_edit_positions(saliency_map, preds)
+
+                # Optimize and generate the new sequence
+                delta = self.update_logits(
+                    og_hidden=og_hidden,
+                    og_logits=og_logits,
+                    hidden_state_mask=mask,
+                    attention_mask=attention_mask,
+                    optimizer=optimizer,
+                    delta=delta
+                )
+
+        h_new = og_hidden + delta.data
+        new_logits, _ = self.get_clf_preds(h_new, attention_mask)
+        
+        return new_logits, h_new
+    
+
+    def sample_guidance(self, tokens, guide_steps, diffusion_steps):
+        x = tokens['input_ids'].to(self.device)
+        attention_mask = tokens['attention_mask'].to(self.device)
+
+        og_pred = self.classifier_model({"input_ids": x, "attention_mask": attention_mask,})
+
+        for _ in range(diffusion_steps):
+            hidden_states = self.diffusion(x, attention_mask, return_hidden=True)
+            logits = self.memdlm_lm(hidden_states)
+            hidden_states = hidden_states.unsqueeze(0) if hidden_states.ndim != 3 else hidden_states
+
+            logits, hidden_states = self.optimized_sampling(logits, hidden_states, attention_mask, guide_steps)
+            logits = self.memdlm_lm(hidden_states)
+            x, _ = self.sample_from_categorical(logits)
+
+        seq = x.squeeze()
+        # _print(seq)
+        # start = 1 if seq[0].item() in self.SPECIAL_TOKEN_IDS else 0
+        # end = -1 if seq[-1].item() in self.SPECIAL_TOKEN_IDS else len(seq)
+        # seq = seq[start:end]
+        # _print(seq)
+
+        _print(seq)
+        seq = seq[1:-1]
+        _print(seq)
+
+
+        optim_tokens = self.tokenizer.decode(seq, skip_special_tokens=True).replace(" ", "")
+        final_pred = self.classifier_model({"embeds": hidden_states, "attention_mask": attention_mask,})
+
+        return optim_tokens, F.sigmoid(og_pred).item(), F.sigmoid(final_pred).item()

src/sampling/pet_generator.py

@@ -0,0 +1,121 @@
+#!/usr/bin/env python3
+
+
+import os
+import torch
+import pandas as pd
+from tqdm import tqdm
+from datetime import datetime
+from omegaconf import OmegaConf
+from transformers import AutoTokenizer, AutoModelForMaskedLM
+
+from src.lm.memdlm.diffusion_module import MembraneDiffusion
+from src.sampling.pet_sampler import PETSampler
+from src.utils.generate_utils import (
+    mask_for_scaffold,
+    calc_blosum_score,
+    calc_ppl,
+    calc_tm_enrich
+)
+
+from src.utils.model_utils import _print
+from src.utils.config_utils import load_config, repo_path
+
+
+config = load_config("solubility.yaml")
+task = config.guidance.get("task", "solubilize")
+
+results_root = repo_path("results", "heme", config.lm.ft_evoflow)
+todays_date = datetime.today().strftime('%Y-%m-%d')
+
+if config.guidance.prior == 'boltzmann':
+    csv_save_path = results_root / task / todays_date / f"{config.guidance.prior}-t={config.guidance.boltzmann_t}_p={config.guidance.top_p}_tau={config.guidance.sampling_t}"
+elif config.guidance.prior == 'lm_probs':
+    csv_save_path = results_root / task / todays_date / f"{config.guidance.prior}_p={config.guidance.top_p}_tau={config.guidance.sampling_t}"
+
+try: os.makedirs(csv_save_path, exist_ok=False)
+except FileExistsError: pass
+
+
+def main():
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+    tokenizer = AutoTokenizer.from_pretrained(config.lm.pretrained_esm)
+    esm_model = AutoModelForMaskedLM.from_pretrained(config.lm.pretrained_esm).eval().to(device)
+
+    diffusion = MembraneDiffusion(config).to(device)
+    state_dict = diffusion.get_state_dict(str(repo_path("checkpoints", config.lm.ft_evoflow, "best_model.ckpt")))
+    diffusion.load_state_dict(state_dict)
+    diffusion.eval().to(device)
+
+    sampler = PETSampler(config, esm_model, tokenizer, diffusion, device)
+
+    # Update this path to your input CSV of scaffold sequences (uppercase = soluble, lowercase = TM).
+    df = pd.read_csv(str(repo_path("results", "heme", "cybtx.csv")))
+    sequences = df['Sequence'].tolist()
+
+    gen_seqs, ppls, blosums, og_tms, gen_tms, delta_tms = [], [], [], [], [], []
+
+
+    for seq in tqdm(sequences, desc='Infilling Sequences'):
+        masked_seq = mask_for_scaffold(seq, generate_type='lowercase', mask_token='<mask>')
+        tokens = tokenizer(masked_seq, return_tensors='pt')
+        input_ids, attn_masks = tokens['input_ids'].to(device), tokens['attention_mask'].to(device)
+        
+        tm_idxs = [i for i in range(len(seq)) if seq[i].islower()] # keep same for blosum calcs, etc
+        soluble_idxs = [i + 1 for i in range(len(seq)) if seq[i].isupper()] # 1 offset for cls token
+
+        infilled_tokens = sampler.optimize_sequence(
+            input_ids=input_ids,
+            attn_masks=attn_masks,
+            soluble_indices=soluble_idxs,
+        )
+        infilled_seq = tokenizer.decode(infilled_tokens).replace(" ", "")[5:-5]
+        
+        try:
+            bl = calc_blosum_score(seq.upper(), infilled_seq, tm_idxs)
+        except:
+            bl = float('inf')
+        
+        try:
+            ppl = calc_ppl(esm_model, tokenizer, infilled_seq, [i for i in range(len(seq))], model_type='esm')
+        except:
+            ppl = float('inf')
+
+        try:
+            og_tm, gen_tm, delta_tm = calc_tm_enrich(seq.upper(), infilled_seq, tm_idxs)
+        except:
+            og_tm, gen_tm, delta_tm = float('inf'), float('inf'), float('inf')
+
+        gen_seqs.append(infilled_seq)
+        ppls.append(ppl)
+        blosums.append(bl)
+        og_tms.append(og_tm)
+        gen_tms.append(gen_tm)
+        delta_tms.append(delta_tm)
+
+        _print(seq)
+        _print(infilled_seq)
+        _print(ppl)
+        _print(bl)
+        _print(og_tm)   
+        _print(gen_tm)
+        _print(delta_tm)
+        _print('\n')
+
+
+    df['MeMDLM Sequence'] = gen_seqs
+    df['MeMDLM PPL'] = ppls
+    df['MeMDLM BLOSUM'] = blosums
+    df['OG TM Enrichment'] = og_tms
+    df['MeMDLM TM Enrichment'] = gen_tms
+    df['Delta TM Enrichment'] = delta_tms
+
+    _print(df)
+    df.to_csv(csv_save_path / "infilled_seqs.csv", index=False)
+
+
+    
+if __name__ == "__main__":
+    main()
+    

src/sampling/{guided_sampler.py → pet_sampler.py}

@@ -4,11 +4,12 @@ import torch
 import torch.nn.functional as F
 
 from src.utils.model_utils import _print
-from src.guidance.solubility_module import SolubilityClassifier
+from src.utils.config_utils import repo_path
+from src.guidance.solubility.solubility_module import SolubilityClassifier
 from src.sampling.unconditional_sampler import UnconditionalSampler
 
 
-class GuidedSampler:
+class PETSampler:
     def __init__(self, config, esm_model, tokenizer, diffusion, device):
         self.config = config
         self.device = device
@@ -18,7 +19,7 @@ class GuidedSampler:
         self.tokenizer = tokenizer
         self.uncond_generator = UnconditionalSampler(self.tokenizer, self.memdlm)
 
-        ckpt_path = os.path.join(f"/home/a03-sgoel/MeMDLM_v2/checkpoints/{config.wandb.name}/best_model.ckpt")
+        ckpt_path = str(repo_path("checkpoints", config.wandb.name, "best_model.ckpt"))
         self.classifier_model = SolubilityClassifier(config)
         state_dict = self.classifier_model.get_state_dict(ckpt_path)
         self.classifier_model.load_state_dict(state_dict)
@@ -31,7 +32,7 @@ class GuidedSampler:
         self.saliency_t = self.config.guidance.saliency_t
         self.sampling_t = self.config.guidance.sampling_t
         self.boltzmann_t = self.config.guidance.boltzmann_t
-    
+        self.task = self.config.guidance.get("task", "solubilize")
 
     def embed_sequence(self, input_ids, attention_masks):
         with torch.no_grad():
@@ -107,7 +108,13 @@ class GuidedSampler:
             bias[other_idxs] = 0.0
 
             sol_scores = torch.sigmoid(solubility_logits)
-            token_bias = sol_scores.unsqueeze(-1) * bias
+            task = self.task
+            if task == "desolubilize":
+                # At soluble positions, bias toward hydrophobic / TM-like residues.
+                token_bias = sol_scores.unsqueeze(-1) * (-bias)
+            else:
+                # At insoluble positions, bias toward hydrophilic residues.
+                token_bias = (1.0 - sol_scores).unsqueeze(-1) * bias
 
             lm_probs = F.softmax(logits_prior / self.sampling_t, dim=-1)
             boltz_weight = torch.exp(token_bias / self.boltzmann_t)
@@ -143,6 +150,10 @@ class GuidedSampler:
 
         # Initialize a mask to store the editable token positions
         edit_mask = torch.ones(seq_len, dtype=torch.bool, device=self.device)
+        
+        # ignore <cls> and <eos> tokens
+        edit_mask[0] = False
+        edit_mask[-1] = False
 
         # Check for any provided soluble residues, otherwise use classifier preds
         if len(soluble_indices) > 0:
@@ -151,10 +162,12 @@ class GuidedSampler:
             solubility_preds = F.sigmoid(solubility_logits)
             edit_mask[solubility_preds > 0.5] = False
 
-        # Find additional TM residues
-        num_conserved = max(1, int(0.1 * edit_mask.sum()))
-        _, topk_idxs = torch.topk(saliency_map, num_conserved)
-        edit_mask[topk_idxs] = False
+        # during solubilization, we also fix the high-saliency TM residues
+        # but for desolubilization, skip this as high-saliency positions are the soluble sites we want to edit
+        if self.task != "desolubilize":
+            num_conserved = max(1, int(0.1 * edit_mask.sum()))
+            _, topk_idxs = torch.topk(saliency_map, num_conserved)
+            edit_mask[topk_idxs] = False
 
         edit_idxs = edit_mask.nonzero(as_tuple=True)[0]
         return edit_idxs
@@ -207,6 +220,7 @@ class GuidedSampler:
         Compute the log probs of the "new" (optimized) token.
         """
         w = torch.sigmoid(saliency_weight * self.alpha)  # Between [0, 1] to ensure valid probs
+        _print(f'w: {w}')
         p_lm = torch.exp(logp_lm)
         p_prior = torch.exp(logp_prior)
         mixed_probs = (1 - w) * p_lm + w * p_prior
@@ -225,8 +239,12 @@ class GuidedSampler:
     def optimize_sequence(self, input_ids, attn_masks, soluble_indices):
         _print(f'soluble idx: {soluble_indices}')
 
+        scaffold_ids = input_ids.squeeze()
+
         # Initialize token ids, logits, and log probs of sequence
         x0, logp_lm, logits_prior = self.denoise_sequence(input_ids, attn_masks)
+        # x0 is resampled from LM logits at all positions, so restore scaffold from input
+        x0[soluble_indices] = scaffold_ids[soluble_indices]
         _print(f'og tokens: {x0}')
         _print(f'og tokens: {x0.shape}')
         _print(f'og log probs: {logp_lm.shape}')
@@ -281,11 +299,12 @@ class GuidedSampler:
         # Sample new tokens        
         x0_prime = torch.distributions.Categorical(logits=logp_lm).sample()
     
-        # Check if any soluble residues have been changed
-        self.check_scaffold(x0, x0_prime, soluble_indices)
+        # Check if any scaffold residues have been changed
+        self.check_scaffold(scaffold_ids, x0_prime, soluble_indices)
 
-        # Preserve the initial sequence scaffold by copying over the soluble tokens
-        x0_prime[soluble_indices] = x0[soluble_indices]
-        self.check_scaffold(x0, x0_prime, soluble_indices)
+        # Preserve scaffold from the _original_ tokenized input, not resampled x0
+        # since the categorical sampling might not sample the original scaffold token
+        x0_prime[soluble_indices] = scaffold_ids[soluble_indices]
+        self.check_scaffold(scaffold_ids, x0_prime, soluble_indices)
 
         return x0_prime

src/sampling/unconditional_generator.py

@@ -14,15 +14,15 @@ from omegaconf import OmegaConf
 from datetime import datetime
 from transformers import AutoTokenizer, AutoModelForMaskedLM
 
-from MeMDLM_v2.src.lm.diffusion_module import MembraneFlow
+from src.lm.memdlm.diffusion_module import MembraneDiffusion
 from src.sampling.unconditional_sampler import UnconditionalSampler
 from src.utils.generate_utils import mask_for_de_novo, calc_ppl
 from src.utils.model_utils import _print
+from src.utils.config_utils import load_config, repo_path
 
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-os.chdir('/home/a03-sgoel/MeMDLM_v2')
-config = OmegaConf.load("./src/configs/lm.yaml")
+config = load_config("lm.yaml")
 
 date = datetime.now().strftime("%Y-%m-%d")
 
@@ -42,24 +42,23 @@ def generate_sequence(prior: str, tokenizer, generator, device):
 
 
 def main():
-    csv_save_path = f'./results/denovo/unconditional/{config.wandb.name}/{date}_tau=3.0_test-set_distribution'
+    csv_save_path = repo_path('results', 'multipass', config.wandb.name, date)
     
     try: os.makedirs(csv_save_path, exist_ok=False)
     except FileExistsError: pass
 
-
     tokenizer = AutoTokenizer.from_pretrained(config.lm.pretrained_evoflow)
     
-    flow = MembraneFlow(config).to(device)
-    state_dict = flow.get_state_dict(f"./checkpoints/{config.wandb.name}/best_model.ckpt")
-    flow.load_state_dict(state_dict)
-    flow.eval()
+    memdlm = MembraneDiffusion(config).to(device)
+    state_dict = memdlm.get_state_dict(str(repo_path("checkpoints", config.wandb.name, "best_model.ckpt")))
+    memdlm.load_state_dict(state_dict)
+    memdlm.eval()
 
     esm_pth = config.lm.pretrained_esm
     esm_model = AutoModelForMaskedLM.from_pretrained(esm_pth).to(device)
     esm_model.eval()
 
-    generator = UnconditionalSampler(tokenizer, flow)
+    generator = UnconditionalSampler(tokenizer, memdlm)
 
     # # Get 100 random sequence lengths to generate
     # seq_lengths = [random.randint(50, 250) for _ in range(5000)]
@@ -69,14 +68,16 @@ def main():
     # seq_lengths = [len(seq) for seq in df['Sequence'].tolist() for _ in range(500)] # generate each length 100 times
     # _print(seq_lengths)
 
-    # Determine lengths from test set distribution
-    df = pd.read_csv("./data/test.csv")
-    seq_lengths = [len(seq) for seq in df['Sequence'].tolist()]
-    length_counts = Counter(seq_lengths) # {L1: freq, L2: freq, ...}
-    total = sum(length_counts.values()) # total number of tokens
-    lengths = np.array(list(length_counts.keys())) # Frequency of each length
-    probs = np.array([length_counts[l] / total for l in lengths])
-    seq_lengths = np.random.choice(lengths, size=len(seq_lengths), p=probs)
+    # # Determine lengths from test set distribution
+    # df = pd.read_csv("./data/test.csv")
+    # seq_lengths = [len(seq) for seq in df['Sequence'].tolist()]
+    # length_counts = Counter(seq_lengths) # {L1: freq, L2: freq, ...}
+    # total = sum(length_counts.values()) # total number of tokens
+    # lengths = np.array(list(length_counts.keys())) # Frequency of each length
+    # probs = np.array([length_counts[l] / total for l in lengths])
+    # seq_lengths = np.random.choice(lengths, size=len(seq_lengths), p=probs)
+
+    seq_lengths = [x for x in range(60, 161) for _ in range(3)]
 
     generation_results = []
     for seq_len in tqdm(seq_lengths, desc=f"Generating sequences: "):
@@ -91,23 +92,23 @@ def main():
             attempts += 1
 
         if len(gen_seq) != seq_len:
-            esm_ppl, flow_ppl = None, None
+            esm_ppl, memdlm_ppl = None, None
         else:
             esm_ppl = calc_ppl(esm_model, tokenizer, gen_seq, [i for i in range(len(gen_seq))], model_type='esm')
-            flow_ppl = calc_ppl(flow, tokenizer, gen_seq, [i for i in range(len(gen_seq))], model_type='flow')
+            memdlm_ppl = calc_ppl(memdlm, tokenizer, gen_seq, [i for i in range(len(gen_seq))], model_type='diffusion')
         
         _print(f'gen seq: {gen_seq}')
         _print(f'esm ppl: {esm_ppl}')
-        _print(f'flow ppl: {flow_ppl}')
+        _print(f'memdlm ppl: {memdlm_ppl}')
         
         seq_res.append(gen_seq)
         seq_res.append(esm_ppl)
-        seq_res.append(flow_ppl)
+        seq_res.append(memdlm_ppl)
 
         generation_results.append(seq_res)
 
-    df = pd.DataFrame(generation_results, columns=['Generated Sequence', 'ESM PPL', 'Flow PPL'])
-    df.to_csv(csv_save_path + "/seqs_with_ppl.csv", index=False)
+    df = pd.DataFrame(generation_results, columns=['Generated Sequence', 'ESM PPL', 'MeMDLM PPL'])
+    df.to_csv(str(csv_save_path / "seqs_with_ppl.csv"), index=False)
 
 
 if __name__ == "__main__":

src/sampling/unconditional_sampler.py

@@ -15,7 +15,7 @@ class UnconditionalSampler:
         self.seed_everything(seed=42)
 
     @torch.inference_mode()
-    def sample_unconditional(self, xt, num_steps, tau=0.7, kappa_fn=lambda t: t, eta=1, alpha=1., banned_token_ids=None, return_logits=None): 
+    def sample_unconditional(self, xt, num_steps, tau=1.0, kappa_fn=lambda t: t, eta=1, alpha=1., banned_token_ids=None, return_logits=None): 
         """
         Stochastic remasking sampling method for iterative refinement of sequences.
 

src/utils/config_utils.py

@@ -0,0 +1,15 @@
+from pathlib import Path
+from omegaconf import OmegaConf
+
+# Repository root (MeMDLM_v2/)
+REPO_ROOT = Path(__file__).resolve().parents[2]
+
+
+def load_config(config_name: str):
+    """Load a YAML config from src/configs/."""
+    return OmegaConf.load(REPO_ROOT / "src" / "configs" / config_name)
+
+
+def repo_path(*parts: str) -> Path:
+    """Build an absolute path relative to the repository root."""
+    return REPO_ROOT.joinpath(*parts)

src/utils/generate_utils.py

@@ -9,13 +9,13 @@ import numpy as np
 from omegaconf import OmegaConf
 from transformers import AutoModelForMaskedLM, AutoModel, AutoTokenizer
 
-from src.lm.memdlm.diffusion_module import MembraneFlow
-from src.lm.dplm.diffusion_module import DPLM
+from src.lm.memdlm.diffusion_module import MembraneDiffusion
 from src.utils.model_utils import get_latents, _print
 from src.sampling.unconditional_sampler import UnconditionalSampler
-from src.lm.dplm.unconditional_sampler import UnconditionalSampler as DPLMUnconditionalSampler
 
-config = OmegaConf.load("/home/a03-sgoel/MeMDLM_v2/src/configs/lm.yaml")
+
+config = OmegaConf.load("/scratch/pranamlab/sgoel/MeMDLM_v2/src/configs/lm.yaml")
+
 
 # -------# Masking #-------- #
 def mask_for_de_novo(sequence_length):
@@ -30,7 +30,7 @@ def mask_for_scaffold(sequence, generate_type, mask_token):
 
 
 # -------# Generation #-------- #
-def memflow_infill_uncond(masked_seq, tokenizer, model: MembraneFlow):
+def memflow_infill_uncond(masked_seq, tokenizer, model: MembraneDiffusion):
     generator = UnconditionalSampler(tokenizer, model) # initialize the generator object
     xt = tokenizer(masked_seq, return_tensors='pt')['input_ids'].to(model.device)
     denoised_tokens = generator.sample_unconditional(xt, config.sampling.n_steps)[0].squeeze()
@@ -68,7 +68,10 @@ def evodiff_infill(motif_seq, tokenizer, model, device, batch_size=1):
     return output[0] #if batch_size==1 else output, og_sample, loc
 
 
-def dplm_infill(masked_seq, tokenizer, model: DPLM, device):
+def dplm_infill(masked_seq, tokenizer, model, device):
+    from src.lm.dplm.diffusion_module import DPLM
+    from src.lm.dplm.unconditional_sampler import UnconditionalSampler as DPLMUnconditionalSampler
+    
     generator = DPLMUnconditionalSampler(tokenizer, model)
     xt = tokenizer(masked_seq, return_tensors='pt')['input_ids'].to(model.device)
     denoised_tokens = generator.sample_unconditional(xt, config.sampling.n_steps)[0].squeeze()
@@ -111,7 +114,7 @@ def calc_ppl(model, tokenizer, generated_sequence, mask_token_indices, model_typ
         with torch.no_grad():
             if model_type == 'esm':
                 loss = model(masked_input, labels=labels).loss.item()
-            elif model_type == 'flow':
+            elif model_type == 'diffusion':
                 logits = model.forward(masked_input, attention_mask=attn_mask)
                 loss = F.cross_entropy(
                     logits.view(-1, logits.size(-1)),
@@ -143,6 +146,28 @@ def calc_blosum_score(og_seq, gen_seq, indices):
     return tot_score / len(indices) if indices else 0
 
 
+# define TM enrichment dict
+TM_ENRICHMENT = {
+    'A': -0.06647, 'C':  0.312396, 'D': -1.52276, 'E': -1.7743,  'F':  0.800983,
+    'G':  0.756636, 'H': -0.44366,  'I':  0.551521, 'K': -1.56399,  'L':  0.299315,
+    'M':  0.279242, 'N': -0.54795,  'P':  0.180039, 'Q': -1.16639,  'R': -1.17808,
+    'S':  0.110367, 'T':  0.228464,  'V':  0.510702, 'W':  0.562891,  'Y':  0.084097
+}
+
+def calc_tm_enrich(og_seq, gen_seq, indices):
+    def _score(seq):
+        chars = [seq[i] for i in indices if i < len(seq)] if indices else list(seq)
+        vals = [TM_ENRICHMENT.get(aa.upper()) for aa in chars]
+        vals = [v for v in vals if v is not None]
+        return float(np.mean(vals)) if vals else float('nan')
+    
+    og_score = _score(og_seq)
+    gen_score = _score(gen_seq)
+    delta_score = gen_score - og_score
+    return og_score, gen_score, delta_score
+
+
+
 def calc_cos_sim(original_sequence, generated_sequence, tokenizer, esm_model, device):
     og_embeddings = get_latents(esm_model, tokenizer, original_sequence.upper(), device)
     new_embeddings = get_latents(esm_model, tokenizer, generated_sequence, device)