peptide/generation.py

import argparse
import math
import torch
import torch.nn as nn
from tqdm import tqdm
from transformers import AutoTokenizer

# --- Model Architecture ---
def modulate(x, shift, scale):
    """
    Modulates the input tensor x with a shift and scale.
    """
    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)

class TimestepEmbedder(nn.Module):
    """
    Embeds a continuous scalar timestep t in [0, 1] into a vector representation.
    """
    def __init__(self, hidden_size):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.Linear(1, hidden_size, bias=True),
            nn.SiLU(),
            nn.Linear(hidden_size, hidden_size, bias=True),
        )

    def forward(self, t):
        # t is shape (batch_size,), needs to be (batch_size, 1) for the Linear layer.
        return self.mlp(t.unsqueeze(-1))

class DiTBlock(nn.Module):
    """
    A single block of the Diffusion Transformer.
    """
    def __init__(self, hidden_size, n_heads):
        super().__init__()
        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
        self.attn = nn.MultiheadAttention(hidden_size, n_heads, batch_first=True)
        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
        self.mlp = nn.Sequential(
            nn.Linear(hidden_size, 4 * hidden_size),
            nn.GELU(),
            nn.Linear(4 * hidden_size, hidden_size)
        )
        self.adaLN_modulation = nn.Sequential(
            nn.SiLU(),
            nn.Linear(hidden_size, 6 * hidden_size, bias=True)
        )

    def forward(self, x, c):
        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=1)
        x_norm1 = modulate(self.norm1(x), shift_msa, scale_msa)
        attn_output, _ = self.attn(x_norm1, x_norm1, x_norm1)
        x = x + gate_msa.unsqueeze(1) * attn_output
        x_norm2 = modulate(self.norm2(x), shift_mlp, scale_mlp)
        mlp_output = self.mlp(x_norm2)
        x = x + gate_mlp.unsqueeze(1) * mlp_output
        return x

class MDLM(nn.Module):
    """
    Masked Diffusion Language Model (MDLM) using a DiT backbone.
    """
    def __init__(self, vocab_size, seq_len, model_dim, n_heads, n_layers):
        super().__init__()
        self.vocab_size = vocab_size
        self.seq_len = seq_len
        self.model_dim = model_dim
        self.mask_token_id = vocab_size # Use vocab_size as the ID for the mask token

        self.token_embedder = nn.Embedding(vocab_size + 1, model_dim) # +1 for the mask token
        self.pos_embedder = nn.Parameter(torch.randn(1, seq_len, model_dim))
        self.time_embedder = TimestepEmbedder(model_dim)

        self.transformer_blocks = nn.ModuleList([
            DiTBlock(model_dim, n_heads) for _ in range(n_layers)
        ])

        self.final_norm = nn.LayerNorm(model_dim)
        self.lm_head = nn.Linear(model_dim, vocab_size)

    def forward(self, x, t):
        seq_len = x.shape[1]
        x_embed = self.token_embedder(x) + self.pos_embedder[:, :seq_len, :]
        t_embed = self.time_embedder(t)
        for block in self.transformer_blocks:
            x_embed = block(x_embed, t_embed)
        x_embed = self.final_norm(x_embed)
        logits = self.lm_head(x_embed)
        return logits

# --- Generation Function ---

def generate_samples(model, device, num_samples, seq_len, steps, temperature):
    """
    Generates samples by starting from a random sequence and progressively refining it.
    """
    model.eval()
    
    # Start with a completely random sequence of tokens
    shape = (num_samples, seq_len)
    x = torch.randint(0, model.vocab_size, shape, dtype=torch.long, device=device)
    
    # Cosine schedule determines how many tokens we *keep* from the previous step.
    # It goes from 0 (keep none) to seq_len (keep all).
    keep_schedule = torch.cos(torch.linspace(math.pi / 2, 0, steps, device=device)) * seq_len
    keep_schedule = torch.round(keep_schedule).long()

    with torch.no_grad():
        progress_bar = tqdm(range(steps), desc="Generating Samples")
        for i in progress_bar:
            # Time `t` should go from 0 (pure noise) up to 1 (pure data)
            t_continuous = torch.full((num_samples,), (i) / steps, device=device)
            
            logits = model(x, t_continuous)
            
            # Apply temperature scaling to control diversity
            scaled_logits = logits / temperature
            probs = torch.nn.functional.softmax(scaled_logits, dim=-1)
            
            # Sample a full new sequence from the model's prediction
            sampled_tokens = torch.multinomial(probs.view(-1, model.vocab_size), 1).view(shape)
            
            # For the last step, the new sample is our final result
            if i == steps - 1:
                x = sampled_tokens
                break

            # Determine which tokens from the *newly sampled sequence* to keep, based on confidence
            confidence = torch.gather(probs, 2, sampled_tokens.unsqueeze(-1)).squeeze(-1)
            
            # Find the indices of the most confident tokens to keep
            num_to_keep = keep_schedule[i]
            _, indices_to_keep = torch.topk(confidence, num_to_keep, largest=True, dim=-1)
            
            # Create a mask for the tokens we are keeping
            keep_mask = torch.zeros_like(x, dtype=torch.bool).scatter_(1, indices_to_keep, True)
            
            # The next sequence `x` is a mix:
            # - Where keep_mask is True, we use the new, confident sampled_tokens.
            # - Where keep_mask is False, we keep the tokens from the previous step `x`.
            x = torch.where(keep_mask, sampled_tokens, x)
            
    return x

# --- Main Execution ---

def main(args):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")

    print(f"Loading checkpoint from {args.checkpoint}...")
    try:
        checkpoint = torch.load(args.checkpoint, map_location=device, weights_only=False)
        model_args = checkpoint['args']
    except FileNotFoundError:
        print(f"Error: Checkpoint file not found at {args.checkpoint}")
        return
    except Exception as e:
        print(f"Error loading checkpoint: {e}")
        return

    print("Initializing model...")
    model = MDLM(
        vocab_size=model_args.vocab_size,
        seq_len=model_args.seq_len,
        model_dim=model_args.model_dim,
        n_heads=model_args.n_heads,
        n_layers=model_args.n_layers
    ).to(device)
    
    model.load_state_dict(checkpoint['model_state_dict'])
    print("Model loaded successfully.")

    gen_len = args.gen_len if args.gen_len is not None else model_args.seq_len
    if gen_len > model_args.seq_len:
        raise ValueError(f"Requested generation length ({gen_len}) is greater than the model's max length ({model_args.seq_len}).")
    print(f"Generating sequences of length {gen_len}.")

    generated_tokens = generate_samples(
        model=model,
        device=device,
        num_samples=args.num_samples,
        seq_len=gen_len,
        steps=args.gen_steps,
        temperature=args.temperature
    )
    
    print("Decoding and saving samples...")
    tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D")
    
    with open(args.output_file, 'w') as f:
        for sample_tokens in generated_tokens:
            sequence = tokenizer.decode(sample_tokens.tolist(), skip_special_tokens=False)
            clean_sequence = sequence.replace(" ", "")[5:-5]
            f.write(clean_sequence + "\n")
            print(clean_sequence)
    
    print(f"Generation complete. {args.num_samples} sequences saved to {args.output_file}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Generate samples from a trained ReDi (MDLM) model starting from random noise.")

    parser.add_argument("--checkpoint", type=str, required=True, help="Path to the model checkpoint file.")
    parser.add_argument("--num_samples", type=int, default=128, help="Number of samples to generate.")
    parser.add_argument("--output_file", type=str, default="./generated_peptides.txt", help="File to save the generated peptide sequences.")
    parser.add_argument("--gen_steps", type=int, default=16, help="Number of steps for the progressive refinement process.")
    parser.add_argument("--gen_len", type=int, default=None, help="Desired length of the generated sequences. Defaults to the model's maximum trained length.")
    parser.add_argument("--temperature", type=float, default=1.0, help="Sampling temperature. >1 increases diversity, <1 decreases it.")

    args = parser.parse_args()
    main(args)