import argparse from typing import List, Callable, Optional, Tuple, Dict, Any import torch import torch.nn.functional as F import yaml from easydict import EasyDict as edict from tqdm import tqdm import time import math from pathlib import Path import os import sys REPO_ROOT = Path(__file__).resolve().parents[1] if str(REPO_ROOT) not in sys.path: sys.path.insert(0, str(REPO_ROOT)) from generate import build_model_and_stuff, tokenize_input_str, detokenize_output from objectives import * from constraints import * from smiles_tokenizer.my_tokenizers import SMILES_SPE_Tokenizer from paths import smiles_vocab_files import pdb import warnings warnings.filterwarnings("ignore", category=FutureWarning) from transformers.utils import logging logging.set_verbosity_error() # --------------------------------------------------------------------------- # small utilities # --------------------------------------------------------------------------- def mix_score(ratio, base_scores, admet_scores): return [r * b + (1.0 - r) * a for r, b, a in zip(ratio, base_scores, admet_scores)] def extract_objective_vector(smiles_tokens, smiles_seqs, objective_models, device): sm_scores = None sm_ratio = None values = [] for obj in objective_models: name, scores = obj(smiles_tokens, smiles_seqs) # list of length B if name == "small_molecule": sm_scores = scores sm_ratio = sm_scores["ratio"] continue if name == "motif" and isinstance(scores, tuple): values.append(torch.tensor(scores[0], device=device, dtype=torch.float32)) values.append(torch.tensor(scores[1], device=device, dtype=torch.float32)) continue if name in ["non_toxicity", "solubility", "permeability", "halflife", "affinity"] and sm_scores is not None: scores = mix_score(sm_ratio, scores, sm_scores[name]) values.append(torch.tensor(scores, device=device, dtype=torch.float32)) return torch.stack(values, dim=1) # (B,m) def compute_scores_print(smiles_tokens, smiles_seqs, objective_models, constraint_models, device, return_scores=False): objective_scores = extract_objective_vector(smiles_tokens, smiles_seqs, objective_models, device) # (B,m) constraint_scores = [] for constraint in constraint_models: constraint_score = constraint(smiles_tokens, smiles_seqs) # list of length B constraint_scores.append(torch.tensor(constraint_score, device=device, dtype=torch.float32)) constraint_scores = torch.stack(constraint_scores, dim=1) # (B, n) # pdb.set_trace() scores = torch.concat([objective_scores, constraint_scores], dim=1) # (B, m+n) print(scores) ratio = analyze_peptide_likeness(smiles_seqs[0])['peptide_likeness'] print("Peptide Likeliness: ", ratio) if return_scores: return scores # --------------------------------------------------------------------------- # edit utilities # --------------------------------------------------------------------------- @torch.no_grad() def _sample_single_edit_batch( x: torch.Tensor, # (B, Lmax) padded lam_ins: torch.Tensor, # (B, Lmax) logits_ins: torch.Tensor, # (B, Lmax, V) lam_del: torch.Tensor, # (B, Lmax) lam_sub: torch.Tensor, # (B, Lmax) logits_sub: torch.Tensor, # (B, Lmax, V) pad_id: int, bos_id: int, eos_id: int, allowed_tokens: Optional[torch.Tensor] = None, # 1D LongTensor of vocab ids delta: float = 1.0, max_len_cap: Optional[int] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: """ Single-edit proposal (batched): - Mask invalid positions/ops (same as multi-edit version) - Choose ONE position i (per sequence) proportional to p_fire_i = 1 - exp(-delta * λ_i) (fallback to λ_i if all p_fire are ~0 but λ has mass) - At that position: choose op ~ proportional to (λ_ins, λ_del, λ_sub) - If ins/sub: sample token from softmax(logits_{ins/sub}[i]) (with allowed_tokens mask) - Apply exactly that one edit. Returns: x_out: (B, Lout) padded base_rate: (B,) relative proposal weight (consistent with your multi-edit log_ratio per fired event) base_rate[b] = expm1(delta * λ_i) * P(op|i) * P(tok|op,i) (tok prob = 1 for delete) """ assert x.dim() == 2, f"x must be (B,Lmax), got {tuple(x.shape)}" device = x.device B, Lmax = x.shape V = logits_ins.shape[-1] eps = 1e-30 if allowed_tokens is not None: if not torch.is_tensor(allowed_tokens): allowed_tokens = torch.tensor(allowed_tokens, device=device, dtype=torch.long) else: allowed_tokens = allowed_tokens.to(device=device, dtype=torch.long) # masks nonpad = (x != pad_id) lengths = nonpad.sum(dim=1) # (B,) is_bos = (x == bos_id) is_eos = (x == eos_id) # mask rates on invalid positions # lam_ins[:] = 0.0 ins_rate = lam_ins.clone() ins_rate = ins_rate.masked_fill(~nonpad, 0.0) ins_rate = ins_rate.masked_fill(is_eos, 0.0) # no insertion at eos del_rate = lam_del.clone() del_rate = del_rate.masked_fill(~nonpad, 0.0) del_rate = del_rate.masked_fill(is_bos | is_eos, 0.0) # no delete bos/eos sub_rate = lam_sub.clone() sub_rate = sub_rate.masked_fill(~nonpad, 0.0) sub_rate = sub_rate.masked_fill(is_bos | is_eos, 0.0) # no sub bos/eos # if at cap, disallow insertions if max_len_cap is not None: at_cap = lengths >= max_len_cap if at_cap.any(): ins_rate = ins_rate.masked_fill(at_cap.unsqueeze(1), 0.0) lam_pos = ins_rate + del_rate + sub_rate # (B, Lmax) # pdb.set_trace() # p_fire per position (masked) a = (delta * lam_pos).clamp_min(0.0) p_fire = (-torch.expm1(-a)).masked_fill(~nonpad, 0.0) # (B, Lmax) # choose one position per sequence pos_mass = p_fire.clone() sum_mass = pos_mass.sum(dim=1, keepdim=True) # (B,1) # fallback if all p_fire are ~0 but lam_pos has mass (common when delta small) fallback = (sum_mass.squeeze(1) <= 1e-12) & (lam_pos.sum(dim=1) > 1e-12) if fallback.any(): pos_mass[fallback] = lam_pos[fallback] sum_mass[fallback] = pos_mass[fallback].sum(dim=1, keepdim=True) # sequences with literally no valid edits no_edit = (sum_mass.squeeze(1) <= 1e-12) # make multinomial happy (won't be used because we'll treat no_edit separately) if no_edit.any(): pos_mass[no_edit, 0] = 1.0 sum_mass[no_edit] = pos_mass[no_edit].sum(dim=1, keepdim=True) pos_probs = pos_mass / sum_mass.clamp_min(1e-12) pos_idx = torch.multinomial(pos_probs, 1).squeeze(1) # (B,) # gather rates at chosen positions b_arange = torch.arange(B, device=device) ins_at = ins_rate[b_arange, pos_idx] # (B,) del_at = del_rate[b_arange, pos_idx] sub_at = sub_rate[b_arange, pos_idx] lam_at = lam_pos[b_arange, pos_idx].clamp_min(1e-12) # (B,) rates3_at = torch.stack([ins_at, del_at, sub_at], dim=1) # (B,3) op_probs_at = rates3_at / rates3_at.sum(dim=1, keepdim=True).clamp_min(1e-12) # sample op for edit rows only op_choice = torch.zeros((B,), device=device, dtype=torch.long) # 0=ins,1=del,2=sub edit_mask = ~no_edit if edit_mask.any(): op_choice[edit_mask] = torch.multinomial(op_probs_at[edit_mask], 1).squeeze(1) ins_mask = edit_mask & (op_choice == 0) del_mask = edit_mask & (op_choice == 1) sub_mask = edit_mask & (op_choice == 2) def _mask_logits_2d(logits_2d: torch.Tensor) -> torch.Tensor: # logits_2d: (K, V) if allowed_tokens is None: return logits_2d add = torch.full_like(logits_2d, -1e9) add[:, allowed_tokens] = 0.0 return logits_2d + add # sample tokens for ins/sub ins_tok = torch.full((B,), pad_id, device=device, dtype=torch.long) sub_tok = torch.full((B,), pad_id, device=device, dtype=torch.long) log_tok = torch.zeros((B,), device=device, dtype=torch.float32) if ins_mask.any(): idx = ins_mask.nonzero(as_tuple=True)[0] logits_sel = logits_ins[idx, pos_idx[idx], :] # (K,V) logits_sel = _mask_logits_2d(logits_sel) q = F.softmax(logits_sel, dim=-1) samp = torch.multinomial(q, 1).squeeze(1) ins_tok[idx] = samp logq = F.log_softmax(logits_sel, dim=-1) log_tok[idx] = logq.gather(1, samp.view(-1, 1)).squeeze(1) if sub_mask.any(): idx = sub_mask.nonzero(as_tuple=True)[0] logits_sel = logits_sub[idx, pos_idx[idx], :] # (K,V) logits_sel = _mask_logits_2d(logits_sel) q = F.softmax(logits_sel, dim=-1) samp = torch.multinomial(q, 1).squeeze(1) sub_tok[idx] = samp logq = F.log_softmax(logits_sel, dim=-1) log_tok[idx] = logq.gather(1, samp.view(-1, 1)).squeeze(1) # base_rate: expm1(delta * λ_i) * P(op|i) * P(tok|op,i) base_log = torch.zeros((B,), device=device, dtype=torch.float32) if edit_mask.any(): a_sel = a[b_arange, pos_idx].to(torch.float32) # (B,) log_expm1 = torch.log(torch.expm1(a_sel).clamp_min(eps)) # (B,) op_prob_sel = op_probs_at.gather(1, op_choice.view(-1, 1)).squeeze(1).clamp_min(eps) log_op = torch.log(op_prob_sel) # delete has token_prob = 1, so log_tok already 0 there base_log = log_expm1 + log_op + log_tok # for no_edit rows, set base_rate=1 (won't be used if cand==x) base_log = torch.where(no_edit, torch.zeros_like(base_log), base_log) base_rate = torch.exp(base_log).clamp_min(0.0) # (B,) # apply single edit new_seqs = [] new_lens = [] for b in range(B): seq = x[b] tokens = seq[seq != pad_id].tolist() if len(tokens) == 0 or no_edit[b].item(): out_tokens = tokens if len(tokens) > 0 else [eos_id] else: i = int(pos_idx[b].item()) i = min(i, len(tokens) - 1) if del_mask[b].item(): # delete token at i (already masked so not BOS/EOS) out_tokens = tokens[:i] + tokens[i+1:] elif sub_mask[b].item(): out_tokens = tokens[:] out_tokens[i] = int(sub_tok[b].item()) else: # insertion out_tokens = tokens[:] out_tokens.insert(i + 1, int(ins_tok[b].item())) if len(out_tokens) == 0 or out_tokens[-1] != eos_id: out_tokens.append(eos_id) if max_len_cap is not None and len(out_tokens) > max_len_cap: out_tokens = out_tokens[:max_len_cap] if out_tokens[-1] != eos_id: out_tokens[-1] = eos_id t = torch.tensor(out_tokens, device=device, dtype=torch.long) new_seqs.append(t) new_lens.append(t.numel()) Lout = max(1, max(new_lens)) x_out = torch.full((B, Lout), pad_id, device=device, dtype=x.dtype) for b, s in enumerate(new_seqs): x_out[b, :s.numel()] = s return x_out, base_rate # --------------------------------------------------------------------------- # ATC + G_T # --------------------------------------------------------------------------- def _augmented_tchebycheff( f_vals: torch.Tensor, w: torch.Tensor, rho: float, z: torch.Tensor, ) -> torch.Tensor: diff = f_vals - z term1 = torch.min(w * diff, dim=1).values term2 = rho * torch.sum(w * diff, dim=1) return term1 + term2 def _G_T( selfies_tokens: torch.Tensor, objective_models: List[Callable[[torch.Tensor], Tuple[str, Any]]], constraint_models: List[Callable[[torch.Tensor], torch.Tensor]], w: torch.Tensor, rho: float, z: torch.Tensor, beta: float, selfies_tokenizer, smiles_tokenizer, ws_for_invalid=False ): device = selfies_tokens.device smiles_seqs = selfies_tokenizer.batch_decode(selfies_tokens.tolist(), return_smiles=True) smiles_tokens = smiles_tokenizer(smiles_seqs, return_tensors='pt', padding=True)['input_ids'].to(device) constraint_results = [] for constraint in constraint_models: res = constraint(smiles_tokens, smiles_seqs) # list of length B constraint_results.append(res) constraint_results = torch.tensor(constraint_results, device=device) survived_seq_indices = (constraint_results == 1).all(dim=0).nonzero(as_tuple=True)[0] survived_seqs = [smiles_seqs[idx] for idx in survived_seq_indices.tolist()] # (B') survived_tokens = smiles_tokens[survived_seq_indices.tolist()] # pdb.set_trace() weighted_sum_full = torch.full((len(smiles_seqs),), float("-inf"), device=device) G_full = torch.full((len(smiles_seqs),), float("-inf"), device=device) # objectives if ws_for_invalid: f_vals = extract_objective_vector(smiles_tokens, smiles_seqs, objective_models, device) weighted_sum_full = torch.sum(w * f_vals, dim=1) u_atc = _augmented_tchebycheff(f_vals, w, rho, z) G = beta * u_atc G_full[survived_seq_indices] = G[survived_seq_indices] else: if survived_seq_indices.numel() > 0: f_vals = extract_objective_vector(survived_tokens, survived_seqs, objective_models, device) # (B', m) u_atc = _augmented_tchebycheff(f_vals, w, rho, z) # (B',) G = beta * u_atc # (B',) weighted_sum = torch.sum(w * f_vals, dim=1) # (B',) G_full[survived_seq_indices] = G weighted_sum_full[survived_seq_indices] = weighted_sum # return full-size tensors (B,) return G_full, weighted_sum_full # --------------------------------------------------------------------------- # rollout # --------------------------------------------------------------------------- @torch.no_grad() def short_rollout_batch( model, x0: torch.Tensor, # (B, Lmax) padded time_grid: torch.Tensor, start_idx: int, pad_id: int, bos_id: int, eos_id: int, allowed_tokens: Optional[torch.Tensor], max_len_cap: Optional[int], num_rollouts: int = 1, num_steps: int =32 ) -> torch.Tensor: """ Returns: xT: (B*num_rollouts, Lmax) Grouping: xT[i*num_rollouts:(i+1)*num_rollouts] corresponds to candidate i. """ device = x0.device B, Lmax = x0.shape # repeat each candidate num_rollouts times (grouped) x = x0.repeat_interleave(num_rollouts, dim=0) # (B*num_rollouts, Lmax) # rollout in batch for j in range(start_idx + 1, time_grid.numel()): t_j = time_grid[j].view(1).to(device) mask = (x != pad_id) lam_ins, logits_ins, lam_del, lam_sub, logits_sub, *_ = model(x_t=x, mask=mask, t=t_j) x, _ = _sample_single_edit_batch( x, lam_ins, logits_ins, lam_del, lam_sub, logits_sub, pad_id, bos_id, eos_id, allowed_tokens, delta=float(1/(num_steps-1)), max_len_cap=max_len_cap, ) return x # --------------------------------------------------------------------------- # finalizer # --------------------------------------------------------------------------- def _finalize_from_last( model, x_last: torch.Tensor, time_grid: torch.Tensor, last_step: int, pad_id: int, bos_id: int, eos_id: int, allowed_tokens: Optional[torch.Tensor], objective_models: List[Callable[[torch.Tensor], Tuple[str, Any]]], constraint_models: List[Callable[[torch.Tensor], torch.Tensor]], w: torch.Tensor, rho: float, ref_z: torch.Tensor, beta_final: float, max_len_cap: Optional[int] = None, num_final_rollouts: int = 50, num_steps: int = 32, selfies_tokenizer=None, smiles_tokenizer=None ) -> torch.Tensor: logG_last, _ = _G_T(x_last, objective_models, constraint_models, w, rho, ref_z, beta_final, selfies_tokenizer, smiles_tokenizer, ws_for_invalid=False) # start_idx = min(last_step, time_grid.numel() - 2) if time_grid.numel() >= 2 else 0 x_Ts = short_rollout_batch(model, x_last, time_grid, last_step, pad_id, bos_id, eos_id, allowed_tokens, max_len_cap, num_final_rollouts, num_steps) logG, _, = _G_T(x_Ts, objective_models, constraint_models, w, rho, ref_z, beta_final, selfies_tokenizer, smiles_tokenizer, ws_for_invalid=False) idx = torch.isfinite(logG).nonzero(as_tuple=True)[0].tolist() if len(idx) == 0 or torch.max(logG) < logG_last: return x_last, logG_last else: best_idx = torch.argmax(logG).item() best_seq = x_Ts[best_idx].unsqueeze(0) return best_seq, logG[best_idx] def pCoMole( model, x0: torch.Tensor, *, pad_id: int, bos_id: int, eos_id: int, allowed_tokens: Optional[torch.Tensor], objective_models: List[Callable[[torch.Tensor], Tuple[str, Any]]], constraint_models: List[Callable[[torch.Tensor], torch.Tensor]], w: torch.Tensor, rho: float, ref_z: torch.Tensor, beta_start: float = 1.0, beta_end: float = 3.0, num_steps: int = 32, num_candidates: int = 8, num_rollouts: int = 4, max_len_cap: Optional[int] = None, device: Optional[torch.device] = None, num_final_rollouts: int = 16, cfg, selfies_tokenizer, smiles_tokenizer, args, diagnostics=None, ) -> torch.Tensor: if device is None: device = x0.device x = x0.clone().to(device) x_input = x0.clone().to(device) time_grid = torch.linspace(0.0, 1.0, steps=num_steps, device=device) last_timestep = 0 best_terminal = None best_terminal_logG = float("-inf") def _sync(): if diagnostics is not None and torch.cuda.is_available(): torch.cuda.synchronize() def _tic(): _sync() return time.perf_counter() def _toc(bucket, t0): if diagnostics is None: return _sync() getattr(diagnostics, bucket).add(time.perf_counter() - t0) with torch.no_grad(): for step in tqdm(range(num_steps - 1)): if diagnostics is not None: diagnostics.num_executed_steps += 1 t = time_grid[step].view(1) frac = step / max(1, (num_steps - 1)) beta_t = beta_start + (beta_end - beta_start) * frac # model forward + candidate proposal t0 = _tic() mask = (x != pad_id) lam_ins, logits_ins, lam_del, lam_sub, logits_sub, lam_total, pi_type = model(x_t=x, mask=mask, t=t) candidates = [x.squeeze(0)] # compute the scores of current sequence with the candidates base_rates = [] for _ in range(num_candidates): cand_seq, base_rate = _sample_single_edit_batch( x, lam_ins, logits_ins, lam_del, lam_sub, logits_sub, pad_id, bos_id, eos_id, allowed_tokens, delta=float(1/(num_steps-1)), max_len_cap=max_len_cap, ) if not torch.equal(cand_seq, x): candidates.append(cand_seq.squeeze(0)) base_rates.append(base_rate) candidates = list(set(candidates)) batch_candidates = torch.nn.utils.rnn.pad_sequence(candidates, batch_first=True, padding_value=pad_id) _toc("proposal", t0) # We only want the survived candidates to improve the objective weights t0 = _tic() cand_logG, cand_ws = _G_T(batch_candidates, objective_models, constraint_models, w, rho, ref_z, beta_t, selfies_tokenizer, smiles_tokenizer, ws_for_invalid=True) _toc("screening", t0) curr_logG = cand_logG[0] curr_ws = cand_ws[0] cand_logG = cand_logG[1:] cand_ws = cand_ws[1:] batch_candidates = batch_candidates[1:, :] if len(batch_candidates) == 0: if diagnostics is not None: diagnostics.num_steps_no_selectable_candidate += 1 continue improve_idx = (cand_ws > curr_ws).nonzero(as_tuple=True)[0] survived_candidates = batch_candidates[improve_idx, :] base_rates = [base_rates[i] for i in improve_idx] # (num_survived_candidates,) if len(improve_idx) == 0: if diagnostics is not None: diagnostics.num_steps_no_selectable_candidate += 1 continue # Keep all the rollout terminal sequences in one batch t0 = _tic() x_Ts = short_rollout_batch(model, survived_candidates, time_grid, step, pad_id, bos_id, eos_id, allowed_tokens, max_len_cap, num_rollouts, num_steps) _toc("rollout", t0) if diagnostics is not None: diagnostics.num_rollouts_simulated += int(survived_candidates.shape[0] * num_rollouts) diagnostics.num_evaluated_candidates += int(survived_candidates.shape[0]) # Constraints are taken into account for the terminal sequences t0 = _tic() logG, _, = _G_T(x_Ts, objective_models, constraint_models, w, rho, ref_z, beta_t, selfies_tokenizer, smiles_tokenizer, ws_for_invalid=False) _toc("terminal_oracle", t0) if diagnostics is not None: logG_reshaped_tmp = logG.reshape(survived_candidates.shape[0], num_rollouts) n_all_inf = int((~torch.isfinite(logG_reshaped_tmp).any(dim=1)).sum().item()) diagnostics.num_candidates_all_rollouts_infeasible += n_all_inf if torch.isfinite(logG).any(): diagnostics.observed_any_feasible_terminal = True # Save the best teminal sequence curr_best_terminal_logG = torch.max(logG) if best_terminal_logG <= curr_best_terminal_logG: best_terminal_idx = torch.argmax(logG) best_terminal = x_Ts[best_terminal_idx].unsqueeze(0) best_terminal_logG = curr_best_terminal_logG best_terminal_smiles_seq = selfies_tokenizer.batch_decode(best_terminal.tolist(), return_smiles=True) best_terminal_smiles_token = smiles_tokenizer(best_terminal_smiles_seq, return_tensors='pt')['input_ids'].to(device) print("\nSaved Best Terminal: ", best_terminal_smiles_seq) print("Saved Best Terminal Length: ", best_terminal_smiles_token.shape[1]) print("Saved Best Terminal logG: ", best_terminal_logG) t0 = _tic() logG = logG.reshape(survived_candidates.shape[0], num_rollouts) log_h_hat = torch.logsumexp(logG, dim=1) - math.log(num_rollouts) # (num_survived_candidates,) idx = (logG.max(dim=1).values > curr_logG).nonzero(as_tuple=True)[0] final_survived_candidates = survived_candidates[idx, :] if len(final_survived_candidates) == 0: _toc("selection", t0) if diagnostics is not None: diagnostics.num_steps_no_selectable_candidate += 1 continue # Doob-like transform log_h_hat = log_h_hat[idx] base_rates_t = torch.tensor([base_rates[i] for i in idx.tolist()], device=device, dtype=torch.float32) log_base = 0.5 * torch.log(base_rates_t.clamp_min(1e-30)) log_weights = log_base + log_h_hat probs = torch.softmax(log_weights, dim=0) selected_idx = torch.multinomial(probs, 1).item() x = final_survived_candidates[selected_idx].unsqueeze(0) _toc("selection", t0) smiles_seq = selfies_tokenizer.batch_decode(x.tolist(), return_smiles=True) smiles_token = smiles_tokenizer(smiles_seq, return_tensors='pt')['input_ids'].to(device) print(smiles_seq[0]) print("Current Length: ", smiles_token.shape[1]) scores = compute_scores_print(smiles_token, smiles_seq, objective_models, constraint_models, device, return_scores=True).squeeze(0) last_timestep = step # finalize t0 = _tic() x_final_rollout, logG_final_rollout = _finalize_from_last( model, x, time_grid, last_timestep, pad_id, bos_id, eos_id, allowed_tokens, objective_models, constraint_models, w, rho, ref_z, beta_end, max_len_cap=max_len_cap, num_final_rollouts=num_final_rollouts, num_steps=num_steps, selfies_tokenizer=selfies_tokenizer, smiles_tokenizer=smiles_tokenizer, ) if diagnostics is not None: _sync() elapsed = time.perf_counter() - t0 diagnostics.rollout.add(elapsed * 0.5) diagnostics.terminal_oracle.add(elapsed * 0.5) if torch.isfinite(logG_final_rollout): diagnostics.observed_any_feasible_terminal = True if logG_final_rollout >= best_terminal_logG: best_terminal = x_final_rollout if diagnostics is not None: if best_terminal is None or not diagnostics.observed_any_feasible_terminal: diagnostics.run_failure = True diagnostics.used_original_input_fallback = True if best_terminal is None: best_terminal = x_input if best_terminal is not None: diagnostics.output_sequence = selfies_tokenizer.batch_decode( best_terminal.tolist(), return_smiles=True )[0] return best_terminal # --------------------------------------------------------------------------- # main # --------------------------------------------------------------------------- def main(): parser = argparse.ArgumentParser() parser.add_argument("--pcomol_config", type=str, required=True) parser.add_argument("--bindevaluator_config", type=str, required=False) parser.add_argument("--ckpt", type=str, required=True) parser.add_argument("--input", type=str, required=True) parser.add_argument("--num_steps", type=int, default=32) parser.add_argument("--max_len_cap", type=int, default=None) parser.add_argument("--num_candidates", type=int, default=10) parser.add_argument("--num_rollouts", type=int, default=5) parser.add_argument("--beta_start", type=float, default=1.0) parser.add_argument("--beta_end", type=float, default=3.0) parser.add_argument("--alpha_start", type=float, default=0.8) parser.add_argument("--alpha_end", type=float, default=0.1) parser.add_argument("--num_final_rollouts", type=int, default=50) parser.add_argument("--objective_weights", type=float, nargs='+') parser.add_argument("--ref_z", type=float, nargs='+') parser.add_argument("--rho", type=float, default=1) parser.add_argument("--target", type=str, default=None) parser.add_argument("--motifs", type=str, default=None) parser.add_argument("--specificity", action='store_true') parser.add_argument("--output_file", type=str, default="outputs/pcomole_peptidomimetics.csv") args = parser.parse_args() if args.output_file: os.makedirs(os.path.dirname(os.path.abspath(args.output_file)) or ".", exist_ok=True) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") with open(args.pcomol_config, "r") as f: cfg = edict(yaml.safe_load(f)) if args.bindevaluator_config: with open(args.bindevaluator_config, "r") as f: bindevaluator_config = edict(yaml.safe_load(f)) editflow, source_dist, selfies_tokenizer, pad_id, bos_id, eos_id, eps_id = build_model_and_stuff(cfg, device) vocab_path, splits_path = smiles_vocab_files() smiles_tokenizer = SMILES_SPE_Tokenizer(str(vocab_path), str(splits_path)) ckpt = torch.load(args.ckpt, map_location=device) editflow.load_state_dict(ckpt["state_dict"], strict=False) model = editflow.model.to(device) model.eval() x0 = tokenize_input_str(args.input, selfies_tokenizer, bos_id, eos_id, device) allowed_tokens = torch.tensor( [tok for tok in source_dist._allowed_tokens if tok not in (eps_id,)], device=device, dtype=torch.long, ) toxicity = Toxicity(device) solubility = Solubility(device) permeability = Permeability(device) halflife = Halflife(device, apply_log1p=True) if args.target: affinity = BindingAffinity(args.target, device) if args.motifs: motif = MotifModel(bindevaluator_config, args.target, args.motifs, device, args.specificity) sm = SmallMolecule(args.target, device) objective_models = [sm, toxicity, solubility, permeability, halflife, affinity, motif] # objective_models = [sm, halflife, affinity, motif] num_objectives = len(objective_models) if args.specificity else len(objective_models) - 1 if not args.objective_weights: objective_weights = torch.tensor([1.0 / num_objectives] * num_objectives).to(device) else: objective_weights = torch.tensor(args.objective_weights).to(device) if not args.ref_z: ref_z = torch.zeros(num_objectives).to(device) else: ref_z = torch.tensor(args.ref_z).to(device) peptidomimetic_hard_constraint = Peptidomimetic() length_soft_constraint = Length(args.input, smiles_tokenizer) disulfide_hard_constraint = DisulfideIntegrity() peptide_hard_constraint = Peptide() no_constraint = Pass() constraint_models = [peptidomimetic_hard_constraint, length_soft_constraint] input_smiles_token = smiles_tokenizer(args.input, return_tensors='pt')['input_ids'].to(device) print("Initial Scores:") compute_scores_print(input_smiles_token, [args.input], objective_models, constraint_models, device, return_scores=False) valid = 0 target_valid = 100 # how many successful designs you want attempt = 0 max_attempts = 200 # safety cap so you don't infinite-loop if it keeps OOM'ing start = time.time() while valid < target_valid and attempt < max_attempts: attempt += 1 try: x_T = pCoMole( model=model, x0=x0, pad_id=pad_id, bos_id=bos_id, eos_id=eos_id, allowed_tokens=allowed_tokens, objective_models=objective_models, constraint_models=constraint_models, w=objective_weights, rho=0.5, ref_z=ref_z, beta_start=args.beta_start, beta_end=args.beta_end, num_steps=args.num_steps, num_candidates=args.num_candidates, num_rollouts=args.num_rollouts, max_len_cap=args.max_len_cap, num_final_rollouts=args.num_final_rollouts, cfg=cfg, selfies_tokenizer=selfies_tokenizer, smiles_tokenizer=smiles_tokenizer, args=args ) out_str = selfies_tokenizer.batch_decode(x_T.tolist())[0] smiles_token = smiles_tokenizer(out_str, return_tensors='pt')['input_ids'].to(device) print("----------------------------") print(f"\nDesigned Sequence: {out_str}\n") print("Final scores:") scores = compute_scores_print( smiles_token, [out_str], objective_models, constraint_models, device, return_scores=True ).squeeze(0) # Only count + save on success valid += 1 with open(args.output_file, 'a') as f: f.write(f"{smiles_token.shape[1]}") for score in scores: f.write(f",{score.item()}") f.write(f",{out_str}\n") except torch.cuda.OutOfMemoryError: # Don't count this round; skip to next attempt print("[WARN] CUDA OOM during pCoMole (or downstream). Skipping this round.") torch.cuda.empty_cache() torch.cuda.ipc_collect() # optional, can help in some cases continue # for _ in range(100): # x_T = pCoMole( # model=model, # x0=x0, # pad_id=pad_id, # bos_id=bos_id, # eos_id=eos_id, # allowed_tokens=allowed_tokens, # objective_models=objective_models, # constraint_models=constraint_models, # w=objective_weights, # rho=0.5, # ref_z=ref_z, # beta_start=args.beta_start, # beta_end=args.beta_end, # num_steps=args.num_steps, # num_candidates=args.num_candidates, # num_rollouts=args.num_rollouts, # max_len_cap=args.max_len_cap, # num_final_rollouts=args.num_final_rollouts, # cfg=cfg, # selfies_tokenizer=selfies_tokenizer, # smiles_tokenizer=smiles_tokenizer # ) # out_str = selfies_tokenizer.batch_decode(x_T.tolist())[0] # smiles_token = smiles_tokenizer(out_str, return_tensors='pt')['input_ids'].to(device) # print("----------------------------") # # print(f"Initial Sequence: {args.input}\n") # # print(f"Initial Scores:") # # compute_scores_print([args.input], objective_models, constraint_models, device) # print(f"\nDesigned Sequence: {out_str}\n") # print("Final scores:") # scores = compute_scores_print(smiles_token, [out_str], objective_models, constraint_models, device, return_scores=True).squeeze(0) # with open(args.output_file, 'a') as f: # f.write(f"{smiles_token.shape[1]}") # for score in scores: # f.write(f",{score.item()}") # f.write(f",{out_str}\n") end = time.time() - start print("Total Time: ", end) print("Average Time Per Sequence: ", end / attempt) if __name__ == "__main__": main()