peptide/rectify_train.py

import argparse
import math
import os
from collections import Counter

import torch
import torch.nn as nn
import torch.nn.functional as F
from datasets import load_from_disk
from torch.optim import AdamW
from torch.optim.lr_scheduler import LambdaLR
from torch.utils.data import DataLoader
from tqdm import tqdm

# --- Model Architecture (Identical to before) ---
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)
        self.apply(self._init_weights)
    def _init_weights(self, module):
        if isinstance(module, (nn.Linear, nn.Embedding)):
            module.weight.data.normal_(mean=0.0, std=0.02)
            if isinstance(module, nn.Linear) and module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.LayerNorm):
            if module.bias is not None:
                module.bias.data.zero_()
            if module.weight is not None:
                module.weight.data.fill_(1.0)
    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

# --- Learning Rate Scheduler ---
def get_lr_scheduler(optimizer, warmup_steps, total_steps, lr_min, lr_max):
    def lr_lambda(current_step):
        if current_step < warmup_steps:
            lr_range = lr_max - lr_min
            lr = lr_min + lr_range * (current_step / warmup_steps)
            return lr / lr_max
        else:
            denominator = total_steps - warmup_steps
            if denominator == 0: progress = 1.0
            else: progress = (current_step - warmup_steps) / denominator
            cosine_decay = 0.5 * (1.0 + math.cos(math.pi * progress))
            lr_range = lr_max - lr_min
            lr = lr_min + lr_range * cosine_decay
            return lr / lr_max
    return LambdaLR(optimizer, lr_lambda)

# --- Training and Validation Functions ---
def train_one_epoch(model, dataloader, optimizer, scheduler, device, epoch, args):
    model.train()
    total_loss = 0.0
    progress_bar = tqdm(dataloader, desc=f"Epoch {epoch+1} [Train]")

    # Check if this is a rectified dataset
    is_rectified_dataset = 'input_ids_x0' in dataloader.dataset.column_names

    for batch in progress_bar:
        optimizer.zero_grad()
        
        if is_rectified_dataset:
            x_0 = torch.tensor(batch['input_ids_x0']).to(device)
            x_1 = torch.tensor(batch['input_ids_x1']).to(device)
        else: # Original dataset
            x_1 = torch.tensor(batch['input_ids']).to(device)
            x_0 = torch.randint(0, model.vocab_size, x_1.shape, device=device)

        batch_size, _ = x_1.shape
        t_continuous = torch.rand(batch_size, device=device)
        mask = torch.rand(x_1.shape, device=device) < t_continuous.view(-1, 1)
        x_t = torch.where(mask, x_1, x_0)
        
        logits = model(x_t, t_continuous)
        loss = F.cross_entropy(logits.view(-1, model.vocab_size), x_1.view(-1), label_smoothing=args.label_smoothing)

        loss.backward()
        optimizer.step()
        scheduler.step()

        total_loss += loss.item()
        progress_bar.set_postfix(loss=loss.item(), lr=scheduler.get_last_lr()[0])

    return total_loss / len(dataloader)

def validate(model, val_dataloader, device, epoch, args):
    model.eval()
    total_val_nll = 0.0
    total_tc = 0.0
    tc_batches = 0
    progress_bar = tqdm(val_dataloader, desc=f"Epoch {epoch+1} [Val]")
    
    is_rectified_dataset = 'input_ids_x0' in val_dataloader.dataset.column_names

    with torch.no_grad():
        for i, batch in enumerate(progress_bar):
            if is_rectified_dataset:
                x_0 = torch.tensor(batch['input_ids_x0']).to(device)
                x_1 = torch.tensor(batch['input_ids_x1']).to(device)
            else: # Original dataset
                x_1 = torch.tensor(batch['input_ids']).to(device)
                x_0 = torch.randint(0, model.vocab_size, x_1.shape, device=device)

            batch_size, seq_len = x_1.shape
            t_continuous = torch.rand(batch_size, device=device)
            mask = torch.rand(x_1.shape, device=device) < t_continuous.view(-1, 1)
            x_t = torch.where(mask, x_1, x_0)
            
            logits = model(x_t, t_continuous)
            val_nll = F.cross_entropy(logits.view(-1, model.vocab_size), x_1.view(-1))
            total_val_nll += val_nll.item()

            if i < args.tc_batches:
                k = args.tc_k_samples
                p_marginal = F.softmax(logits, dim=-1)
                sampled_x1 = torch.multinomial(p_marginal.view(-1, model.vocab_size), k, replacement=True).view(batch_size, seq_len, k)
                
                kl_divs = []
                for b in range(batch_size):
                    sample_tuples = [tuple(s.tolist()) for s in sampled_x1[b].T]
                    joint_counts = Counter(sample_tuples)
                    p_joint_est = {k: v / len(sample_tuples) for k, v in joint_counts.items()}
                    kl_sum = 0
                    for seq_tuple, p_j in p_joint_est.items():
                        log_p_marginal_prod = 0
                        for pos, token_id in enumerate(seq_tuple):
                            log_p_marginal_prod += torch.log(p_marginal[b, pos, token_id] + 1e-9)
                        kl_sum += p_j * (math.log(p_j + 1e-9) - log_p_marginal_prod)
                    kl_divs.append(kl_sum)
                total_tc += sum(kl_divs) / len(kl_divs)
                tc_batches += 1

    avg_val_nll = total_val_nll / len(val_dataloader)
    perplexity = math.exp(avg_val_nll)
    avg_tc = total_tc / tc_batches if tc_batches > 0 else 0
    return avg_val_nll, perplexity, avg_tc

# --- Main Execution ---
def main(args):
    run_name = f"rectify_version{args.version}_lr{args.learning_rate}_ls{args.label_smoothing}"
    
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")
    
    args.checkpoint_dir = os.path.join(args.checkpoint_dir, run_name)
    print(f"Saving to {args.checkpoint_dir}")
    os.makedirs(args.checkpoint_dir, exist_ok=True)

    print("Loading datasets...")
    train_dataset = load_from_disk(args.train_dataset_path)
    val_dataset = load_from_disk(args.val_dataset_path)
    train_dataloader = DataLoader(train_dataset, batch_size=None, shuffle=False)
    val_dataloader = DataLoader(val_dataset, batch_size=None, shuffle=False)

    print("Initializing model...")
    model = MDLM(args.vocab_size, args.seq_len, args.model_dim, args.n_heads, args.n_layers).to(device)
    
    if args.resume_from_checkpoint:
        print(f"Loading weights from: {args.resume_from_checkpoint}")
        try:
            checkpoint = torch.load(args.resume_from_checkpoint, map_location=device, weights_only=False)
            model.load_state_dict(checkpoint['model_state_dict'])
            print("Model weights loaded successfully for fine-tuning.")
        except Exception as e:
            print(f"Error loading checkpoint for resume: {e}")
            return

    print(f"Model initialized with {sum(p.numel() for p in model.parameters()):,} parameters.")
    
    optimizer = AdamW(model.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay)
    num_training_steps = args.epochs * len(train_dataloader)
    warmup_steps = int(num_training_steps * 0.1)
    scheduler = get_lr_scheduler(optimizer, warmup_steps, num_training_steps, args.min_learning_rate, args.learning_rate)

    best_val_nll = float('inf')

    print("Starting training...")
    for epoch in range(args.epochs):
        train_loss = train_one_epoch(model, train_dataloader, optimizer, scheduler, device, epoch, args)
        val_nll, perplexity, tc_error = validate(model, val_dataloader, device, epoch, args)

        print(f"Epoch {epoch+1}/{args.epochs} -> Train Loss: {train_loss:.4f}, Val NLL: {val_nll:.4f}, Val PPL: {perplexity:.2f}, TC: {tc_error:.4f}")

        if val_nll < best_val_nll:
            best_val_nll = val_nll
            best_checkpoint_path = os.path.join(args.checkpoint_dir, "best.pt")
            torch.save({
                'epoch': epoch + 1, 'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(), 'scheduler_state_dict': scheduler.state_dict(),
                'val_nll': val_nll, 'tc_error': tc_error, 'args': args
            }, best_checkpoint_path)
            print(f"New best checkpoint saved to {best_checkpoint_path}")

    print("Training complete.")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Train or fine-tune a ReDi (MDLM) model.")

    parser.add_argument("--train_dataset_path", type=str, required=True)
    parser.add_argument("--val_dataset_path", type=str, required=True)
    parser.add_argument("--model_dim", type=int, default=1024)
    parser.add_argument("--n_heads", type=int, default=8)
    parser.add_argument("--n_layers", type=int, default=6)
    parser.add_argument("--vocab_size", type=int, default=24)
    parser.add_argument("--seq_len", type=int, default=100)
    parser.add_argument("--epochs", type=int, default=50)
    parser.add_argument("--learning_rate", type=float, default=1e-4)
    parser.add_argument("--min_learning_rate", type=float, default=1e-5)
    parser.add_argument("--weight_decay", type=float, default=1e-5)
    parser.add_argument("--label_smoothing", type=float, default=0)
    parser.add_argument("--tc_batches", type=int, default=5)
    parser.add_argument("--tc_k_samples", type=int, default=10)
    parser.add_argument("--checkpoint_dir", type=str, default="./checkpoints")
    parser.add_argument("--version", type=str, default=1, help="Rectification iteration number.")
    parser.add_argument("--resume_from_checkpoint", type=str, default=None, help="Path to a checkpoint to load model weights from for fine-tuning.")

    args = parser.parse_args()
    main(args)