Move the PDC release into CupOFMoCA-PDC/
06cc59e verified 12 days ago
29.5 kB
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)