| import argparse |
| import os |
| import sys |
| import time |
| from dataclasses import dataclass |
| from typing import Dict, List, Optional |
|
|
| import numpy as np |
| ROOT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) |
| if ROOT_DIR not in sys.path: |
| sys.path.insert(0, ROOT_DIR) |
|
|
| import torch |
| from hydra import compose, initialize_config_dir |
| from hydra.core.global_hydra import GlobalHydra |
|
|
| from diffusion import Diffusion |
| from scoring.scoring_functions import ScoringFunctions |
| from scoring.functions.binding import MultiTargetBindingAffinity |
| from td3b.direction_oracle import DirectionalOracle, resolve_device |
| from td3b.data_utils import peptide_seq_to_smiles, smiles_token_length |
|
|
| from baselines.baselines import ( |
| RewardInputs, |
| RewardWrapper, |
| classifier_guidance, |
| peptune_mctg_sampling, |
| sequential_monte_carlo, |
| twisted_diffusion_sampler, |
| unguided_sampling, |
| ) |
|
|
|
|
| AMINO_ACIDS = "ACDEFGHIKLMNPQRSTVWY" |
|
|
|
|
| @dataclass |
| class ProteinTokenizer: |
| aa_to_id: Dict[str, int] |
| pad_id: int = 0 |
|
|
| @classmethod |
| def default(cls) -> "ProteinTokenizer": |
| aa_to_id = {aa: idx + 1 for idx, aa in enumerate(AMINO_ACIDS)} |
| return cls(aa_to_id=aa_to_id, pad_id=0) |
|
|
| def encode(self, seq: str) -> torch.Tensor: |
| ids = [self.aa_to_id.get(aa, self.pad_id) for aa in seq] |
| return torch.tensor([ids], dtype=torch.long) |
|
|
|
|
| def load_base_model( |
| ckpt_path: str, |
| device: str, |
| config_name: str = "peptune_config.yaml", |
| ) -> Diffusion: |
| GlobalHydra.instance().clear() |
| config_dir = os.path.join(os.path.dirname(__file__), "..", "configs") |
| initialize_config_dir(config_dir=config_dir, job_name="load_model") |
| cfg = compose(config_name=config_name) |
| try: |
| model = Diffusion.load_from_checkpoint( |
| ckpt_path, |
| config=cfg, |
| mode="eval", |
| device=device, |
| map_location=device, |
| ) |
| model.eval() |
| return model |
| except Exception as exc: |
| print(f"[load_base_model] Lightning load failed, falling back to raw state_dict: {exc}") |
|
|
| checkpoint = torch.load(ckpt_path, map_location=device, weights_only=False) |
| if isinstance(checkpoint, dict): |
| if "model_state_dict" in checkpoint: |
| state_dict = checkpoint["model_state_dict"] |
| elif "state_dict" in checkpoint: |
| state_dict = checkpoint["state_dict"] |
| else: |
| state_dict = checkpoint |
| else: |
| raise ValueError(f"Unsupported checkpoint format: {type(checkpoint)}") |
|
|
| model = Diffusion( |
| config=cfg, |
| mode="eval", |
| device=device, |
| ) |
| missing, unexpected = model.load_state_dict(state_dict, strict=False) |
| if missing: |
| print(f"[load_base_model] Missing keys: {len(missing)}") |
| if unexpected: |
| print(f"[load_base_model] Unexpected keys: {len(unexpected)}") |
| model.eval() |
| model.to(device) |
| return model |
|
|
|
|
| def load_reward_models( |
| prot_seq: Optional[str], |
| device: str, |
| base_model: Optional[Diffusion] = None, |
| base_path: Optional[str] = None, |
| multi_target: bool = False, |
| score_func_names: Optional[List[str]] = None, |
| ): |
| if multi_target: |
| if base_model is None or base_path is None: |
| raise ValueError("base_model and base_path are required for multi-target affinity.") |
| return MultiTargetBindingAffinity( |
| tokenizer=base_model.tokenizer, |
| base_path=base_path, |
| device=device, |
| emb_model=base_model.backbone, |
| ) |
| if score_func_names is None: |
| score_func_names = [ |
| "binding_affinity1", |
| "solubility", |
| "hemolysis", |
| "nonfouling", |
| "permeability", |
| ] |
| if prot_seq is None: |
| raise ValueError("prot_seq is required for single-target scoring.") |
| return ScoringFunctions(score_func_names, prot_seqs=[prot_seq], device=device) |
|
|
|
|
| def load_direction_oracle(args, device: str) -> DirectionalOracle: |
| 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, |
| ) |
| oracle.eval() |
| return oracle |
|
|
|
|
| def run_baseline( |
| baseline: str, |
| base_model: Diffusion, |
| reward_fn: RewardWrapper, |
| batch_size: int, |
| seq_length: int, |
| num_steps: int, |
| guidance_scale: float, |
| alpha: float, |
| guidance_steps: Optional[int], |
| mcts_iterations: int, |
| num_children: int, |
| sample_prob_weight: float, |
| invalid_penalty: float, |
| pareto_max_size: Optional[int], |
| ) -> Dict[str, torch.Tensor]: |
| baseline = baseline.lower() |
| if baseline == "cg": |
| return classifier_guidance( |
| base_model, |
| reward_fn, |
| batch_size=batch_size, |
| seq_length=seq_length, |
| num_steps=num_steps, |
| guidance_scale=guidance_scale, |
| guidance_steps=guidance_steps, |
| ) |
| if baseline == "unguided": |
| return unguided_sampling( |
| base_model, |
| batch_size=batch_size, |
| seq_length=seq_length, |
| num_steps=num_steps, |
| ) |
| if baseline == "smc": |
| return sequential_monte_carlo( |
| base_model, |
| reward_fn, |
| batch_size=batch_size, |
| seq_length=seq_length, |
| num_steps=num_steps, |
| alpha=alpha, |
| ) |
| if baseline == "tds": |
| return twisted_diffusion_sampler( |
| base_model, |
| reward_fn, |
| batch_size=batch_size, |
| seq_length=seq_length, |
| num_steps=num_steps, |
| guidance_scale=guidance_scale, |
| alpha=alpha, |
| guidance_steps=guidance_steps, |
| ) |
| if baseline == "peptune": |
| return peptune_mctg_sampling( |
| base_model, |
| reward_fn, |
| batch_size=batch_size, |
| seq_length=seq_length, |
| num_steps=num_steps, |
| mcts_iterations=mcts_iterations, |
| num_children=num_children, |
| alpha=alpha, |
| sample_prob_weight=sample_prob_weight, |
| invalid_penalty=invalid_penalty, |
| pareto_max_size=pareto_max_size, |
| ) |
| raise ValueError(f"Unknown baseline: {baseline}") |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--ckpt_path", type=str, required=True) |
| parser.add_argument("--device", type=str, default="cuda:0") |
| parser.add_argument("--baseline", type=str, default="cg", choices=["cg", "smc", "tds", "unguided", "peptune"]) |
| parser.add_argument("--prot_seq", type=str, default=None) |
| parser.add_argument("--targets_csv", type=str, default=None) |
| parser.add_argument("--d_star", type=float, default=1.0) |
| parser.add_argument("--batch_size", type=int, default=32) |
| parser.add_argument("--seq_length", type=int, default=200) |
| parser.add_argument("--binder_seq", type=str, default=None) |
| parser.add_argument("--num_steps", type=int, default=128) |
| parser.add_argument("--guidance_scale", type=float, default=1.0) |
| parser.add_argument("--alpha", type=float, default=0.1) |
| parser.add_argument("--reward_alpha", type=float, default=None) |
| parser.add_argument("--mcts_iterations", type=int, default=20) |
| parser.add_argument("--num_children", type=int, default=24) |
| parser.add_argument("--sample_prob_weight", type=float, default=0.1) |
| parser.add_argument("--invalid_penalty", type=float, default=1.0) |
| parser.add_argument("--pareto_max_size", type=int, default=None) |
| parser.add_argument("--guidance_steps", type=int, default=None) |
| parser.add_argument("--fast_direction", action="store_true", default=False) |
| parser.add_argument("--num_batches", type=int, default=1) |
| parser.add_argument("--output_dir", type=str, default=None) |
| parser.add_argument("--shard_id", type=int, default=None) |
| parser.add_argument("--num_shards", type=int, default=None) |
| parser.add_argument("--direction_oracle_ckpt", type=str, default=None) |
| parser.add_argument("--direction_oracle_tr2d2_checkpoint", type=str, default=None) |
| parser.add_argument("--direction_oracle_tokenizer_vocab", type=str, default=None) |
| parser.add_argument("--direction_oracle_tokenizer_splits", type=str, default=None) |
| parser.add_argument("--direction_oracle_esm_name", type=str, default="facebook/esm2_t33_650M_UR50D") |
| parser.add_argument("--direction_oracle_esm_cache_dir", type=str, default=None) |
| parser.add_argument("--direction_oracle_esm_local_files_only", action="store_true", default=False) |
| parser.add_argument("--direction_oracle_max_ligand_length", type=int, default=768) |
| parser.add_argument("--direction_oracle_max_protein_length", type=int, default=1024) |
| parser.add_argument("--direction_oracle_d_model", type=int, default=256) |
| parser.add_argument("--direction_oracle_n_heads", type=int, default=4) |
| parser.add_argument("--direction_oracle_n_self_attn_layers", type=int, default=1) |
| parser.add_argument("--direction_oracle_n_bmca_layers", type=int, default=2) |
| parser.add_argument("--direction_oracle_dropout", type=float, default=0.3) |
| args = parser.parse_args() |
|
|
| rank_env = os.environ.get("LOCAL_RANK") |
| world_env = os.environ.get("WORLD_SIZE") |
| if rank_env is not None or world_env is not None: |
| rank = int(rank_env or 0) |
| world_size = int(world_env or 1) |
| else: |
| rank = int(args.shard_id) if args.shard_id is not None else 0 |
| world_size = int(args.num_shards) if args.num_shards is not None else 1 |
| if world_size < 1: |
| world_size = 1 |
| if world_size > 1 and str(args.device).lower() in {"cuda", "cuda:0", "auto"}: |
| args.device = f"cuda:{rank}" |
|
|
| resolved_device = resolve_device(args.device) |
| args.device = str(resolved_device) |
|
|
| tr2d2_root = ROOT_DIR |
| if args.direction_oracle_ckpt is None: |
| args.direction_oracle_ckpt = os.path.join( |
| tr2d2_root, "direction_oracle.pt" |
| ) |
| if args.direction_oracle_tr2d2_checkpoint is None: |
| args.direction_oracle_tr2d2_checkpoint = os.path.join( |
| tr2d2_root, "pretrained", "peptune-pretrained.ckpt" |
| ) |
| if args.direction_oracle_tokenizer_vocab is None: |
| args.direction_oracle_tokenizer_vocab = os.path.join( |
| tr2d2_root, "tokenizer", "new_vocab.txt" |
| ) |
| if args.direction_oracle_tokenizer_splits is None: |
| args.direction_oracle_tokenizer_splits = os.path.join( |
| tr2d2_root, "tokenizer", "new_splits.txt" |
| ) |
|
|
| if args.targets_csv is None and args.prot_seq is None: |
| 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, "..")) |
| multi_target = args.targets_csv is not None |
| scoring_fn = load_reward_models( |
| args.prot_seq if not multi_target else None, |
| args.device, |
| base_model=base_model, |
| base_path=base_path, |
| multi_target=multi_target, |
| ) |
| direction_oracle = load_direction_oracle(args, args.device) |
| reward_alpha = args.reward_alpha if args.reward_alpha is not None else args.alpha |
|
|
| if args.targets_csv: |
| import pandas as pd |
|
|
| df = pd.read_csv(args.targets_csv) |
| if "Target_Sequence" not in df.columns: |
| raise ValueError("targets_csv must contain a 'Target_Sequence' column.") |
| if "Ligand_Sequence" not in df.columns: |
| raise ValueError("targets_csv must contain a 'Ligand_Sequence' column.") |
|
|
| targets = [] |
| for row_idx, row in df.iterrows(): |
| target_seq = str(row["Target_Sequence"]) if pd.notna(row["Target_Sequence"]) else None |
| if not target_seq: |
| continue |
| binder_seq = row["Ligand_Sequence"] |
| if pd.isna(binder_seq): |
| binder_seq = None |
| else: |
| binder_seq = str(binder_seq) |
| if binder_seq.strip() == "": |
| binder_seq = None |
| targets.append( |
| { |
| "target_seq": target_seq, |
| "binder_seq": binder_seq, |
| "row_index": int(row_idx), |
| } |
| ) |
| else: |
| targets = [{"target_seq": args.prot_seq, "binder_seq": args.binder_seq, "row_index": 0}] |
|
|
| if world_size > 1: |
| targets = [item for idx, item in enumerate(targets) if idx % world_size == rank] |
| print(f"[shard] rank {rank}/{world_size}: {len(targets)} targets") |
|
|
| output_dir = args.output_dir |
| if output_dir is None: |
| output_dir = os.path.join(os.path.dirname(__file__), "outputs") |
| os.makedirs(output_dir, exist_ok=True) |
|
|
| from utils.app import PeptideAnalyzer |
|
|
| analyzer = PeptideAnalyzer() |
| all_rows = [] |
| batch_rows = [] |
| metrics_rows = [] |
| def resolve_seq_length(binder_seq: Optional[str]) -> int: |
| if not binder_seq: |
| return args.seq_length |
| try: |
| smiles = peptide_seq_to_smiles(binder_seq) |
| if not smiles: |
| return args.seq_length |
| if base_model.tokenizer is None: |
| return len(smiles) |
| return smiles_token_length(smiles, base_model.tokenizer) |
| except Exception as exc: |
| print(f"Warning: failed to derive seq_length from binder_seq; using {args.seq_length}. Error: {exc}") |
| return args.seq_length |
|
|
| for target_idx, target_info in enumerate(targets): |
| target_seq = target_info["target_seq"] |
| binder_seq = target_info.get("binder_seq") |
| row_index = target_info.get("row_index", target_idx) |
| seq_length = resolve_seq_length(binder_seq) |
| protein_tokens = direction_oracle.encode_protein(target_seq) |
| for direction_name, d_star in [("agonist", 1.0), ("antagonist", -1.0)]: |
|
|
| reward_inputs = RewardInputs( |
| protein_tokens=protein_tokens, |
| d_star=d_star, |
| protein_seq=target_seq, |
| ) |
| reward_fn = RewardWrapper( |
| scoring_fn=scoring_fn, |
| direction_oracle=direction_oracle, |
| base_model=base_model, |
| tokenizer=base_model.tokenizer, |
| reward_inputs=reward_inputs, |
| device=torch.device(args.device), |
| fast_direction=args.fast_direction, |
| reward_alpha=reward_alpha, |
| ) |
|
|
| num_batches = 1 if multi_target else args.num_batches |
| for batch_idx in range(num_batches): |
| start = time.perf_counter() |
| result = run_baseline( |
| args.baseline, |
| base_model, |
| reward_fn, |
| batch_size=args.batch_size, |
| seq_length=seq_length, |
| num_steps=args.num_steps, |
| guidance_scale=args.guidance_scale, |
| alpha=args.alpha, |
| guidance_steps=args.guidance_steps, |
| mcts_iterations=args.mcts_iterations, |
| num_children=args.num_children, |
| sample_prob_weight=args.sample_prob_weight, |
| invalid_penalty=args.invalid_penalty, |
| pareto_max_size=args.pareto_max_size, |
| ) |
| elapsed = time.perf_counter() - start |
|
|
| scores = reward_fn.evaluate_tokens( |
| result["tokens"], |
| torch.ones_like(result["tokens"], device=result["tokens"].device), |
| ) |
| sequences = scores["sequences"] |
| affinity = scores["affinity"].detach().cpu().numpy() |
| direction = scores["direction"].detach().cpu().numpy() |
| gated_reward = scores["gated_reward"].detach().cpu().numpy() |
| valid_mask = np.array([analyzer.is_peptide(seq) for seq in sequences], dtype=np.float32) |
| valid_fraction = float(valid_mask.mean()) if len(valid_mask) else 0.0 |
| consistency = d_star * (direction - 0.5) |
| if d_star > 0: |
| direction_correct = (direction >= 0.5).astype(np.float32) |
| else: |
| direction_correct = (direction < 0.5).astype(np.float32) |
| success = direction_correct * valid_mask |
| direction_mean = float(np.mean(direction)) |
| direction_std = float(np.std(direction)) |
| affinity_mean = float(np.mean(affinity)) |
| affinity_std = float(np.std(affinity)) |
| consistency_mean = float(np.mean(consistency)) |
| consistency_std = float(np.std(consistency)) |
| gated_reward_mean = float(np.mean(gated_reward)) |
| gated_reward_std = float(np.std(gated_reward)) |
| direction_acc_mean = float(np.mean(direction_correct)) |
| direction_acc_std = float(np.std(direction_correct)) |
| success_rate_mean = float(np.mean(success)) |
| success_rate_std = float(np.std(success)) |
| batch_metrics = { |
| "direction_mean": direction_mean, |
| "direction_std": direction_std, |
| "affinity_mean": affinity_mean, |
| "affinity_std": affinity_std, |
| "consistency_mean": consistency_mean, |
| "consistency_std": consistency_std, |
| "gated_reward_mean": gated_reward_mean, |
| "gated_reward_std": gated_reward_std, |
| "direction_accuracy_mean": direction_acc_mean, |
| "direction_accuracy_std": direction_acc_std, |
| "valid_fraction": valid_fraction, |
| "success_rate_mean": success_rate_mean, |
| "success_rate_std": success_rate_std, |
| } |
|
|
| for i, seq in enumerate(sequences): |
| all_rows.append( |
| { |
| "rank": rank, |
| "sequence": seq, |
| "affinity": float(affinity[i]), |
| "direction": float(direction[i]), |
| "d_star": float(d_star), |
| "direction_name": direction_name, |
| "target_seq": target_seq, |
| "target_index": target_idx, |
| "row_index": row_index, |
| "binder_seq": binder_seq, |
| "seq_length": seq_length, |
| "gated_reward": float(gated_reward[i]), |
| "consistency_reward": float(consistency[i]), |
| "direction_accuracy": float(direction_correct[i]), |
| "valid": float(valid_mask[i]), |
| "success": float(success[i]), |
| "batch_index": batch_idx, |
| "batch_time_sec": elapsed, |
| **batch_metrics, |
| } |
| ) |
| batch_rows.append( |
| { |
| "rank": rank, |
| "batch_index": batch_idx, |
| "batch_time_sec": elapsed, |
| "target_index": target_idx, |
| "row_index": row_index, |
| "binder_seq": binder_seq, |
| "seq_length": seq_length, |
| "direction_name": direction_name, |
| } |
| ) |
| metrics_rows.append( |
| { |
| "rank": rank, |
| "target_index": target_idx, |
| "target_seq": target_seq, |
| "row_index": row_index, |
| "binder_seq": binder_seq, |
| "seq_length": seq_length, |
| "direction_name": direction_name, |
| "d_star": float(d_star), |
| "batch_index": batch_idx, |
| "num_samples": len(sequences), |
| **batch_metrics, |
| } |
| ) |
| print( |
| f"Target {target_idx} dir {direction_name}: " |
| f"generated {len(sequences)} sequences in {elapsed:.3f}s" |
| ) |
|
|
| import pandas as pd |
|
|
| if world_size > 1: |
| output_csv = os.path.join(output_dir, f"{args.baseline}_samples_rank{rank}.csv") |
| batch_csv = os.path.join(output_dir, f"batch_times_rank{rank}.csv") |
| metrics_csv = os.path.join(output_dir, f"{args.baseline}_metrics_rank{rank}.csv") |
| else: |
| output_csv = os.path.join(output_dir, f"{args.baseline}_samples.csv") |
| batch_csv = os.path.join(output_dir, "batch_times.csv") |
| metrics_csv = os.path.join(output_dir, f"{args.baseline}_metrics.csv") |
| pd.DataFrame(all_rows).to_csv(output_csv, index=False) |
| pd.DataFrame(batch_rows).to_csv(batch_csv, index=False) |
| pd.DataFrame(metrics_rows).to_csv(metrics_csv, index=False) |
|
|
| print(f"Saved samples to {output_csv}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
