Complete rename of lectures_02_03 to lecture_3

lectures_02_03/GUIDANCE_NOTES.md

@@ -1,156 +0,0 @@
-# ESM-2 flow matching and diffusion guidance
-
-Two independent teaching scripts, each with encoding, normalization, a small
-conditional generator, a time-conditioned reward predictor, training, three
-guidance examples, and constrained sequence decoding. Neither imports the other.
-
-## Run
-
-Install Python 3.10+ and the dependencies in your own environment:
-
-```sh
-pip install torch transformers==4.57.6
-python esm2_flow_guidance.py --epochs 200 --samples 8
-python esm2_diffusion_guidance.py --epochs 200 --samples 8
-```
-
-Keep both scripts and `esm2_example.csv` together. Their default paths are
-relative to the scripts, so the commands also work with absolute script paths.
-The first run downloads the public `facebook/esm2_t6_8M_UR50D` checkpoint into
-`.esm2_cache`. CUDA is used when available; otherwise the scripts use CPU.
-These examples were executed with PyTorch 2.9.1 and Transformers 4.57.6.
-
-## Motivation
-
-We will generate protein-sequence representations with a flow and a diffusion
-model, then guide generation toward specified properties. Using frozen ESM-2
-residue embeddings, we can compare classifier-free conditioning with
-single-objective and weighted multi-objective reward steering in the same latent
-space. After generation, we decode the latents into amino acids and enforce an
-explicit residue-count constraint. This separates learning the sequence
-distribution, expressing preferences, and guaranteeing a discrete output rule.
-
-## Where the rewards come from
-
-The bundled dataset contains 64 synthetic sequences of length 24. These are
-teaching examples, not natural proteins or experimentally validated peptides.
-
-For sequence s of length L, `composition_proxies()` computes:
-
-- r1 = (number of K and R minus number of D and E) / L.
-  This is a side-chain charge-count proxy, not a complete pH-dependent charge model.
-- r2 = number of residues in `DEHKNQRST` / L.
-  This is a defined polar/charged composition fraction, not measured solubility.
-- The bundled class label c is 1 when r1 > 0, and 0 otherwise.
-
-CSV columns are `sequence,c,r1,r2`. Both classes 0 and 1 must be present. If r1
-and r2 are omitted, the scripts compute the composition proxies directly.
-Supplied r1/r2 columns override the proxy calculation, allowing measured
-objectives or external predictor outputs. Orient each objective so larger is
-better before supplying it. For an undesirable quantity, negate it first.
-
-All sequences in one run must have the same length, between 1 and 128, using
-the 20 canonical amino acids. The default hard minimum is 12, so shorter
-sequences require a smaller `--min-polar` value.
-
-## Normalization and reward prediction
-
-- Preserve one ESM-2 embedding per residue, giving [N, L, 320]. Remove BOS/EOS.
-- Standardize each latent feature using its mean and standard deviation over
-  training sequences and residue positions.
-- Standardize each property separately: r_tilde = (r - mean) / std.
-- Save these statistics. Do not refit them on generated, validation, or test data.
-- Train a small predictor on intermediate latents and their clean-sequence
-  standardized property labels. Its two outputs estimate expected endpoint
-  rewards from the current state and time.
-
-The examples use all 64 sequences for training and do not claim held-out quality.
-For research, split and cluster data before fitting statistics or networks, then
-assess prediction and sequence reconstruction on held-out sequences.
-
-## Three sampling modes in each script
-
-1. CFG: `w=2`, `eta=0`. Combine unconditional and class-conditioned predictions
-   as F_uncond + w * (F_cond - F_uncond). Here w=0 is unconditional and w=1 is
-   ordinary conditional sampling. Drop the class to null index 2 in 20% of
-   training examples. CFG does not use the external reward predictor.
-2. Single objective: `w=0`, `eta=1`, `lambdas=(1,0)`.
-3. Multiple objectives: `w=0`, `eta=1`, `lambdas=(0.7,0.3)`.
-
-The weights are normalized to sum to one. Lambda controls the relative
-tradeoff in standardized property units; eta controls overall steering strength.
-Set both w and eta positive to combine CFG and reward steering. Edit the three
-calls in `main()` to change the demonstration settings.
-
-Flow training uses Z_0=noise, Z_1=data, a straight conditional path, and target
-velocity Z_1-Z_0. Sampling integrates from t=0 to t=1 with Euler steps. Reward
-steering adds kappa(t) times the reward gradient to the velocity, with the chosen
-schedule kappa(t)=4*eta*t*(1-t). This is heuristic velocity steering, not a claim
-of exact sampling from a reward-tilted density.
-
-DDPM training uses Z_0=data and predicts the Gaussian noise used to construct
-Z_k. Sampling runs k=1000 down to 1. The small noise predictor includes the
-Gaussian-reference skip sqrt(1-alpha_bar_k)*Z_k and learns an additive correction
-scaled by sqrt(alpha_bar_k). This is a parameterization of the noise predictor;
-the target and DDPM equations remain noise-prediction equations. It lets the
-small network pass through high-dimensional noise without reconstructing every
-coordinate through its narrow hidden layer.
-
-For DDPM steering, epsilon_guided = epsilon_CFG - eta *
-sqrt(1-alpha_bar_k) * gradient(R_lambda). This follows the score-to-noise
-conversion s = -epsilon / sqrt(1-alpha_bar_k). The posterior standard deviation
-used for the reverse random increment is a different quantity.
-
-In both scripts, gradient ascent uses a learned expected-reward predictor.
-It is not the exact log conditional likelihood or log exponential-reward
-expectation needed for exact conditional or reward-tilted sampling. The
-generator and reward weights are frozen during sampling; gradients are enabled
-only for the current latent. No gradients through a discrete sequence scorer
-are required. A research implementation must validate surrogate quality and
-recheck the true objectives after decoding.
-
-## One decoder at the end
-
-The identical `decode()` function appears in both files to keep each script
-standalone. Both samplers use this same final decoding operation.
-
-1. Undo latent standardization.
-2. Apply ESM-2's frozen language-model head and retain the 20 amino-acid logits.
-3. Take the highest-logit residue at each position.
-4. If fewer than M positions contain a residue in `DEHKNQRST`, replace exactly
-   the missing number using the lowest logit-cost changes to that set.
-
-For the default M=12, every decoded 24-residue output contains at least 12
-members of the selected set. `--min-polar 0` removes the constraint. The procedure
-maximizes the sum of the fixed per-position logits subject to this minimum
-count. It does not enforce the constraint along the latent trajectory, preserve
-an exact conditioned generative distribution, or guarantee physical solubility.
-
-The language-model head is a simple available decoder, not a mathematically
-exact inverse of ESM-2. Generated latents can be off the encoding manifold.
-For serious protein generation, validate decoding and consider training a
-dedicated sequence decoder. Property improvements in the surrogate need not
-survive discretization or the constrained substitutions.
-
-## Outputs and checks
-
-Each script writes `cfg.fasta`, `single.fasta`, `multi.fasta`, and `results.pt`
-to its own output directory. The tensor file contains standardized generated
-latents, both networks' state dictionaries, normalization statistics, ESM model
-name, latent shape, and constraint settings. The terminal prints example
-sequences and re-evaluated mean composition proxies.
-
-Both examples have been executed through actual ESM-2 encoding, 200 training
-epochs, all three guidance modes, and final decoding. Checks included finite
-outputs, reward input gradients, normalization, schedule indexing, and the
-hard constraint. For small test cases, constrained decoding was compared with
-exhaustive search over polar/nonpolar assignments. These checks establish code
-mechanics, not biological validity or reliable property optimization.
-
-## Primary references
-
-- ESM model: https://huggingface.co/facebook/esm2_t6_8M_UR50D
-- ESM implementation: https://huggingface.co/docs/transformers/model_doc/esm
-- Flow Matching: https://arxiv.org/abs/2210.02747
-- DDPM: https://arxiv.org/abs/2006.11239
-- Classifier-Free Diffusion Guidance: https://arxiv.org/abs/2207.12598

lectures_02_03/README.md

@@ -1,310 +0,0 @@
-# Lectures 2 and 3 · Continuous Generative Models
-
-These examples accompany Lectures 2 and 3 of CIS 6270. We implement flow matching and DDPM diffusion using frozen ESM-2
-residue embeddings, with shared property definitions and normalization to
-compare guidance across the two sampling processes.
-
-Each example includes data loading, latent and property normalization, a small
-conditional model, training, classifier-free guidance, single-objective reward
-steering, weighted multi-objective steering, and final sequence decoding with a
-hard residue-count constraint.
-
-## Examples
-
-| Example | What the model learns | Sampling |
-| --- | --- | --- |
-| [Flow matching](esm2_flow_guidance.py) | A velocity field between noise and clean ESM-2 latents | Euler integration from noise to data |
-| [Diffusion](esm2_diffusion_guidance.py) | The noise added to clean ESM-2 latents | A 1,000-step DDPM reverse chain |
-| [Mathematical and implementation notes](GUIDANCE_NOTES.md) | Normalization, guidance conventions, decoding, and limitations | Companion reading for both scripts |
-
-Each script contains the complete training and sampling implementation,
-including the final sequence decoder, and can run independently.
-
-## Quick start
-
-Use **Python 3.11**, or another compatible Python version at least 3.10. A fresh
-virtual environment is recommended. The dependencies are pinned to the versions
-used to test these examples: PyTorch 2.9.1 and Transformers 4.57.6.
-
-### 1. Download and install
-
-```bash
-git clone https://huggingface.co/ChatterjeeLab/CIS6270
-cd CIS6270
-
-python3 -m venv .venv
-source .venv/bin/activate
-python -m pip install --upgrade pip
-python -m pip install -r requirements.txt
-```
-
-On Windows PowerShell, use `python -m venv .venv` and activate with
-`.venv\Scripts\Activate.ps1`.
-
-### 2. Train and sample
-
-Run either example, or both:
-
-```bash
-python lectures_02_03/esm2_flow_guidance.py --epochs 200 --samples 8
-python lectures_02_03/esm2_diffusion_guidance.py --epochs 200 --samples 8
-```
-
-Each command trains a generator and a property predictor from scratch, then
-generates sequences using three guidance settings. The first run downloads the
-public [ESM-2 8M checkpoint](https://huggingface.co/facebook/esm2_t6_8M_UR50D)
-into `lectures_02_03/.esm2_cache/`, which subsequent runs reuse. Both the repository
-and checkpoint are publicly accessible.
-
-The scripts use CUDA when available and CPU otherwise. Apple Silicon currently
-uses the CPU path. Runtime depends on hardware and the first-run download.
-
-For a short end-to-end installation check:
-
-```bash
-python lectures_02_03/esm2_flow_guidance.py --epochs 2 --samples 2 --output outputs/flow_smoke
-python lectures_02_03/esm2_diffusion_guidance.py --epochs 2 --samples 2 --output outputs/diffusion_smoke
-```
-
-The two-epoch runs exercise encoding, training, sampling, and decoding.
-
-### 3. Find the generated sequences
-
-By default, outputs are written beside the scripts:
-
-```text
-lectures_02_03/
-├── esm2_flow_outputs/
-│   ├── cfg.fasta
-│   ├── single.fasta
-│   ├── multi.fasta
-│   └── results.pt
-└── esm2_diffusion_outputs/
-    ├── cfg.fasta
-    ├── single.fasta
-    ├── multi.fasta
-    └── results.pt
-```
-
-The FASTA files contain decoded sequences. `results.pt` contains generated
-standardized latents, generator and reward-predictor state dictionaries,
-normalization statistics, the ESM checkpoint name, latent dimensions, and
-constraint settings. The terminal also prints re-evaluated composition
-properties of the decoded sequences.
-
-Each run initializes new generator and reward-predictor parameters. Repeating
-a command with the same output directory replaces its saved files; choose a
-new `--output` directory to retain results from separate experiments.
-
-## Implementation
-
-1. **Encode sequences.** Use frozen ESM-2 to obtain one 320-dimensional vector
-   per residue, retaining position information and removing BOS/EOS tokens.
-2. **Normalize.** Standardize latent coordinates and each property using
-   training-set statistics.
-3. **Train.** Learn a conditional generative field and a time-conditioned
-   predictor of the clean sequence's standardized properties.
-4. **Guide generation.** Compare CFG, one reward, and a weighted reward sum.
-5. **Decode once at the end.** Undo latent normalization, apply the frozen ESM-2
-   language-model head, and enforce the selected residue-count constraint.
-6. **Re-evaluate.** Calculate the composition properties on the actual decoded
-   sequences so comparisons between guidance settings reflect the final
-   amino-acid outputs.
-
-## Sequence data and property labels
-
-For the teaching examples, we use [64 synthetic sequences of length
-24](esm2_example.csv). We calculate both property labels directly
-from residue counts and assign the conditioning class according to the first
-property:
-
-| Column | Definition | Interpretation |
-| --- | --- | --- |
-| `sequence` | A canonical amino-acid sequence | Input to frozen ESM-2 |
-| `r1` | `(count(K) + count(R) - count(D) - count(E)) / length` | A simple charge-count proxy |
-| `r2` | `count(residues in DEHKNQRST) / length` | A defined polar/charged fraction |
-| `c` | `1` if `r1 > 0`, otherwise `0` | The binary conditioning class |
-
-### Use your own data
-
-Supply a CSV with `sequence,c,r1,r2` columns:
-
-```csv
-sequence,c,r1,r2
-QYWSDSWWESQMMSPWYPMSPLSV,0,-0.08333333,0.41666667
-CKSEFQPPHLMGHDFFACEMRNFK,0,0.00000000,0.45833333
-WPYGEHMLADNNVVKKRLQQWCFI,1,0.04166667,0.41666667
-KVVGALPIESFYTAKMESIAVEVI,0,-0.04166667,0.33333333
-```
-
-```bash
-python lectures_02_03/esm2_flow_guidance.py --data my_sequences.csv --output outputs/my_flow
-python lectures_02_03/esm2_diffusion_guidance.py --data my_sequences.csv --output outputs/my_diffusion
-```
-
-- Include at least four sequences and both class labels, 0 and 1.
-- Use one fixed sequence length, at most 128 residues, and the 20 canonical
-  amino acids to match the fixed-length latent tensors in both models.
-- Supply finite property values and orient each objective so **larger is better**.
-  For a quantity to minimize, negate it before writing the CSV.
-- If both `r1` and `r2` are omitted, `composition_proxies()` calculates the
-  two example properties directly. Supplied reward columns override them.
-- Experimental measurements or scores from a separate predictor can fill the
-  reward columns. We fit a differentiable surrogate to these labels and
-  evaluate its gradients with respect to the intermediate latent during sampling.
-- The printed composition proxies always retain their defined count-based
-  meaning, even when custom reward labels are supplied. Re-evaluate custom
-  objectives with the corresponding assay or scorer after decoding.
-
-The bundled demonstration uses all 64 sequences for training. For studies with
-held-out evaluation, partition sequences by cluster before fitting normalization
-statistics or model parameters.
-
-## Property normalization and scalarization
-
-Each objective is standardized independently:
-
-```text
-r̃ₘ(s) = [rₘ(s) − μₘ] / max(σₘ, 10⁻⁶)
-```
-
-Here μₘ and σₘ are the training-set mean and standard deviation of property m.
-Standardization expresses each property in units of its training-set variation,
-so the scalarization weights specify relative preferences on a common scale.
-Reuse the saved training statistics for new examples to maintain that scale
-throughout evaluation and generation.
-
-The reward predictor estimates standardized clean-sequence properties from an
-intermediate latent and its time. During sampling, we combine its predictions:
-
-```text
-Rλ(z,t) = λ₁ r̂₁(z,t) + λ₂ r̂₂(z,t)
-λ₁ ≥ 0, λ₂ ≥ 0, λ₁ + λ₂ = 1
-```
-
-The nonnegative tradeoff weights λ are normalized to sum to one. Overall
-steering strength η is a separate parameter.
-
-## Guidance configurations
-
-Both scripts run these settings in `main()`:
-
-| Output | CFG strength `w` | Reward strength `eta` | Property weights `lambdas` |
-| --- | ---: | ---: | --- |
-| `cfg.fasta` | 2.0, class 1 | 0.0 | Unused |
-| `single.fasta` | 0.0 | 1.0 | `(1.0, 0.0)` |
-| `multi.fasta` | 0.0 | 1.0 | `(0.7, 0.3)` |
-
-**Classifier-free guidance.** During training, we replace 20% of class labels
-with null class index 2. At sampling time, we combine the conditional and
-unconditional predictions as `F_uncond + w * (F_cond - F_uncond)`.
-Thus `w=0` is unconditional, `w=1` is ordinary conditional sampling, and `w>1` amplifies the conditional
-difference. F is a velocity for flow matching and predicted noise for DDPM.
-
-**Reward steering.** Freeze both networks, enable gradients only for the
-current latent, calculate the scalarized reward, and differentiate it with
-respect to that latent. Set `(1, 0)` for the first objective alone or change
-the weights to express a tradeoff. This correction uses the property predictor,
-while CFG uses the conditional and unconditional generative predictions.
-
-To change these settings, edit the three `sample(...)` calls in `main()`.
-To combine CFG with reward steering, set both strengths:
-
-```python
-latent = sample(model, reward_model, n=8, c=1,
-                w=2.0, eta=1.0, lambdas=(0.7, 0.3))
-```
-
-Each call resets the sampling seed to compare methods using the same random
-draws for the same batch size.
-
-### Guidance in the sampling dynamics
-
-| | Flow matching | DDPM diffusion |
-| --- | --- | --- |
-| Clean-data endpoint | Z₁ | Z₀ |
-| Training target | Velocity Z₁ − Z₀ | Injected Gaussian noise ε |
-| Generation direction | t = 0 → 1 | k = 1000 → 1 |
-| Reward correction | Add κ(t)∇Rλ to velocity | Subtract η√(1−ᾱₖ)∇Rλ from predicted noise |
-| Numerical step | Euler ODE update | Reverse mean plus posterior Gaussian noise |
-
-For the flow, we choose `κ(t)=4ηt(1−t)` to taper reward steering near the noise
-and data endpoints. For DDPM, we convert the reward-gradient score correction
-into a noise-prediction correction using `s=−ε/√(1−ᾱₖ)`. Both implementations
-use heuristic gradient steering based on learned estimates of endpoint properties.
-
-The DDPM predictor includes a Gaussian-reference skip plus a learned correction
-so the small MLP can carry high-dimensional noise. The training target remains
-noise, and the sampler uses the standard DDPM reverse equations.
-
-## Constrained sequence decoding
-
-The default decoder requires **at least 12 of the 24 output residues** to
-belong to:
-
-```text
-𝒫 = {D, E, H, K, N, Q, R, S, T}
-```
-
-We begin with the highest-logit residue at each position and count the members
-of 𝒫. When the count falls below the specified minimum, we replace the required
-number of residues using the lowest-cost substitutions into 𝒫. This discrete
-decoding procedure maximizes the sum of the fixed per-position logits subject
-to the minimum-count constraint.
-
-```bash
-# Require at least 16 selected residues.
-python lectures_02_03/esm2_flow_guidance.py --min-polar 16
-
-# Remove the minimum-count constraint.
-python lectures_02_03/esm2_diffusion_guidance.py --min-polar 0
-```
-
-Set `--min-polar` between zero and the sequence length to specify the minimum
-number of selected polar or charged residues in each decoded sequence.
-
-## Verification
-
-Run the offline unit checks after installing dependencies:
-
-```bash
-python -m unittest discover -s tests -v
-```
-
-The five unit tests cover dataset annotations, scalarization weights, reward
-input gradients, DDPM schedule indexing, and constrained decoding. For the
-decoder test, we compare the selected sequences with exhaustive solutions on
-small examples. All unit tests run locally with the installed dependencies.
-
-We also checked both scripts end to end using the ESM-2 checkpoint, 200 training
-epochs, and all three guidance modes. All 48 decoded sequences in that run
-satisfied the specified minimum residue count.
-
-## Troubleshooting
-
-- **Download failure:** the first run needs internet access to the ESM-2
-  checkpoint. Cached subsequent runs can use `HF_HUB_OFFLINE=1` once all files
-  have been downloaded.
-- **Dependency conflicts:** use a clean virtual environment and the pinned
-  requirements. For a specific CUDA build, follow the official
-  [PyTorch installation instructions](https://pytorch.org/get-started/locally/).
-- **Input validation error:** check equal lengths, canonical residues, both
-  class labels, finite rewards, and a minimum count no greater than the length.
-- **Weak or repetitive samples:** inspect training convergence, evaluate
-  reconstruction through the ESM-2 head, and assess property-predictor accuracy
-  on held-out sequences before adjusting guidance strength. Compare guidance
-  settings using properties recalculated after sequence decoding.
-
-## References
-
-- [ESM-2 checkpoint](https://huggingface.co/facebook/esm2_t6_8M_UR50D)
-  and [Transformers ESM documentation](https://huggingface.co/docs/transformers/model_doc/esm).
-- [Flow Matching for Generative Modeling](https://arxiv.org/abs/2210.02747).
-- [Denoising Diffusion Probabilistic Models](https://arxiv.org/abs/2006.11239).
-- [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
-
-## License
-
-Repository code is distributed under the [MIT License](https://huggingface.co/ChatterjeeLab/CIS6270/blob/main/LICENSE), matching the
-repository's license setting. ESM-2 weights are downloaded separately and remain
-subject to their original distribution terms.

lectures_02_03/esm2_diffusion_guidance.py

@@ -1,264 +0,0 @@
-"""Standalone ESM-2 diffusion guidance example for CIS 6270.
-Input CSV: sequence,c,r1,r2. Sequences have one fixed length; c is 0 or 1.
-Both objectives are oriented so larger is better. The bundled CSV is synthetic.
-Run: python esm2_diffusion_guidance.py --data esm2_example.csv --epochs 200
-Install: pip install torch transformers==4.57.6
-Outputs: guided residue latents, model weights, and decoded amino-acid sequences.
-"""
-import argparse
-import csv
-from pathlib import Path
-
-import torch
-from torch import nn
-import torch.nn.functional as F
-from torch.utils.data import DataLoader, TensorDataset
-from transformers import AutoTokenizer, EsmForMaskedLM
-
-ROOT = Path(__file__).resolve().parent
-DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-ESM_NAME = "facebook/esm2_t6_8M_UR50D"
-AMINO_ACIDS = "ACDEFGHIKLMNPQRSTVWY"
-POLAR_RESIDUES = "DEHKNQRST"  # Operational polar/charged set; not a solubility assay.
-BATCH_SIZE, HIDDEN, LEARNING_RATE = 16, 128, 1e-3
-CONDITION_DROP = 0.2  # Drop c during training so the same model learns the null condition.
-
-
-def composition_proxies(sequences):
-    # Transparent teaching rewards; these are not measured activity or solubility.
-    return torch.tensor([
-        [(sum(a in "KR" for a in s) - sum(a in "DE" for a in s)) / len(s),
-         sum(a in POLAR_RESIDUES for a in s) / len(s)]
-        for s in sequences
-    ], dtype=torch.float32)
-
-
-# 1. Load annotated sequences and encode frozen ESM-2 residue vectors.
-@torch.no_grad()
-def load_data(path):
-    with Path(path).open(newline="") as handle:
-        rows = list(csv.DictReader(handle))
-    sequences = [row["sequence"].strip().upper() for row in rows]
-    if len(rows) < 4 or any(not s or set(s) - set(AMINO_ACIDS) for s in sequences):
-        raise ValueError("Supply at least four sequences using the 20 standard amino acids")
-    lengths = {len(s) for s in sequences}
-    if len(lengths) != 1 or max(lengths) > 128:
-        raise ValueError("This compact example requires one fixed sequence length, at most 128")
-    c = torch.tensor([int(row["c"]) for row in rows], dtype=torch.long)
-    # With sequence,c only, compute example rewards directly from residue counts.
-    # Optional r1/r2 columns override them with supplied labels, such as measurements.
-    if {"r1", "r2"}.issubset(rows[0]):
-        r = torch.tensor([[float(row["r1"]), float(row["r2"])] for row in rows])
-    else:
-        r = composition_proxies(sequences)
-    if set(c.tolist()) != {0, 1} or not torch.isfinite(r).all():
-        raise ValueError("Include both c=0 and c=1, with finite r1 and r2")
-
-    tokenizer = AutoTokenizer.from_pretrained(ESM_NAME, cache_dir=ROOT / ".esm2_cache")
-    esm = EsmForMaskedLM.from_pretrained(
-        ESM_NAME, cache_dir=ROOT / ".esm2_cache", use_safetensors=True
-    ).to(DEVICE).eval().requires_grad_(False)
-    encoded = []
-    for start in range(0, len(sequences), BATCH_SIZE):
-        tokens = tokenizer(sequences[start:start + BATCH_SIZE], return_tensors="pt")
-        tokens = {key: value.to(DEVICE) for key, value in tokens.items()}
-        hidden = esm.esm(**tokens).last_hidden_state
-        encoded.append(hidden[:, 1:-1].cpu())  # Remove BOS/EOS; retain all L residue positions.
-    z = torch.cat(encoded)                    # [N, L, D], with D=320 for this checkpoint.
-    z_mean = z.mean((0, 1), keepdim=True)
-    z_std = z.std((0, 1), correction=0, keepdim=True).clamp_min(1e-4)
-    r_mean, r_std = r.mean(0), r.std(0, correction=0).clamp_min(1e-6)
-    dataset = TensorDataset((z - z_mean) / z_std, c, (r - r_mean) / r_std)
-    stats = {"z_mean": z_mean, "z_std": z_std, "r_mean": r_mean, "r_std": r_std}
-    return dataset, esm, tokenizer, stats
-
-
-# 2. Small model classes: flattening lets each output depend on the entire sequence.
-class DiffusionModel(nn.Module):
-    def __init__(self, length, dim):
-        super().__init__()
-        self.length, self.dim = length, dim
-        self.time = nn.Sequential(nn.Linear(1, 32), nn.SiLU(), nn.Linear(32, 32))
-        self.condition = nn.Embedding(3, 16)    # Indices 0,1 are classes; 2 is the null class.
-        self.net = nn.Sequential(
-            nn.Linear(length * dim + 48, HIDDEN), nn.SiLU(),
-            nn.Linear(HIDDEN, HIDDEN), nn.SiLU(),
-            nn.Linear(HIDDEN, length * dim),
-        )
-
-    def forward(self, z, t, c):
-        time = self.time(t[:, None])
-        inputs = torch.cat([z.flatten(1), time, self.condition(c)], dim=1)
-        k = (t * K).round().long().clamp(0, K)
-        a = alpha_bars[k, None, None]
-        return (1 - a).sqrt() * z + a.sqrt() * self.net(inputs).reshape_as(z)
-        # Gaussian-reference noise prediction plus a learned correction.
-        # The skip carries all coordinates; the correction vanishes near pure noise.
-
-
-class RewardModel(nn.Module):
-    def __init__(self, length, dim):
-        super().__init__()
-        self.net = nn.Sequential(
-            nn.Linear(length * dim + 1, HIDDEN), nn.SiLU(),
-            nn.Linear(HIDDEN, HIDDEN), nn.SiLU(), nn.Linear(HIDDEN, 2),
-        )
-
-    def forward(self, z, t):
-        return self.net(torch.cat([z.flatten(1), t[:, None]], dim=1))
-        # Two predicted standardized endpoint objectives, given the current latent and time.
-
-
-# The same DDPM schedule as Section 7; k=0 denotes the clean latent.
-K = 1000
-betas = torch.cat([torch.zeros(1), torch.linspace(1e-4, 0.02, K)]).to(DEVICE)
-alphas = 1.0 - betas
-alpha_bars = alphas.cumprod(0)
-previous = torch.cat([torch.ones(1, device=DEVICE), alpha_bars[:-1]])
-posterior_variances = betas * (1 - previous) / (1 - alpha_bars).clamp_min(1e-20)
-
-
-# 3. DDPM: corrupt Z_0 at step k and regress the exact noise epsilon.
-def train(dataset, epochs):
-    _, length, dim = dataset.tensors[0].shape
-    loader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
-    model = DiffusionModel(length, dim).to(DEVICE)
-    reward_model = RewardModel(length, dim).to(DEVICE)
-    optimizer = torch.optim.Adam(list(model.parameters()) + list(reward_model.parameters()), lr=LEARNING_RATE)
-    for epoch in range(epochs):
-        total = 0.0
-        for z0, c, r_tilde in loader:
-            z0, c, r_tilde = z0.to(DEVICE), c.to(DEVICE), r_tilde.to(DEVICE)
-            k = torch.randint(1, K + 1, (len(z0),), device=DEVICE)
-            t = k.float() / K                  # Normalized k is only the network's time input.
-            a = alpha_bars[k, None, None]
-            epsilon = torch.randn_like(z0)
-            zk = a.sqrt() * z0 + (1 - a).sqrt() * epsilon
-            dropped = c.masked_fill(torch.rand(len(c), device=DEVICE) < CONDITION_DROP, 2)
-            loss_diffusion = F.mse_loss(model(zk, t, dropped), epsilon)
-            loss_reward = F.mse_loss(reward_model(zk, t), r_tilde)
-            loss = loss_diffusion + loss_reward
-            optimizer.zero_grad(set_to_none=True)
-            loss.backward()
-            optimizer.step()
-            total += loss.item()
-        if (epoch + 1) % 50 == 0 or epoch + 1 == epochs:
-            print(f"epoch {epoch+1}: combined training loss {total / len(loader):.4f}")
-    for network in (model, reward_model):
-        network.eval().requires_grad_(False)
-        for parameter in network.parameters():
-            parameter.grad = None
-    return model, reward_model
-
-
-# 4. R_lambda = sum_m lambda_m * r_tilde_m; differentiate the current latent only.
-def reward_gradient(reward_model, z, t, lambdas):
-    with torch.enable_grad():                  # Re-enable input gradients inside sampling.
-        state = z.detach().requires_grad_(True)
-        R_lambda = (reward_model(state, t) * lambdas).sum(dim=1)
-        grad = torch.autograd.grad(R_lambda.sum(), state)[0]
-    return grad.detach()                      # Do not retain graphs across sampling steps.
-
-
-def normalize_weights(lambdas):
-    values = torch.as_tensor(lambdas, dtype=torch.float32, device=DEVICE)
-    if values.shape != (2,) or not torch.isfinite(values).all() or (values < 0).any() or values.sum() <= 0:
-        raise ValueError("Use two finite, nonnegative weights with a positive sum")
-    return values / values.sum()              # lambda controls tradeoffs, eta controls strength.
-
-
-# 5. Sampling: CFG first; optional reward gradient then changes predicted noise.
-@torch.no_grad()
-def sample(model, reward_model, n=8, c=1, w=0.0, eta=0.0, lambdas=(1.0, 0.0)):
-    if c not in (0, 1) or w < 0 or eta < 0:
-        raise ValueError("Use c=0/1 and nonnegative w and eta")
-    lambdas = normalize_weights(lambdas)
-    torch.manual_seed(123)                     # Same starting and reverse noise across comparisons.
-    z = torch.randn(n, model.length, model.dim, device=DEVICE)
-    null = torch.full((n,), 2, dtype=torch.long, device=DEVICE)
-    condition = torch.full((n,), c, dtype=torch.long, device=DEVICE)
-    for k in range(K, 0, -1):
-        t = torch.full((n,), k / K, device=DEVICE)
-        eps_uncond = model(z, t, null)
-        eps = eps_uncond
-        if w != 0:
-            eps = eps_uncond + w * (model(z, t, condition) - eps_uncond)
-        sigma = (1 - alpha_bars[k]).sqrt()       # FORWARD corruption standard deviation.
-        if eta != 0:
-            grad_R = reward_gradient(reward_model, z, t, lambdas)
-            eps = eps - eta * sigma * grad_R    # s_guided=s_theta+eta*grad R; s=-epsilon/sigma.
-        mean = (z - betas[k] * eps / sigma) / alphas[k].sqrt()
-        z = mean + posterior_variances[k].sqrt() * torch.randn_like(z) if k > 1 else mean
-    return z                                   # No fresh noise at k=1.
-
-
-# 6. One decoder, used only AFTER the flow or diffusion trajectory is complete.
-@torch.no_grad()
-def decode(z, esm, tokenizer, stats, min_polar=12):
-    latent = z * stats["z_std"].to(z.device) + stats["z_mean"].to(z.device)
-    logits = esm.lm_head(latent)               # Frozen head produces one vocabulary distribution per residue.
-    aa_ids = torch.tensor(tokenizer.convert_tokens_to_ids(list(AMINO_ACIDS)), device=z.device)
-    logits = logits.index_select(-1, aa_ids)   # Restrict the vocabulary to the 20 amino acids.
-    if not 0 <= min_polar <= z.shape[1]:
-        raise ValueError("min_polar must lie between zero and the sequence length")
-    polar = torch.tensor([a in POLAR_RESIDUES for a in AMINO_ACIDS], device=z.device)
-    polar_ids = polar.nonzero().flatten()
-    best_scores, choices = logits.max(dim=-1)  # Start from unrestricted amino-acid argmax.
-    polar_scores, local = logits[..., polar_ids].max(dim=-1)
-    polar_choices = polar_ids[local]          # Best polar/charged amino acid at each position.
-    for i in range(len(z)):
-        already_polar = polar[choices[i]]
-        missing = max(0, min_polar - int(already_polar.sum()))
-        if missing:
-            cost = (best_scores[i] - polar_scores[i]).masked_fill(already_polar, float("inf"))
-            positions = cost.topk(missing, largest=False).indices
-            choices[i, positions] = polar_choices[i, positions]
-    return ["".join(AMINO_ACIDS[i] for i in row) for row in choices.cpu().tolist()]
-    # Exact maximum-logit decode subject to at least min_polar selected residues.
-    # This enforces composition, not measured solubility; no gradient through argmax.
-
-
-# 7. Run CFG, single-objective steering, and scalarized multi-objective steering in order.
-def main():
-    parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument("--data", type=Path, default=ROOT / "esm2_example.csv")
-    parser.add_argument("--epochs", type=int, default=200)
-    parser.add_argument("--samples", type=int, default=8)
-    parser.add_argument("--min-polar", type=int, default=12)
-    parser.add_argument("--output", type=Path, default=ROOT / "esm2_diffusion_outputs")
-    args = parser.parse_args()
-    if args.epochs < 1 or args.samples < 1:
-        parser.error("epochs and samples must be positive")
-    torch.manual_seed(7)
-    if DEVICE.type == "cpu":
-        torch.set_num_threads(2)
-    dataset, esm, tokenizer, stats = load_data(args.data)
-    if not 0 <= args.min_polar <= dataset.tensors[0].shape[1]:
-        parser.error("min-polar must lie between zero and the sequence length")
-    model, reward_model = train(dataset, args.epochs)
-    outputs = {
-        "cfg": sample(model, reward_model, args.samples, c=1, w=2.0),
-        "single": sample(model, reward_model, args.samples, eta=1.0, lambdas=(1.0, 0.0)),
-        "multi": sample(model, reward_model, args.samples, eta=1.0, lambdas=(0.7, 0.3)),
-    }
-    args.output.mkdir(parents=True, exist_ok=True)
-    for name, latent in outputs.items():
-        if not torch.isfinite(latent).all():
-            raise RuntimeError(f"Nonfinite {name} output; reduce guidance or check training")
-        sequences = decode(latent, esm, tokenizer, stats, args.min_polar)
-        assert all(sum(a in POLAR_RESIDUES for a in s) >= args.min_polar for s in sequences)
-        fasta = "".join(f">{name}_{i+1}\n{seq}\n" for i, seq in enumerate(sequences))
-        (args.output / f"{name}.fasta").write_text(fasta)
-        count = sum(a in POLAR_RESIDUES for a in sequences[0])
-        print(f"{name}: {sequences[0]}  polar/charged residues={count}")
-        print("  decoded mean composition proxies:", composition_proxies(sequences).mean(0).tolist())
-    torch.save({"standardized_latents": {k: v.cpu() for k, v in outputs.items()},
-                "model": model.state_dict(), "reward_model": reward_model.state_dict(),
-                "stats": stats, "esm_name": ESM_NAME,
-                "length": model.length, "dim": model.dim, "min_polar": args.min_polar,
-                "polar_residues": POLAR_RESIDUES}, args.output / "results.pt")
-    print(f"Saved latent tensors and FASTA sequences to {args.output}")
-
-
-if __name__ == "__main__":
-    main()

lectures_02_03/esm2_example.csv

@@ -1,65 +0,0 @@
-sequence,c,r1,r2
-QYWSDSWWESQMMSPWYPMSPLSV,0,-0.08333333,0.41666667
-CKSEFQPPHLMGHDFFACEMRNFK,0,0.00000000,0.45833333
-WPYGEHMLADNNVVKKRLQQWCFI,1,0.04166667,0.41666667
-KVVGALPIESFYTAKMESIAVEVI,0,-0.04166667,0.33333333
-MEFHAYSLGPASRWSSRHGYFTNL,1,0.04166667,0.45833333
-TCYLFCLAQIVMFGRAYVGGYWAK,1,0.08333333,0.16666667
-VFSHIVACCEIETRDTCVWNHMYM,0,-0.08333333,0.41666667
-QRNQTEIRTLIYMPYGNKTYLRGS,1,0.12500000,0.54166667
-WDNQRIKTALAGHTGVDSGWMHNF,0,0.00000000,0.50000000
-PPVCRGWQMPRAALWVMRCRFPRQ,1,0.20833333,0.29166667
-HPNAKNYGETIFFGLWGNDRKNLV,1,0.04166667,0.45833333
-CMKEIYESFLDETTIWVIHMIAAR,0,-0.08333333,0.41666667
-RWDTGMWGKDCIHRPIQKKDPTWA,1,0.08333333,0.50000000
-WCCQLMAYPPAKMGHPTFDHPLPK,1,0.04166667,0.29166667
-QPVRDFVLEVGQSILESHSLKVPS,0,-0.04166667,0.50000000
-GREMRGFQKTPNATKKEGTYKPAP,1,0.16666667,0.54166667
-ENEWNGSRHYWCVRWLGTLKILTF,1,0.04166667,0.45833333
-AVSQIAPHARQQASSQENKHGETT,0,0.00000000,0.66666667
-NADTREKMPSNDLSWWLYNMFMPA,0,-0.04166667,0.45833333
-SLWFAIMPYVAKYVIYGHVEPTAT,0,0.00000000,0.25000000
-YGWHYTLLMWYDERCGFLGRGRAQ,1,0.04166667,0.33333333
-EEQEKNLDDMNSLDFSQNLSLFEY,0,-0.25000000,0.66666667
-ETWDHYPQYKCVKCKNNHCPHVAY,1,0.04166667,0.50000000
-YHIMRKVPSTNEFRPLAHNNCTLQ,1,0.08333333,0.54166667
-DFRMNQYGSKDYPKTTGAREHLFM,1,0.04166667,0.54166667
-IQWWSDVQTMQCMADEGMHHQYFV,0,-0.12500000,0.45833333
-WAQMMMVRQIDCRNFIRHTVYSNF,1,0.08333333,0.45833333
-DMTFGHSSDAILAKVYDSKNCPQL,0,-0.04166667,0.50000000
-DRSICNKIPPSQTIPFMGPFISQN,1,0.04166667,0.45833333
-QNHLSDAEIKYCIKWNVLNGKMLP,1,0.04166667,0.45833333
-RQYCPEPSWSQHTNKPACMKVEEC,0,0.00000000,0.54166667
-IWWWHVNCQCRDDHEHCPLCYAPI,0,-0.08333333,0.37500000
-SDKSVTHIMERDIWPIIAWLRDRV,0,0.00000000,0.50000000
-WSDAWDVYIWILAMNCEVGRMYYC,0,-0.08333333,0.25000000
-VAKELIRNYGDEFCWLVNSSMKED,0,-0.08333333,0.50000000
-VTWANPHINITFCDPMHIIRIQCQ,0,0.00000000,0.41666667
-GYYWVTHLVHMVDIWFSDMPIHYE,0,-0.12500000,0.33333333
-TRALRHDDVWVDYYFIWTALTYQF,0,-0.04166667,0.41666667
-CGNGSDFDTFIQMWICACGMCIDK,0,-0.08333333,0.33333333
-NGQMACTQENWAWSVFHFEWWCPK,0,-0.04166667,0.41666667
-GSRETKHEVTADCPFQWNGRALKK,1,0.08333333,0.58333333
-CHGKYTFTFGMEIRDMPAWYGHFV,0,0.00000000,0.33333333
-GSYDNPLAIKWQMVCNSSRMDWTM,0,0.00000000,0.45833333
-FRIAQIRSKIGYLQNIFVMSLLLR,1,0.16666667,0.37500000
-SCCDWRQMQDRVGPAMMKEWACGE,0,-0.04166667,0.41666667
-FFSELMKCHYHYCYAWRRPIAKDW,1,0.08333333,0.37500000
-AWGPVLDFNIQSDIFYPCRTDMGL,0,-0.08333333,0.33333333
-PHTCRRIMMMEQDLTNFPLAYCTQ,0,0.00000000,0.45833333
-DQINDGSAHENQYCQNPSDPFVWK,0,-0.12500000,0.58333333
-QMRSHFQYGCGWMITPMRFPKNMA,1,0.12500000,0.37500000
-EECYALAGDYDFKVRIQHAQCTWQ,0,-0.08333333,0.45833333
-ITKQMWVNDVWDTEVRHHSFMTLF,0,-0.04166667,0.54166667
-EHAKYVKNSNCYNKENNVDLFYCN,0,0.00000000,0.58333333
-SYYVPKTKQRLTTSRCGLKGVVKA,1,0.25000000,0.50000000
-TDAKGWHYNWHWPLHQFCHLPVHA,0,0.00000000,0.41666667
-YIVRMSDSFHSGHDHDKKNLIMYH,0,0.00000000,0.58333333
-CTHVQNPNSIETNCHEHYDFGTSM,0,-0.12500000,0.62500000
-RWWYQNHLMCYNSKDACYHPWINM,1,0.04166667,0.41666667
-RKSNGGKSLIQNAGITCVKWCYGV,1,0.16666667,0.41666667
-WWNQITTPHWCKWWASHWMQYSRW,1,0.08333333,0.45833333
-IDTICWQVCILVQESDAKDVFEGN,0,-0.16666667,0.45833333
-NFHEWVLWELTGQWRVWSRQHSKE,0,0.00000000,0.58333333
-VTYCMMAKFAGMDYEQKICVHVAC,0,0.00000000,0.29166667
-KVKTTLKVEREAPPYTRMSLSAWT,1,0.12500000,0.54166667

lectures_02_03/esm2_flow_guidance.py

@@ -1,251 +0,0 @@
-"""Standalone ESM-2 flow matching guidance example for CIS 6270.
-Input CSV: sequence,c,r1,r2. Sequences have one fixed length; c is 0 or 1.
-Both objectives are oriented so larger is better. The bundled CSV is synthetic.
-Run: python esm2_flow_guidance.py --data esm2_example.csv --epochs 200
-Install: pip install torch transformers==4.57.6
-Outputs: guided residue latents, model weights, and decoded amino-acid sequences.
-"""
-import argparse
-import csv
-from pathlib import Path
-
-import torch
-from torch import nn
-import torch.nn.functional as F
-from torch.utils.data import DataLoader, TensorDataset
-from transformers import AutoTokenizer, EsmForMaskedLM
-
-ROOT = Path(__file__).resolve().parent
-DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-ESM_NAME = "facebook/esm2_t6_8M_UR50D"
-AMINO_ACIDS = "ACDEFGHIKLMNPQRSTVWY"
-POLAR_RESIDUES = "DEHKNQRST"  # Operational polar/charged set; not a solubility assay.
-BATCH_SIZE, HIDDEN, LEARNING_RATE = 16, 128, 1e-3
-CONDITION_DROP = 0.2  # Drop c during training so the same model learns the null condition.
-
-
-def composition_proxies(sequences):
-    # Transparent teaching rewards; these are not measured activity or solubility.
-    return torch.tensor([
-        [(sum(a in "KR" for a in s) - sum(a in "DE" for a in s)) / len(s),
-         sum(a in POLAR_RESIDUES for a in s) / len(s)]
-        for s in sequences
-    ], dtype=torch.float32)
-
-
-# 1. Load annotated sequences and encode frozen ESM-2 residue vectors.
-@torch.no_grad()
-def load_data(path):
-    with Path(path).open(newline="") as handle:
-        rows = list(csv.DictReader(handle))
-    sequences = [row["sequence"].strip().upper() for row in rows]
-    if len(rows) < 4 or any(not s or set(s) - set(AMINO_ACIDS) for s in sequences):
-        raise ValueError("Supply at least four sequences using the 20 standard amino acids")
-    lengths = {len(s) for s in sequences}
-    if len(lengths) != 1 or max(lengths) > 128:
-        raise ValueError("This compact example requires one fixed sequence length, at most 128")
-    c = torch.tensor([int(row["c"]) for row in rows], dtype=torch.long)
-    # With sequence,c only, compute example rewards directly from residue counts.
-    # Optional r1/r2 columns override them with supplied labels, such as measurements.
-    if {"r1", "r2"}.issubset(rows[0]):
-        r = torch.tensor([[float(row["r1"]), float(row["r2"])] for row in rows])
-    else:
-        r = composition_proxies(sequences)
-    if set(c.tolist()) != {0, 1} or not torch.isfinite(r).all():
-        raise ValueError("Include both c=0 and c=1, with finite r1 and r2")
-
-    tokenizer = AutoTokenizer.from_pretrained(ESM_NAME, cache_dir=ROOT / ".esm2_cache")
-    esm = EsmForMaskedLM.from_pretrained(
-        ESM_NAME, cache_dir=ROOT / ".esm2_cache", use_safetensors=True
-    ).to(DEVICE).eval().requires_grad_(False)
-    encoded = []
-    for start in range(0, len(sequences), BATCH_SIZE):
-        tokens = tokenizer(sequences[start:start + BATCH_SIZE], return_tensors="pt")
-        tokens = {key: value.to(DEVICE) for key, value in tokens.items()}
-        hidden = esm.esm(**tokens).last_hidden_state
-        encoded.append(hidden[:, 1:-1].cpu())  # Remove BOS/EOS; retain all L residue positions.
-    z = torch.cat(encoded)                    # [N, L, D], with D=320 for this checkpoint.
-    z_mean = z.mean((0, 1), keepdim=True)
-    z_std = z.std((0, 1), correction=0, keepdim=True).clamp_min(1e-4)
-    r_mean, r_std = r.mean(0), r.std(0, correction=0).clamp_min(1e-6)
-    dataset = TensorDataset((z - z_mean) / z_std, c, (r - r_mean) / r_std)
-    stats = {"z_mean": z_mean, "z_std": z_std, "r_mean": r_mean, "r_std": r_std}
-    return dataset, esm, tokenizer, stats
-
-
-# 2. Small model classes: flattening lets each output depend on the entire sequence.
-class FlowModel(nn.Module):
-    def __init__(self, length, dim):
-        super().__init__()
-        self.length, self.dim = length, dim
-        self.time = nn.Sequential(nn.Linear(1, 32), nn.SiLU(), nn.Linear(32, 32))
-        self.skip = nn.Linear(32, 1)           # Preserve full-dimensional state/noise through a time gate.
-        self.condition = nn.Embedding(3, 16)    # Indices 0,1 are classes; 2 is the null class.
-        self.net = nn.Sequential(
-            nn.Linear(length * dim + 48, HIDDEN), nn.SiLU(),
-            nn.Linear(HIDDEN, HIDDEN), nn.SiLU(),
-            nn.Linear(HIDDEN, length * dim),
-        )
-
-    def forward(self, z, t, c):
-        time = self.time(t[:, None])
-        inputs = torch.cat([z.flatten(1), time, self.condition(c)], dim=1)
-        return self.skip(time)[:, :, None] * z + self.net(inputs).reshape_as(z)
-        # A time-dependent linear part plus a learned nonlinear velocity correction.
-
-
-class RewardModel(nn.Module):
-    def __init__(self, length, dim):
-        super().__init__()
-        self.net = nn.Sequential(
-            nn.Linear(length * dim + 1, HIDDEN), nn.SiLU(),
-            nn.Linear(HIDDEN, HIDDEN), nn.SiLU(), nn.Linear(HIDDEN, 2),
-        )
-
-    def forward(self, z, t):
-        return self.net(torch.cat([z.flatten(1), t[:, None]], dim=1))
-        # Two predicted standardized endpoint objectives, given the current latent and time.
-
-
-# 3. Flow matching: Z_t=(1-t)Z_0+tZ_1; target velocity is Z_1-Z_0.
-def train(dataset, epochs):
-    _, length, dim = dataset.tensors[0].shape
-    loader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
-    model = FlowModel(length, dim).to(DEVICE)
-    reward_model = RewardModel(length, dim).to(DEVICE)
-    optimizer = torch.optim.Adam(list(model.parameters()) + list(reward_model.parameters()), lr=LEARNING_RATE)
-    for epoch in range(epochs):
-        total = 0.0
-        for z1, c, r_tilde in loader:
-            z1, c, r_tilde = z1.to(DEVICE), c.to(DEVICE), r_tilde.to(DEVICE)
-            z0 = torch.randn_like(z1)          # Flow convention: Z_0=noise, Z_1=clean latent.
-            t = torch.rand(len(z1), device=DEVICE)
-            zt = (1 - t[:, None, None]) * z0 + t[:, None, None] * z1
-            dropped = c.masked_fill(torch.rand(len(c), device=DEVICE) < CONDITION_DROP, 2)
-            loss_flow = F.mse_loss(model(zt, t, dropped), z1 - z0)
-            loss_reward = F.mse_loss(reward_model(zt, t), r_tilde)
-            loss = loss_flow + loss_reward    # Independent parameter sets; one optimizer suffices.
-            optimizer.zero_grad(set_to_none=True)
-            loss.backward()
-            optimizer.step()
-            total += loss.item()
-        if (epoch + 1) % 50 == 0 or epoch + 1 == epochs:
-            print(f"epoch {epoch+1}: combined training loss {total / len(loader):.4f}")
-    for network in (model, reward_model):
-        network.eval().requires_grad_(False)   # Freeze weights; still allow gradients of input z.
-        for parameter in network.parameters():
-            parameter.grad = None
-    return model, reward_model
-
-
-# 4. R_lambda = sum_m lambda_m * r_tilde_m; differentiate the current latent only.
-def reward_gradient(reward_model, z, t, lambdas):
-    with torch.enable_grad():                  # Re-enable input gradients inside sampling.
-        state = z.detach().requires_grad_(True)
-        R_lambda = (reward_model(state, t) * lambdas).sum(dim=1)
-        grad = torch.autograd.grad(R_lambda.sum(), state)[0]
-    return grad.detach()                      # Do not retain graphs across sampling steps.
-
-
-def normalize_weights(lambdas):
-    values = torch.as_tensor(lambdas, dtype=torch.float32, device=DEVICE)
-    if values.shape != (2,) or not torch.isfinite(values).all() or (values < 0).any() or values.sum() <= 0:
-        raise ValueError("Use two finite, nonnegative weights with a positive sum")
-    return values / values.sum()              # lambda controls tradeoffs, eta controls strength.
-
-
-# 5. Sampling: CFG first; optional reward gradient then changes the velocity.
-@torch.no_grad()
-def sample(model, reward_model, n=8, c=1, w=0.0, eta=0.0, lambdas=(1.0, 0.0), steps=200):
-    if c not in (0, 1) or steps < 1 or w < 0 or eta < 0:
-        raise ValueError("Use c=0/1, positive steps, and nonnegative w and eta")
-    lambdas = normalize_weights(lambdas)
-    torch.manual_seed(123)                     # Compare methods from the same initial noise.
-    z = torch.randn(n, model.length, model.dim, device=DEVICE)
-    null = torch.full((n,), 2, dtype=torch.long, device=DEVICE)
-    condition = torch.full((n,), c, dtype=torch.long, device=DEVICE)
-    dt = 1.0 / steps
-    for step in range(steps):
-        t = torch.full((n,), step * dt, device=DEVICE)
-        v_uncond = model(z, t, null)
-        v = v_uncond
-        if w != 0:
-            v = v_uncond + w * (model(z, t, condition) - v_uncond)
-        if eta != 0:
-            grad_R = reward_gradient(reward_model, z, t, lambdas)
-            kappa = eta * 4 * t[:, None, None] * (1 - t[:, None, None])
-            v = v + kappa * grad_R             # Chosen steering rule; not exact conditional transport.
-        z = z + dt * v                         # Integrate forward from t=0 to t=1.
-    return z
-
-
-# 6. One decoder, used only AFTER the flow or diffusion trajectory is complete.
-@torch.no_grad()
-def decode(z, esm, tokenizer, stats, min_polar=12):
-    latent = z * stats["z_std"].to(z.device) + stats["z_mean"].to(z.device)
-    logits = esm.lm_head(latent)               # Frozen head produces one vocabulary distribution per residue.
-    aa_ids = torch.tensor(tokenizer.convert_tokens_to_ids(list(AMINO_ACIDS)), device=z.device)
-    logits = logits.index_select(-1, aa_ids)   # Restrict the vocabulary to the 20 amino acids.
-    if not 0 <= min_polar <= z.shape[1]:
-        raise ValueError("min_polar must lie between zero and the sequence length")
-    polar = torch.tensor([a in POLAR_RESIDUES for a in AMINO_ACIDS], device=z.device)
-    polar_ids = polar.nonzero().flatten()
-    best_scores, choices = logits.max(dim=-1)  # Start from unrestricted amino-acid argmax.
-    polar_scores, local = logits[..., polar_ids].max(dim=-1)
-    polar_choices = polar_ids[local]          # Best polar/charged amino acid at each position.
-    for i in range(len(z)):
-        already_polar = polar[choices[i]]
-        missing = max(0, min_polar - int(already_polar.sum()))
-        if missing:
-            cost = (best_scores[i] - polar_scores[i]).masked_fill(already_polar, float("inf"))
-            positions = cost.topk(missing, largest=False).indices
-            choices[i, positions] = polar_choices[i, positions]
-    return ["".join(AMINO_ACIDS[i] for i in row) for row in choices.cpu().tolist()]
-    # Exact maximum-logit decode subject to at least min_polar selected residues.
-    # This enforces composition, not measured solubility; no gradient through argmax.
-
-
-# 7. Run CFG, single-objective steering, and scalarized multi-objective steering in order.
-def main():
-    parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument("--data", type=Path, default=ROOT / "esm2_example.csv")
-    parser.add_argument("--epochs", type=int, default=200)
-    parser.add_argument("--samples", type=int, default=8)
-    parser.add_argument("--min-polar", type=int, default=12)
-    parser.add_argument("--output", type=Path, default=ROOT / "esm2_flow_outputs")
-    args = parser.parse_args()
-    if args.epochs < 1 or args.samples < 1:
-        parser.error("epochs and samples must be positive")
-    torch.manual_seed(7)
-    if DEVICE.type == "cpu":
-        torch.set_num_threads(2)
-    dataset, esm, tokenizer, stats = load_data(args.data)
-    if not 0 <= args.min_polar <= dataset.tensors[0].shape[1]:
-        parser.error("min-polar must lie between zero and the sequence length")
-    model, reward_model = train(dataset, args.epochs)
-    outputs = {
-        "cfg": sample(model, reward_model, args.samples, c=1, w=2.0),
-        "single": sample(model, reward_model, args.samples, eta=1.0, lambdas=(1.0, 0.0)),
-        "multi": sample(model, reward_model, args.samples, eta=1.0, lambdas=(0.7, 0.3)),
-    }
-    args.output.mkdir(parents=True, exist_ok=True)
-    for name, latent in outputs.items():
-        if not torch.isfinite(latent).all():
-            raise RuntimeError(f"Nonfinite {name} output; reduce guidance or check training")
-        sequences = decode(latent, esm, tokenizer, stats, args.min_polar)
-        assert all(sum(a in POLAR_RESIDUES for a in s) >= args.min_polar for s in sequences)
-        fasta = "".join(f">{name}_{i+1}\n{seq}\n" for i, seq in enumerate(sequences))
-        (args.output / f"{name}.fasta").write_text(fasta)
-        count = sum(a in POLAR_RESIDUES for a in sequences[0])
-        print(f"{name}: {sequences[0]}  polar/charged residues={count}")
-        print("  decoded mean composition proxies:", composition_proxies(sequences).mean(0).tolist())
-    torch.save({"standardized_latents": {k: v.cpu() for k, v in outputs.items()},
-                "model": model.state_dict(), "reward_model": reward_model.state_dict(),
-                "stats": stats, "esm_name": ESM_NAME,
-                "length": model.length, "dim": model.dim, "min_polar": args.min_polar,
-                "polar_residues": POLAR_RESIDUES}, args.output / "results.pt")
-    print(f"Saved latent tensors and FASTA sequences to {args.output}")
-
-
-if __name__ == "__main__":
-    main()