#!/usr/bin/env python3 """ TD3B Inference Script Generate directional binders for target proteins using a finetuned TD3B model. Usage: python inference.py \ --ckpt_path checkpoints/td3b.ckpt \ --val_csv data/test.csv \ --save_path results/ \ --seed 42 """ import argparse import os import sys import logging from typing import Dict, List, Tuple import numpy as np import pandas as pd 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 models.diffusion import Diffusion from configs.finetune_config import ( DiffusionConfig, RoFormerConfig, NoiseConfig, TrainingConfig, SamplingConfig, EvalConfig, OptimConfig, MCTSConfig, ) 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") # ─── Defaults ───────────────────────────────────────────────────────────────── DEFAULTS = dict( seq_length=200, sampling_eps=1e-3, total_num_steps=128, hidden_dim=768, num_layers=8, num_heads=8, alpha=0.1, min_affinity_threshold=0.0, sigmoid_temperature=0.1, num_pool=32, val_samples_per_target=8, ) def load_model(ckpt_path: str, device: torch.device): """Load finetuned TD3B model from checkpoint.""" ckpt = torch.load(ckpt_path, map_location=device, weights_only=False) state_dict = ckpt.get("model_state_dict") or ckpt.get("state_dict") or ckpt config = ckpt.get("config") or {} tokenizer = load_tokenizer(ROOT_DIR) cfg = DiffusionConfig( roformer=RoFormerConfig( hidden_size=config.get("hidden_dim", 768), n_layers=config.get("num_layers", 8), n_heads=config.get("num_heads", 8), ), 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=device).to(device) 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 def sample_sequences(model, batch_size: int, seq_length: int, num_steps: int, eps: float = 1e-5): """Sample sequences from the diffusion model.""" x = model.sample_prior(batch_size, seq_length).to(model.device, dtype=torch.long) timesteps = torch.linspace(1, eps, num_steps + 1, device=model.device) dt = torch.tensor((1 - eps) / num_steps, device=model.device) for i in range(num_steps): t = timesteps[i] * torch.ones(x.shape[0], 1, device=model.device) _, x = model.single_reverse_step(x, t=t, dt=dt) x = x.to(model.device) # Remove remaining masks mask_pos = (x == model.mask_index) if mask_pos.any(): t = timesteps[-2] * torch.ones(x.shape[0], 1, device=model.device) _, x = model.single_noise_removal(x, t=t, dt=dt) x = x.to(model.device) return x def score_sequences(reward_model, sequences: List[str]): """Score sequences with the TD3B reward function.""" result = reward_model(sequences) if isinstance(result, tuple): rewards, info = result return ( np.asarray(rewards), np.asarray(info.get("affinities", rewards)), np.asarray(info.get("directions", np.zeros_like(rewards))), np.asarray(info.get("confidences", np.ones_like(rewards))), ) rewards = np.asarray(result) return rewards, rewards, np.zeros_like(rewards), np.ones_like(rewards) def main(): parser = argparse.ArgumentParser(description="TD3B Inference") parser.add_argument("--ckpt_path", type=str, required=True, help="Path to TD3B checkpoint") parser.add_argument("--val_csv", type=str, required=True, help="CSV with Target_Sequence, Ligand_Sequence, label columns") parser.add_argument("--save_path", type=str, default="results", help="Output directory") parser.add_argument("--device", type=str, default="cuda:0") parser.add_argument("--seed", type=int, default=42) parser.add_argument("--num_pool", type=int, default=32, help="Pool size for candidate generation") parser.add_argument("--val_samples_per_target", type=int, default=8, help="Samples to keep per target-direction") 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 device = torch.device(args.device if torch.cuda.is_available() else "cpu") torch.manual_seed(args.seed) np.random.seed(args.seed) os.makedirs(args.save_path, exist_ok=True) analyzer = PeptideAnalyzer() # Load model logger.info(f"Loading model from {args.ckpt_path}") model, tokenizer = load_model(args.ckpt_path, device) # Load targets logger.info(f"Loading targets from {args.val_csv}") 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": max(1, min(int(seq_len) if seq_len else DEFAULTS["seq_length"], args.max_seq_length)), }) # 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 (reuses the preloaded oracle; only re-encodes the protein) try: reward_model = create_td3b_reward_function( affinity_predictor=target_affinity, directional_oracle=directional_oracle, target_protein_seq=target["target_seq"], target_direction=d_name, peptide_tokenizer=tokenizer, device=device, 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}") continue # Generate pool of candidates target_length = target.get("seq_length", 200) 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). 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) 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.array([], dtype=int) # no valid finite-reward candidate → save nothing # Save resampled samples (already gated to valid + finite above) for i in chosen: records.append({ "target": target["target_seq"][:20], "target_uid": target["target_uid"], "sequence": sequences[i], "target_direction": d_star, "direction_name": d_name, "is_valid": True, "affinity": float(affinities[i]), "gated_reward": float(gated_rewards[i]), "direction_oracle": float(directions[i]), "direction_accuracy": float(direction_accuracy[i]), }) # Save results out_df = pd.DataFrame(records) out_path = os.path.join(args.save_path, f"td3b_results_seed{args.seed}.csv") out_df.to_csv(out_path, index=False) # Print summary if len(out_df) > 0: dp = out_df[out_df["target_direction"] == 1.0] dm = out_df[out_df["target_direction"] == -1.0] logger.info(f"\n{'='*60}") logger.info(f"Results saved to {out_path} ({len(out_df)} valid samples)") logger.info(f" Aff(d*=+1) = {dp['affinity'].mean():.2f}" if len(dp) else " No agonist samples") logger.info(f" Aff(d*=-1) = {dm['affinity'].mean():.2f}" if len(dm) else " No antagonist samples") logger.info(f" DA(d*=+1) = {dp['direction_accuracy'].mean():.3f}" if len(dp) else "") logger.info(f" DA(d*=-1) = {dm['direction_accuracy'].mean():.3f}" if len(dm) else "") logger.info(f" Gated Reward = {out_df['gated_reward'].mean():.2f}") 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__": main()