File size: 5,904 Bytes
295b1cd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 | import argparse
from pathlib import Path
import torch
import torch.nn.functional as F
from tqdm import tqdm
# Import necessary classes from your training script
from smiles_train import MDLMLightningModule, PeptideAnalyzer
from smiles_tokenizer.my_tokenizers import SMILES_SPE_Tokenizer
import pdb
def generate_smiles(model, tokenizer, args):
"""
Generates peptide SMILES strings using the trained MDLM model
with a forward (t=0 to t=1) flow matching process.
Args:
model (MDLMLightningModule): The trained PyTorch Lightning model.
tokenizer (SMILES_SPE_Tokenizer): The tokenizer used for training.
args (argparse.Namespace): Command-line arguments containing sampling parameters.
Returns:
list[str]: A list of generated SMILES strings.
float: The validity rate of the generated SMILES.
"""
print("Starting SMILES generation with forward flow matching (t=0 to t=1)...")
model.eval()
device = args.device
# 1. Start with a tensor of random tokens (pure noise at t=0)
x = torch.randint(
0,
model.model.vocab_size,
(args.n_samples, args.seq_len),
device=device
)
# 2. Define the time schedule for the forward process (0.0 to 1.0)
time_steps = torch.linspace(0.0, 1.0, args.n_steps + 1, device=device)
# 3. Iteratively follow the flow from noise to data
with torch.no_grad():
for i in tqdm(range(args.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.n_samples,), t_curr, device=device)
# Get the model's prediction for the final clean sequence (at t=1)
logits = model(x, t_tensor)
logits = logits / args.temperature
pred_x1 = torch.argmax(logits, dim=-1)
# On the last step, the result is the final prediction
if i == args.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=device) < noise_prob
# Generate new random tokens for the noise positions
noise = torch.randint(
0,
model.model.vocab_size,
x.shape,
device=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 = tokenizer.batch_decode(x)
# 5. Analyze the validity of the generated sequences
peptide_analyzer = PeptideAnalyzer()
valid_count = 0
valid_smiles = []
for seq in generated_sequences:
if peptide_analyzer.is_peptide(seq):
valid_count += 1
valid_smiles.append(seq)
validity_rate = valid_count / len(generated_sequences)
print(f"\nGeneration complete. Validity rate: {validity_rate:.2%}")
return valid_smiles, validity_rate
def main():
parser = argparse.ArgumentParser(description="Sample from a trained ReDi model.")
# --- Required Arguments ---
parser.add_argument("--checkpoint_path", type=str, required=True, help="Path to the model checkpoint (.ckpt file).")
# --- Sampling Arguments ---
parser.add_argument("--n_samples", type=int, default=16, help="Number of SMILES strings to generate.")
parser.add_argument("--seq_len", type=int, default=256, help="Maximum sequence length for generated SMILES.")
parser.add_argument("--n_steps", type=int, default=100, help="Number of denoising steps for sampling.")
parser.add_argument("--temperature", type=float, default=1.0, help="Sampling temperature. Higher values increase diversity.")
# --- Environment Arguments ---
parser.add_argument("--vocab_path", type=str, default='/scratch/pranamlab/tong/ReDi_discrete/smiles/smiles_tokenizer/new_vocab.txt', help="Path to tokenizer vocabulary file.")
parser.add_argument("--splits_path", type=str, default='/scratch/pranamlab/tong/ReDi_discrete/smiles/smiles_tokenizer/new_splits.txt', help="Path to tokenizer splits file.")
parser.add_argument("--output_file", type=str, default="generated_smiles.txt", help="File to save the valid generated SMILES.")
args = parser.parse_args()
# Set up device
device = "cuda" if torch.cuda.is_available() else "cpu"
args.device = device
print(f"Using device: {device}")
# --- Load Model and Tokenizer ---
print("Loading tokenizer...")
tokenizer = SMILES_SPE_Tokenizer(args.vocab_path, args.splits_path)
print(f"Loading model from checkpoint: {args.checkpoint_path}")
# Load hyperparameters from the checkpoint to ensure model architecture matches
checkpoint = torch.load(args.checkpoint_path, map_location=device, weights_only=False)
model_hparams = checkpoint["hyper_parameters"]["args"]
# Instantiate the model with the loaded hyperparameters
model = MDLMLightningModule.load_from_checkpoint(
args.checkpoint_path,
args=model_hparams,
tokenizer=tokenizer,
map_location=device,
strict=False # Recommended if you have updated the code since training
)
model.to(device)
# --- Generate SMILES ---
valid_smiles, validity_rate = generate_smiles(model, tokenizer, args)
# pdb.set_trace()
with open('./v0_samples_200.csv', 'a') as f:
for smiles in valid_smiles:
# print(smiles)
f.write(smiles + '\n')
print(validity_rate)
if __name__ == "__main__":
main() |