#!/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()