peptide/new_coupling.py

import argparse
import math
import os
from collections import defaultdict

import torch
import torch.nn as nn
from tqdm import tqdm
from datasets import Dataset, DatasetDict

# --- Model Architecture (Must match the trained model) ---
def modulate(x, shift, scale):
    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)

class TimestepEmbedder(nn.Module):
    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):
        return self.mlp(t.unsqueeze(-1))

class DiTBlock(nn.Module):
    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):
    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
        self.token_embedder = nn.Embedding(vocab_size + 1, model_dim)
        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 & Utility Functions ---

def generate_x1_from_x0(model, device, x0_batch, steps, temperature):
    model.eval()
    x = x0_batch.clone()
    num_samples, seq_len = x.shape
    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():
        for i in range(steps):
            t_continuous = torch.full((num_samples,), 1.0 - (i / steps), device=device)
            logits = model(x, t_continuous)
            scaled_logits = logits / temperature
            probs = torch.nn.functional.softmax(scaled_logits, dim=-1)
            sampled_tokens = torch.multinomial(probs.view(-1, model.vocab_size), 1).view(x.shape)
            if i == steps - 1:
                x = sampled_tokens
                break
            confidence = torch.gather(probs, 2, sampled_tokens.unsqueeze(-1)).squeeze(-1)
            num_to_keep = keep_schedule[i]
            _, indices_to_keep = torch.topk(confidence, num_to_keep, largest=True, dim=-1)
            keep_mask = torch.zeros_like(x, dtype=torch.bool).scatter_(1, indices_to_keep, True)
            x = torch.where(keep_mask, sampled_tokens, x)
    return x

def is_sample_valid(sample_x1):
    """
    Checks if special tokens [0, 1, 2, 3] appear in the middle of the sequence.
    """
    middle_sequence = sample_x1[1:-1]
    invalid_tokens = {0, 1, 2, 3}
    for token in middle_sequence:
        if token in invalid_tokens:
            return False
    return True

def create_prebatched_dataset(dataset, max_tokens_per_batch=500):
    """
    Groups samples into batches and restructures the dataset.
    Each row in the new dataset is a complete batch.
    """
    # Group samples by their length
    data_by_length = defaultdict(list)
    for sample in dataset:
        length = len(sample['input_ids_x1'])
        data_by_length[length].append(sample)

    # Create the actual batches
    batched_data = {'input_ids_x0': [], 'input_ids_x1': []}
    for length, samples in data_by_length.items():
        samples_per_batch = max(1, max_tokens_per_batch // length)
        for i in range(0, len(samples), samples_per_batch):
            batch_samples = samples[i:i + samples_per_batch]
            
            batch_x0 = [s['input_ids_x0'] for s in batch_samples]
            batch_x1 = [s['input_ids_x1'] for s in batch_samples]
            
            batched_data['input_ids_x0'].append(batch_x0)
            batched_data['input_ids_x1'].append(batch_x1)
            
    return Dataset.from_dict(batched_data)

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

    all_x0 = []
    all_x1 = []

    # 1. Generate samples for each length
    for length in range(args.min_len, args.max_len + 1):
        print(f"Generating {args.samples_per_len} valid samples for length {length}...")
        valid_samples_count = 0
        pbar = tqdm(total=args.samples_per_len)
        while valid_samples_count < args.samples_per_len:
            remaining = args.samples_per_len - valid_samples_count
            batch_size = min(args.batch_size, remaining)
            
            shape = (batch_size, length)
            x0_batch = torch.randint(0, model.vocab_size, shape, dtype=torch.long, device=device)
            x1_batch = generate_x1_from_x0(model, device, x0_batch, args.gen_steps, args.temperature)

            # 2. Perform sanity check on each sample
            for x0, x1 in zip(x0_batch, x1_batch):
                if is_sample_valid(x1.tolist()):
                    all_x0.append(x0.cpu().tolist())
                    all_x1.append(x1.cpu().tolist())
                    valid_samples_count += 1
                    pbar.update(1)
                    if valid_samples_count >= args.samples_per_len:
                        break
        pbar.close()

    # 3. Create dataset and split
    print("Splitting dataset...")
    rectified_data = {'input_ids_x0': all_x0, 'input_ids_x1': all_x1}
    dataset = Dataset.from_dict(rectified_data)
    train_test_split = dataset.train_test_split(test_size=0.2, seed=42)
    valid_test_split = train_test_split['test'].train_test_split(test_size=0.5, seed=42)
    final_dataset_dict = DatasetDict({
        'train': train_test_split['train'],
        'validation': valid_test_split['train'],
        'test': valid_test_split['test']
    })

    # 4. Pre-batch each split
    print("Pre-batching splits...")
    batched_dataset_dict = DatasetDict()
    for split_name, split_dataset in final_dataset_dict.items():
        print(f"Processing {split_name} split...")
        batched_dataset_dict[split_name] = create_prebatched_dataset(split_dataset)

    # 5. Save the final dataset
    output_path = f"{args.output_path}/v{args.version}"
    print(f"Saving new batched dataset to {output_path}...")
    batched_dataset_dict.save_to_disk(output_path)
    
    print("Rectification complete.")
    print(f"Train on this by updating your training script's dataset path to '{output_path}'.")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Generate a rectified dataset with variable lengths and pre-batching.")

    parser.add_argument("--checkpoint", type=str, required=True)
    parser.add_argument("--output_path", type=str, default="./rectified_datasets")
    parser.add_argument("--version", type=str, default='1')
    parser.add_argument("--samples_per_len", type=int, default=10000)
    parser.add_argument("--min_len", type=int, default=6)
    parser.add_argument("--max_len", type=int, default=49)
    parser.add_argument("--gen_steps", type=int, default=16)
    parser.add_argument("--temperature", type=float, default=1.0)
    parser.add_argument("--batch_size", type=int, default=128)

    args = parser.parse_args()
    main(args)