TD3B-dev runnability fixes + finetune-on-target + validity samplers

.gitignore

@@ -1,2 +1,8 @@
 __pycache__/
 *.pyc
+results/
+*.egg-info/
+.ipynb_checkpoints/
+/checkpoints
+/data
+scoring/functions/classifiers/binding-affinity.pt

CLAUDE.md

@@ -0,0 +1,250 @@
+# CLAUDE.md
+
+This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
+
+## What this is
+
+TD3B (Transition-Directed Discrete Diffusion for Allosteric Binder Generation) is a sequence-based
+generative framework that designs peptide binders with a **specified direction** — agonist or
+antagonist — for a target protein. It extends **TR2-D2** (a masked discrete-diffusion peptide
+generator: MDLM backbone + MCTS amortized finetuning) with three additions:
+
+1. a **Direction Oracle** `f_φ` that predicts agonist vs. antagonist behavior,
+2. a **soft binding-affinity gate** `g_ψ`, and
+3. a **gated reward** `R = g_ψ · σ(d*·(f_φ − 0.5)/τ)` that steers generation toward direction `d* ∈ {+1 agonist, −1 antagonist}`.
+
+Finetuning distills MCTS-discovered high-reward sequences into the diffusion policy via a total loss
+`L = L_WDCE + λ·L_ctr + β·L_KL` (weighted denoising CE + directional contrastive + KL-to-reference).
+
+Paper: arXiv:2605.09810 (LMRL Workshop, ICLR 2026).
+
+## Repo layout: dev vs. OSS, and where the artifacts are
+
+- This repo (`ChatterjeeLab/TD3B-dev`) is the **full/dev code**. The clean code-only public release is
+  `ChatterjeeLab/TD3B` on HuggingFace. Keep changes runnable against that OSS release.
+- **No checkpoints or data are in git.** Trained checkpoints, train/test CSVs, and generated binders
+  ship as a single ~3.4 GB archive `td3b_dev_artifacts.zip` on Google Drive (link in README). Unzip at
+  repo root to populate `checkpoints/`, `data/`, `scoring/functions/classifiers/`, `generated_binders/`.
+- `.gitattributes` routes all weight/data extensions (`*.ckpt *.pt *.pth *.npy *.csv? *.zip ...`) through
+  Git LFS. Only three XGBoost classifier JSONs (`hemolysis`, `nonfouling`, `solubility`) are in-repo;
+  `binding-affinity.pt` and `permeability-xgboost.json` come from the archive.
+
+## Environment & commands
+
+```bash
+conda env create -f env.yml   # creates env "td3b" (python 3.10, pytorch-cuda 12.1)
+conda activate td3b
+pip install -e .              # installs the `td3b` package (setup.py)
+```
+
+Core deps: PyTorch + Lightning 2.5.5, HuggingFace `transformers` 4.56.2, `fair-esm` 2.0.0 (ESM2),
+`rdkit`, `SmilesPE`, `xgboost`, `wandb`, `hydra-core`. There is **no test suite, no linter config, and
+no Makefile** — do not look for `pytest`/`tox`/`ruff`. Verify changes by actually running inference.
+
+**Inference** (primary OSS entry point — generate directional binders):
+```bash
+python inference.py \
+    --ckpt_path checkpoints/td3b.ckpt \
+    --val_csv data/test.csv \
+    --save_path results/ \
+    --seed 42 --num_pool 32 --val_samples_per_target 8 --resample_alpha 0.1
+```
+For each target row it generates for **both** directions (agonist `d*=+1`, antagonist `d*=-1`), scores
+with the oracle + affinity gate, applies softmax(reward/α) weighted resampling (Algorithm 2), keeps only
+RDKit-valid peptides, and writes `results/td3b_results_seed{seed}.csv`.
+
+**Training** (multi-target): edit paths in `launch_multi_target.sh` (`BASE_PATH`, checkpoints, data,
+oracle, and `WANDB_ENTITY` — blank by default; set your own), then `bash launch_multi_target.sh`. It calls
+`finetune_multi_target.py`. Key knobs live in the launch script: `CONTRASTIVE_WEIGHT`(λ), `KL_BETA`(β),
+`SIGMOID_TEMPERATURE`(τ), `NUM_ITER`/`NUM_CHILDREN` (MCTS), `TARGETS_PER_MCTS`(K), and the cadence flags
+below.
+
+**Baselines** (CG, SMC, TDS, PepTune, Unguided): `cd baselines && ./run.sh <csv> <baseline> <device>`.
+Multi-GPU via a 5th/6th arg (`torchrun`). Note this script loads `../pretrained/peptune-pretrained.ckpt`,
+**not** `checkpoints/pretrained.ckpt` (see path landmines below).
+
+**Demo**: `notebooks/TD3B_Inference_Demo.ipynb` (Colab T4).
+
+## Architecture — the cross-file big picture
+
+The pieces below only make sense together; reading any one file in isolation misses the flow.
+
+### Diffusion backbone — `models/`
+- `Diffusion(L.LightningModule)` in `models/diffusion.py` is the MDLM core (absorbing-state / masked
+  discrete diffusion). The denoiser is `models/roformer.py::Roformer`, a thin wrapper over HuggingFace
+  `RoFormerForMaskedLM` (rotary embeddings). Tokens are **SMILES** via `tokenizer/my_tokenizers.py::SMILES_SPE_Tokenizer`.
+- Absorbing state = the tokenizer's `[MASK]` id (`self.mask_index`), not a constant. Generation starts
+  fully masked (`sample_prior` → all-mask) and reverse-diffuses. SUBS parameterization
+  (`subs_parameterization`) forbids predicting MASK and pins already-unmasked positions ("carry-over
+  unmasking"). Callers use `single_reverse_step` / `single_noise_removal` (final step guarantees no
+  surviving MASK); MCTS uses the `batch_mcts_reverse_step` / `mcts_reverse_step` variants which also
+  return per-step policy vs. pretrained log-probs for the importance log-ratio `log_rnd`.
+- Config is built by `configs/finetune_config.py::DiffusionConfig`, a **shim** that synthesizes
+  duck-typed attribute objects (`type(...)()`) for backward compat. It fixes `parameterization='subs'`,
+  `T=0` (continuous-time MDLM loss), `time_conditioning=False`, and `max_position_embeddings=1035`.
+  It is a *partial* interface — only the finetune/eval/MCTS fields exist; pure-training paths
+  (`antithetic_sampling`, `noise.state_dependent`, `vocab`, `model.length`) are absent and would
+  `AttributeError` under it.
+- `models/noise_schedule.py`: `get_noise(config)` supports geometric/loglinear/cosine/linear, but every
+  reverse/sampling step **asserts `loglinear`**. A second hardcoded `LogPolyNoise` (cubic) masks
+  peptide-bond tokens more slowly.
+
+### Directional reward — `td3b/td3b_scoring.py`
+`TD3BRewardFunction.__call__(List[str] of peptide SMILES) -> (rewards, info)`. Internals:
+`g_ψ` = `scoring/functions/binding.py::BindingAffinity` (magnitude); `f_φ` = the `DirectionalOracle`
+(direction prob ∈[0,1] + confidence κ); reward = `g_ψ · σ(d*·(f_φ−0.5)/τ)`. `create_td3b_reward_function`
+is the factory that builds/loads the oracle, caches the encoded protein tokens, maps `'agonist'/'antagonist'`
+→ `d* = +1/−1`, and returns the configured reward. `TD3BConfidenceWeighting` provides the
+confidence-weighted importance weights used by MCTS (`w = κ·exp(S/α)`).
+
+### Direction Oracle — `td3b/direction_oracle.py`
+`DirectionalOracle` wraps `ESM_TR2D2_GPCRClassifier`: **frozen ESM2** (`facebook/esm2_t33_650M_UR50D`,
+1280-d — downloads from HF unless `esm_cache_dir`/`esm_local_files_only` set) encodes the **protein**;
+a **frozen TR2-D2 RoFormer** encodes the **ligand/peptide SMILES**; both project to `d_model=256`, pass
+1 self-attention layer each, then **2 stacked bidirectional cross-attention (BMCA) layers**, mean-pool,
+concat, MLP → 2 logits. `predict_with_confidence` returns `f_φ = p_agonist` and `κ = max(softmax)`.
+Loading needs four assets: the oracle `.pt`, the TR2-D2 ligand checkpoint, and the SMILES tokenizer
+vocab+splits. The RoFormer config (768/8/8/1035) is hardcoded and must match the checkpoint.
+
+### Losses — `td3b/td3b_losses.py`
+`TD3BTotalLoss` = `L_WDCE + λ·L_ctr + β·L_KL` (λ=`contrastive_weight`, β=`kl_beta`, both default 0.1).
+`L_WDCE` is computed **externally** by `training/finetune_utils.py::loss_wdce` and passed in — it is the
+policy-distillation term that reweights MCTS samples by `softmax(log_rnd)`. `L_ctr` is `ContrastiveLoss`
+(margin, default) or `InfoNCELoss` over agonist/antagonist embeddings from `extract_embeddings_from_mdlm`
+(reaches into `model.backbone.model`, RoFormer last hidden state, **must not** be under `no_grad`).
+`L_KL` is per-position categorical KL to a **frozen deepcopy reference model** (the pretrained weights).
+
+### MCTS — `mcts/peptide_mcts.py` + `td3b/td3b_mcts.py`
+Base `MCTS` does Pareto/multi-objective tree search: root is fully masked; `select` descends via a
+non-dominated (not scalar-UCB) set with `rd.choice`; `expand` samples `num_children` one-step
+unmaskings, rolls each to a full sequence, filters with `PeptideAnalyzer.is_peptide`, scores valid ones,
+and maintains a **Pareto buffer** of finished trajectories (each storing `x_final`, `log_rnd`,
+reward, score vector). `TD3B_MCTS` subclasses it: injects the gated `TD3BRewardFunction`, pads the (N,2)
+directional score vector to (N,5) so the base Pareto machinery works, folds confidence into `log_rnd`,
+and `forward`/`consolidateBuffer` return **7** values (adding `directional_labels`, `confidences`).
+
+### Training loop — `finetune_multi_target.py`
+Self-contained inline loop (it does **not** call `td3b/td3b_finetune.py::td3b_finetune`, which is a
+legacy/unused single-target loop). Per epoch it alternates:
+- **MCTS generation phase** (every `resample_every_n_step` epochs, default 10): for each sampled target ×
+  each direction, build a per-(target,direction) reward, run a fresh `TD3B_MCTS`, push Pareto survivors
+  into a replay buffer with `directional_label` **forced to the intended `d*`** (not the oracle guess).
+- **Gradient phase** (every epoch): shuffle the buffer, pad variable-length `x` to max-len with
+  `mask_index` + build an attention mask, then WDCE + KL every batch, and contrastive only when a batch
+  mixes both directions.
+
+Three independent epoch-indexed cadence knobs: `resample_targets_every` (redraw K targets),
+`resample_every_n_step` (MCTS phase), `reset_every_n_step` (reset vs. reuse the search tree).
+`add_td3b_sampling_to_model(policy)` monkey-patches `sample_finetuned_td3b` onto the Diffusion instance —
+required before validation/eval works (that method is not native to `Diffusion`).
+
+### Data schema — `td3b/data_utils.py`
+CSVs use columns `Target_Sequence`, `Ligand_Sequence` (peptide AA string), and `label`
+(`agonist`/`antagonist`, mapped to `d*`); `TD3BDataset` also reads `Action`, `Target_UniProt_ID`,
+`Ligand_UniProt_ID`. Binders are converted AA→SMILES via RDKit `MolFromSequence`. The multi-target
+script uses its own in-file `TargetDataset` (groups by target, stores per-direction **median binder
+length** used to set generation length), reading only `Target_Sequence`/`Ligand_Sequence`/`label`.
+`inference.py` also reads `Target_UniProt_ID`, `Ligand_SMILES` (for length).
+
+### New entry points (added 2026-07-12) — `finetune_on_target.py`, `generate_valid.py`
+- **`finetune_on_target.py` (Function A)** — user-facing "bring your own target" wrapper; it does **not** reimplement training. It normalizes `--target_seq` (repeatable) / `--targets_csv` into a temp CSV (seeding a placeholder poly-G binder of `--binder_length` residues for any missing direction so `TargetDataset` can set a length prior), then **subprocess-invokes `finetune_multi_target.py`** with `--targets_per_mcts=<#targets>` + `--resample_targets_every 1` (finetune on ONLY those targets), and finally **generates in-process** reusing `inference.load_model`/`sample_sequences`/`score_sequences` + `create_td3b_reward_function` + Algorithm-2 resampling. `--direction` restricts only what is generated (finetuning always searches both). Writes `results/finetune_on_target/binders_<dir>_validity-<on|off>_seed<seed>.csv` + the finetuned ckpt under `results/<run_name>_<ts>/`. Paths default to repo root; missing heavy artifacts fail fast with a Google-Drive hint.
+- **`generate_valid.py` + `sampling_strategies.py` (Function B)** — sampling-time validity boosters (no retraining) that reuse the model's own `sample_prior`/`single_reverse_step`/`single_noise_removal`/`forward` and change only token SELECTION (temperature → top-k → softmax → top-p), plus a **remask self-correction loop** (remask the lowest-confidence K% of invalid sequences and re-denoise for R rounds) and a **best-of-N** validity-guided rejection wrapper — no diffusion math is reimplemented. `sampling_strategies.generate(...)` dispatches `baseline, more_steps, top_p(=nucleus), top_k, low_temp, remask, best_of_n, nucleus_remask` (default `nucleus_remask`). `generate_valid.py` loads the real ckpt via `inference.load_model`, or falls back to `build_random_model` (random-init `Diffusion`, CPU dev) when `--ckpt_path` is absent. Output: valid-only CSV (`idx,sequence,n_chars`) + a printed valid-yield summary.
+- **`inference.py --sampler`** (opt-in; default `baseline` = original behavior, byte-for-byte) selects a
+  Function-B sampling strategy for the candidate pool, with pass-through knobs
+  (`--num_steps --top_p --top_k --temperature --remask_rounds --remask_frac --best_of_n`). Non-baseline
+  strategies dispatch to `sampling_strategies.generate`. Function B is also available standalone via
+  `generate_valid.py`.
+
+## Landmines
+
+A batch of path/wiring bugs that stopped the OSS release from running was fixed on 2026-07-09.
+`python inference.py --help` and `python finetune_multi_target.py --help` now import cleanly (verified
+in the `tr2d2-pep` conda env); a full generation run still needs the Google-Drive artifacts + a GPU.
+
+**Fixed — do NOT reintroduce:**
+- `td3b/td3b_finetune.py`: the `from plotting import ...` (no such module) is now `try/except`-guarded,
+  and `loss_wdce` is imported **lazily inside `td3b_finetune()`** to break a `finetune_utils ↔ td3b`
+  circular import. Keep both — `td3b/__init__.py` eagerly imports `td3b_finetune`, and `finetune_utils`
+  imports the `td3b` package at module load, so any module-level `from training.finetune_utils import …`
+  in `td3b_finetune.py` re-creates the cycle (it was inference.py's first crash, before plotting).
+- Stale `tr2d2-pep/` prefix stripped from the tokenizer loader (`finetune_utils.load_tokenizer`), every
+  classifier loader (`scoring/functions/*.py`, `scoring/scoring_functions.py`, `binding.py`), the
+  `Diffusion` fallback tokenizer, and the training results dir. Assets now resolve to the README layout
+  (`tokenizer/`, `scoring/functions/classifiers/`, `results/`).
+- `inference.py` reward wiring rewritten: build `MultiTargetBindingAffinity` + `DirectionalOracle` once,
+  wrap each target with `TargetSpecificBindingAffinity`, call `create_td3b_reward_function`. (It used to
+  call a non-existent `create_reward_function` signature whose `TypeError` was swallowed → empty CSV.)
+- Demo notebook: `from models.diffusion import Diffusion` (was `from diffusion import …`), clone URL now
+  points at the OSS HF repo with `GIT_LFS_SKIP_SMUDGE=1`, `total_memory` typo fixed.
+- Added `__init__.py` to the 8 dirs `find_packages()` missed; fixed corrupted `configs/peptune_config.yaml`
+  key (`batchinohup ng` → `batching`).
+
+**Fixed — round 2 (runtime hardening, 2026-07-10; verified dynamically in `tr2d2-pep`):**
+- `inference.py` Algorithm-2 resampling: now gates candidates by `finite-reward AND valid-peptide`
+  **before** sampling and draws **without replacement** (`k=min(val_samples_per_target, n_eligible)`).
+  Previously `replacement=True` + the peaked softmax produced duplicate rows (inflated counts, skewed
+  means) and validity was filtered only afterward (could save 0 samples despite valid candidates).
+- Checkpoint-load guards (silent-random-weights class): `inference.py::load_model` now raises if the
+  backbone loaded **no** weights and warns on partial loads; `direction_oracle.py::TR2D2RoFormerEncoder`
+  now unwraps `model_state_dict`/`state_dict` and raises if **zero** RoFormer keys matched (was silently
+  leaving the ligand encoder random); `binding.py` tolerates raw/`state_dict`/`model_state_dict` ckpt
+  containers; `_load_state_dict_flexible` loudly flags missing **non-ESM** (trained) keys. All guards
+  fail only on the impossible-for-a-valid-checkpoint case, so they can't break a good load.
+
+**Must NOT change — tied to pretrained weights:** `max_position_embeddings=1035`,
+`parameterization='subs'`, `T=0`, `time_conditioning=False`. `Diffusion.forward` hard-fails if `seq_len > 1035`.
+
+**Remaining known quirks (not inference blockers):**
+- **Checkpoint generates valid SMILES only at SHORT length (verified end-to-end 2026-07-12 on real weights
+  from `/data1/hanqun/TD3B/checkpoints`).** Valid-peptide yield vs generation length (SMILES tokens),
+  `is_peptide` over 32–64 samples: L=40 → 28% baseline / 59% `best_of_n`; L=60 → ~5–30%; **L≥150 → ~0%
+  regardless of step count (128/256/512) or best_of_n.** Real 25–30-residue binders are 150–210 tokens, i.e.
+  outside the valid regime — that is why a naive `inference.py` run writes an empty CSV. Fix applied:
+  `inference.py` now derives length from the binder's *token* count (not char count; reads the `SMILES`
+  column, not the nonexistent `Ligand_SMILES`), adds `--seq_length`/`--max_seq_length`, and hints on empty
+  output. Working run: `--seq_length 50 --sampler best_of_n --best_of_n 6` → 16 valid oracle-scored binders,
+  antagonist direction-accuracy ≈ 1.0. This is a model-capability limit, not a code bug.
+- **ESM2 network fetch:** `facebook/esm2_t33_650M_UR50D` is downloaded at inference twice — by the oracle
+  (`transformers`) and by `BindingAffinity`/`MultiTargetBindingAffinity` (`fair-esm`, no offline flag).
+  Offline runs need warm HF + torch-hub caches.
+- **Training default paths** (`finetune_multi_target.py:551`, factory fallback `td3b_scoring.py:345`) still
+  resolve under `{base_path}/tr2d2-pep/...` with a wrong default oracle filename
+  (`best_model_tr2d2_gpcr_fixed.pt` vs. shipped `direction_oracle.pt`). `launch_multi_target.sh` overrides
+  all of these with explicit `checkpoints/` args, so training-via-launch-script works; a bare
+  `python finetune_multi_target.py` does not.
+- `baselines/run.sh` loads `pretrained/peptune-pretrained.ckpt` (not `checkpoints/pretrained.ckpt`) and
+  defaults `CSV_PATH="To Be Added"` — pass real args.
+- `td3b/td3b_finetune.py::td3b_finetune()` is legacy/unused (superseded by `finetune_multi_target.py`);
+  still writes to `{base}/TR2-D2/tr2d2-pep/results/...` and needs the optional `plotting` module.
+- `TD3BDataset`/`load_td3b_data` (oracle-training path) require CSV columns `Action`, `Ligand_UniProt_ID`
+  that the inference/finetune CSVs don't carry.
+- Dead/inert flags: `--contrastive_type` (loss always uses `'margin'`), `--num_epoch_for_sampling`,
+  `min_affinity_threshold=0.0` (down-weight branch never fires), `use_confidence_weighting` (stored but
+  never applied in `compute_gated_reward`).
+- **`PeptideAnalyzer.is_peptide` (`utils/app.py`) false-positives** on atom-only / pure-amino-acid-letter
+  strings (e.g. `"CCNCCF"`, `"cccccc"` → True) because it checks AA-letter membership before RDKit. Harmless
+  for a well-trained model (emits real peptide SMILES) but can inflate `valid_mask` for an untrained/early
+  checkpoint. Shared by MCTS/baselines — do not change its semantics without checking all callers.
+- Oracle confidence κ is actually ∈[0.5,1] (max of a 2-class softmax), not the documented [0,1]; unused at
+  inference. `TD3BConfidenceWeighting.compute_importance_weights` uses raw `exp(reward/α)` which can overflow
+  for large affinities — off the inference path (MCTS/training only).
+- **Device coupling:** `models/roformer.py` and helpers pick their own device independently of `--device`
+  (some hardcode `cuda:0`); `resolve_device` in `scoring_functions.py`/`direction_oracle.py` silently falls
+  back to `cuda:0`/CPU. Mixed-device errors are easy to create.
+
+**New entry points (2026-07-12) — usage caveats:**
+- **The validity toggle is a FILTER, not a reward term.** `finetune_on_target.py --validity_reward {on,off}`
+  (default `on`) never changes the reward formula (`affinity × direction`); it toggles the
+  `PeptideAnalyzer.is_peptide` gate on **both** halves — the finetune-side MCTS expansion (forwarded to
+  `finetune_multi_target.py --validity_reward` → `args.enforce_validity`) and the generation-side Algorithm-2
+  resampling. `off` retains invalid samples (each output row still records `is_valid`).
+  `--finetune_validity_hook {on,off}` decouples the finetune-side gate from the generation-side toggle.
+- **Function B benchmarks were on a RANDOM-INIT model** → only *relative* trends (which strategy helps most at
+  long length) are meaningful; **absolute** valid-yields need the real checkpoint. `generate_valid.py`
+  **silently falls back to random init** when `--ckpt_path` is missing/omitted (it prints a NOTE, but the
+  numbers are garbage) �� always pass a real `--ckpt_path` for reportable yields.
+- `finetune_on_target.py` runs the finetune half as a **subprocess** (default `WANDB_MODE=disabled`) and locates
+  the produced checkpoint by diffing `results/<run_name>_*` dirs before/after (newest `model_final.ckpt`, else
+  newest `model_epoch_*.ckpt`); a renamed/failed run dir would break that discovery. Its own oracle/path defaults
+  are passed through explicitly so the subprocess never falls back to `finetune_multi_target.py`'s legacy
+  `{base}/tr2d2-pep/...` defaults.

README.md

@@ -11,6 +11,8 @@
 
 TD3B is a sequence-based generative framework that designs peptide binders with specified agonist or antagonist behavior. It combines a Direction Oracle, a soft binding-affinity gate, and amortized fine-tuning of a pre-trained discrete diffusion model (MDLM).
 
+> **Development release.** This repository (`ChatterjeeLab/TD3B-dev`) carries the full code. The heavy artifacts — trained checkpoints, training/test data, and generated binders — are distributed as a single archive on Google Drive (see [Data and Checkpoints](#data-and-checkpoints)). For the clean code-only release see [`ChatterjeeLab/TD3B`](https://huggingface.co/ChatterjeeLab/TD3B).
+
 ## Installation
 
 ```bash
@@ -29,23 +31,35 @@ To run it: download the notebook from this repository, open [Google Colab](https
 
 ## Data and Checkpoints
 
-The checkpoints ship with this repository. The expected layout is:
+The checkpoints, datasets, and generated binders are bundled in a single archive on Google Drive:
+
+**Download:** [`td3b_dev_artifacts.zip`](https://drive.google.com/file/d/1EU0-pcF-UoEb0OEMuO4qhodZ5RjMU4n9/view?usp=sharing) (~3.4 GB)
+
+Unzip it at the repository root to restore the full layout:
+
+```bash
+unzip td3b_dev_artifacts.zip -d .
+```
 
 ```
-TD3B/
+TD3B-dev/
 ├── checkpoints/
-│   ├── pretrained.ckpt          # Pre-trained MDLM weights
-│   ├── td3b.ckpt                # Fine-tuned TD3B model
-│   └── direction_oracle.pt      # Direction Oracle weights
+│   ├── pretrained.ckpt              # Pre-trained MDLM weights (1.4 GB)
+│   ├── td3b.ckpt                    # Fine-tuned TD3B model (231 MB)
+│   └── direction_oracle.pt          # Direction Oracle weights (2.85 GB)
 ├── scoring/functions/classifiers/
-│   ├── binding-affinity.pt
-│   ├── hemolysis-xgboost.json
-│   ├── nonfouling-xgboost.json
-│   ├── permeability-xgboost.json
-│   └── solubility-xgboost.json
-└── tokenizer/
-    ├── new_vocab.txt
-    └── new_splits.txt
+│   ├── binding-affinity.pt          # (from archive)
+│   ├── permeability-xgboost.json    # (from archive)
+│   ├── hemolysis-xgboost.json       # (in repo)
+│   ├── nonfouling-xgboost.json      # (in repo)
+│   └── solubility-xgboost.json      # (in repo)
+├── data/
+│   ├── train.csv                    # Training set (target-binder pairs)
+│   ├── test.csv                     # Test set
+│   └── td3b_data_new.csv            # Full labeled target/ligand pairs (agonist/antagonist)
+└── generated_binders/
+    ├── agonist/                     # 1106 generated agonist binders (.pdb + .trb), 249 targets
+    └── antagonist.tar.gz            # Generated antagonist binders
 ```
 
 ## Code Structure
@@ -53,6 +67,9 @@ TD3B/
 ```
 TD3B/
 ├── inference.py                 # Generate binders (main inference entry point)
+├── finetune_on_target.py        # Finetune on your own target(s) + generate (Function A)
+├── generate_valid.py            # Validity-boosting sampling CLI (Function B)
+├── sampling_strategies.py       # Validity-boosting sampler library (Function B)
 ├── finetune_multi_target.py     # Multi-target TD3B training
 ├── launch_multi_target.sh       # Training launcher script
 ├── models/
@@ -95,21 +112,88 @@ python inference.py \
 
 This generates 32 candidates per (target, direction), scores them with the Direction Oracle and affinity predictor, applies Algorithm 2 weighted resampling, and saves only valid peptide samples.
 
-Output: `results/td3b_results_seed42.csv` with columns: target, sequence, direction, affinity, gated_reward, direction_oracle, direction_accuracy.
+> **Generation length matters.** The length (in SMILES tokens) comes from `--seq_length` if given, else it is
+> derived from each row's reference binder (correctly tokenized). This released checkpoint reliably produces
+> **valid** SMILES only at **short lengths (≲100 tokens)**; long binders (e.g. 25–30-residue peptides ≈ 150–210
+> tokens) yield few or no valid samples regardless of step count. For non-empty output on such targets, pass a
+> shorter `--seq_length` and a validity-boosting sampler (see below), e.g.:
+> ```bash
+> python inference.py --ckpt_path checkpoints/td3b.ckpt --val_csv data/test.csv \
+>     --seq_length 50 --sampler best_of_n --best_of_n 6 --num_pool 32 --val_samples_per_target 4
+> ```
+> On a real run this yields valid, oracle-scored binders in both directions (antagonist direction accuracy ≈ 1.0).
+
+Output: `results/td3b_results_seed42.csv` with columns: target, target_uid, sequence, direction_name, target_direction, is_valid, affinity, gated_reward, direction_oracle, direction_accuracy.
+
+## Finetune on your own target(s)
+
+`finetune_on_target.py` takes protein target(s) you supply, finetunes the pretrained TD3B policy on **only those target(s)**, then generates directional (agonist/antagonist) binders for them. It reuses the existing machinery: the finetune half subprocess-invokes `finetune_multi_target.py` (with `K` set to the number of targets), and the generation half runs in-process with the same reward, oracle, and Algorithm 2 resampling as `inference.py`.
+
+```bash
+python finetune_on_target.py \
+    --target_seq MKTAYIAKQRQISFVKSHFSRQLEERLGLIEVQ \
+    --target_seq GSHMKELVLALYDYQEKSPREVTMKKGDILTLL \
+    --direction both \
+    --validity_reward on \
+    --num_epochs 20 \
+    --num_pool 32 \
+    --gen_samples_per_target 8 \
+    --device cuda:0 --seed 42
+```
+
+Provide targets by repeating `--target_seq` and/or via `--targets_csv` (a CSV with a `Target_Sequence` column; optional `Ligand_Sequence`/`label` rows seed the per-direction length prior):
+
+```bash
+python finetune_on_target.py --targets_csv my_targets.csv --direction agonist --binder_length 20
+```
+
+Key flags:
+- `--target_seq SEQ` (repeatable) and/or `--targets_csv` — provide at least one target.
+- `--direction {agonist,antagonist,both}` (default `both`) — which binders to **generate**. Finetuning always searches both directions.
+- `--validity_reward {on,off}` (default `on`) — validity gate applied to **both** halves. `on` keeps the `is_peptide` filter (invalid MCTS children are zero-rewarded during finetuning; only valid peptides are eligible for resampling during generation). `off` drops the gate on both halves (invalid samples are retained, each row keeps `is_valid`); the reward itself is always affinity × direction — validity is only ever a filter, never part of the reward formula. `--finetune_validity_hook {on,off}` overrides just the finetune-side gate.
+- `--binder_length` (default 20) — placeholder binder length (residues) used to seed the generation-length prior for any target with no known binder.
+- `--skip_finetune` (generate directly from `--td3b_checkpoint`) / `--skip_generate` (finetune only).
+- Training knobs: `--num_epochs` (20), `--num_iter` (10) / `--num_children` (16) (MCTS), `--learning_rate` (3e-4), `--seq_length` (200). Generation knobs: `--num_pool` (32), `--gen_samples_per_target` (8), `--resample_alpha` (0.1), `--total_num_steps` (128).
+- Paths default to the repo root; override `--pretrained_checkpoint`, `--direction_oracle_ckpt`, `--output_dir`, `--device` (`auto`/`cpu`/`cuda:N`), `--seed` as needed.
+
+Output: `results/finetune_on_target/binders_<direction>_validity-<on|off>_seed<seed>.csv` with columns target, sequence, direction_name, target_direction, is_valid, affinity, gated_reward, direction_oracle, direction_accuracy, gen_length. The finetuned checkpoint is written under `results/<run_name>_<timestamp>/` (used automatically for the generation half).
+
+## Higher-validity sampling for long binders
+
+As the target length grows, the fraction of decoded SMILES that pass the RDKit peptide check drops. `generate_valid.py` (backed by `sampling_strategies.py`) applies **sampling-time** strategies — no retraining — to raise the valid-peptide yield, especially at large length, and reports the valid fraction. The strategies only change token selection (temperature / top-k / top-p) and add a remask self-correction loop and best-of-N rejection on top of the model's existing diffusion primitives.
+
+```bash
+python generate_valid.py \
+    --ckpt_path checkpoints/td3b.ckpt \
+    --length 400 --num_samples 64 \
+    --strategy nucleus_remask \
+    --device cuda:0 --seed 42 \
+    --save_path results/valid_len400.csv
+```
+
+Key flags:
+- `--strategy` (default `nucleus_remask`) — one of `baseline`, `more_steps`, `top_p` (a.k.a. `nucleus`), `top_k`, `low_temp`, `remask`, `best_of_n`, `nucleus_remask`. On random-init benchmarks the self-correcting strategies (`remask`, `best_of_n`, `nucleus_remask`) give the largest relative validity gains at length ≥ 200; `nucleus_remask` is the recommended default.
+- `--length` (default 200, in tokens) and `--num_samples` (default 64).
+- `--ckpt_path` — TD3B checkpoint. If omitted or missing, a **random-init** model is used (exercises the sampling machinery only; yields are meaningless — pass a real checkpoint for real numbers).
+- Per-strategy overrides (else the preset default is used): `--top_p`, `--top_k`, `--temperature`, `--steps_per_token`, `--remask_rounds`, `--remask_frac`, `--remask_steps`, `--best_of_n`, `--num_steps`.
+
+Output: prints the valid yield (valid / `num_samples`) with per-round counts, and writes the valid sequences to `--save_path` (default `results/valid_<strategy>_len<L>.csv`) with columns idx, sequence, n_chars.
 
 ## Training
 
 ### Multi-target TD3B
 
-1. Edit `launch_multi_target.sh` — set paths to checkpoints, data, and oracle:
+1. Edit `launch_multi_target.sh` only if needed — `BASE_PATH` auto-detects the repo root (this script's own directory), and the checkpoint, data, and oracle paths derive from it:
 
 ```bash
-BASE_PATH="/path/to/TD3B"
+BASE_PATH="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"  # auto-detects the repo root
 PRETRAINED_CHECKPOINT="${BASE_PATH}/checkpoints/pretrained.ckpt"
 TRAIN_CSV="${BASE_PATH}/data/train.csv"
 ORACLE_CKPT="${BASE_PATH}/checkpoints/direction_oracle.pt"
 ```
 
+`finetune_multi_target.py --base_path` also defaults to the repo root, so you only need to override these when your checkpoints or data live elsewhere.
+
 2. Launch training:
 
 ```bash
@@ -138,10 +222,11 @@ bash run.sh --baseline tds --device cuda:0
 
 ```bibtex
 @inproceedings{
-  cao2026tdb,
-  title={{TD3B: Transition-Directed Discrete Diffusion for Allosteric Binder Generation}},
-  author={Hanqun Cao and Aastha Pal and Sophia Tang and Yinuo Zhang and Jingjie Zhang and Pheng-Ann Heng and Pranam Chatterjee},
-  booktitle={Forty-third International Conference on Machine Learning (Spotlight)},
-  year={2026}
+cao2026tdb,
+title={{TD}3B: Transition-Directed Discrete Diffusion for Allosteric Binder Generation},
+author={Hanqun Cao and Aastha Pal and Sophia Tang and Yinuo Zhang and Jingjie Zhang and Pheng-Ann Heng and Pranam Chatterjee},
+booktitle={Learning Meaningful Representations of Life (LMRL) Workshop at ICLR 2026},
+year={2026},
+url={https://openreview.net/forum?id=uMUicLylVp}
 }
 ```

baselines/run_mcts_tr2d2.py

@@ -48,7 +48,7 @@ def _extract_state_and_config(ckpt: Dict[str, Any]) -> Dict[str, Any]:
 
 def _build_args(cfg: Dict[str, Any], cli: argparse.Namespace) -> argparse.Namespace:
     defaults = {
-        "base_path": "To Be Added",
+        "base_path": ROOT_DIR,
         "seq_length": 200,
         "sampling_eps": 1e-3,
         "total_num_steps": 128,

baselines/run_validation_td3b.py

@@ -48,7 +48,7 @@ def _extract_state_and_config(ckpt: Dict[str, Any]) -> Dict[str, Any]:
 
 def _build_args(cfg: Dict[str, Any], cli: argparse.Namespace) -> argparse.Namespace:
     defaults = {
-        "base_path": "To Be Added",
+        "base_path": ROOT_DIR,
         "seq_length": 200,
         "sampling_eps": 1e-3,
         "total_num_steps": 128,

baselines/sampling_setup.py

@@ -300,7 +300,7 @@ def main():
         raise ValueError("--prot_seq is required when --targets_csv is not provided.")
 
     base_model = load_base_model(args.ckpt_path, args.device)
-    base_path = os.path.abspath(os.path.join(ROOT_DIR, ".."))
+    base_path = ROOT_DIR
     multi_target = args.targets_csv is not None
     scoring_fn = load_reward_models(
         args.prot_seq if not multi_target else None,

configs/peptune_config.yaml

@@ -36,7 +36,7 @@ lr_scheduler:
 data:
   train: To Be Added
   valid: To Be Added
-  batchinohup ng: wrapping # padding / wrapping
+  batching: wrapping # padding / wrapping
 
 loader:
   global_batch_size: 64

finetune_multi_target.py

@@ -397,8 +397,9 @@ def parse_args():
 
     # Paths
     path_group = parser.add_argument_group('Paths')
-    path_group.add_argument('--base_path', type=str, required=True,
-                           help='Base path for TR2-D2 project')
+    path_group.add_argument('--base_path', type=str,
+                           default=os.path.dirname(os.path.abspath(__file__)),
+                           help='Base path for TR2-D2 project (defaults to the repo root)')
     path_group.add_argument('--train_csv', type=str, required=True,
                            help='Path to training CSV file')
     path_group.add_argument('--val_csv', type=str, default=None,
@@ -455,6 +456,12 @@ def parse_args():
                            help='Temperature for importance weighting')
     mcts_group.add_argument('--exploration', type=float, default=1.0,
                            help='UCB exploration constant')
+    mcts_group.add_argument('--validity_reward', choices=['on', 'off'], default='on',
+                           help="Validity gate during MCTS expansion. 'on' (default): "
+                                "invalid decoded peptides are gated as zero-reward nodes "
+                                "(PeptideAnalyzer.is_peptide filter). 'off': disable the gate "
+                                "so invalid peptides are kept and scored on the pure "
+                                "affinity x direction reward.")
 
     # TD3B loss hyperparameters
     loss_group = parser.add_argument_group('TD3B Loss')
@@ -510,8 +517,8 @@ def parse_args():
     log_group = parser.add_argument_group('Logging')
     log_group.add_argument('--wandb_project', type=str, default='TD3B-multi-target',
                           help='W&B project name')
-    log_group.add_argument('--wandb_entity', type=str, default='phos_zj',
-                          help='W&B entity name')
+    log_group.add_argument('--wandb_entity', type=str, default=None,
+                          help='W&B entity name (default: your personal/default entity)')
 
     # Directional oracle
     oracle_group = parser.add_argument_group('Directional Oracle')
@@ -570,6 +577,9 @@ def parse_args():
     args.time_conditioning = False
     args.num_obj = 5  # Must match padded score vector size
     args.scalarization = "sum"
+    # Validity gate toggle -> consumed by MCTS (mcts/peptide_mcts.py) via the
+    # args threading through create_mcts_instance/create_td3b_mcts/TD3B_MCTS.
+    args.enforce_validity = (args.validity_reward == "on")
 
     # Create save path
     args.save_path = create_output_directory(

finetune_on_target.py

@@ -0,0 +1,701 @@
+#!/usr/bin/env python3
+"""
+TD3B — Finetune-on-Target entry point (BUILD FUNCTION A).
+
+User-facing workflow: the user supplies one or more protein *targets*; the system
+(1) FINETUNES the pretrained TD3B diffusion policy on ONLY those target(s), then
+(2) GENERATES directional (agonist/antagonist) peptide binders for them — with a
+toggle to include or exclude the validity reward.
+
+This script does NOT reimplement the training loop. It REUSES the existing
+machinery:
+
+  * FINETUNE half  -> runs ``finetune_multi_target.py`` as a subprocess on a
+    temporary CSV built from the provided target(s), with K (``--targets_per_mcts``)
+    set to the number of provided targets. That script builds
+    ``MultiTargetBindingAffinity`` + ``DirectionalOracle``, runs ``TD3B_MCTS`` per
+    (target, direction) to search candidate binders, and distills them into the
+    diffusion policy via WDCE + KL + contrastive losses.
+
+  * GENERATE half  -> runs IN-PROCESS, reusing ``inference.py``'s
+    ``load_model`` / ``sample_sequences`` / ``score_sequences`` and
+    ``td3b.td3b_scoring.create_td3b_reward_function`` (gated reward
+    R = g_psi * sigma(d*(f_phi - 0.5)/tau)) plus the Algorithm-2 weighted
+    resampling. Running this half in-process (rather than shelling out to
+    ``inference.py``) is what lets us honor ``--validity_reward off`` at the
+    generation/scoring layer — see the VALIDITY TOGGLE section below.
+
+--------------------------------------------------------------------------------
+VALIDITY TOGGLE  (``--validity_reward {on,off}``)
+--------------------------------------------------------------------------------
+Validity never enters the reward *formula* itself (the reward is always
+affinity x direction). It enters TD3B only as a FILTER: ``PeptideAnalyzer.is_peptide``
+gates MCTS expansion during finetuning and gates the Algorithm-2 resampling during
+generation. This flag now controls BOTH halves:
+
+  * ``on``  (default): keep the validity gate on both halves. During finetuning,
+    invalid decoded children are gated as zero-reward MCTS nodes; during
+    generation only valid peptides are eligible for resampling (identical to
+    ``inference.py``).
+  * ``off``: drop the validity gate on both halves. FINETUNE side — forwarded as
+    ``--validity_reward off`` to ``finetune_multi_target.py``, which disables the
+    ``is_peptide`` gate in MCTS expansion (``mcts/peptide_mcts.py``) so invalid
+    peptides are kept and scored on the pure affinity x direction reward instead
+    of being zero-rewarded. GENERATION side — every finite-reward candidate is
+    eligible for resampling and invalid samples are retained in the output (each
+    row still records ``is_valid`` for inspection). The reward stays pure
+    affinity x direction throughout.
+
+Use ``--finetune_validity_hook {on,off}`` to override the finetune-side gate
+independently of the generation-side toggle (default: follow ``--validity_reward``).
+
+--------------------------------------------------------------------------------
+Also note (reuse-without-editing limitations):
+  * ``finetune_multi_target.py`` ALWAYS searches both agonist and antagonist per
+    target in its MCTS phase. The training CSV therefore always seeds a length
+    prior for BOTH directions (real binder length if provided, else a placeholder
+    of ``--binder_length`` residues). ``--direction`` restricts only what the
+    GENERATE half emits, not what finetuning trains on.
+
+Base paths DEFAULT to the repo root (auto-detected from ``__file__``); nothing is
+hardcoded to an absolute path. Heavy artifacts (checkpoints / ESM / binding
+predictor) are validated up-front with a clear error if missing.
+"""
+
+import argparse
+import glob
+import logging
+import os
+import subprocess
+import sys
+import tempfile
+from typing import Dict, List, Optional, Tuple
+
+import numpy as np
+import pandas as pd
+
+# Repo root — every default path is derived from this, never hardcoded absolute.
+REPO_ROOT = os.path.dirname(os.path.abspath(__file__))
+if REPO_ROOT not in sys.path:
+    sys.path.insert(0, REPO_ROOT)
+
+logger = logging.getLogger("finetune_on_target")
+logging.basicConfig(level=logging.INFO,
+                    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+
+AA_SET = set("ACDEFGHIKLMNPQRSTVWY")
+DIRECTIONS = {"agonist": 1.0, "antagonist": -1.0}
+GDRIVE_HINT = ("These heavy artifacts ship separately (see README 'Data and "
+               "Checkpoints' — ~3.4 GB on Google Drive). Unzip them at the repo "
+               "root to restore checkpoints/ and scoring/functions/classifiers/.")
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Target-set construction
+# ───────────────────────────────────────────────────────────��─────────────────
+def _looks_like_protein(seq: str) -> bool:
+    if not isinstance(seq, str):
+        return False
+    s = seq.strip().upper()
+    return len(s) >= 2 and all(c in AA_SET for c in s)
+
+
+def make_placeholder_binder(binder_length: int, mode: str = "polyG",
+                            rng: Optional[np.random.Generator] = None) -> str:
+    """Synthesize a placeholder binder purely to seed the per-direction LENGTH
+    prior (never scored as a real binder). Poly-glycine by default."""
+    n = max(2, int(binder_length))
+    if mode == "random":
+        rng = rng or np.random.default_rng(0)
+        return "".join(rng.choice(list(AA_SET), size=n))
+    return "G" * n
+
+
+def build_target_training_frame(
+    target_seqs: List[str],
+    targets_csv: Optional[str],
+    binder_length: int,
+    placeholder_mode: str = "polyG",
+    seed: int = 42,
+) -> Tuple[pd.DataFrame, List[str]]:
+    """Build the normalized training frame consumed by ``finetune_multi_target``.
+
+    Columns: ``Target_Sequence``, ``Ligand_Sequence``, ``label``.
+
+    For every unique target we ensure BOTH directions have at least one row so the
+    downstream ``TargetDataset`` can compute a per-direction median binder length.
+    Real (Ligand_Sequence, label) rows from the CSV are kept as-is; any missing
+    direction is seeded with a placeholder binder of ``binder_length`` residues.
+
+    Returns (frame, unique_target_order).
+    """
+    rng = np.random.default_rng(seed)
+    # target_seq -> {'agonist': [binders...], 'antagonist': [...]}
+    per_target: Dict[str, Dict[str, List[str]]] = {}
+    order: List[str] = []
+
+    def _ensure(t: str) -> Dict[str, List[str]]:
+        if t not in per_target:
+            per_target[t] = {"agonist": [], "antagonist": []}
+            order.append(t)
+        return per_target[t]
+
+    # 1) Explicit --target_seq entries (no known binders).
+    for t in target_seqs or []:
+        t = t.strip()
+        if not t:
+            continue
+        if not _looks_like_protein(t):
+            logger.warning("Target sequence does not look like a protein (non-AA "
+                           "characters); using it anyway: %.40s...", t)
+        _ensure(t)
+
+    # 2) Optional CSV with Target_Sequence (+ optional Ligand_Sequence/label).
+    if targets_csv:
+        if not os.path.isfile(targets_csv):
+            raise FileNotFoundError(f"--targets_csv not found: {targets_csv}")
+        df = pd.read_csv(targets_csv)
+        if "Target_Sequence" not in df.columns:
+            raise ValueError(f"{targets_csv} must contain a 'Target_Sequence' column "
+                             f"(found: {list(df.columns)})")
+        for _, row in df.iterrows():
+            t = str(row["Target_Sequence"]).strip()
+            if not t:
+                continue
+            bucket = _ensure(t)
+            binder = row.get("Ligand_Sequence")
+            label = str(row.get("label", "")).strip().lower()
+            if isinstance(binder, str) and binder.strip() and label in DIRECTIONS:
+                bucket[label].append(binder.strip())
+
+    if not order:
+        raise ValueError("No targets provided. Pass at least one --target_seq or a "
+                         "--targets_csv with a 'Target_Sequence' column.")
+
+    # 3) Materialize rows, seeding placeholders for any empty direction.
+    rows = []
+    for t in order:
+        for direction in ("agonist", "antagonist"):
+            binders = per_target[t][direction]
+            if not binders:
+                binders = [make_placeholder_binder(binder_length, placeholder_mode, rng)]
+            for b in binders:
+                rows.append({"Target_Sequence": t, "Ligand_Sequence": b, "label": direction})
+
+    frame = pd.DataFrame(rows, columns=["Target_Sequence", "Ligand_Sequence", "label"])
+    return frame, order
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Checkpoint / artifact validation
+# ─────────────────────────────────────────────────────────────────────────────
+def _require_files(pairs: List[Tuple[str, str]], stage: str) -> None:
+    """Raise a single clear error listing every missing artifact for ``stage``."""
+    missing = [(label, path) for label, path in pairs if not os.path.isfile(path)]
+    if missing:
+        lines = "\n".join(f"  - {label}: {path}" for label, path in missing)
+        raise FileNotFoundError(
+            f"Cannot run the {stage} stage — required artifact(s) are missing:\n"
+            f"{lines}\n{GDRIVE_HINT}"
+        )
+
+
+def finetune_artifact_checks(args) -> List[Tuple[str, str]]:
+    return [
+        ("pretrained_checkpoint", args.pretrained_checkpoint),
+        ("direction_oracle_ckpt", args.direction_oracle_ckpt),
+        ("direction_oracle_tr2d2_checkpoint", args.direction_oracle_tr2d2_checkpoint),
+        ("direction_oracle_tokenizer_vocab", args.direction_oracle_tokenizer_vocab),
+        ("direction_oracle_tokenizer_splits", args.direction_oracle_tokenizer_splits),
+        ("binding_affinity_predictor",
+         os.path.join(args.base_path, "scoring", "functions", "classifiers", "binding-affinity.pt")),
+    ]
+
+
+def generate_artifact_checks(args, ckpt_path: str) -> List[Tuple[str, str]]:
+    return [
+        ("td3b_checkpoint", ckpt_path),
+        ("direction_oracle_ckpt", args.direction_oracle_ckpt),
+        ("direction_oracle_tr2d2_checkpoint", args.direction_oracle_tr2d2_checkpoint),
+        ("binding_affinity_predictor",
+         os.path.join(args.base_path, "scoring", "functions", "classifiers", "binding-affinity.pt")),
+    ]
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# FINETUNE half  (subprocess reuse of finetune_multi_target.py)
+# ─────────────────────────────────────────────────────────────────────────────
+def run_finetune(args, train_csv: str, num_targets: int) -> str:
+    """Finetune on the provided target(s) by invoking finetune_multi_target.py.
+
+    Returns the path to the finetuned checkpoint (model_final.ckpt, else the
+    newest model_epoch_*.ckpt) discovered in the run's output directory.
+    """
+    _require_files(finetune_artifact_checks(args), "finetune")
+
+    # Validity gate forwarded to the finetune-side MCTS. Follows --validity_reward
+    # unless --finetune_validity_hook explicitly overrides it (decoupling the
+    # finetune-side gate from the generation-side toggle).
+    finetune_validity = args.finetune_validity_hook or args.validity_reward
+
+    results_root = os.path.join(args.base_path, "results")
+    os.makedirs(results_root, exist_ok=True)
+    # finetune_multi_target appends a timestamp to run_name -> record pre-existing
+    # dirs so we can identify the NEW one it creates.
+    before = set(glob.glob(os.path.join(results_root, f"{args.run_name}_*")))
+
+    # K = number of provided targets (train on ONLY these).
+    cmd = [
+        sys.executable, os.path.join(args.base_path, "finetune_multi_target.py"),
+        "--base_path", args.base_path,
+        "--train_csv", train_csv,
+        "--pretrained_checkpoint", args.pretrained_checkpoint,
+        "--run_name", args.run_name,
+        "--device", args.device,
+        "--targets_per_mcts", str(num_targets),
+        "--resample_targets_every", "1",
+        "--num_epochs", str(args.num_epochs),
+        "--learning_rate", str(args.learning_rate),
+        "--train_batch_size", str(args.train_batch_size),
+        "--gradient_accumulation_steps", str(args.gradient_accumulation_steps),
+        "--resample_every_n_step", str(args.resample_every_n_step),
+        "--save_every_n_epochs", str(args.save_every_n_epochs),
+        "--reset_every_n_step", str(max(1, args.num_epochs)),  # reset tree once, at epoch 0
+        "--num_iter", str(args.num_iter),
+        "--num_children", str(args.num_children),
+        "--buffer_size", str(args.buffer_size),
+        "--validity_reward", finetune_validity,  # forwards the validity gate to finetune-side MCTS
+        "--alpha", str(args.alpha),
+        "--min_affinity_threshold", str(args.min_affinity_threshold),
+        "--sigmoid_temperature", str(args.sigmoid_temperature),
+        "--seq_length", str(args.seq_length),
+        "--wandb_project", args.wandb_project,
+        # oracle wiring (pass our repo-root defaults through explicitly so the
+        # subprocess does not fall back to its own legacy default paths)
+        "--direction_oracle_ckpt", args.direction_oracle_ckpt,
+        "--direction_oracle_tr2d2_checkpoint", args.direction_oracle_tr2d2_checkpoint,
+        "--direction_oracle_tokenizer_vocab", args.direction_oracle_tokenizer_vocab,
+        "--direction_oracle_tokenizer_splits", args.direction_oracle_tokenizer_splits,
+        "--direction_oracle_esm_name", args.direction_oracle_esm_name,
+        "--direction_oracle_d_model", str(args.direction_oracle_d_model),
+        "--direction_oracle_n_heads", str(args.direction_oracle_n_heads),
+        "--direction_oracle_n_self_attn_layers", str(args.direction_oracle_n_self_attn_layers),
+        "--direction_oracle_n_bmca_layers", str(args.direction_oracle_n_bmca_layers),
+        "--direction_oracle_dropout", str(args.direction_oracle_dropout),
+    ]
+    if args.direction_oracle_esm_cache_dir:
+        cmd += ["--direction_oracle_esm_cache_dir", args.direction_oracle_esm_cache_dir]
+    if args.direction_oracle_esm_local_files_only:
+        cmd += ["--direction_oracle_esm_local_files_only"]
+    if args.grad_clip:
+        cmd += ["--grad_clip", "--gradnorm_clip", str(args.gradnorm_clip)]
+
+    env = dict(os.environ)
+    env.setdefault("WANDB_MODE", args.wandb_mode)  # default: disabled (offline, no prompt)
+
+    logger.info("Launching finetune subprocess (K=%d target(s)):\n  %s",
+                num_targets, " ".join(cmd))
+    proc = subprocess.run(cmd, cwd=args.base_path, env=env)
+    if proc.returncode != 0:
+        raise RuntimeError(
+            f"finetune_multi_target.py exited with code {proc.returncode}. "
+            f"See the subprocess log above for the underlying cause (e.g. a missing "
+            f"checkpoint, OOM, or oracle/ESM load failure)."
+        )
+
+    # Locate the checkpoint produced by this run.
+    after = set(glob.glob(os.path.join(results_root, f"{args.run_name}_*")))
+    new_dirs = sorted(after - before, key=os.path.getmtime)
+    search_dirs = new_dirs or sorted(after, key=os.path.getmtime)
+    if not search_dirs:
+        raise RuntimeError(
+            f"Finetuning finished but no results dir matching "
+            f"{results_root}/{args.run_name}_* was found; cannot locate the "
+            f"finetuned checkpoint.")
+    run_dir = search_dirs[-1]
+    final_ckpt = os.path.join(run_dir, "model_final.ckpt")
+    if os.path.isfile(final_ckpt):
+        ckpt = final_ckpt
+    else:
+        epoch_ckpts = sorted(glob.glob(os.path.join(run_dir, "model_epoch_*.ckpt")),
+                             key=os.path.getmtime)
+        if not epoch_ckpts:
+            raise RuntimeError(
+                f"Finetuning finished but no checkpoint (model_final.ckpt or "
+                f"model_epoch_*.ckpt) was found in {run_dir}.")
+        ckpt = epoch_ckpts[-1]
+    logger.info("Finetuned checkpoint: %s", ckpt)
+    return ckpt
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# GENERATE half  (in-process reuse of inference.py + td3b scoring)
+# ─────────────────────────────────────────────────────────────────────────────
+# These three thin seams are what the smoke test monkeypatches to inject a
+# random-init tiny model + stub reward (so the generation control flow can run on
+# CPU without the 3.4 GB artifacts / ESM).
+def build_generation_model(ckpt_path: str, device):
+    """Load the finetuned TD3B model. Reuses inference.load_model (which refuses
+    to run on an all-random backbone, guarding against silent garbage output)."""
+    from inference import load_model
+    return load_model(ckpt_path, device)
+
+
+def build_reward_components(args, device, tokenizer, model):
+    """Build the shared (expensive) reward components ONCE: the multi-target
+    affinity predictor (ESM2 + binding-affinity.pt) and the Direction Oracle."""
+    from scoring.functions.binding import MultiTargetBindingAffinity
+    from td3b.direction_oracle import DirectionalOracle
+
+    multi_affinity = MultiTargetBindingAffinity(
+        tokenizer=tokenizer, base_path=args.base_path, device=device,
+        emb_model=model.backbone,
+    )
+    directional_oracle = DirectionalOracle(
+        model_ckpt=args.direction_oracle_ckpt,
+        tr2d2_checkpoint=args.direction_oracle_tr2d2_checkpoint,
+        tokenizer_vocab=args.direction_oracle_tokenizer_vocab,
+        tokenizer_splits=args.direction_oracle_tokenizer_splits,
+        esm_name=args.direction_oracle_esm_name,
+        d_model=args.direction_oracle_d_model,
+        n_heads=args.direction_oracle_n_heads,
+        n_self_attn_layers=args.direction_oracle_n_self_attn_layers,
+        n_bmca_layers=args.direction_oracle_n_bmca_layers,
+        dropout=args.direction_oracle_dropout,
+        max_ligand_length=args.direction_oracle_max_ligand_length,
+        max_protein_length=args.direction_oracle_max_protein_length,
+        device=device,
+        esm_cache_dir=args.direction_oracle_esm_cache_dir,
+        esm_local_files_only=args.direction_oracle_esm_local_files_only,
+    )
+    directional_oracle.eval()
+    return multi_affinity, directional_oracle
+
+
+def make_reward_function(multi_affinity, directional_oracle, target_seq, direction,
+                         tokenizer, device, args):
+    """Bind the shared reward components to (target, direction). Reuses the exact
+    gated reward from td3b.td3b_scoring."""
+    from scoring.functions.binding import TargetSpecificBindingAffinity
+    from td3b.td3b_scoring import create_td3b_reward_function
+
+    target_affinity = TargetSpecificBindingAffinity(multi_affinity, target_seq)
+    return create_td3b_reward_function(
+        affinity_predictor=target_affinity,
+        directional_oracle=directional_oracle,
+        target_protein_seq=target_seq,
+        target_direction=direction,
+        peptide_tokenizer=tokenizer,
+        device=device,
+        min_affinity_threshold=args.min_affinity_threshold,
+        temperature=args.sigmoid_temperature,
+    )
+
+
+def generate_binders(args, model, tokenizer, multi_affinity, directional_oracle,
+                     targets, length_provider, device) -> pd.DataFrame:
+    """Sample -> score -> Algorithm-2 resample, for every (target, requested
+    direction). Honors the validity toggle at the resampling gate.
+
+    Args:
+        targets: ordered list of unique target protein sequences.
+        length_provider: fn(target_seq, direction_name) -> int generation length.
+    Returns a results DataFrame (also written to disk by the caller).
+    """
+    import torch
+    from inference import sample_sequences, score_sequences
+    from utils.app import PeptideAnalyzer
+
+    analyzer = PeptideAnalyzer()
+    validity_on = (args.validity_reward == "on")
+    directions = (["agonist", "antagonist"] if args.direction == "both" else [args.direction])
+    columns = ["target", "target_full", "sequence", "direction_name", "target_direction",
+               "is_valid", "validity_reward", "affinity", "gated_reward",
+               "direction_oracle", "direction_accuracy", "gen_length"]
+    records = []
+
+    for tidx, target_seq in enumerate(targets):
+        for d_name in directions:
+            d_star = DIRECTIONS[d_name]
+            logger.info("[%d/%d] target=%.20s... direction=%s",
+                        tidx + 1, len(targets), target_seq, d_name)
+
+            reward_model = make_reward_function(
+                multi_affinity, directional_oracle, target_seq, d_name,
+                tokenizer, device, args)
+
+            gen_len = int(max(1, length_provider(target_seq, d_name)))
+            x_pool = sample_sequences(model, args.num_pool, gen_len, args.total_num_steps)
+            sequences = tokenizer.batch_decode(x_pool)
+
+            valid_mask = np.array([analyzer.is_peptide(s) for s in sequences], dtype=bool)
+            gated_rewards, affinities, oracle_dirs, _conf = score_sequences(reward_model, sequences)
+            direction_accuracy = ((oracle_dirs > 0.5).astype(float) if d_star > 0
+                                  else (oracle_dirs < 0.5).astype(float))
+
+            # ── VALIDITY TOGGLE ───────────────────────────────────────────────
+            # on : eligible = finite reward AND valid peptide (== inference.py)
+            # off: eligible = finite reward only (validity gate skipped; invalid
+            #      samples retained; reward stays pure affinity x direction)
+            finite = np.isfinite(gated_rewards)
+            eligible = (finite & valid_mask) if validity_on else finite
+
+            if eligible.any():
+                rewards_t = torch.as_tensor(gated_rewards[eligible], device=device,
+                                            dtype=torch.float32)
+                alpha = max(args.resample_alpha, 1e-6)
+                weights = torch.softmax(rewards_t / alpha, dim=0)
+                k = min(args.gen_samples_per_target, int(eligible.sum()))
+                idx = torch.multinomial(weights, num_samples=k, replacement=False)
+                chosen = np.where(eligible)[0][idx.cpu().numpy()]
+            else:
+                logger.warning("No eligible candidates for target=%.20s dir=%s "
+                               "(validity_reward=%s).", target_seq, d_name, args.validity_reward)
+                chosen = np.array([], dtype=int)
+
+            for i in chosen:
+                records.append({
+                    "target": target_seq[:20],
+                    "target_full": target_seq,
+                    "sequence": sequences[i],
+                    "direction_name": d_name,
+                    "target_direction": d_star,
+                    "is_valid": bool(valid_mask[i]),
+                    "validity_reward": args.validity_reward,
+                    "affinity": float(affinities[i]),
+                    "gated_reward": float(gated_rewards[i]),
+                    "direction_oracle": float(oracle_dirs[i]),
+                    "direction_accuracy": float(direction_accuracy[i]),
+                    "gen_length": gen_len,
+                })
+
+    # Always return the defined columns so an empty result still writes a header
+    # row (clearer than a 0-byte file — never silently produce empty output).
+    return pd.DataFrame(records, columns=columns)
+
+
+def run_generation(args, ckpt_path: str, train_csv: str, targets: List[str]) -> str:
+    """Orchestrate the generate half and write the results CSV. Returns its path."""
+    import torch
+    _require_files(generate_artifact_checks(args, ckpt_path), "generate")
+
+    device = torch.device(args.device if (args.device != "auto") else
+                          ("cuda:0" if torch.cuda.is_available() else "cpu"))
+    if device.type == "cuda" and not torch.cuda.is_available():
+        logger.warning("CUDA requested but unavailable; using CPU.")
+        device = torch.device("cpu")
+
+    torch.manual_seed(args.seed)
+    np.random.seed(args.seed)
+
+    logger.info("Loading finetuned model: %s", ckpt_path)
+    model, tokenizer = build_generation_model(ckpt_path, device)
+
+    logger.info("Building reward components (affinity predictor + Direction Oracle)...")
+    multi_affinity, directional_oracle = build_reward_components(args, device, tokenizer, model)
+
+    # Per-direction generation length prior, reusing finetune_multi_target's
+    # TargetDataset (SMILES-token median length) for consistency with training.
+    length_provider = build_length_provider(train_csv, tokenizer, args.binder_length)
+
+    df = generate_binders(args, model, tokenizer, multi_affinity, directional_oracle,
+                          targets, length_provider, device)
+
+    os.makedirs(args.output_dir, exist_ok=True)
+    out_path = os.path.join(
+        args.output_dir,
+        f"binders_{args.direction}_validity-{args.validity_reward}_seed{args.seed}.csv")
+    df.to_csv(out_path, index=False)
+
+    if len(df):
+        logger.info("=" * 60)
+        logger.info("Wrote %d binder(s) -> %s", len(df), out_path)
+        for d_name in (["agonist", "antagonist"] if args.direction == "both" else [args.direction]):
+            sub = df[df["direction_name"] == d_name]
+            if len(sub):
+                logger.info("  %-10s n=%d  affinity=%.3f  gated=%.3f  dir_acc=%.3f  valid_frac=%.3f",
+                            d_name, len(sub), sub["affinity"].mean(), sub["gated_reward"].mean(),
+                            sub["direction_accuracy"].mean(), sub["is_valid"].mean())
+        logger.info("=" * 60)
+    else:
+        logger.warning("No binders generated (empty results). Wrote header-only CSV -> %s", out_path)
+    return out_path
+
+
+def build_length_provider(train_csv: str, tokenizer, fallback_binder_length: int):
+    """Return fn(target_seq, direction_name) -> generation length, reusing
+    finetune_multi_target.TargetDataset so the generated length matches the length
+    prior used during training. Falls back to a token-length estimate of a
+    placeholder of ``fallback_binder_length`` residues if a target is unknown."""
+    from finetune_multi_target import TargetDataset
+    from td3b.data_utils import peptide_seq_to_smiles, smiles_token_length
+
+    dataset = TargetDataset(train_csv, tokenizer=tokenizer)
+    fallback_len = smiles_token_length(
+        peptide_seq_to_smiles(make_placeholder_binder(fallback_binder_length)), tokenizer)
+
+    def provider(target_seq: str, direction_name: str) -> int:
+        try:
+            return dataset.get_sequence_length(target_seq, direction_name)
+        except KeyError:
+            return fallback_len
+
+    return provider
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# CLI
+# ─────────────────────────────────────────────────────────────────────────────
+def parse_args(argv=None):
+    p = argparse.ArgumentParser(
+        description="Finetune TD3B on user-provided protein target(s), then "
+                    "generate directional binders (with a validity-reward toggle).",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+
+    tgt = p.add_argument_group("Targets (provide at least one)")
+    tgt.add_argument("--target_seq", action="append", default=None, metavar="SEQ",
+                     help="Target protein amino-acid sequence. Repeatable for multiple targets.")
+    tgt.add_argument("--targets_csv", type=str, default=None,
+                     help="CSV with at least 'Target_Sequence' (optional 'Ligand_Sequence'/'label').")
+
+    ctrl = p.add_argument_group("Control")
+    ctrl.add_argument("--direction", choices=["agonist", "antagonist", "both"], default="both",
+                      help="Which directional binders to GENERATE (finetuning always trains both).")
+    ctrl.add_argument("--validity_reward", choices=["on", "off"], default="on",
+                      help="on: keep the is_peptide validity gate during BOTH finetuning "
+                           "(MCTS expansion) and generation; off: pure affinity x direction, "
+                           "no validity gate in either phase (see module docstring).")
+    ctrl.add_argument("--binder_length", type=int, default=20,
+                      help="Placeholder binder length (residues) used to seed the length prior "
+                           "when a target has no known binder.")
+    ctrl.add_argument("--placeholder_mode", choices=["polyG", "random"], default="polyG",
+                      help="Placeholder binder composition (length-prior seed only).")
+
+    stages = p.add_argument_group("Stages")
+    stages.add_argument("--skip_finetune", action="store_true",
+                        help="Skip finetuning and generate from --td3b_checkpoint directly.")
+    stages.add_argument("--skip_generate", action="store_true",
+                        help="Finetune only; do not generate.")
+
+    train = p.add_argument_group("Training knobs (small sane defaults)")
+    train.add_argument("--num_epochs", type=int, default=20)
+    train.add_argument("--num_iter", type=int, default=10, help="MCTS iterations per resample.")
+    train.add_argument("--num_children", type=int, default=16, help="Children per MCTS expansion.")
+    train.add_argument("--learning_rate", type=float, default=3e-4)
+    train.add_argument("--train_batch_size", type=int, default=8)
+    train.add_argument("--gradient_accumulation_steps", type=int, default=1)
+    train.add_argument("--resample_every_n_step", type=int, default=10)
+    train.add_argument("--save_every_n_epochs", type=int, default=20)
+    train.add_argument("--buffer_size", type=int, default=50)
+    train.add_argument("--alpha", type=float, default=0.1, help="Importance-weight temperature.")
+    train.add_argument("--min_affinity_threshold", type=float, default=0.0)
+    train.add_argument("--sigmoid_temperature", type=float, default=0.1)
+    train.add_argument("--seq_length", type=int, default=200, help="Max sequence length.")
+    train.add_argument("--total_num_steps", type=int, default=128, help="Diffusion steps for generation.")
+    train.add_argument("--grad_clip", action="store_true")
+    train.add_argument("--gradnorm_clip", type=float, default=1.0)
+    train.add_argument("--finetune_validity_hook", choices=["on", "off"], default=None,
+                       help="Optional override for the FINETUNE-side validity gate forwarded to "
+                            "finetune_multi_target.py's --validity_reward. Default (unset): follow "
+                            "--validity_reward. Set to decouple the finetune-side gate from the "
+                            "generation-side toggle.")
+
+    gen = p.add_argument_group("Generation knobs")
+    gen.add_argument("--num_pool", type=int, default=32, help="Candidates sampled per (target, direction).")
+    gen.add_argument("--gen_samples_per_target", type=int, default=8, help="Binders kept per (target, direction).")
+    gen.add_argument("--resample_alpha", type=float, default=0.1, help="Algorithm-2 resampling temperature.")
+
+    paths = p.add_argument_group("Paths (default to repo root, auto-detected)")
+    paths.add_argument("--base_path", type=str, default=REPO_ROOT,
+                       help="Repo root; source of scoring/, tokenizer/, checkpoints/.")
+    paths.add_argument("--pretrained_checkpoint", type=str,
+                       default=os.path.join(REPO_ROOT, "checkpoints", "pretrained.ckpt"))
+    paths.add_argument("--td3b_checkpoint", type=str, default=None,
+                       help="Finetuned checkpoint for --skip_finetune generation. If omitted and "
+                            "finetuning ran, the produced checkpoint is used.")
+    paths.add_argument("--output_dir", type=str,
+                       default=os.path.join(REPO_ROOT, "results", "finetune_on_target"))
+    paths.add_argument("--run_name", type=str, default="finetune_on_target")
+    paths.add_argument("--device", type=str, default="auto", help="'auto', 'cpu', 'cuda', 'cuda:N'.")
+    paths.add_argument("--seed", type=int, default=42)
+    paths.add_argument("--wandb_project", type=str, default="TD3B-finetune-on-target")
+    paths.add_argument("--wandb_mode", type=str, default="disabled",
+                       help="WANDB_MODE for the finetune subprocess (disabled/offline/online).")
+
+    orc = p.add_argument_group("Directional oracle")
+    orc.add_argument("--direction_oracle_ckpt", type=str,
+                     default=os.path.join(REPO_ROOT, "checkpoints", "direction_oracle.pt"))
+    orc.add_argument("--direction_oracle_tr2d2_checkpoint", type=str,
+                     default=os.path.join(REPO_ROOT, "checkpoints", "pretrained.ckpt"))
+    orc.add_argument("--direction_oracle_tokenizer_vocab", type=str,
+                     default=os.path.join(REPO_ROOT, "tokenizer", "new_vocab.txt"))
+    orc.add_argument("--direction_oracle_tokenizer_splits", type=str,
+                     default=os.path.join(REPO_ROOT, "tokenizer", "new_splits.txt"))
+    orc.add_argument("--direction_oracle_esm_name", type=str, default="facebook/esm2_t33_650M_UR50D")
+    orc.add_argument("--direction_oracle_esm_cache_dir", type=str, default=None)
+    orc.add_argument("--direction_oracle_esm_local_files_only", action="store_true")
+    orc.add_argument("--direction_oracle_max_ligand_length", type=int, default=768)
+    orc.add_argument("--direction_oracle_max_protein_length", type=int, default=1024)
+    orc.add_argument("--direction_oracle_d_model", type=int, default=256)
+    orc.add_argument("--direction_oracle_n_heads", type=int, default=4)
+    orc.add_argument("--direction_oracle_n_self_attn_layers", type=int, default=1)
+    orc.add_argument("--direction_oracle_n_bmca_layers", type=int, default=2)
+    orc.add_argument("--direction_oracle_dropout", type=float, default=0.3)
+
+    args = p.parse_args(argv)
+
+    if not args.target_seq and not args.targets_csv:
+        p.error("Provide at least one --target_seq or a --targets_csv.")
+    if args.skip_finetune and args.skip_generate:
+        p.error("--skip_finetune and --skip_generate together leave nothing to do.")
+    # Normalize base-path-derived defaults if the user overrode --base_path only.
+    return args
+
+
+def write_temp_training_csv(frame: pd.DataFrame, output_dir: str) -> str:
+    os.makedirs(output_dir, exist_ok=True)
+    fd, path = tempfile.mkstemp(prefix="td3b_target_train_", suffix=".csv", dir=output_dir)
+    os.close(fd)
+    frame.to_csv(path, index=False)
+    return path
+
+
+def main(argv=None):
+    args = parse_args(argv)
+    logger.info("=" * 80)
+    logger.info("TD3B finetune-on-target  |  direction=%s  validity_reward=%s",
+                args.direction, args.validity_reward)
+    logger.info("=" * 80)
+
+    # 1) Build + normalize the target training set (seeds length priors).
+    frame, targets = build_target_training_frame(
+        args.target_seq, args.targets_csv, args.binder_length,
+        placeholder_mode=args.placeholder_mode, seed=args.seed)
+    logger.info("Prepared %d unique target(s); %d training row(s).", len(targets), len(frame))
+
+    train_csv = write_temp_training_csv(frame, args.output_dir)
+    logger.info("Temp training CSV: %s", train_csv)
+
+    if args.finetune_validity_hook is not None:
+        logger.info("--finetune_validity_hook=%s overrides the finetune-side validity gate "
+                    "(generation-side stays validity_reward=%s).",
+                    args.finetune_validity_hook, args.validity_reward)
+
+    # 2) FINETUNE half.
+    ckpt_path = args.td3b_checkpoint
+    if not args.skip_finetune:
+        ckpt_path = run_finetune(args, train_csv, num_targets=len(targets))
+    elif not ckpt_path:
+        raise ValueError("--skip_finetune requires --td3b_checkpoint (nothing to generate from).")
+
+    # 3) GENERATE half.
+    if not args.skip_generate:
+        out_path = run_generation(args, ckpt_path, train_csv, targets)
+        logger.info("Done. Results: %s", out_path)
+    else:
+        logger.info("Done (finetune only). Checkpoint: %s", ckpt_path)
+
+
+if __name__ == "__main__":
+    main()

generate_valid.py

@@ -0,0 +1,153 @@
+#!/usr/bin/env python3
+"""
+generate_valid.py -- standalone CLI to generate peptide SMILES with a chosen
+validity-boosting sampling strategy (see ``sampling_strategies.py``) and report the
+fraction that pass ``utils.app.PeptideAnalyzer.is_peptide``.
+
+Examples
+--------
+# Real checkpoint, long peptides, nucleus + remask self-correction (recommended default):
+python generate_valid.py \
+    --ckpt_path checkpoints/td3b.ckpt \
+    --length 400 --num_samples 64 \
+    --strategy nucleus_remask \
+    --device cuda:0 --seed 42 \
+    --save_path results/valid_len400.csv
+
+# No checkpoint available -> RANDOM-init model on CPU (development / API smoke test;
+# absolute yields are garbage, only the sampling machinery is exercised):
+python generate_valid.py --length 200 --num_samples 32 --strategy remask --device cpu
+
+Strategies: baseline, more_steps, top_p (nucleus), top_k, low_temp, remask,
+best_of_n, nucleus_remask.  Per-strategy knobs below override the preset defaults.
+"""
+import argparse
+import csv
+import logging
+import os
+import sys
+
+import numpy as np
+import torch
+
+ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
+if ROOT_DIR not in sys.path:
+    sys.path.insert(0, ROOT_DIR)
+
+from sampling_strategies import generate, build_random_model, available_strategies
+from utils.app import PeptideAnalyzer
+
+logger = logging.getLogger("generate_valid")
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+
+
+def _load_model(ckpt_path, device, base_path, hidden_size, n_layers, n_heads):
+    """Load the real checkpoint via ``inference.load_model`` when available; otherwise
+    fall back to a RANDOM-init model (imports of ``inference`` are done lazily so the
+    random path has no heavy dependencies)."""
+    if ckpt_path and os.path.isfile(ckpt_path):
+        logger.info("Loading real checkpoint from %s", ckpt_path)
+        from inference import load_model  # reuse the canonical loader
+        model, tokenizer = load_model(ckpt_path, device)
+        return model, tokenizer, False
+    if ckpt_path:
+        logger.warning("Checkpoint %s not found -- falling back to RANDOM-init model.", ckpt_path)
+    else:
+        logger.warning("No --ckpt_path given -- using RANDOM-init model "
+                       "(yields are meaningless; API/mechanism check only).")
+    model, tokenizer = build_random_model(
+        device=device, hidden_size=hidden_size, n_layers=n_layers,
+        n_heads=n_heads, base_path=base_path)
+    return model, tokenizer, True
+
+
+def build_parser():
+    p = argparse.ArgumentParser(description="Generate valid peptides with a chosen sampling strategy.")
+    p.add_argument("--ckpt_path", type=str, default=None,
+                   help="Path to TD3B checkpoint. If missing/omitted, a random-init model is used.")
+    p.add_argument("--base_path", type=str, default=ROOT_DIR, help="Repo root (for tokenizer files).")
+    p.add_argument("--length", type=int, default=200, help="Target sequence length (tokens).")
+    p.add_argument("--num_samples", type=int, default=64, help="Number of sequences to generate.")
+    p.add_argument("--strategy", type=str, default="nucleus_remask",
+                   choices=available_strategies(), help="Sampling strategy.")
+    p.add_argument("--device", type=str, default="cuda:0")
+    p.add_argument("--seed", type=int, default=42)
+    p.add_argument("--save_path", type=str, default=None,
+                   help="CSV path to save VALID sequences (default: results/valid_<strategy>_len<L>.csv).")
+    # strategy knobs (None -> use the strategy preset default)
+    p.add_argument("--num_steps", type=int, default=128, help="Base reverse-diffusion steps.")
+    p.add_argument("--eps", type=float, default=1e-5)
+    p.add_argument("--temperature", type=float, default=None, help="<1 sharpens logits (low_temp).")
+    p.add_argument("--top_p", type=float, default=None, help="Nucleus mass in (0,1].")
+    p.add_argument("--top_k", type=int, default=None, help="Top-k tokens per position.")
+    p.add_argument("--steps_per_token", type=float, default=None,
+                   help="more_steps: num_steps = max(num_steps, round(steps_per_token*length)).")
+    p.add_argument("--remask_rounds", type=int, default=None, help="Self-correction rounds.")
+    p.add_argument("--remask_frac", type=float, default=None, help="Fraction of lowest-conf tokens to remask.")
+    p.add_argument("--remask_steps", type=int, default=None, help="Re-denoise steps per remask round.")
+    p.add_argument("--best_of_n", type=int, default=None, help="Oversample N per slot, keep first valid.")
+    # random-fallback architecture (ignored when a real checkpoint loads)
+    p.add_argument("--hidden_size", type=int, default=768)
+    p.add_argument("--n_layers", type=int, default=8)
+    p.add_argument("--n_heads", type=int, default=8)
+    return p
+
+
+def main():
+    args = build_parser().parse_args()
+
+    torch.manual_seed(args.seed)
+    np.random.seed(args.seed)
+    device = torch.device(args.device if (args.device.startswith("cpu") or torch.cuda.is_available()) else "cpu")
+
+    model, tokenizer, is_random = _load_model(
+        args.ckpt_path, device, args.base_path, args.hidden_size, args.n_layers, args.n_heads)
+    analyzer = PeptideAnalyzer()
+
+    logger.info("Generating %d sequences of length %d with strategy=%s on %s",
+                args.num_samples, args.length, args.strategy, device)
+
+    tokens, sequences, valid_mask, stats = generate(
+        model, tokenizer, analyzer,
+        batch_size=args.num_samples, length=args.length, strategy=args.strategy,
+        num_steps=args.num_steps, eps=args.eps,
+        temperature=args.temperature, top_p=args.top_p, top_k=args.top_k,
+        steps_per_token=args.steps_per_token,
+        remask_rounds=args.remask_rounds, remask_frac=args.remask_frac,
+        remask_steps=args.remask_steps, best_of_n=args.best_of_n,
+        verbose=False,
+    )
+
+    valid_seqs = [s for s, v in zip(sequences, valid_mask) if v]
+
+    save_path = args.save_path or os.path.join(
+        args.base_path, "results", f"valid_{args.strategy}_len{args.length}.csv")
+    os.makedirs(os.path.dirname(os.path.abspath(save_path)), exist_ok=True)
+    with open(save_path, "w", newline="") as f:
+        w = csv.writer(f)
+        w.writerow(["idx", "sequence", "n_chars"])
+        for i, s in enumerate(valid_seqs):
+            w.writerow([i, s, len(s)])
+
+    print("\n" + "=" * 66)
+    print(f"  strategy         : {stats['strategy']}")
+    print(f"  length           : {stats['length']}")
+    print(f"  num_samples      : {stats['batch_size']}")
+    print(f"  num_steps        : {stats['num_steps']}"
+          f"   (temp={stats['temperature']}, top_p={stats['top_p']}, top_k={stats['top_k']})")
+    print(f"  remask           : rounds={stats['remask_rounds']} frac={stats['remask_frac']} "
+          f"steps={stats['remask_steps']}   best_of_n={stats['best_of_n']}")
+    if len(stats.get("round_valid_counts", [])) > 1:
+        print(f"  valid per round  : {stats['round_valid_counts']}  (round 0 = before remask)")
+    print(f"  VALID YIELD      : {stats['valid_count']}/{stats['batch_size']} "
+          f"= {stats['valid_rate']:.1%}")
+    print(f"  wall time        : {stats['wall_time_s']}s")
+    print(f"  saved valid seqs : {save_path}  ({len(valid_seqs)} rows)")
+    if is_random:
+        print("  NOTE: RANDOM-init model -- yields are meaningless; rerun with --ckpt_path "
+              "<real.ckpt> for real numbers.")
+    print("=" * 66)
+
+
+if __name__ == "__main__":
+    main()

inference.py

@@ -29,10 +29,13 @@ from configs.finetune_config import (
     DiffusionConfig, RoFormerConfig, NoiseConfig,
     TrainingConfig, SamplingConfig, EvalConfig, OptimConfig, MCTSConfig,
 )
-from training.finetune_utils import load_tokenizer, create_reward_function
+from training.finetune_utils import load_tokenizer
 from td3b.direction_oracle import DirectionalOracle
 from td3b.td3b_scoring import create_td3b_reward_function
+from scoring.functions.binding import MultiTargetBindingAffinity, TargetSpecificBindingAffinity
+from td3b.data_utils import peptide_seq_to_smiles, smiles_token_length
 from utils.app import PeptideAnalyzer
+import sampling_strategies
 
 logger = logging.getLogger(__name__)
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
@@ -76,7 +79,22 @@ def load_model(ckpt_path: str, device: torch.device):
     )
 
     model = Diffusion(config=cfg, tokenizer=tokenizer, device=device).to(device)
-    model.load_state_dict(state_dict, strict=False)
+    incompatible = model.load_state_dict(state_dict, strict=False)
+    # A correct TD3B checkpoint round-trips to 0 missing backbone keys; strict=False otherwise
+    # silently leaves the denoiser at random init and generates garbage with no error.
+    backbone_missing = [k for k in incompatible.missing_keys if k.startswith("backbone.")]
+    total_backbone = sum(1 for k in model.state_dict() if k.startswith("backbone."))
+    if backbone_missing:
+        msg = (f"{len(backbone_missing)}/{total_backbone} backbone weights did not load from "
+               f"{ckpt_path} (e.g. {backbone_missing[:3]}) — the generator would run on "
+               f"partially-random weights.")
+        if len(backbone_missing) == total_backbone:
+            raise RuntimeError(msg + " The backbone loaded NO weights; refusing to run — "
+                               "check the checkpoint format/key namespace.")
+        logger.warning(msg)
+    if incompatible.unexpected_keys:
+        logger.warning("Checkpoint had %d unexpected keys (e.g. %s)",
+                       len(incompatible.unexpected_keys), incompatible.unexpected_keys[:3])
     model.eval()
     model.tokenizer = tokenizer
     return model, tokenizer
@@ -130,6 +148,30 @@ def main():
     parser.add_argument("--resample_alpha", type=float, default=0.1, help="Temperature for weighted resampling")
     parser.add_argument("--direction_oracle_ckpt", type=str, default=None)
     parser.add_argument("--direction_oracle_tr2d2_checkpoint", type=str, default=None)
+    # ─── Opt-in validity-boosting sampler (default "baseline" == original behavior) ───
+    parser.add_argument("--sampler", type=str, default="baseline",
+                        choices=sampling_strategies.available_strategies(),
+                        help="Candidate-pool sampling strategy (default: baseline == original behavior)")
+    parser.add_argument("--num_steps", type=int, default=128,
+                        help="Reverse-diffusion steps for non-baseline samplers")
+    parser.add_argument("--top_p", type=float, default=None,
+                        help="Nucleus cumulative mass for top_p/nucleus/nucleus_remask samplers")
+    parser.add_argument("--top_k", type=int, default=None,
+                        help="Top-k cutoff on clean-token logits for the top_k sampler")
+    parser.add_argument("--temperature", type=float, default=None,
+                        help="Softmax temperature on clean-token logits (e.g. low_temp)")
+    parser.add_argument("--remask_rounds", type=int, default=None,
+                        help="Remask self-correction rounds (remask/nucleus_remask samplers)")
+    parser.add_argument("--remask_frac", type=float, default=None,
+                        help="Fraction of lowest-confidence tokens remasked each round")
+    parser.add_argument("--best_of_n", type=int, default=None,
+                        help="Best-of-N validity-guided oversampling per slot (best_of_n sampler)")
+    parser.add_argument("--seq_length", type=int, default=None,
+                        help="Generation length in SMILES tokens (override). If unset, derived from the "
+                             "reference binder. NOTE: this checkpoint only yields valid SMILES at short "
+                             "lengths (~<100 tokens) — set this (e.g. 50) for non-empty output on long binders.")
+    parser.add_argument("--max_seq_length", type=int, default=200,
+                        help="Cap on the derived generation length (SMILES tokens)")
     args = parser.parse_args()
 
     # Setup
@@ -149,36 +191,68 @@ def main():
     df = pd.read_csv(args.val_csv)
     targets = []
     for _, row in df.iterrows():
+        # Generation length in SMILES *tokens*: explicit --seq_length override, else the
+        # reference binder's tokenized length (from a SMILES/Ligand_SMILES column, or derived
+        # from Ligand_Sequence), else the default. Capped at --max_seq_length.
+        if args.seq_length is not None:
+            seq_len = args.seq_length
+        else:
+            smi = row.get("SMILES") if isinstance(row.get("SMILES"), str) else row.get("Ligand_SMILES")
+            if not isinstance(smi, str):
+                lig = row.get("Ligand_Sequence")
+                smi = peptide_seq_to_smiles(lig) if isinstance(lig, str) else None
+            seq_len = smiles_token_length(smi, tokenizer) if isinstance(smi, str) else DEFAULTS["seq_length"]
         targets.append({
             "target_seq": row["Target_Sequence"],
             "target_uid": row.get("Target_UniProt_ID", ""),
             "binder_seq": row.get("Ligand_Sequence", ""),
             "label": row.get("label", ""),
-            "seq_length": min(len(row.get("Ligand_SMILES", "x" * 200)), 200),
+            "seq_length": max(1, min(int(seq_len) if seq_len else DEFAULTS["seq_length"], args.max_seq_length)),
         })
 
-    # Build reward function for each target
+    # Build the shared reward components ONCE. The affinity predictor (ESM2) and the
+    # Direction Oracle (ESM2 + TR2-D2) are expensive to load, so we build them a single
+    # time and rebind the per-target protein via lightweight wrappers inside the loop.
     logger.info("Building reward functions...")
     oracle_ckpt = args.direction_oracle_ckpt or os.path.join(ROOT_DIR, "checkpoints", "direction_oracle.pt")
     oracle_tr2d2 = args.direction_oracle_tr2d2_checkpoint or os.path.join(ROOT_DIR, "checkpoints", "pretrained.ckpt")
+    vocab_path = os.path.join(ROOT_DIR, "tokenizer", "new_vocab.txt")
+    splits_path = os.path.join(ROOT_DIR, "tokenizer", "new_splits.txt")
+
+    multi_affinity = MultiTargetBindingAffinity(
+        tokenizer=tokenizer,
+        base_path=ROOT_DIR,
+        device=device,
+        emb_model=model.backbone,
+    )
+    directional_oracle = DirectionalOracle(
+        model_ckpt=oracle_ckpt,
+        tr2d2_checkpoint=oracle_tr2d2,
+        tokenizer_vocab=vocab_path,
+        tokenizer_splits=splits_path,
+        device=device,
+    )
+    directional_oracle.eval()
 
     records = []
 
     for tidx, target in enumerate(targets):
+        # Bind the shared affinity predictor to this target's protein sequence.
+        target_affinity = TargetSpecificBindingAffinity(multi_affinity, target["target_seq"])
         for d_star, d_name in [(1.0, "agonist"), (-1.0, "antagonist")]:
             logger.info(f"[{tidx+1}/{len(targets)}] Target {target['target_uid']} direction={d_name}")
 
-            # Create reward function
+            # Create reward function (reuses the preloaded oracle; only re-encodes the protein)
             try:
-                reward_model = create_reward_function(
-                    base_path=ROOT_DIR,
-                    tokenizer=tokenizer,
+                reward_model = create_td3b_reward_function(
+                    affinity_predictor=target_affinity,
+                    directional_oracle=directional_oracle,
                     target_protein_seq=target["target_seq"],
-                    target_direction="agonist" if d_star > 0 else "antagonist",
+                    target_direction=d_name,
+                    peptide_tokenizer=tokenizer,
                     device=device,
-                    emb_model=model.backbone,
-                    directional_oracle_checkpoint=oracle_ckpt,
-                    direction_oracle_tr2d2_checkpoint=oracle_tr2d2,
+                    min_affinity_threshold=DEFAULTS["min_affinity_threshold"],
+                    temperature=DEFAULTS["sigmoid_temperature"],
                 )
             except Exception as e:
                 logger.warning(f"Failed to create reward for {target['target_uid']}: {e}")
@@ -186,35 +260,48 @@ def main():
 
             # Generate pool of candidates
             target_length = target.get("seq_length", 200)
-            x_pool = sample_sequences(model, args.num_pool, target_length, 128)
-            sequences = tokenizer.batch_decode(x_pool)
-
-            # Check validity
-            valid_mask = np.array([analyzer.is_peptide(seq) for seq in sequences])
+            if args.sampler == "baseline":
+                # Original behavior, unchanged (byte-for-byte).
+                x_pool = sample_sequences(model, args.num_pool, target_length, 128)
+                sequences = tokenizer.batch_decode(x_pool)
+                # Check validity
+                valid_mask = np.array([analyzer.is_peptide(seq) for seq in sequences])
+            else:
+                # Opt-in validity-boosting sampler. Reuses its decode + is_peptide validity.
+                # None-valued knobs fall back to each strategy's preset inside generate().
+                x_pool, sequences, valid_mask, _sampler_stats = sampling_strategies.generate(
+                    model, tokenizer, analyzer,
+                    batch_size=args.num_pool, length=target_length,
+                    strategy=args.sampler,
+                    num_steps=args.num_steps, top_p=args.top_p, top_k=args.top_k,
+                    temperature=args.temperature, remask_rounds=args.remask_rounds,
+                    remask_frac=args.remask_frac, best_of_n=args.best_of_n,
+                )
 
             # Score all
             gated_rewards, affinities, directions, confidences = score_sequences(reward_model, sequences)
             direction_accuracy = ((directions > 0.5).astype(float) if d_star > 0
                                   else (directions < 0.5).astype(float))
 
-            # Weighted resampling (Algorithm 2)
-            finite = np.isfinite(gated_rewards)
-            if finite.any():
-                rewards_t = torch.as_tensor(gated_rewards[finite], device=device)
+            # Weighted resampling (Algorithm 2). Resample only among candidates that are BOTH
+            # finite-reward AND valid peptides, and draw WITHOUT replacement so the output holds
+            # up to val_samples_per_target DISTINCT binders. (With replacement + the peaked
+            # softmax at alpha=0.1, multinomial repeatedly draws the argmax → duplicate rows that
+            # inflate the sample count and skew every mean metric; filtering validity only AFTER
+            # resampling could also drop every draw and save 0 samples despite valid candidates.)
+            eligible = np.isfinite(gated_rewards) & valid_mask.astype(bool)
+            if eligible.any():
+                rewards_t = torch.as_tensor(gated_rewards[eligible], device=device, dtype=torch.float32)
                 alpha = max(args.resample_alpha, 1e-6)
                 weights = torch.softmax(rewards_t / alpha, dim=0)
-                idx = torch.multinomial(weights, num_samples=args.val_samples_per_target, replacement=True)
-                valid_idx = np.where(finite)[0]
-                chosen = valid_idx[idx.cpu().numpy()]
+                k = min(args.val_samples_per_target, int(eligible.sum()))
+                idx = torch.multinomial(weights, num_samples=k, replacement=False)
+                chosen = np.where(eligible)[0][idx.cpu().numpy()]
             else:
-                chosen = np.arange(min(args.val_samples_per_target, len(sequences)))
+                chosen = np.array([], dtype=int)  # no valid finite-reward candidate → save nothing
 
-            # Save only VALID resampled samples
+            # Save resampled samples (already gated to valid + finite above)
             for i in chosen:
-                is_valid = bool(valid_mask[i]) if valid_mask.size else False
-                if not is_valid:
-                    continue  # Skip invalid samples
-
                 records.append({
                     "target": target["target_seq"][:20],
                     "target_uid": target["target_uid"],
@@ -247,6 +334,9 @@ def main():
         logger.info(f"{'='*60}")
     else:
         logger.warning("No valid samples generated.")
+        logger.warning("Hint: this checkpoint only produces valid SMILES at short lengths "
+                       "(~<100 tokens). Re-run with a shorter --seq_length (e.g. 50) and/or "
+                       "--sampler best_of_n (or nucleus_remask) to raise the valid yield.")
 
 
 if __name__ == "__main__":

launch_multi_target.sh

@@ -7,8 +7,9 @@
 # Configuration
 # ============================================================================
 
-# Paths — update these to your local paths
-BASE_PATH="/path/to/TD3B"
+# Paths — BASE_PATH auto-detects the repo root (this script's own directory);
+# the paths below derive from it, so they usually need no editing.
+BASE_PATH="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
 PRETRAINED_CHECKPOINT="${BASE_PATH}/checkpoints/pretrained.ckpt"
 TRAIN_CSV="${BASE_PATH}/data/train.csv"
 VAL_CSV="${BASE_PATH}/data/test.csv"  # Optional: create validation split
@@ -75,7 +76,7 @@ fi
 
 # W&B (optional)
 WANDB_PROJECT="tr2d2-multi-target"
-WANDB_ENTITY="phos_zj"
+WANDB_ENTITY=""  # set to your W&B entity/team (empty = your default entity)
 
 # ============================================================================
 # Launch Training

mcts/peptide_mcts.py

@@ -225,13 +225,22 @@ class MCTS:
         rootNode=None,
         reward_func=None,
         num_obj=None,
+        enforce_validity=None,
     ):
         self.timer = StepTimer(policy_model.device)
-        
+
         self.device = policy_model.device
-        
+
         self.args = args
         self.config = config
+        # Validity gate toggle for MCTS expansion. Default (None) consults
+        # args.enforce_validity (threaded from --validity_reward), falling back
+        # to True so DEFAULT behavior is unchanged. When False, invalid decoded
+        # children are NOT zero-rewarded/filtered -- expansion proceeds on the
+        # pure affinity x direction reward.
+        if enforce_validity is None:
+            enforce_validity = getattr(args, "enforce_validity", True)
+        self.enforce_validity = bool(enforce_validity)
         self.noise = noise_schedule.get_noise(config)
         self.time_conditioning = args.time_conditioning
         
@@ -573,7 +582,11 @@ class MCTS:
                 childSeq = childSequences[i]
                 
                 # check if the peptide is valid
-                if self.analyzer.is_peptide(childSeq):
+                # When enforce_validity is False, skip the is_peptide gate so
+                # invalid decoded children are kept and scored on the pure
+                # affinity x direction reward instead of being added as
+                # zero-reward nodes. (short-circuits before calling is_peptide)
+                if (not self.enforce_validity) or self.analyzer.is_peptide(childSeq):
                     valid_x_children.append(x_children[i])
                     valid_x_final.append(x_rollout[i])
                     validSequences.append(childSeq)

models/diffusion.py

@@ -147,8 +147,8 @@ class Diffusion(L.LightningModule):
                 
         # PeptideCLM tokenizer 
         if tokenizer is None:
-            self.tokenizer = SMILES_SPE_Tokenizer(f'{base_path}/tr2d2-pep/tokenizer/new_vocab.txt',
-                                    f'{base_path}/tr2d2-pep/tokenizer/new_splits.txt')
+            self.tokenizer = SMILES_SPE_Tokenizer(f'{base_path}/tokenizer/new_vocab.txt',
+                                    f'{base_path}/tokenizer/new_splits.txt')
         else:
             self.tokenizer = tokenizer
             

notebooks/TD3B_Inference_Demo.ipynb

@@ -1,457 +1,409 @@
 {
-  "nbformat": 4,
-  "nbformat_minor": 0,
-  "metadata": {
-    "colab": {
-      "provenance": [],
-      "gpuType": "T4"
-    },
-    "kernelspec": {
-      "name": "python3",
-      "display_name": "Python 3"
-    },
-    "language_info": {
-      "name": "python"
-    },
-    "accelerator": "GPU"
+ "nbformat": 4,
+ "nbformat_minor": 0,
+ "metadata": {
+  "colab": {
+   "provenance": [],
+   "gpuType": "T4"
   },
-  "cells": [
-    {
-      "cell_type": "markdown",
-      "source": [
-        "# TD3B: Transition-Directed Discrete Diffusion for Allosteric Binder Generation\n",
-        "\n",
-        "This notebook demonstrates **TD3B inference** — generating peptide binders with specified agonist or antagonist behavior for GPCR targets.\n",
-        "\n",
-        "**What TD3B does:**\n",
-        "- Takes a target protein sequence + desired direction (agonist / antagonist)\n",
-        "- Generates peptide binder sequences using a finetuned discrete diffusion model\n",
-        "- Scores them with a Direction Oracle and binding affinity predictor\n",
-        "- Returns the best candidates via weighted resampling (Algorithm 2)\n",
-        "\n",
-        "**Requirements:** GPU runtime (T4 or better). Click **Runtime → Change runtime type → GPU**."
-      ],
-      "metadata": {}
-    },
-    {
-      "cell_type": "markdown",
-      "source": [
-        "## 1. Setup"
-      ],
-      "metadata": {}
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "# Install dependencies\n",
-        "!pip install -q torch torchvision --index-url https://download.pytorch.org/whl/cu121\n",
-        "!pip install -q transformers fair-esm SmilesPE rdkit-pypi scipy pandas numpy xgboost pytorch-lightning lightning hydra-core loguru timm huggingface_hub"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "# Clone TD3B repository and download checkpoints from HuggingFace\n",
-        "!git clone https://github.com/chq1155/TD3B_ICML.git TD3B\n",
-        "%cd TD3B\n",
-        "\n",
-        "from huggingface_hub import hf_hub_download\n",
-        "import os\n",
-        "\n",
-        "REPO_ID = \"ChatterjeeLab/TD3B\"\n",
-        "os.makedirs(\"checkpoints\", exist_ok=True)\n",
-        "os.makedirs(\"data\", exist_ok=True)\n",
-        "\n",
-        "# Download checkpoints (this may take a few minutes)\n",
-        "for fname in [\"checkpoints/td3b.ckpt\", \"checkpoints/pretrained.ckpt\",\n",
-        "              \"checkpoints/direction_oracle.pt\",\n",
-        "              \"scoring/functions/classifiers/binding-affinity.pt\",\n",
-        "              \"data/test.csv\", \"data/train.csv\"]:\n",
-        "    print(f\"Downloading {fname}...\")\n",
-        "    hf_hub_download(repo_id=REPO_ID, filename=fname, local_dir=\".\")\n",
-        "\n",
-        "print(\"\\nAll files downloaded!\")\n",
-        "!ls -lh checkpoints/"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "source": [
-        "## 2. Load Model and Oracle"
-      ],
-      "metadata": {}
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "import sys\n",
-        "sys.path.insert(0, \".\")\n",
-        "\n",
-        "import torch\n",
-        "import numpy as np\n",
-        "import pandas as pd\n",
-        "\n",
-        "from diffusion import Diffusion\n",
-        "from configs.finetune_config import (\n",
-        "    DiffusionConfig, RoFormerConfig, NoiseConfig,\n",
-        "    TrainingConfig, SamplingConfig, EvalConfig, OptimConfig, MCTSConfig,\n",
-        ")\n",
-        "from tokenizer.my_tokenizers import SMILES_SPE_Tokenizer\n",
-        "from td3b.direction_oracle import DirectionalOracle\n",
-        "from td3b.td3b_scoring import TD3BRewardFunction, create_td3b_reward_function\n",
-        "from scoring.functions.binding import BindingAffinity\n",
-        "from utils.app import PeptideAnalyzer\n",
-        "\n",
-        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
-        "print(f\"Using device: {device}\")\n",
-        "if torch.cuda.is_available():\n",
-        "    print(f\"GPU: {torch.cuda.get_device_name(0)}\")\n",
-        "    print(f\"Memory: {torch.cuda.get_device_properties(0).total_mem / 1e9:.1f} GB\")"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "# Load tokenizer\n",
-        "tokenizer = SMILES_SPE_Tokenizer(\"tokenizer/new_vocab.txt\", \"tokenizer/new_splits.txt\")\n",
-        "print(f\"Tokenizer vocab size: {len(tokenizer)}\")\n",
-        "\n",
-        "# Load diffusion model\n",
-        "print(\"\\nLoading TD3B model...\")\n",
-        "cfg = DiffusionConfig(\n",
-        "    roformer=RoFormerConfig(hidden_size=768, n_layers=8, n_heads=8),\n",
-        "    noise=NoiseConfig(),\n",
-        "    training=TrainingConfig(sampling_eps=1e-3),\n",
-        "    sampling=SamplingConfig(steps=128, sampling_eps=1e-3),\n",
-        "    eval_cfg=EvalConfig(), optim=OptimConfig(lr=3e-4), mcts=MCTSConfig(),\n",
-        ")\n",
-        "model = Diffusion(config=cfg, tokenizer=tokenizer, device=device).to(device)\n",
-        "\n",
-        "ckpt = torch.load(\"checkpoints/td3b.ckpt\", map_location=device, weights_only=False)\n",
-        "state_dict = ckpt.get(\"model_state_dict\") or ckpt.get(\"state_dict\") or ckpt\n",
-        "model.load_state_dict(state_dict, strict=False)\n",
-        "model.eval()\n",
-        "model.tokenizer = tokenizer\n",
-        "print(\"TD3B model loaded!\")\n",
-        "\n",
-        "# Load Direction Oracle\n",
-        "print(\"\\nLoading Direction Oracle...\")\n",
-        "oracle = DirectionalOracle(\n",
-        "    model_ckpt=\"checkpoints/direction_oracle.pt\",\n",
-        "    tr2d2_checkpoint=\"checkpoints/pretrained.ckpt\",\n",
-        "    tokenizer_vocab=\"tokenizer/new_vocab.txt\",\n",
-        "    tokenizer_splits=\"tokenizer/new_splits.txt\",\n",
-        "    device=device,\n",
-        ")\n",
-        "oracle.eval()\n",
-        "print(\"Direction Oracle loaded!\")\n",
-        "\n",
-        "# Load Affinity Predictor\n",
-        "print(\"\\nLoading Affinity Predictor...\")\n",
-        "analyzer = PeptideAnalyzer()\n",
-        "print(\"\\nAll models loaded!\")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "source": [
-        "## 3. Define Helper Functions"
-      ],
-      "metadata": {}
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "def sample_sequences(model, batch_size, seq_length, num_steps=128, eps=1e-5):\n",
-        "    \"\"\"Sample sequences from the diffusion model.\"\"\"\n",
-        "    x = model.sample_prior(batch_size, seq_length).to(model.device, dtype=torch.long)\n",
-        "    timesteps = torch.linspace(1, eps, num_steps + 1, device=model.device)\n",
-        "    dt = torch.tensor((1 - eps) / num_steps, device=model.device)\n",
-        "\n",
-        "    for i in range(num_steps):\n",
-        "        t = timesteps[i] * torch.ones(x.shape[0], 1, device=model.device)\n",
-        "        _, x = model.single_reverse_step(x, t=t, dt=dt)\n",
-        "        x = x.to(model.device)\n",
-        "\n",
-        "    mask_pos = (x == model.mask_index)\n",
-        "    if mask_pos.any():\n",
-        "        t = timesteps[-2] * torch.ones(x.shape[0], 1, device=model.device)\n",
-        "        _, x = model.single_noise_removal(x, t=t, dt=dt)\n",
-        "    return x\n",
-        "\n",
-        "\n",
-        "def generate_binders(target_seq, direction=\"agonist\", num_pool=32,\n",
-        "                     num_keep=8, alpha=0.1, seq_length=200):\n",
-        "    \"\"\"\n",
-        "    Generate directional binders for a target protein.\n",
-        "    \n",
-        "    Args:\n",
-        "        target_seq: Target protein amino acid sequence\n",
-        "        direction: 'agonist' or 'antagonist'\n",
-        "        num_pool: Number of candidates to generate\n",
-        "        num_keep: Number of final samples after resampling\n",
-        "        alpha: Temperature for weighted resampling\n",
-        "        seq_length: Binder sequence length (in SMILES tokens)\n",
-        "    \n",
-        "    Returns:\n",
-        "        DataFrame with generated binders and scores\n",
-        "    \"\"\"\n",
-        "    d_star = 1.0 if direction == \"agonist\" else -1.0\n",
-        "    \n",
-        "    # Build reward function\n",
-        "    affinity_pred = BindingAffinity(\n",
-        "        prot_seq=target_seq, tokenizer=tokenizer,\n",
-        "        base_path=\".\", device=device, emb_model=model.backbone\n",
-        "    )\n",
-        "    reward_fn = create_td3b_reward_function(\n",
-        "        affinity_predictor=affinity_pred,\n",
-        "        directional_oracle=oracle,\n",
-        "        target_protein_seq=target_seq,\n",
-        "        target_direction=direction,\n",
-        "        peptide_tokenizer=tokenizer,\n",
-        "        device=device,\n",
-        "    )\n",
-        "    \n",
-        "    # Generate candidates\n",
-        "    with torch.no_grad():\n",
-        "        x_pool = sample_sequences(model, num_pool, seq_length)\n",
-        "    sequences = tokenizer.batch_decode(x_pool)\n",
-        "    \n",
-        "    # Score all\n",
-        "    rewards, info = reward_fn(sequences)\n",
-        "    affinities = info[\"affinities\"]\n",
-        "    directions = info[\"directions\"]\n",
-        "    \n",
-        "    # Weighted resampling (Algorithm 2)\n",
-        "    rewards_t = torch.as_tensor(rewards, device=device)\n",
-        "    weights = torch.softmax(rewards_t / max(alpha, 1e-6), dim=0)\n",
-        "    idx = torch.multinomial(weights, num_samples=num_keep, replacement=True)\n",
-        "    chosen = idx.cpu().numpy()\n",
-        "    \n",
-        "    # Filter to valid peptides only\n",
-        "    results = []\n",
-        "    for i in chosen:\n",
-        "        is_valid = analyzer.is_peptide(sequences[i])\n",
-        "        da = float(directions[i] > 0.5) if d_star > 0 else float(directions[i] < 0.5)\n",
-        "        results.append({\n",
-        "            \"sequence\": sequences[i],\n",
-        "            \"direction\": direction,\n",
-        "            \"is_valid\": is_valid,\n",
-        "            \"affinity\": float(affinities[i]),\n",
-        "            \"gated_reward\": float(rewards[i]),\n",
-        "            \"p_agonist\": float(directions[i]),\n",
-        "            \"direction_accuracy\": da,\n",
-        "        })\n",
-        "    \n",
-        "    df = pd.DataFrame(results)\n",
-        "    return df"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "source": [
-        "## 4. Generate Binders\n",
-        "\n",
-        "Let's generate **agonist** and **antagonist** binders for a test target and compare the Direction Oracle predictions."
-      ],
-      "metadata": {}
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "# Load test targets\n",
-        "test_df = pd.read_csv(\"data/test.csv\")\n",
-        "print(f\"Test set: {len(test_df)} target-binder pairs\")\n",
-        "\n",
-        "# Pick first target for demo\n",
-        "target_row = test_df.iloc[0]\n",
-        "TARGET_SEQ = target_row[\"Target_Sequence\"]\n",
-        "TARGET_UID = target_row[\"Target_UniProt_ID\"]\n",
-        "print(f\"\\nTarget: {TARGET_UID}\")\n",
-        "print(f\"Sequence length: {len(TARGET_SEQ)} aa\")\n",
-        "print(f\"Sequence: {TARGET_SEQ[:60]}...\")"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "%%time\n",
-        "# Generate AGONIST binders\n",
-        "print(\"Generating agonist binders (d*=+1)...\")\n",
-        "torch.manual_seed(42)\n",
-        "np.random.seed(42)\n",
-        "df_agonist = generate_binders(TARGET_SEQ, direction=\"agonist\", num_pool=32, num_keep=8)\n",
-        "\n",
-        "print(f\"\\nGenerated {len(df_agonist)} samples ({df_agonist['is_valid'].sum()} valid)\")\n",
-        "print(f\"Mean p(agonist): {df_agonist['p_agonist'].mean():.3f}\")\n",
-        "print(f\"Mean affinity: {df_agonist['affinity'].mean():.2f}\")\n",
-        "print(f\"Mean gated reward: {df_agonist['gated_reward'].mean():.2f}\")"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "%%time\n",
-        "# Generate ANTAGONIST binders\n",
-        "print(\"Generating antagonist binders (d*=-1)...\")\n",
-        "torch.manual_seed(42)\n",
-        "np.random.seed(42)\n",
-        "df_antagonist = generate_binders(TARGET_SEQ, direction=\"antagonist\", num_pool=32, num_keep=8)\n",
-        "\n",
-        "print(f\"\\nGenerated {len(df_antagonist)} samples ({df_antagonist['is_valid'].sum()} valid)\")\n",
-        "print(f\"Mean p(agonist): {df_antagonist['p_agonist'].mean():.3f}\")\n",
-        "print(f\"Mean affinity: {df_antagonist['affinity'].mean():.2f}\")\n",
-        "print(f\"Mean gated reward: {df_antagonist['gated_reward'].mean():.2f}\")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "source": [
-        "## 5. Compare Directional Control"
-      ],
-      "metadata": {}
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "import matplotlib.pyplot as plt\n",
-        "\n",
-        "fig, axes = plt.subplots(1, 3, figsize=(15, 4))\n",
-        "\n",
-        "# Plot 1: Direction Oracle p(agonist)\n",
-        "axes[0].hist(df_agonist[\"p_agonist\"], bins=20, alpha=0.7, label=\"d*=+1 (agonist)\", color=\"#e74c3c\")\n",
-        "axes[0].hist(df_antagonist[\"p_agonist\"], bins=20, alpha=0.7, label=\"d*=-1 (antagonist)\", color=\"#3498db\")\n",
-        "axes[0].axvline(0.5, color=\"gray\", linestyle=\"--\", label=\"threshold\")\n",
-        "axes[0].set_xlabel(\"p(agonist)\")\n",
-        "axes[0].set_ylabel(\"Count\")\n",
-        "axes[0].set_title(\"Direction Oracle Predictions\")\n",
-        "axes[0].legend()\n",
-        "\n",
-        "# Plot 2: Binding Affinity\n",
-        "axes[1].hist(df_agonist[\"affinity\"], bins=20, alpha=0.7, label=\"Agonist\", color=\"#e74c3c\")\n",
-        "axes[1].hist(df_antagonist[\"affinity\"], bins=20, alpha=0.7, label=\"Antagonist\", color=\"#3498db\")\n",
-        "axes[1].set_xlabel(\"Predicted Binding Affinity\")\n",
-        "axes[1].set_ylabel(\"Count\")\n",
-        "axes[1].set_title(\"Binding Affinity Distribution\")\n",
-        "axes[1].legend()\n",
-        "\n",
-        "# Plot 3: Gated Reward\n",
-        "axes[2].hist(df_agonist[\"gated_reward\"], bins=20, alpha=0.7, label=\"Agonist\", color=\"#e74c3c\")\n",
-        "axes[2].hist(df_antagonist[\"gated_reward\"], bins=20, alpha=0.7, label=\"Antagonist\", color=\"#3498db\")\n",
-        "axes[2].set_xlabel(\"Gated Reward\")\n",
-        "axes[2].set_ylabel(\"Count\")\n",
-        "axes[2].set_title(\"Gated Reward Distribution\")\n",
-        "axes[2].legend()\n",
-        "\n",
-        "plt.tight_layout()\n",
-        "plt.savefig(\"td3b_results.png\", dpi=150, bbox_inches=\"tight\")\n",
-        "plt.show()\n",
-        "\n",
-        "print(\"\\nSummary:\")\n",
-        "print(f\"  Agonist mode:    p(agonist)={df_agonist['p_agonist'].mean():.3f}  Affinity={df_agonist['affinity'].mean():.2f}  Gated={df_agonist['gated_reward'].mean():.2f}\")\n",
-        "print(f\"  Antagonist mode: p(agonist)={df_antagonist['p_agonist'].mean():.3f}  Affinity={df_antagonist['affinity'].mean():.2f}  Gated={df_antagonist['gated_reward'].mean():.2f}\")\n",
-        "print(f\"  Directional gap: Δp = {df_agonist['p_agonist'].mean() - df_antagonist['p_agonist'].mean():.3f}\")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "source": [
-        "## 6. Run on Multiple Targets\n",
-        "\n",
-        "Generate binders for the first 5 test targets and compute aggregate metrics."
-      ],
-      "metadata": {}
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "N_TARGETS = 5  # Number of targets to evaluate (increase for full benchmark)\n",
-        "\n",
-        "all_results = []\n",
-        "targets = test_df.drop_duplicates(\"Target_UniProt_ID\").head(N_TARGETS)\n",
-        "\n",
-        "for i, (_, row) in enumerate(targets.iterrows()):\n",
-        "    uid = row[\"Target_UniProt_ID\"]\n",
-        "    seq = row[\"Target_Sequence\"]\n",
-        "    print(f\"[{i+1}/{N_TARGETS}] {uid} (len={len(seq)})\")\n",
-        "    \n",
-        "    for direction in [\"agonist\", \"antagonist\"]:\n",
-        "        torch.manual_seed(42)\n",
-        "        np.random.seed(42)\n",
-        "        df = generate_binders(seq, direction=direction, num_pool=32, num_keep=8)\n",
-        "        df[\"target_uid\"] = uid\n",
-        "        all_results.append(df)\n",
-        "        \n",
-        "        d_star = 1.0 if direction == \"agonist\" else -1.0\n",
-        "        da = df[\"direction_accuracy\"].mean()\n",
-        "        print(f\"    {direction:>10s}: DA={da:.2f}  Aff={df['affinity'].mean():.2f}  Gated={df['gated_reward'].mean():.2f}  valid={df['is_valid'].sum()}/{len(df)}\")\n",
-        "\n",
-        "combined = pd.concat(all_results, ignore_index=True)\n",
-        "\n",
-        "print(f\"\\n{'='*60}\")\n",
-        "print(f\"AGGREGATE METRICS ({N_TARGETS} targets)\")\n",
-        "print(f\"{'='*60}\")\n",
-        "for d_name, d_val in [(\"Agonist (d*=+1)\", \"agonist\"), (\"Antagonist (d*=-1)\", \"antagonist\")]:\n",
-        "    sub = combined[combined[\"direction\"] == d_val]\n",
-        "    valid = sub[sub[\"is_valid\"] == True]\n",
-        "    print(f\"  {d_name}:\")\n",
-        "    print(f\"    Affinity: {sub['affinity'].mean():.2f}\")\n",
-        "    print(f\"    Direction Accuracy: {sub['direction_accuracy'].mean():.3f}\")\n",
-        "    print(f\"    Gated Reward (all): {sub['gated_reward'].mean():.2f}\")\n",
-        "    if len(valid) > 0:\n",
-        "        print(f\"    Gated Reward (valid): {valid['gated_reward'].mean():.2f}\")\n",
-        "    print(f\"    Valid: {sub['is_valid'].sum()}/{len(sub)}\")\n",
-        "\n",
-        "# Save\n",
-        "combined.to_csv(\"td3b_demo_results.csv\", index=False)\n",
-        "print(f\"\\nResults saved to td3b_demo_results.csv\")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "source": [
-        "## Citation\n",
-        "\n",
-        "```bibtex\n",
-        "@article{cao2026td3b,\n",
-        "  title={TD3B: Transition-Directed Discrete Diffusion for Allosteric Binder Generation},\n",
-        "  author={Cao, Hanqun and Pal, Aastha and Tang, Sophia and Zhang, Yinuo and Zhang, Jingjie and Heng, Pheng Ann and Chatterjee, Pranam},\n",
-        "  journal={arXiv preprint arXiv:2605.09810},\n",
-        "  year={2026}",
-        "}\n",
-        "```"
-      ],
-      "metadata": {}
-    }
-  ]
-}
+  "kernelspec": {
+   "name": "python3",
+   "display_name": "Python 3"
+  },
+  "language_info": {
+   "name": "python"
+  },
+  "accelerator": "GPU"
+ },
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "source": [
+    "# TD3B: Transition-Directed Discrete Diffusion for Allosteric Binder Generation\n",
+    "\n",
+    "This notebook demonstrates **TD3B inference** — generating peptide binders with specified agonist or antagonist behavior for GPCR targets.\n",
+    "\n",
+    "**What TD3B does:**\n",
+    "- Takes a target protein sequence + desired direction (agonist / antagonist)\n",
+    "- Generates peptide binder sequences using a finetuned discrete diffusion model\n",
+    "- Scores them with a Direction Oracle and binding affinity predictor\n",
+    "- Returns the best candidates via weighted resampling (Algorithm 2)\n",
+    "\n",
+    "**Requirements:** GPU runtime (T4 or better). Click **Runtime → Change runtime type → GPU**."
+   ],
+   "metadata": {}
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "## 1. Setup"
+   ],
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Install dependencies\n",
+    "!pip install -q torch torchvision --index-url https://download.pytorch.org/whl/cu121\n",
+    "!pip install -q transformers fair-esm SmilesPE rdkit-pypi scipy pandas numpy xgboost pytorch-lightning lightning hydra-core loguru timm huggingface_hub"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": "# Clone the TD3B code and download checkpoints from HuggingFace.\n# Skip LFS blobs on clone (the large checkpoints are pulled explicitly below).\n!GIT_LFS_SKIP_SMUDGE=1 git clone https://huggingface.co/ChatterjeeLab/TD3B TD3B\n%cd TD3B\n\nfrom huggingface_hub import hf_hub_download\nimport os\n\nREPO_ID = \"ChatterjeeLab/TD3B\"\nos.makedirs(\"checkpoints\", exist_ok=True)\nos.makedirs(\"data\", exist_ok=True)\n\n# Download checkpoints (this may take a few minutes)\nfor fname in [\"checkpoints/td3b.ckpt\", \"checkpoints/pretrained.ckpt\",\n              \"checkpoints/direction_oracle.pt\",\n              \"scoring/functions/classifiers/binding-affinity.pt\",\n              \"data/test.csv\", \"data/train.csv\"]:\n    print(f\"Downloading {fname}...\")\n    hf_hub_download(repo_id=REPO_ID, filename=fname, local_dir=\".\")\n\nprint(\"\\nAll files downloaded!\")\n!ls -lh checkpoints/"
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "## 2. Load Model and Oracle"
+   ],
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": "import sys\nsys.path.insert(0, \".\")\n\nimport torch\nimport numpy as np\nimport pandas as pd\n\nfrom models.diffusion import Diffusion\nfrom configs.finetune_config import (\n    DiffusionConfig, RoFormerConfig, NoiseConfig,\n    TrainingConfig, SamplingConfig, EvalConfig, OptimConfig, MCTSConfig,\n)\nfrom tokenizer.my_tokenizers import SMILES_SPE_Tokenizer\nfrom td3b.direction_oracle import DirectionalOracle\nfrom td3b.td3b_scoring import TD3BRewardFunction, create_td3b_reward_function\nfrom scoring.functions.binding import BindingAffinity\nfrom utils.app import PeptideAnalyzer\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(f\"Using device: {device}\")\nif torch.cuda.is_available():\n    print(f\"GPU: {torch.cuda.get_device_name(0)}\")\n    print(f\"Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB\")"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Load tokenizer\n",
+    "tokenizer = SMILES_SPE_Tokenizer(\"tokenizer/new_vocab.txt\", \"tokenizer/new_splits.txt\")\n",
+    "print(f\"Tokenizer vocab size: {len(tokenizer)}\")\n",
+    "\n",
+    "# Load diffusion model\n",
+    "print(\"\\nLoading TD3B model...\")\n",
+    "cfg = DiffusionConfig(\n",
+    "    roformer=RoFormerConfig(hidden_size=768, n_layers=8, n_heads=8),\n",
+    "    noise=NoiseConfig(),\n",
+    "    training=TrainingConfig(sampling_eps=1e-3),\n",
+    "    sampling=SamplingConfig(steps=128, sampling_eps=1e-3),\n",
+    "    eval_cfg=EvalConfig(), optim=OptimConfig(lr=3e-4), mcts=MCTSConfig(),\n",
+    ")\n",
+    "model = Diffusion(config=cfg, tokenizer=tokenizer, device=device).to(device)\n",
+    "\n",
+    "ckpt = torch.load(\"checkpoints/td3b.ckpt\", map_location=device, weights_only=False)\n",
+    "state_dict = ckpt.get(\"model_state_dict\") or ckpt.get(\"state_dict\") or ckpt\n",
+    "model.load_state_dict(state_dict, strict=False)\n",
+    "model.eval()\n",
+    "model.tokenizer = tokenizer\n",
+    "print(\"TD3B model loaded!\")\n",
+    "\n",
+    "# Load Direction Oracle\n",
+    "print(\"\\nLoading Direction Oracle...\")\n",
+    "oracle = DirectionalOracle(\n",
+    "    model_ckpt=\"checkpoints/direction_oracle.pt\",\n",
+    "    tr2d2_checkpoint=\"checkpoints/pretrained.ckpt\",\n",
+    "    tokenizer_vocab=\"tokenizer/new_vocab.txt\",\n",
+    "    tokenizer_splits=\"tokenizer/new_splits.txt\",\n",
+    "    device=device,\n",
+    ")\n",
+    "oracle.eval()\n",
+    "print(\"Direction Oracle loaded!\")\n",
+    "\n",
+    "# Load Affinity Predictor\n",
+    "print(\"\\nLoading Affinity Predictor...\")\n",
+    "analyzer = PeptideAnalyzer()\n",
+    "print(\"\\nAll models loaded!\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "## 3. Define Helper Functions"
+   ],
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def sample_sequences(model, batch_size, seq_length, num_steps=128, eps=1e-5):\n",
+    "    \"\"\"Sample sequences from the diffusion model.\"\"\"\n",
+    "    x = model.sample_prior(batch_size, seq_length).to(model.device, dtype=torch.long)\n",
+    "    timesteps = torch.linspace(1, eps, num_steps + 1, device=model.device)\n",
+    "    dt = torch.tensor((1 - eps) / num_steps, device=model.device)\n",
+    "\n",
+    "    for i in range(num_steps):\n",
+    "        t = timesteps[i] * torch.ones(x.shape[0], 1, device=model.device)\n",
+    "        _, x = model.single_reverse_step(x, t=t, dt=dt)\n",
+    "        x = x.to(model.device)\n",
+    "\n",
+    "    mask_pos = (x == model.mask_index)\n",
+    "    if mask_pos.any():\n",
+    "        t = timesteps[-2] * torch.ones(x.shape[0], 1, device=model.device)\n",
+    "        _, x = model.single_noise_removal(x, t=t, dt=dt)\n",
+    "    return x\n",
+    "\n",
+    "\n",
+    "def generate_binders(target_seq, direction=\"agonist\", num_pool=32,\n",
+    "                     num_keep=8, alpha=0.1, seq_length=200):\n",
+    "    \"\"\"\n",
+    "    Generate directional binders for a target protein.\n",
+    "    \n",
+    "    Args:\n",
+    "        target_seq: Target protein amino acid sequence\n",
+    "        direction: 'agonist' or 'antagonist'\n",
+    "        num_pool: Number of candidates to generate\n",
+    "        num_keep: Number of final samples after resampling\n",
+    "        alpha: Temperature for weighted resampling\n",
+    "        seq_length: Binder sequence length (in SMILES tokens)\n",
+    "    \n",
+    "    Returns:\n",
+    "        DataFrame with generated binders and scores\n",
+    "    \"\"\"\n",
+    "    d_star = 1.0 if direction == \"agonist\" else -1.0\n",
+    "    \n",
+    "    # Build reward function\n",
+    "    affinity_pred = BindingAffinity(\n",
+    "        prot_seq=target_seq, tokenizer=tokenizer,\n",
+    "        base_path=\".\", device=device, emb_model=model.backbone\n",
+    "    )\n",
+    "    reward_fn = create_td3b_reward_function(\n",
+    "        affinity_predictor=affinity_pred,\n",
+    "        directional_oracle=oracle,\n",
+    "        target_protein_seq=target_seq,\n",
+    "        target_direction=direction,\n",
+    "        peptide_tokenizer=tokenizer,\n",
+    "        device=device,\n",
+    "    )\n",
+    "    \n",
+    "    # Generate candidates\n",
+    "    with torch.no_grad():\n",
+    "        x_pool = sample_sequences(model, num_pool, seq_length)\n",
+    "    sequences = tokenizer.batch_decode(x_pool)\n",
+    "    \n",
+    "    # Score all\n",
+    "    rewards, info = reward_fn(sequences)\n",
+    "    affinities = info[\"affinities\"]\n",
+    "    directions = info[\"directions\"]\n",
+    "    \n",
+    "    # Weighted resampling (Algorithm 2)\n",
+    "    rewards_t = torch.as_tensor(rewards, device=device)\n",
+    "    weights = torch.softmax(rewards_t / max(alpha, 1e-6), dim=0)\n",
+    "    idx = torch.multinomial(weights, num_samples=num_keep, replacement=True)\n",
+    "    chosen = idx.cpu().numpy()\n",
+    "    \n",
+    "    # Filter to valid peptides only\n",
+    "    results = []\n",
+    "    for i in chosen:\n",
+    "        is_valid = analyzer.is_peptide(sequences[i])\n",
+    "        da = float(directions[i] > 0.5) if d_star > 0 else float(directions[i] < 0.5)\n",
+    "        results.append({\n",
+    "            \"sequence\": sequences[i],\n",
+    "            \"direction\": direction,\n",
+    "            \"is_valid\": is_valid,\n",
+    "            \"affinity\": float(affinities[i]),\n",
+    "            \"gated_reward\": float(rewards[i]),\n",
+    "            \"p_agonist\": float(directions[i]),\n",
+    "            \"direction_accuracy\": da,\n",
+    "        })\n",
+    "    \n",
+    "    df = pd.DataFrame(results)\n",
+    "    return df"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "## 4. Generate Binders\n",
+    "\n",
+    "Let's generate **agonist** and **antagonist** binders for a test target and compare the Direction Oracle predictions."
+   ],
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Load test targets\n",
+    "test_df = pd.read_csv(\"data/test.csv\")\n",
+    "print(f\"Test set: {len(test_df)} target-binder pairs\")\n",
+    "\n",
+    "# Pick first target for demo\n",
+    "target_row = test_df.iloc[0]\n",
+    "TARGET_SEQ = target_row[\"Target_Sequence\"]\n",
+    "TARGET_UID = target_row[\"Target_UniProt_ID\"]\n",
+    "print(f\"\\nTarget: {TARGET_UID}\")\n",
+    "print(f\"Sequence length: {len(TARGET_SEQ)} aa\")\n",
+    "print(f\"Sequence: {TARGET_SEQ[:60]}...\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%time\n",
+    "# Generate AGONIST binders\n",
+    "print(\"Generating agonist binders (d*=+1)...\")\n",
+    "torch.manual_seed(42)\n",
+    "np.random.seed(42)\n",
+    "df_agonist = generate_binders(TARGET_SEQ, direction=\"agonist\", num_pool=32, num_keep=8)\n",
+    "\n",
+    "print(f\"\\nGenerated {len(df_agonist)} samples ({df_agonist['is_valid'].sum()} valid)\")\n",
+    "print(f\"Mean p(agonist): {df_agonist['p_agonist'].mean():.3f}\")\n",
+    "print(f\"Mean affinity: {df_agonist['affinity'].mean():.2f}\")\n",
+    "print(f\"Mean gated reward: {df_agonist['gated_reward'].mean():.2f}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%time\n",
+    "# Generate ANTAGONIST binders\n",
+    "print(\"Generating antagonist binders (d*=-1)...\")\n",
+    "torch.manual_seed(42)\n",
+    "np.random.seed(42)\n",
+    "df_antagonist = generate_binders(TARGET_SEQ, direction=\"antagonist\", num_pool=32, num_keep=8)\n",
+    "\n",
+    "print(f\"\\nGenerated {len(df_antagonist)} samples ({df_antagonist['is_valid'].sum()} valid)\")\n",
+    "print(f\"Mean p(agonist): {df_antagonist['p_agonist'].mean():.3f}\")\n",
+    "print(f\"Mean affinity: {df_antagonist['affinity'].mean():.2f}\")\n",
+    "print(f\"Mean gated reward: {df_antagonist['gated_reward'].mean():.2f}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "## 5. Compare Directional Control"
+   ],
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "fig, axes = plt.subplots(1, 3, figsize=(15, 4))\n",
+    "\n",
+    "# Plot 1: Direction Oracle p(agonist)\n",
+    "axes[0].hist(df_agonist[\"p_agonist\"], bins=20, alpha=0.7, label=\"d*=+1 (agonist)\", color=\"#e74c3c\")\n",
+    "axes[0].hist(df_antagonist[\"p_agonist\"], bins=20, alpha=0.7, label=\"d*=-1 (antagonist)\", color=\"#3498db\")\n",
+    "axes[0].axvline(0.5, color=\"gray\", linestyle=\"--\", label=\"threshold\")\n",
+    "axes[0].set_xlabel(\"p(agonist)\")\n",
+    "axes[0].set_ylabel(\"Count\")\n",
+    "axes[0].set_title(\"Direction Oracle Predictions\")\n",
+    "axes[0].legend()\n",
+    "\n",
+    "# Plot 2: Binding Affinity\n",
+    "axes[1].hist(df_agonist[\"affinity\"], bins=20, alpha=0.7, label=\"Agonist\", color=\"#e74c3c\")\n",
+    "axes[1].hist(df_antagonist[\"affinity\"], bins=20, alpha=0.7, label=\"Antagonist\", color=\"#3498db\")\n",
+    "axes[1].set_xlabel(\"Predicted Binding Affinity\")\n",
+    "axes[1].set_ylabel(\"Count\")\n",
+    "axes[1].set_title(\"Binding Affinity Distribution\")\n",
+    "axes[1].legend()\n",
+    "\n",
+    "# Plot 3: Gated Reward\n",
+    "axes[2].hist(df_agonist[\"gated_reward\"], bins=20, alpha=0.7, label=\"Agonist\", color=\"#e74c3c\")\n",
+    "axes[2].hist(df_antagonist[\"gated_reward\"], bins=20, alpha=0.7, label=\"Antagonist\", color=\"#3498db\")\n",
+    "axes[2].set_xlabel(\"Gated Reward\")\n",
+    "axes[2].set_ylabel(\"Count\")\n",
+    "axes[2].set_title(\"Gated Reward Distribution\")\n",
+    "axes[2].legend()\n",
+    "\n",
+    "plt.tight_layout()\n",
+    "plt.savefig(\"td3b_results.png\", dpi=150, bbox_inches=\"tight\")\n",
+    "plt.show()\n",
+    "\n",
+    "print(\"\\nSummary:\")\n",
+    "print(f\"  Agonist mode:    p(agonist)={df_agonist['p_agonist'].mean():.3f}  Affinity={df_agonist['affinity'].mean():.2f}  Gated={df_agonist['gated_reward'].mean():.2f}\")\n",
+    "print(f\"  Antagonist mode: p(agonist)={df_antagonist['p_agonist'].mean():.3f}  Affinity={df_antagonist['affinity'].mean():.2f}  Gated={df_antagonist['gated_reward'].mean():.2f}\")\n",
+    "print(f\"  Directional gap: Δp = {df_agonist['p_agonist'].mean() - df_antagonist['p_agonist'].mean():.3f}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "## 6. Run on Multiple Targets\n",
+    "\n",
+    "Generate binders for the first 5 test targets and compute aggregate metrics."
+   ],
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "N_TARGETS = 5  # Number of targets to evaluate (increase for full benchmark)\n",
+    "\n",
+    "all_results = []\n",
+    "targets = test_df.drop_duplicates(\"Target_UniProt_ID\").head(N_TARGETS)\n",
+    "\n",
+    "for i, (_, row) in enumerate(targets.iterrows()):\n",
+    "    uid = row[\"Target_UniProt_ID\"]\n",
+    "    seq = row[\"Target_Sequence\"]\n",
+    "    print(f\"[{i+1}/{N_TARGETS}] {uid} (len={len(seq)})\")\n",
+    "    \n",
+    "    for direction in [\"agonist\", \"antagonist\"]:\n",
+    "        torch.manual_seed(42)\n",
+    "        np.random.seed(42)\n",
+    "        df = generate_binders(seq, direction=direction, num_pool=32, num_keep=8)\n",
+    "        df[\"target_uid\"] = uid\n",
+    "        all_results.append(df)\n",
+    "        \n",
+    "        d_star = 1.0 if direction == \"agonist\" else -1.0\n",
+    "        da = df[\"direction_accuracy\"].mean()\n",
+    "        print(f\"    {direction:>10s}: DA={da:.2f}  Aff={df['affinity'].mean():.2f}  Gated={df['gated_reward'].mean():.2f}  valid={df['is_valid'].sum()}/{len(df)}\")\n",
+    "\n",
+    "combined = pd.concat(all_results, ignore_index=True)\n",
+    "\n",
+    "print(f\"\\n{'='*60}\")\n",
+    "print(f\"AGGREGATE METRICS ({N_TARGETS} targets)\")\n",
+    "print(f\"{'='*60}\")\n",
+    "for d_name, d_val in [(\"Agonist (d*=+1)\", \"agonist\"), (\"Antagonist (d*=-1)\", \"antagonist\")]:\n",
+    "    sub = combined[combined[\"direction\"] == d_val]\n",
+    "    valid = sub[sub[\"is_valid\"] == True]\n",
+    "    print(f\"  {d_name}:\")\n",
+    "    print(f\"    Affinity: {sub['affinity'].mean():.2f}\")\n",
+    "    print(f\"    Direction Accuracy: {sub['direction_accuracy'].mean():.3f}\")\n",
+    "    print(f\"    Gated Reward (all): {sub['gated_reward'].mean():.2f}\")\n",
+    "    if len(valid) > 0:\n",
+    "        print(f\"    Gated Reward (valid): {valid['gated_reward'].mean():.2f}\")\n",
+    "    print(f\"    Valid: {sub['is_valid'].sum()}/{len(sub)}\")\n",
+    "\n",
+    "# Save\n",
+    "combined.to_csv(\"td3b_demo_results.csv\", index=False)\n",
+    "print(f\"\\nResults saved to td3b_demo_results.csv\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "## Citation\n",
+    "\n",
+    "```bibtex\n",
+    "@article{caotd3b,\n",
+    "  title={TD3B: Transition-Directed Discrete Diffusion for Allosteric Binder Generation},\n",
+    "  author={Cao, Hanqun and Pal, Aastha and Tang, Sophia and Zhang, Yinuo and Zhang, Jingjie and Heng, Pheng-Ann and Chatterjee, Pranam}\n",
+    "}\n",
+    "```"
+   ],
+   "metadata": {}
+  }
+ ]
+}

sampling_strategies.py

@@ -0,0 +1,461 @@
+#!/usr/bin/env python3
+"""
+sampling_strategies.py  --  TD3B / MDLM validity-boosting samplers (FUNCTION B).
+
+Goal
+----
+As the target peptide length grows, the fraction of decoded SMILES that pass
+``utils.app.PeptideAnalyzer.is_peptide`` (RDKit) collapses.  This module supplies a
+library of *sampling-time* techniques (no retraining) that raise that valid-yield,
+especially at long length.
+
+All samplers REUSE the model's existing masked-diffusion primitives:
+    * ``Diffusion.sample_prior``        -> fully masked start
+    * ``Diffusion.single_reverse_step`` -> one reverse (denoising) step  (returns log_p, x_next)
+    * ``Diffusion.single_noise_removal``-> final step guaranteeing no surviving [MASK]
+    * ``Diffusion.forward``             -> per-position log p(x0)  (SUBS-parameterised logits)
+
+We DO NOT reimplement the diffusion math.  The reverse-posterior construction
+(``q_xs = p_x0 * (t - (t-dt))`` with the mask-stay probability, and the carry-over
+copy_flag) is left entirely to ``single_reverse_step`` / ``single_noise_removal``.
+We only change **token SELECTION** -- i.e. we transform the model's clean-token
+distribution ``p_x0`` (temperature / top-k / top-p) that we feed back into
+``single_reverse_step(..., p_x0=...)`` -- and add a **remask self-correction loop**
+and a **best-of-N validity-guided rejection** wrapper on top.
+
+Common entry point
+------------------
+    generate(model, tokenizer, analyzer, batch_size, length, strategy="baseline", **kw)
+        -> (tokens, sequences, valid_mask, stats)
+
+Strategies (``strategy=``):
+    baseline    : reverse diffusion identical to inference.sample_sequences (num_steps=128)
+    more_steps  : same, but num_steps scales with length (steps_per_token)
+    top_p       : nucleus -- restrict p_x0 to the smallest set with cumulative mass >= p
+    nucleus     : alias of top_p
+    top_k       : restrict p_x0 to the k most-probable clean tokens
+    low_temp    : temperature < 1 on the clean-token logits (sharper, more valid)
+    remask      : self-correction -- generate, then remask the lowest-confidence K%%
+                  of tokens of INVALID sequences and re-denoise, for R rounds
+    best_of_n   : oversample N per slot, keep the first valid decode
+    nucleus_remask : nucleus + remask (recommended default for long length)
+
+Any preset can be combined with explicit kwargs, e.g.
+    generate(..., strategy="remask", top_p=0.9, remask_rounds=4)
+
+The library is model-agnostic: it works with a RANDOM-init ``Diffusion`` (for CPU
+development / relative benchmarking) and, unchanged, with the real checkpoint.
+"""
+
+from __future__ import annotations
+
+import time
+from typing import Callable, Dict, List, Optional, Tuple
+
+import numpy as np
+import torch
+
+# --------------------------------------------------------------------------- #
+#  Distribution transforms  (token SELECTION only -- no diffusion math here)   #
+# --------------------------------------------------------------------------- #
+_NEG_INF = -1e9
+
+
+def _apply_top_k(logits: torch.Tensor, k: int) -> torch.Tensor:
+    """Keep the k largest logits per position, push the rest to -inf."""
+    if k is None or k <= 0 or k >= logits.shape[-1]:
+        return logits
+    kth = torch.topk(logits, k, dim=-1).values[..., -1, None]  # (..., 1) k-th largest
+    return logits.masked_fill(logits < kth, _NEG_INF)
+
+
+def _apply_top_p(probs: torch.Tensor, p: float) -> torch.Tensor:
+    """Nucleus filter: zero out the low-mass tail so the kept set has cumulative
+    mass >= p (always keeps at least the argmax), then renormalise."""
+    if p is None or p >= 1.0:
+        return probs
+    sorted_probs, sorted_idx = torch.sort(probs, dim=-1, descending=True)
+    cumsum = sorted_probs.cumsum(dim=-1)
+    # a token is DROPPED if the mass strictly *before* it already reached p
+    drop_sorted = (cumsum - sorted_probs) > p
+    sorted_probs = sorted_probs.masked_fill(drop_sorted, 0.0)
+    new_probs = torch.zeros_like(probs).scatter_(-1, sorted_idx, sorted_probs)
+    return new_probs / new_probs.sum(dim=-1, keepdim=True).clamp_min(1e-12)
+
+
+def make_transform(temperature: float = 1.0,
+                   top_p: Optional[float] = None,
+                   top_k: Optional[int] = None) -> Optional[Callable[[torch.Tensor], torch.Tensor]]:
+    """Build a callable ``log_p (B,L,V) -> p_x0 (B,L,V)`` implementing
+    temperature -> top_k (logit space) -> softmax -> top_p (prob space).
+
+    Returns ``None`` when the transform is the identity (baseline path), so the
+    caller can take the cheaper single-``single_reverse_step`` route.
+    """
+    temp_on = temperature is not None and abs(temperature - 1.0) > 1e-8
+    if not temp_on and top_p is None and top_k is None:
+        return None
+
+    def _transform(log_p: torch.Tensor) -> torch.Tensor:
+        v = log_p
+        if temp_on:
+            v = v / float(temperature)
+        if top_k is not None:
+            v = _apply_top_k(v, int(top_k))
+        probs = torch.softmax(v, dim=-1)
+        if top_p is not None:
+            probs = _apply_top_p(probs, float(top_p))
+        return probs
+
+    return _transform
+
+
+# --------------------------------------------------------------------------- #
+#  Core reverse-diffusion loop (delegates all diffusion math to the model)     #
+# --------------------------------------------------------------------------- #
+def _reverse_step(model, x: torch.Tensor, t: torch.Tensor, dt: torch.Tensor,
+                  attn_mask: torch.Tensor,
+                  transform: Optional[Callable]) -> Tuple[torch.Tensor, torch.Tensor]:
+    """One reverse step.  Returns (log_p, x_next).
+
+    * transform is None  -> plain ``single_reverse_step`` (one forward, returns log_p).
+    * transform given    -> one ``forward`` to get log p_x0, transform it, then hand the
+      modified p_x0 to ``single_reverse_step`` which still builds the reverse posterior
+      and samples exactly as in the original math.
+    """
+    if transform is None:
+        log_p, x_next = model.single_reverse_step(x, t=t, dt=dt, attn_mask=attn_mask)
+        return log_p, x_next
+    # lower-level logits (SUBS-parameterised); sigma comes from the model's own schedule
+    sigma_t, _ = model.noise(t)
+    log_p = model.forward(x, attn_mask=attn_mask, sigma=sigma_t)
+    p_x0 = transform(log_p)
+    _, x_next = model.single_reverse_step(x, t=t, dt=dt, p_x0=p_x0, attn_mask=attn_mask)
+    return log_p, x_next
+
+
+def _update_confidence(conf: torch.Tensor, x_prev: torch.Tensor, x_next: torch.Tensor,
+                       log_p: torch.Tensor, mask_index: int) -> torch.Tensor:
+    """Record, for every position that transitions mask->token at this step, the model's
+    probability of the chosen token.  Because of carry-over unmasking each position is
+    written exactly once (when it first unmasks); higher conf == model more certain."""
+    if log_p is None:
+        return conf
+    newly = (x_prev == mask_index) & (x_next != mask_index)
+    if not newly.any():
+        return conf
+    probs = log_p.exp()
+    chosen = x_next.clamp(0, probs.shape[-1] - 1).unsqueeze(-1)
+    step_conf = probs.gather(-1, chosen).squeeze(-1)
+    return torch.where(newly, step_conf, conf)
+
+
+def _reverse_diffusion(model, x: torch.Tensor, num_steps: int,
+                       t_start: float = 1.0, eps: float = 1e-5,
+                       transform: Optional[Callable] = None,
+                       attn_mask: Optional[torch.Tensor] = None,
+                       conf: Optional[torch.Tensor] = None,
+                       noise_removal: bool = True) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Reverse-diffuse ``x`` from time ``t_start`` down to ``eps`` over ``num_steps``.
+
+    Mirrors ``inference.sample_sequences`` (linspace timesteps, final noise-removal to
+    guarantee no surviving [MASK]) but supports (a) an arbitrary starting time -- so the
+    remask loop can resume from a partially-masked state -- and (b) a p_x0 ``transform``.
+    Already-unmasked positions are preserved by the model's copy_flag (carry-over).
+    Returns (x, confidence) with confidence in [0, 1] per position.
+    """
+    device = model.device
+    x = x.to(device, dtype=torch.long)
+    if attn_mask is None:
+        attn_mask = torch.ones_like(x, device=device, dtype=torch.long)
+    if conf is None:
+        conf = torch.zeros(x.shape, device=device, dtype=torch.float32)
+    mask_index = model.mask_index
+
+    num_steps = max(int(num_steps), 1)
+    timesteps = torch.linspace(t_start, eps, num_steps + 1, device=device)
+    dt = torch.tensor((t_start - eps) / num_steps, device=device)
+
+    for i in range(num_steps):
+        t = timesteps[i] * torch.ones(x.shape[0], 1, device=device)
+        log_p, x_next = _reverse_step(model, x, t, dt, attn_mask, transform)
+        conf = _update_confidence(conf, x, x_next, log_p, mask_index)
+        x = x_next.to(device)
+
+    if noise_removal and (x == mask_index).any():
+        t = timesteps[-2] * torch.ones(x.shape[0], 1, device=device)
+        log_p, x_next = model.single_noise_removal(x, t=t, dt=dt, attn_mask=attn_mask)
+        conf = _update_confidence(conf, x, x_next, log_p, mask_index)
+        x = x_next.to(device)
+
+    return x, conf
+
+
+# --------------------------------------------------------------------------- #
+#  Decode / validity helpers                                                    #
+# --------------------------------------------------------------------------- #
+def decode_and_validate(tokenizer, analyzer, x: torch.Tensor) -> Tuple[List[str], np.ndarray]:
+    """Decode token ids to SMILES and test each with ``analyzer.is_peptide``."""
+    sequences = tokenizer.batch_decode(x)
+    valid = np.fromiter((bool(analyzer.is_peptide(s)) for s in sequences),
+                        dtype=bool, count=len(sequences))
+    return sequences, valid
+
+
+# --------------------------------------------------------------------------- #
+#  Remask self-correction (the key technique for long sequences)               #
+# --------------------------------------------------------------------------- #
+def _remask_and_redenoise(model, x: torch.Tensor, conf: torch.Tensor,
+                          rows: torch.Tensor, remask_frac: float, remask_steps: int,
+                          transform: Optional[Callable], eps: float,
+                          attn_mask: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    """For the selected ``rows`` (typically the invalid sequences): set the lowest-
+    confidence ``remask_frac`` of their positions back to [MASK], then re-denoise ONLY
+    those rows from t=remask_frac down to eps.  Returns updated (x, conf) copies.
+
+    The high-confidence tokens are kept (carry-over), so re-denoising only re-samples
+    the positions the model was least sure about -- a cheap, targeted second attempt.
+    """
+    x = x.clone()
+    conf = conf.clone()
+    sub_x = x[rows].clone()
+    sub_conf = conf[rows].clone()
+    n, L = sub_x.shape
+    k = max(1, int(round(remask_frac * L)))
+    k = min(k, L)
+
+    # lowest-confidence k positions per row -> back to MASK
+    low_idx = torch.topk(sub_conf, k, dim=-1, largest=False).indices  # (n, k)
+    sub_x.scatter_(1, low_idx, model.mask_index)
+    sub_conf.scatter_(1, low_idx, 0.0)  # will be recomputed when re-unmasked
+
+    # resume masked-diffusion from a time consistent with the masked fraction (~k/L)
+    t_start = float(min(max(k / L, 2 * eps), 1.0))
+    sub_attn = attn_mask[rows]
+    sub_x, sub_conf = _reverse_diffusion(
+        model, sub_x, num_steps=remask_steps, t_start=t_start, eps=eps,
+        transform=transform, attn_mask=sub_attn, conf=sub_conf, noise_removal=True)
+
+    x[rows] = sub_x
+    conf[rows] = sub_conf
+    return x, conf
+
+
+# --------------------------------------------------------------------------- #
+#  Single generation pass (+ optional remask rounds)                           #
+# --------------------------------------------------------------------------- #
+@torch.no_grad()
+def _generate_once(model, tokenizer, analyzer, batch_size: int, length: int,
+                   num_steps: int, eps: float,
+                   transform: Optional[Callable],
+                   remask_rounds: int, remask_frac: float, remask_steps: int
+                   ) -> Tuple[torch.Tensor, List[str], np.ndarray, torch.Tensor, List[int]]:
+    device = model.device
+    x = model.sample_prior(batch_size, length).to(device, dtype=torch.long)
+    attn_mask = torch.ones_like(x, device=device, dtype=torch.long)
+
+    x, conf = _reverse_diffusion(model, x, num_steps=num_steps, t_start=1.0, eps=eps,
+                                 transform=transform, attn_mask=attn_mask,
+                                 noise_removal=True)
+    sequences, valid = decode_and_validate(tokenizer, analyzer, x)
+    round_valid = [int(valid.sum())]  # valid count after each stage (round 0 == initial)
+
+    for _ in range(int(remask_rounds)):
+        if valid.all():
+            break
+        rows = torch.from_numpy(np.where(~valid)[0]).to(device=device, dtype=torch.long)
+        x, conf = _remask_and_redenoise(model, x, conf, rows, remask_frac, remask_steps,
+                                        transform, eps, attn_mask)
+        # re-decode / re-validate only the rows we touched
+        new_seqs = tokenizer.batch_decode(x[rows])
+        for j, gi in enumerate(rows.tolist()):
+            sequences[gi] = new_seqs[j]
+            valid[gi] = bool(analyzer.is_peptide(new_seqs[j]))
+        round_valid.append(int(valid.sum()))
+
+    return x, sequences, valid, conf, round_valid
+
+
+# --------------------------------------------------------------------------- #
+#  Best-of-N validity-guided rejection wrapper                                 #
+# --------------------------------------------------------------------------- #
+@torch.no_grad()
+def _generate_best_of_n(model, tokenizer, analyzer, batch_size: int, length: int,
+                        num_steps: int, eps: float, transform: Optional[Callable],
+                        remask_rounds: int, remask_frac: float, remask_steps: int,
+                        best_of_n: int
+                        ) -> Tuple[torch.Tensor, List[str], np.ndarray, Dict]:
+    """Draw up to ``best_of_n`` independent candidates per slot; keep the first valid
+    decode for each slot (fall back to the last draw if none is valid)."""
+    device = model.device
+    tokens = None
+    sequences: List[str] = [""] * batch_size
+    valid = np.zeros(batch_size, dtype=bool)
+    draws_used = 0
+    for n in range(max(int(best_of_n), 1)):
+        draws_used = n + 1
+        xg, seqs, vmask, _conf, _rv = _generate_once(
+            model, tokenizer, analyzer, batch_size, length, num_steps, eps,
+            transform, remask_rounds, remask_frac, remask_steps)
+        if tokens is None:
+            tokens = xg.clone()
+            for i in range(batch_size):
+                sequences[i] = seqs[i]
+            valid = vmask.copy()
+        else:
+            # accept this draw only for slots not yet valid
+            take = (~valid) & vmask
+            if take.any():
+                idx = np.where(take)[0]
+                tokens[idx] = xg[idx]
+                for i in idx:
+                    sequences[i] = seqs[i]
+                valid[idx] = True
+            # for still-invalid slots, keep the freshest candidate (so tokens stay consistent)
+            still = np.where(~valid)[0]
+            if len(still):
+                tokens[still] = xg[still]
+                for i in still:
+                    sequences[i] = seqs[i]
+        if valid.all():
+            break
+    stats = {"best_of_n_draws_used": draws_used}
+    return tokens, sequences, valid, stats
+
+
+# --------------------------------------------------------------------------- #
+#  Public dispatch                                                              #
+# --------------------------------------------------------------------------- #
+STRATEGY_PRESETS: Dict[str, Dict] = {
+    "baseline":       dict(),
+    "more_steps":     dict(steps_per_token=1.0),
+    "top_p":          dict(top_p=0.9),
+    "nucleus":        dict(top_p=0.9),
+    "top_k":          dict(top_k=20),
+    "low_temp":       dict(temperature=0.7),
+    "remask":         dict(remask_rounds=3, remask_frac=0.25, remask_steps=32),
+    "best_of_n":      dict(best_of_n=4),
+    "nucleus_remask": dict(top_p=0.9, remask_rounds=3, remask_frac=0.25, remask_steps=32),
+}
+
+
+def available_strategies() -> List[str]:
+    return list(STRATEGY_PRESETS.keys())
+
+
+@torch.no_grad()
+def generate(model, tokenizer, analyzer, batch_size: int, length: int,
+             strategy: str = "baseline",
+             num_steps: int = 128, eps: float = 1e-5,
+             temperature: Optional[float] = None,
+             top_p: Optional[float] = None,
+             top_k: Optional[int] = None,
+             steps_per_token: Optional[float] = None,
+             remask_rounds: Optional[int] = None,
+             remask_frac: Optional[float] = None,
+             remask_steps: Optional[int] = None,
+             best_of_n: Optional[int] = None,
+             verbose: bool = False,
+             ) -> Tuple[torch.Tensor, List[str], np.ndarray, Dict]:
+    """Generate ``batch_size`` peptides of ``length`` tokens with the chosen strategy.
+
+    Returns
+    -------
+    tokens     : LongTensor (batch_size, length)  final token ids
+    sequences  : list[str]                         decoded SMILES
+    valid_mask : np.ndarray[bool] (batch_size,)    analyzer.is_peptide per sequence
+    stats      : dict                              metrics (valid_rate, timing, knobs, ...)
+    """
+    if strategy not in STRATEGY_PRESETS:
+        raise ValueError(f"unknown strategy {strategy!r}; choose from {available_strategies()}")
+
+    # preset provides defaults; any explicitly-passed (non-None) kwarg overrides it
+    cfg = dict(temperature=1.0, top_p=None, top_k=None, steps_per_token=None,
+               remask_rounds=0, remask_frac=0.25, remask_steps=32, best_of_n=1)
+    cfg.update(STRATEGY_PRESETS[strategy])
+    explicit = dict(temperature=temperature, top_p=top_p, top_k=top_k,
+                    steps_per_token=steps_per_token, remask_rounds=remask_rounds,
+                    remask_frac=remask_frac, remask_steps=remask_steps, best_of_n=best_of_n)
+    for key, val in explicit.items():
+        if val is not None:
+            cfg[key] = val
+
+    # steps scale with length for `more_steps` (>= base num_steps)
+    eff_steps = int(num_steps)
+    if cfg["steps_per_token"] is not None:
+        eff_steps = max(eff_steps, int(round(cfg["steps_per_token"] * length)))
+
+    transform = make_transform(cfg["temperature"], cfg["top_p"], cfg["top_k"])
+
+    t0 = time.time()
+    if int(cfg["best_of_n"]) > 1:
+        tokens, sequences, valid, extra = _generate_best_of_n(
+            model, tokenizer, analyzer, batch_size, length, eff_steps, eps, transform,
+            int(cfg["remask_rounds"]), float(cfg["remask_frac"]), int(cfg["remask_steps"]),
+            int(cfg["best_of_n"]))
+        round_valid = [int(valid.sum())]
+    else:
+        tokens, sequences, valid, _conf, round_valid = _generate_once(
+            model, tokenizer, analyzer, batch_size, length, eff_steps, eps, transform,
+            int(cfg["remask_rounds"]), float(cfg["remask_frac"]), int(cfg["remask_steps"]))
+        extra = {}
+    wall = time.time() - t0
+
+    valid_count = int(valid.sum())
+    stats = {
+        "strategy": strategy,
+        "length": int(length),
+        "batch_size": int(batch_size),
+        "num_steps": int(eff_steps),
+        "temperature": float(cfg["temperature"]),
+        "top_p": cfg["top_p"],
+        "top_k": cfg["top_k"],
+        "remask_rounds": int(cfg["remask_rounds"]),
+        "remask_frac": float(cfg["remask_frac"]),
+        "remask_steps": int(cfg["remask_steps"]),
+        "best_of_n": int(cfg["best_of_n"]),
+        "valid_count": valid_count,
+        "valid_rate": valid_count / max(batch_size, 1),
+        "round_valid_counts": round_valid,   # valid count after each remask round (shows mechanism)
+        "wall_time_s": round(wall, 2),
+    }
+    stats.update(extra)
+    if verbose:
+        print(f"[{strategy}] L={length} steps={eff_steps} "
+              f"valid={valid_count}/{batch_size} ({stats['valid_rate']:.1%}) "
+              f"rounds={round_valid} {wall:.1f}s")
+    return tokens, sequences, valid, stats
+
+
+# --------------------------------------------------------------------------- #
+#  Model construction helpers                                                   #
+# --------------------------------------------------------------------------- #
+def build_random_model(device="cpu", hidden_size: int = 768, n_layers: int = 8,
+                       n_heads: int = 8, tokenizer=None, base_path: Optional[str] = None):
+    """Construct a RANDOM-init ``Diffusion`` (no checkpoint) for CPU development /
+    relative benchmarking.  Same architecture/config path as ``inference.load_model``
+    minus the weight load, so samplers exercised here transfer unchanged to the real
+    checkpoint.  Reduce hidden_size/n_layers for fast CPU benchmarks (yields are
+    garbage either way with random init -- only relative trends are meaningful).
+    """
+    import os
+    from configs.finetune_config import (DiffusionConfig, RoFormerConfig, NoiseConfig,
+                                          TrainingConfig, SamplingConfig, EvalConfig,
+                                          OptimConfig, MCTSConfig)
+    from models.diffusion import Diffusion
+    if tokenizer is None:
+        from training.finetune_utils import load_tokenizer
+        # default to this module's own directory (repo root) so it works from any cwd
+        if base_path is None:
+            base_path = os.path.dirname(os.path.abspath(__file__))
+        tokenizer = load_tokenizer(base_path)
+    dev = torch.device(device)
+    cfg = DiffusionConfig(
+        roformer=RoFormerConfig(hidden_size=hidden_size, n_layers=n_layers, n_heads=n_heads),
+        noise=NoiseConfig(), training=TrainingConfig(sampling_eps=1e-3),
+        sampling=SamplingConfig(steps=128, sampling_eps=1e-3), eval_cfg=EvalConfig(),
+        optim=OptimConfig(lr=3e-4), mcts=MCTSConfig())
+    model = Diffusion(config=cfg, tokenizer=tokenizer, device=dev).to(dev)
+    model.eval()
+    model.tokenizer = tokenizer
+    return model, tokenizer

scoring/functions/binding.py

@@ -154,10 +154,12 @@ class BindingAffinity:
             self.max_pep_len = getattr(self.pep_model.config, "max_position_embeddings", None)
 
         self.model = ImprovedBindingPredictor().to(self.device)
-        checkpoint = torch.load(f'{base_path}/tr2d2-pep/scoring/functions/classifiers/binding-affinity.pt',
+        checkpoint = torch.load(f'{base_path}/scoring/functions/classifiers/binding-affinity.pt',
                                 map_location=self.device,
                                 weights_only=False)
-        self.model.load_state_dict(checkpoint['model_state_dict'])
+        _bind_sd = checkpoint.get('model_state_dict', checkpoint.get('state_dict', checkpoint)) \
+            if isinstance(checkpoint, dict) else checkpoint
+        self.model.load_state_dict(_bind_sd)
         
         self.model.eval()
         
@@ -268,10 +270,12 @@ class MultiTargetBindingAffinity:
 
         # Binding affinity prediction model
         self.model = ImprovedBindingPredictor().to(self.device)
-        checkpoint = torch.load(f'{base_path}/tr2d2-pep/scoring/functions/classifiers/binding-affinity.pt',
+        checkpoint = torch.load(f'{base_path}/scoring/functions/classifiers/binding-affinity.pt',
                                 map_location=self.device,
                                 weights_only=False)
-        self.model.load_state_dict(checkpoint['model_state_dict'])
+        _bind_sd = checkpoint.get('model_state_dict', checkpoint.get('state_dict', checkpoint)) \
+            if isinstance(checkpoint, dict) else checkpoint
+        self.model.load_state_dict(_bind_sd)
         self.model.eval()
 
         # Protein (ESM) model

scoring/functions/hemolysis.py

@@ -15,7 +15,7 @@ class Hemolysis:
     
     def __init__(self, tokenizer, base_path, device=None, emb_model=None):
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") if device is None else device
-        self.predictor = xgb.Booster(model_file=f'{base_path}/tr2d2-pep/scoring/functions/classifiers/hemolysis-xgboost.json')
+        self.predictor = xgb.Booster(model_file=f'{base_path}/scoring/functions/classifiers/hemolysis-xgboost.json')
         self.emb_model = emb_model if emb_model is not None else AutoModelForMaskedLM.from_pretrained('aaronfeller/PeptideCLM-23M-all').roformer.to(device).eval()
         self.tokenizer = tokenizer
     

scoring/functions/nonfouling.py

@@ -18,7 +18,7 @@ class Nonfouling:
     
     def __init__(self, tokenizer, base_path, device=None, emb_model=None):
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") if device is None else device
-        self.predictor = xgb.Booster(model_file=f'{base_path}/tr2d2-pep/scoring/functions/classifiers/nonfouling-xgboost.json')
+        self.predictor = xgb.Booster(model_file=f'{base_path}/scoring/functions/classifiers/nonfouling-xgboost.json')
         self.emb_model = emb_model if emb_model is not None else AutoModelForMaskedLM.from_pretrained('aaronfeller/PeptideCLM-23M-all').roformer.to(device).eval()
         self.tokenizer = tokenizer
     

scoring/functions/permeability.py

@@ -100,7 +100,7 @@ class Permeability:
     
     def __init__(self, tokenizer, base_path, device=None, emb_model=None):
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") if device is None else device
-        self.predictor = xgb.Booster(model_file=f'{base_path}/tr2d2-pep/scoring/functions/classifiers/permeability-xgboost.json')
+        self.predictor = xgb.Booster(model_file=f'{base_path}/scoring/functions/classifiers/permeability-xgboost.json')
         if emb_model is not None:
             self.emb_model = emb_model.to(self.device).eval()
         else:

scoring/functions/solubility.py

@@ -14,7 +14,7 @@ warnings.filterwarnings("ignore", category=FutureWarning)
 class Solubility:
     def __init__(self, tokenizer, base_path, device=None, emb_model=None):
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") if device is None else device
-        self.predictor = xgb.Booster(model_file=f'{base_path}/tr2d2-pep/scoring/functions/classifiers/solubility-xgboost.json')
+        self.predictor = xgb.Booster(model_file=f'{base_path}/scoring/functions/classifiers/solubility-xgboost.json')
         if emb_model is not None:
             self.emb_model = emb_model.to(self.device).eval()
         else:

scoring/scoring_functions.py

@@ -47,8 +47,8 @@ class ScoringFunctions:
         emb_model = AutoModelForMaskedLM.from_pretrained(
             'aaronfeller/PeptideCLM-23M-all'
         ).roformer.to(device).eval()
-        tokenizer = SMILES_SPE_Tokenizer(f'{base_path}/tr2d2-pep/tokenizer/new_vocab.txt',
-                                        f'{base_path}/tr2d2-pep/tokenizer/new_splits.txt')
+        tokenizer = SMILES_SPE_Tokenizer(f'{base_path}/tokenizer/new_vocab.txt',
+                                        f'{base_path}/tokenizer/new_splits.txt')
         prot_seqs = prot_seqs if prot_seqs is not None else []
         
         if score_func_names is None:

td3b/direction_oracle.py

@@ -114,14 +114,28 @@ class TR2D2RoFormerEncoder(nn.Module):
         if checkpoint_path:
             print(f"  Loading TR2-D2 checkpoint...")
             ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False)
-            state_dict = ckpt.get("state_dict", ckpt)
+            # Unwrap common checkpoint containers (raw / "state_dict" / "model_state_dict").
+            state_dict = ckpt
+            for _key in ("state_dict", "model_state_dict"):
+                if isinstance(state_dict, dict) and _key in state_dict:
+                    state_dict = state_dict[_key]
+                    break
             roformer_state = {
                 k.replace("model.", "").replace("backbone.", ""): v
                 for k, v in state_dict.items()
                 if "roformer" in k or "encoder" in k or "backbone" in k
             }
-            self.encoder.model.load_state_dict(roformer_state, strict=False)
-            print("  TR2-D2 checkpoint loaded")
+            if not roformer_state:
+                raise RuntimeError(
+                    f"No RoFormer/encoder keys matched in TR2-D2 checkpoint {checkpoint_path} "
+                    f"(top-level keys: {list(state_dict.keys())[:8]}). The ligand encoder would be "
+                    f"left randomly initialized, producing meaningless direction scores."
+                )
+            incompatible = self.encoder.model.load_state_dict(roformer_state, strict=False)
+            if incompatible.missing_keys:
+                print(f"  [warn] TR2-D2 encoder: {len(incompatible.missing_keys)} keys not in "
+                      f"checkpoint (e.g. {incompatible.missing_keys[:3]})")
+            print(f"  TR2-D2 checkpoint loaded ({len(roformer_state)} tensors)")
 
         for p in self.encoder.parameters():
             p.requires_grad = False
@@ -289,6 +303,14 @@ def _load_state_dict_flexible(model: nn.Module, state_dict: Dict, strict: bool =
             len(state_dict)
         )
         incompatible = model.load_state_dict(filtered, strict=False)
+        # ESM keys are expected to be absent (ESM2 is re-loaded pretrained in __init__); any
+        # missing NON-esm key is a trained head left at random init — surface it loudly.
+        trained_missing = [k for k in incompatible.missing_keys if not k.startswith("esm")]
+        if trained_missing:
+            logger.warning(
+                "%d TRAINED (non-ESM) oracle keys did not load and remain randomly initialized: %s",
+                len(trained_missing), trained_missing[:10]
+            )
         if incompatible.missing_keys:
             logger.warning("Missing keys (first 10): %s", incompatible.missing_keys[:10])
         if incompatible.unexpected_keys:

td3b/td3b_finetune.py

@@ -7,11 +7,19 @@ import numpy as np
 import torch
 import wandb
 import os
-from training.finetune_utils import loss_wdce
+# NOTE: `loss_wdce` is imported lazily inside td3b_finetune() (see below) to avoid a
+# circular import: training.finetune_utils imports the td3b package at module load,
+# so importing it here at module level would deadlock the finetune_utils <-> td3b cycle.
 from .td3b_losses import TD3BTotalLoss, extract_embeddings_from_mdlm
 from tqdm import tqdm
 import pandas as pd
-from plotting import plot_data_with_distribution_seaborn, plot_data
+try:
+    from plotting import plot_data_with_distribution_seaborn, plot_data
+except ImportError:
+    # `plotting` is an optional dev-only helper used solely by the legacy
+    # single-target loop below. Guard the import so that `import td3b` (and the
+    # inference/training entry points that depend on it) work without it.
+    plot_data_with_distribution_seaborn = plot_data = None
 
 
 def td3b_finetune(
@@ -234,7 +242,8 @@ def td3b_finetune(
                 directional_labels = directional_labels_saved
                 confidences = confidences_saved
 
-        # Compute WDCE loss
+        # Compute WDCE loss (lazy import breaks the finetune_utils <-> td3b cycle)
+        from training.finetune_utils import loss_wdce
         wdce_loss = loss_wdce(
             policy_model,
             log_rnd,

td3b/td3b_mcts.py

@@ -26,7 +26,8 @@ class TD3B_MCTS(BaseMCTS):
         mask_index: int,
         buffer_size: int = 100,
         noise=None,
-        tokenizer=None
+        tokenizer=None,
+        enforce_validity=None,
     ):
         """
         Args:
@@ -38,6 +39,10 @@ class TD3B_MCTS(BaseMCTS):
             buffer_size: Maximum buffer size
             noise: Noise schedule
             tokenizer: Peptide tokenizer
+            enforce_validity: Validity gate toggle for MCTS expansion. None
+                (default) consults args.enforce_validity, falling back to True
+                (unchanged behavior). False disables the is_peptide gate so
+                invalid peptides are kept/scored on pure affinity x direction.
         """
         # Initialize base MCTS (will set self.rewardFunc later)
         # Note: base MCTS expects 'policy_model' not 'diffusion_model'
@@ -56,7 +61,8 @@ class TD3B_MCTS(BaseMCTS):
             config=config,
             policy_model=diffusion_model,
             pretrained=diffusion_model,  # Use same model
-            score_func_names=['affinity', 'gated_reward', 'placeholder1', 'placeholder2', 'placeholder3']  # 5 objectives
+            score_func_names=['affinity', 'gated_reward', 'placeholder1', 'placeholder2', 'placeholder3'],  # 5 objectives
+            enforce_validity=enforce_validity,
         )
 
         # Set TD3B-specific attributes
@@ -284,7 +290,10 @@ def create_td3b_mcts(
         diffusion_model: MDLM model
         td3b_reward_function: TD3BRewardFunction instance
         alpha: Temperature for importance weighting
-        **kwargs: Additional MCTS arguments
+        **kwargs: Additional MCTS arguments (forwarded verbatim to TD3B_MCTS).
+            This includes ``enforce_validity`` (validity gate toggle); when not
+            passed here it defaults to None inside TD3B_MCTS and is resolved from
+            ``args.enforce_validity`` (i.e. the --validity_reward flag).
 
     Returns:
         mcts: TD3B_MCTS instance

training/finetune_utils.py

@@ -192,8 +192,8 @@ def load_tokenizer(base_path: str) -> SMILES_SPE_Tokenizer:
         >>> tokenizer = load_tokenizer('To Be Added')
     """
     base_path = Path(base_path)
-    vocab_path = base_path / "tr2d2-pep" / "tokenizer" / "new_vocab.txt"
-    spe_path = base_path / "tr2d2-pep" / "tokenizer" / "new_splits.txt"
+    vocab_path = base_path / "tokenizer" / "new_vocab.txt"
+    spe_path = base_path / "tokenizer" / "new_splits.txt"
 
     if not vocab_path.exists():
         raise FileNotFoundError(f"Vocabulary file not found: {vocab_path}")
@@ -340,7 +340,7 @@ def create_output_directory(base_path: str, run_name: str, add_timestamp: bool =
     else:
         dir_name = run_name
 
-    output_dir = os.path.join(base_path, "tr2d2-pep", "results", dir_name)
+    output_dir = os.path.join(base_path, "results", dir_name)
     os.makedirs(output_dir, exist_ok=True)
 
     logger.info("Created output directory: %s", output_dir)