import argparse import re import random from collections import Counter import csv import json import os from datetime import datetime import torch import torch.nn as nn import torch.nn.functional as F from tqdm import tqdm from smiles.models.train import MDLMLightningModule from smiles.smiles_tokenizer.my_tokenizers import SMILES_SPE_Tokenizer import numpy as np from smiles.evaluation.peptide_analyzer import PeptideAnalyzer from smiles.evaluation.pdc_analyzer import PDCAnalyzer from smiles.scoring.classifiers import Conjugation, BindingAffinity class GenerationTracker: """ Tracks detailed validity information during generation """ def __init__(self, output_dir='./generation_logs'): self.output_dir = output_dir os.makedirs(output_dir, exist_ok=True) self.run_id = datetime.now().strftime("%Y%m%d_%H%M%S") self.step_logs = [] self.final_results = [] def log_step(self, step, scores, detailed_results=None, improved_indices=None): """Log information for a generation step""" step_info = { 'step': step, 'scores': scores.cpu().numpy().tolist() if isinstance(scores, torch.Tensor) else scores, 'num_valid': int(sum(1 for s in scores if s > 0)) if hasattr(scores, '__iter__') else 0, 'improved_indices': improved_indices or [] } if detailed_results: step_info['detailed'] = [] for i, result in enumerate(detailed_results): step_info['detailed'].append({ 'sample_idx': i, 'is_valid_peptide': result.get('is_valid_peptide', False), 'has_conjugation_site': result.get('has_conjugation_site', False), 'conjugation_success': result.get('conjugation_success', False), 'pdc_valid': result.get('pdc_valid', False), 'num_conjugates': len(result.get('conjugates', [])), 'error': result.get('error') }) self.step_logs.append(step_info) def log_final(self, sequences, detailed_results_with_props): """Log final results with properties""" for i, (seq, result) in enumerate(zip(sequences, detailed_results_with_props)): final_info = { 'sample_idx': i, 'peptide_smiles': seq, 'is_valid_peptide': result.get('is_valid_peptide', False), 'has_conjugation_site': result.get('has_conjugation_site', False), 'num_conjugates': len(result.get('conjugates', [])), 'num_valid_pdcs': len(result.get('valid_pdcs', [])), 'best_pdc': None, 'properties': result.get('properties', {}) } if result.get('best_pdc'): final_info['best_pdc'] = { 'smiles': result['best_pdc'].get('smiles', ''), 'checks': result['best_pdc'].get('checks', {}) } self.final_results.append(final_info) def save(self, filename_prefix='generation'): """Save all logs to files""" # Save step logs step_log_path = os.path.join(self.output_dir, f"{filename_prefix}_{self.run_id}_steps.json") with open(step_log_path, 'w') as f: json.dump(self.step_logs, f, indent=2, default=str) # Save final results final_path = os.path.join(self.output_dir, f"{filename_prefix}_{self.run_id}_final.json") with open(final_path, 'w') as f: json.dump(self.final_results, f, indent=2, default=str) # Save summary CSV csv_path = os.path.join(self.output_dir, f"{filename_prefix}_{self.run_id}_summary.csv") with open(csv_path, 'w', newline='') as f: writer = csv.writer(f) writer.writerow([ 'sample_idx', 'peptide_smiles', 'is_valid_peptide', 'has_conjugation_site', 'num_conjugates', 'num_valid_pdcs', 'best_pdc_smiles', 'mf', 'polararea', 'xlogp', 'heavycnt', 'rotbonds' ]) for result in self.final_results: props = result.get('properties', {}) writer.writerow([ result['sample_idx'], result['peptide_smiles'], result['is_valid_peptide'], result['has_conjugation_site'], result['num_conjugates'], result['num_valid_pdcs'], result['best_pdc']['smiles'] if result.get('best_pdc') else '', props.get('mf', ''), props.get('polararea', ''), props.get('xlogp', ''), props.get('heavycnt', ''), props.get('rotbonds', '') ]) print(f"\n=== Generation logs saved ===") print(f" Step logs: {step_log_path}") print(f" Final results: {final_path}") print(f" Summary CSV: {csv_path}") return step_log_path, final_path, csv_path def print_step_summary(self, step, scores, detailed_results=None): """Print a summary for the current step""" if isinstance(scores, torch.Tensor): scores_list = scores.cpu().numpy().tolist() else: scores_list = list(scores) num_valid = sum(1 for s in scores_list if s > 0) print(f"\n--- Step {step} Summary ---") print(f" Valid: {num_valid}/{len(scores_list)}") if detailed_results: valid_peptide = sum(1 for r in detailed_results if r.get('is_valid_peptide')) has_site = sum(1 for r in detailed_results if r.get('has_conjugation_site')) conj_success = sum(1 for r in detailed_results if r.get('conjugation_success')) pdc_valid = sum(1 for r in detailed_results if r.get('pdc_valid')) print(f" Breakdown:") print(f" Valid peptides: {valid_peptide}") print(f" Has conjugation site: {has_site}") print(f" Conjugation success: {conj_success}") print(f" PDC valid: {pdc_valid}") def print_final_summary(self): """Print final summary""" print("\n" + "=" * 60) print("FINAL GENERATION SUMMARY") print("=" * 60) total = len(self.final_results) valid_pdcs = sum(1 for r in self.final_results if r.get('num_valid_pdcs', 0) > 0) print(f"\nTotal sequences: {total}") print(f"With valid PDCs: {valid_pdcs}") print("\nDetailed results:") for result in self.final_results: status = "" if result.get('num_valid_pdcs', 0) > 0 else "" print(f"\n [{status}] Sample {result['sample_idx']}:") print(f" Peptide: {result['peptide_smiles'][:60]}...") print(f" Valid peptide: {result['is_valid_peptide']}") print(f" Has conjugation site: {result['has_conjugation_site']}") print(f" Conjugates: {result['num_conjugates']}") print(f" Valid PDCs: {result['num_valid_pdcs']}") if result.get('properties'): props = result['properties'] print(f" Properties:") for k, v in props.items(): print(f" {k}: {v}") def peptide_bond_mask(smiles_list): # Initialize the batch mask batch_size = len(smiles_list) max_seq_length = 1035 #max(len(smiles) for smiles in smiles_list) # Find the longest SMILES mask = torch.zeros((batch_size, max_seq_length), dtype=torch.int) # Mask filled with zeros bond_patterns = [ (r'OC\(=O\)', 'ester'), (r'N\(C\)C\(=O\)', 'n_methyl'), (r'N[12]C\(=O\)', 'peptide'), # Pro peptide bonds (r'NC\(=O\)', 'peptide'), # Regular peptide bonds (r'C\(=O\)N\(C\)', 'n_methyl'), (r'C\(=O\)N[12]?', 'peptide') ] for batch_idx, smiles in enumerate(smiles_list): positions = [] used = set() # Identify bonds for pattern, bond_type in bond_patterns: for match in re.finditer(pattern, smiles): if not any(p in range(match.start(), match.end()) for p in used): positions.append({ 'start': match.start(), 'end': match.end(), 'type': bond_type, 'pattern': match.group() }) used.update(range(match.start(), match.end())) # Update the mask for the current SMILES for pos in positions: mask[batch_idx, pos['start']:pos['end']] = 1 return mask def peptide_token_mask(smiles_list, token_lists): # Initialize the batch mask batch_size = len(smiles_list) token_seq_length = max(len(tokens) for tokens in token_lists) # Find the longest tokenized sequence tokenized_masks = torch.zeros((batch_size, token_seq_length), dtype=torch.int) # Mask filled with zeros atomwise_masks = peptide_bond_mask(smiles_list) for batch_idx, atomwise_mask in enumerate(atomwise_masks): token_seq = token_lists[batch_idx] atom_idx = 0 for token_idx, token in enumerate(token_seq): if token_idx != 0 and token_idx != len(token_seq) - 1: if torch.sum(atomwise_mask[atom_idx:atom_idx+len(token)]) >= 1: tokenized_masks[batch_idx][token_idx] = 1 atom_idx += len(token) return tokenized_masks class MOGGenerator: def __init__(self, model, device, objectives, args, tracker=None): self.model = model self.device = device self.objectives = objectives self.args = args self.num_objectives = len(objectives) self.peptide_analyzer = PeptideAnalyzer() self.pdc_analyzer = PDCAnalyzer(device) self.invalid = [0, 1, 2, 3, 4, 585, 586] self.tracker = tracker # Find the Conjugation objective for detailed tracking self.conjugation_obj = None for obj in objectives: if isinstance(obj, Conjugation): self.conjugation_obj = obj break def _generate_single_x0(self, n_steps=16, temperature=1.0): """Generate a single batch of x0 sequences.""" self.model.eval() # 1. Start with a tensor of random tokens (pure noise at t=0) x = torch.randint( 0, self.model.tokenizer.vocab_size, (self.args.num_samples, self.args.gen_len), device=self.device ) # 2. Define the time schedule for the forward process (0.0 to 1.0) time_steps = torch.linspace(0.0, 1.0, n_steps + 1, device=self.device) # 3. Iteratively follow the flow from noise to data with torch.no_grad(): for i in tqdm(range(n_steps), desc="Flow Matching Steps", leave=False): t_curr = time_steps[i] t_next = time_steps[i+1] # Prepare the current timestep tensor for the model t_tensor = torch.full((self.args.num_samples,), t_curr, device=self.device) # Get the model's prediction for the final clean sequence (at t=1) logits = self.model(x, t_tensor) logits = logits / temperature logits[:, :, 586] = -1000 pred_x1 = torch.argmax(logits, dim=-1) # On the last step, the result is the final prediction if i == n_steps - 1: x = pred_x1 break # --- Construct the next state x_{t_next} --- # The probability of a token being noise at time t_next is (1 - t_next). noise_prob = 1.0 - t_next mask = torch.rand(x.shape, device=self.device) < noise_prob # Generate new random tokens for the noise positions noise = torch.randint( 0, self.model.tokenizer.vocab_size, x.shape, device=self.device ) # Combine the final prediction with noise to form the next intermediate state x = torch.where(mask, noise, pred_x1) return x def _check_conjugation_validity(self, x): """ Check if sequences can be conjugated. Returns a boolean tensor indicating validity for each sample. """ if self.conjugation_obj is None: # Fall back to peptide validity check sequences = self.model.tokenizer.batch_decode(x) return torch.tensor([self.peptide_analyzer.is_peptide(seq) for seq in sequences], device=self.device) # Use conjugation objective - score > 0 means valid scores = self.conjugation_obj(x) return scores > 0 def generate_x0(self, n_steps=16, temperature=1.0, require_valid=False, max_retries=10): """ Generate initial x0 sequences. Args: n_steps: Number of flow matching steps temperature: Sampling temperature require_valid: If True, regenerate invalid samples until all are valid conjugatable peptides max_retries: Maximum number of regeneration attempts per sample """ print("Starting initial SMILES generation...") x = self._generate_single_x0(n_steps, temperature) if require_valid: print("Checking conjugation validity and regenerating invalid samples...") for retry in range(max_retries): validity = self._check_conjugation_validity(x) num_valid = validity.sum().item() num_invalid = self.args.num_samples - num_valid if num_invalid == 0: print(f"All {self.args.num_samples} samples are valid conjugatable peptides!") break print(f"Retry {retry + 1}/{max_retries}: {num_valid}/{self.args.num_samples} valid, regenerating {num_invalid} invalid samples...") # Get indices of invalid samples invalid_indices = torch.where(~validity)[0] # Generate new candidates for invalid samples # Temporarily adjust num_samples to generate replacements original_num_samples = self.args.num_samples self.args.num_samples = num_invalid new_x = self._generate_single_x0(n_steps, temperature) self.args.num_samples = original_num_samples # Replace invalid samples with new ones x[invalid_indices] = new_x # Final validity check final_validity = self._check_conjugation_validity(x) num_final_valid = final_validity.sum().item() if num_final_valid < self.args.num_samples: print(f"Warning: After {max_retries} retries, only {num_final_valid}/{self.args.num_samples} samples are valid.") else: print(f"Success: All {self.args.num_samples} samples are valid conjugatable peptides.") # 4. Decode the final token IDs into SMILES strings generated_sequences = self.model.tokenizer.batch_decode(x) # 5. Analyze the validity of the generated sequences validities = [] for seq in generated_sequences: validities.append(self.peptide_analyzer.is_peptide(seq)) print(seq) print(f"Initial Sequence Validity: {validities}") return x def _get_scores(self, x_batch): """Calculates the normalized scores for a batch of sequences.""" scores = [] for obj_func in self.objectives: scores.append(obj_func(x_batch.to(self.device))) return torch.stack(scores, dim=0).to(self.device) def _barker_g(self, u): """Barker balancing function.""" return u / (1 + u) def validity(self, x): sampled_sequences = self.model.tokenizer.batch_decode(x) return [1.0 if self.peptide_analyzer.is_peptide(seq) else 0.0 for seq in sampled_sequences] def generate(self): """Main generation loop with tracking.""" shape = (self.args.num_samples, self.args.gen_len) x = self.generate_x0( require_valid=self.args.require_valid_x0, max_retries=self.args.max_x0_retries ) print(x) if args.weights is None: # The first weight is for conjugation (peptide + PDC validity) # We need to ensure the SMILES sequences are valid peptides with valid PDCs weights = torch.tensor([1] + [1/(self.num_objectives-1)] * (self.num_objectives-1), device=self.device).view(-1,1) else: weights = torch.tensor(self.args.weights, device=self.device).view(-1, 1) if len(weights) != self.num_objectives: raise ValueError("Number of weights must match number of objectives.") print(f"Weights: {weights}") with torch.no_grad(): for t in tqdm(range(self.args.optimization_steps), desc="MOG Generation"): improved = False improved_indices = [] # Anneal guidance strength eta_t = self.args.eta_min + (self.args.eta_max - self.args.eta_min) * (t / (self.args.optimization_steps - 1)) # Choose a random position to mutate mut_idx = random.randint(1, self.args.gen_len-2) # Determine the generation timestep # We cycle through the timesteps to ensure all are visited generation_step = t % self.args.optimization_steps time_t = torch.full((self.args.num_samples,), (generation_step / self.args.optimization_steps), device=self.device) # Get proposal distribution from ReDi model for the chosen position logits = self.model(x, time_t) probs = F.softmax(logits, dim=-1) pos_probs = probs[:, mut_idx, :] pos_probs[:, x[:, mut_idx]] = 0 # We don't evalute the same token pos_probs[:, self.invalid] = 0 # Prune candidate vocabulary using top-p sampling sorted_probs, sorted_indices = torch.sort(pos_probs, descending=True) cumulative_probs = torch.cumsum(sorted_probs, dim=-1) remove_mask = cumulative_probs > self.args.top_p remove_mask[..., 1:] = remove_mask[..., :-1].clone() remove_mask[..., 0] = 0 # Get the set of candidate tokens for each sample in the batch candidate_tokens_list = [] for i in range(self.args.num_samples): sample_mask = remove_mask[i] candidates = sorted_indices[i, ~sample_mask] candidate_tokens_list.append(candidates) # Get current scores current_scores = self._get_scores(x) w_current = torch.exp(eta_t * torch.min(weights * current_scores, dim=0).values) if t == 0: print(f"Initial Scores: {current_scores}") # Log initial step with detailed results if self.tracker and self.conjugation_obj: detailed = self.conjugation_obj.get_last_detailed_results() self.tracker.log_step(t, current_scores[0], detailed) self.tracker.print_step_summary(t, current_scores[0], detailed) # Evaluate all candidate tokens for each sample final_proposal_tokens = [] for i in range(self.args.num_samples): candidates = candidate_tokens_list[i] candidates = torch.tensor([token for token in candidates if token not in self.invalid], device=candidates.device) if len(candidates) >= 200: candidates = candidates[:200] num_candidates = len(candidates) # Create a batch of proposed sequences for the current sample x_prop_batch = x[i].repeat(num_candidates, 1) x_prop_batch[:, mut_idx] = candidates # Evaluate all proposals proposal_scores = self._get_scores(x_prop_batch) proposal_s_omega = torch.min(weights * proposal_scores, dim=0).values w_proposal = torch.exp(eta_t * proposal_s_omega) # Get ReDi probabilities for the candidates redi_probs = pos_probs[i, candidates] # Calculate unnormalized guided probabilities tilde_q = redi_probs * self._barker_g(w_proposal / w_current[i]) # Normalize and sample the final token final_probs = tilde_q / (torch.sum(tilde_q) + 1e-9) index = torch.multinomial(final_probs, 1).item() if (random.uniform(0,1) > 1 and proposal_scores[:,index][0] == 1) or torch.sum(weights.squeeze(1) * proposal_scores[:, index]) >= torch.sum(weights.squeeze(1) * current_scores[:,i]): final_token = candidates[index] print(f"Previous Weighted Sum: {torch.sum(weights.squeeze(1) * current_scores[:,i])}") print(f"Previous Scores: {current_scores[:,i]}") print(f"New Weighted Sum: {torch.sum(weights.squeeze(1) * proposal_scores[:, index])}") print(f"New Scores: {proposal_scores[:,index]}") improved = True improved_indices.append(i) else: final_token = x[i][mut_idx] final_proposal_tokens.append(final_token) # Update the sequences with the chosen tokens x[torch.arange(self.args.num_samples), mut_idx] = torch.stack(final_proposal_tokens) if improved: print(self.model.tokenizer.batch_decode(x)) # Log step with tracking (every 4 steps to reduce overhead) if self.tracker and self.conjugation_obj and (t % 4 == 0 or t == self.args.optimization_steps - 1): # Re-evaluate for detailed logging step_scores = self._get_scores(x) detailed = self.conjugation_obj.get_last_detailed_results() self.tracker.log_step(t, step_scores[0], detailed, improved_indices) # Print summary for final step if t == self.args.optimization_steps - 1: self.tracker.print_step_summary(t, step_scores[0], detailed) return x # --- Main Execution --- def main(args): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Using device: {device}") target = args.target tokenizer = SMILES_SPE_Tokenizer('./smiles/smiles_tokenizer/new_vocab.txt', './smiles/smiles_tokenizer/new_splits.txt') # Initialize Conjugation model with specified builder type conjugation_model = Conjugation(tokenizer, device=device, builder_type=args.builder_type) # hemolysis_model = Hemolysis(device) # nonfouling_model = Nonfouling(device) # solubility_model = Solubility(device) # permeability_model = Permeability(device) affinity_model = BindingAffinity(target, device) # List of all objective functions OBJECTIVE_FUNCTIONS = [conjugation_model, affinity_model] # Initialize tracker for logging tracker = GenerationTracker(output_dir=args.output_dir) # --- Set random seeds for reproducibility --- if args.seed is not None: import random import numpy as np random.seed(args.seed) np.random.seed(args.seed) torch.manual_seed(args.seed) if torch.cuda.is_available(): torch.cuda.manual_seed(args.seed) torch.cuda.manual_seed_all(args.seed) print(f"Random seed set to: {args.seed}") # --- Load Model --- print(f"Loading model from checkpoint: {args.checkpoint}") checkpoint = torch.load(args.checkpoint, map_location=device, weights_only=False) model_hparams = checkpoint["hyper_parameters"]["args"] model = MDLMLightningModule.load_from_checkpoint( args.checkpoint, args=model_hparams, tokenizer=tokenizer, map_location=device, strict=False ) model.to(device) print("Model loaded successfully.") mog_generator = MOGGenerator(model, device, OBJECTIVE_FUNCTIONS, args, tracker=tracker) all_sequences = [] all_final_results = [] for batch_idx in range(args.num_batches): print(f"\n{'='*60}") print(f"BATCH {batch_idx + 1}/{args.num_batches}") print(f"{'='*60}") generated_tokens = mog_generator.generate() final_scores = mog_generator._get_scores(generated_tokens.detach()).detach().cpu().numpy() sequence_str = tokenizer.batch_decode(generated_tokens) print("\n--- Batch Results ---") print(f"Sequences: {sequence_str}") print(f"Scores: {final_scores}") all_sequences.extend(sequence_str) # Get detailed results with properties for final sequences detailed_with_props = conjugation_model.get_detailed_scores_with_properties(generated_tokens) tracker.log_final(sequence_str, detailed_with_props) all_final_results.extend(detailed_with_props) # Print and save final summary tracker.print_final_summary() tracker.save(filename_prefix='pdc_generation') # Also save sequences to simple text file with open(args.output_file, 'w') as f: for seq in all_sequences: f.write(seq + '\n') print(f"\nSequences saved to: {args.output_file}") print("\nGeneration complete.") if __name__ == "__main__": parser = argparse.ArgumentParser(description="Multi-Objective Generation with LBP-MOG-ReDi (Single Mutation).") parser.add_argument("--checkpoint", type=str, required=True, help="Path to the trained ReDi model checkpoint.") parser.add_argument("--num_samples", type=int, default=10, help="Number of samples to generate.") parser.add_argument("--num_batches", type=int, default=10, help="Number of batches to generate.") parser.add_argument("--output_file", type=str, default="./smiles.txt", help="File to save the generated sequences.") parser.add_argument("--output_dir", type=str, default="./generation_logs", help="Directory to save detailed generation logs.") parser.add_argument("--gen_len", type=int, default=50, help="Length of the sequences to generate.") parser.add_argument("--optimization_steps", type=int, default=16, help="Number of optimization steps per generation.") parser.add_argument("--weights", type=float, nargs='+', required=False, help="Weights for the objectives (e.g., 0.5 0.5).") parser.add_argument("--eta_min", type=float, default=1.0, help="Minimum guidance strength for annealing.") parser.add_argument("--eta_max", type=float, default=20.0, help="Maximum guidance strength for annealing.") parser.add_argument("--top_p", type=float, default=0.9, help="Top-p for pruning candidate tokens.") # PDC builder parser.add_argument("--builder_type", type=str, default="AoaDau", choices=["AoaDau", "ValAlaDau", "ValCitDau", "GlutarylDau", "Dau", "SSCpt", "VcMMAE", "DSPACpt", "MalHexCpt", "DSECCpt", "Cpt", "DSPASN38", "DSEBSN38", "DSEPSN38", "SN38"], help="PDC builder type for conjugation. " "'Dau' samples linker by GNRHR+Dau frequency. " "'Cpt' samples linker by ERBB2+Cpt frequency. " "'SN38' samples linker by NRP1+SN38 frequency.") parser.add_argument("--seed", type=int, default=None, help="Random seed for reproducibility.") parser.add_argument("--require_valid_x0", action="store_true", help="If set, regenerate x0 samples until they are valid conjugatable peptides.") parser.add_argument("--max_x0_retries", type=int, default=10, help="Maximum retries when require_valid_x0 is enabled.") parser.add_argument("--target", type=str, required=True, help="Target protein sequence for binding affinity.") args = parser.parse_args() main(args)