smiles/moo.py

import argparse
import re
import random
from collections import Counter
import csv
import torch
import torch.nn as nn
import torch.nn.functional as F
from tqdm import tqdm
from train import MDLMLightningModule, PeptideAnalyzer
from smiles_tokenizer.my_tokenizers import SMILES_SPE_Tokenizer
import numpy as np

from smiles_classifiers import Analyzer, Hemolysis, Nonfouling, Solubility, BindingAffinity

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):
        self.model = model
        self.device = device
        self.objectives = objectives
        self.args = args
        self.num_objectives = len(objectives)
        self.peptide_analyzer = PeptideAnalyzer()
        self.invalid = [0, 1, 2, 3, 4, 585, 586]

    def generate_x0(self, n_steps=16, temperature=1.0):
        print("Starting initial SMILES generation...")
        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,
            (args.num_samples, 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"):
                t_curr = time_steps[i]
                t_next = time_steps[i+1]

                # Prepare the current timestep tensor for the model
                t_tensor = torch.full((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)

        # 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."""
        shape = (self.args.num_samples, self.args.gen_len)
        x = self.generate_x0()
        print(x)

        if args.weights is None:
            # The first weight is for peptide analyzer
            # We need to ensure the SMILES sequences are valid peptides
            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
                # 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}")

                # 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
                    
                    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))

        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('/scratch/pranamlab/tong/ReDi_discrete/smiles/smiles_tokenizer/new_vocab.txt', '/scratch/pranamlab/tong/ReDi_discrete/smiles/smiles_tokenizer/new_splits.txt')
    
    analyzer_model = Analyzer(tokenizer)
    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 = [analyzer_model, hemolysis_model, nonfouling_model, solubility_model, affinity_model]

    # --- 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)
    
    for _ in range(args.num_batches):
        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(sequence_str)
        print(final_scores)

    print("Generation 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 samples to generate.")
    parser.add_argument("--output_file", type=str, default="./smiles.txt", help="File to save the generated sequences.")
    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 passes over the sequence.")
    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.")

    parser.add_argument("--target", type=str, required=True)
    args = parser.parse_args()
    main(args)