code/contrastive/contrastive_trainer_v51_FINAL.py

#!/usr/bin/env python3
"""
V5.1-FIXED Contrastive Trainer V6 CORRECT

Uses EXACTLY the correct MultimodalGlycanBERTConfig arguments.
"""

import os
import sys
import pickle
import random
import logging
from pathlib import Path

import yaml
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR
from torch.cuda.amp import autocast, GradScaler
from tqdm import tqdm

sys.path.insert(0, str(Path(__file__).parents[2]))
from bert_training_v4.model.multimodal_glycan_bert_v3 import (
    MultimodalGlycanBERT, 
    MultimodalGlycanBERTConfig
)

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


class InfoNCELoss(nn.Module):
    def __init__(self, temperature: float = 0.1):
        super().__init__()
        self.temperature = temperature
    
    def forward(self, anchor, positive, negatives):
        anchor = F.normalize(anchor, dim=-1)
        positive = F.normalize(positive, dim=-1)
        negatives = F.normalize(negatives, dim=-1)
        
        pos_sim = (anchor * positive).sum(-1) / self.temperature
        neg_sim = torch.bmm(negatives, anchor.unsqueeze(-1)).squeeze(-1) / self.temperature
        
        all_logits = torch.cat([pos_sim.unsqueeze(-1), neg_sim], dim=-1)
        return (-pos_sim + torch.logsumexp(all_logits, dim=-1)).mean()


class ContrastiveDataset(Dataset):
    def __init__(self, positives, negatives, n_neg=5, max_len=256):
        self.positives = positives
        self.negatives = negatives
        self.n_neg = n_neg
        self.max_len = max_len
        
    def __len__(self):
        return len(self.positives)
    
    def _prepare(self, token_ids):
        token_ids = token_ids[:self.max_len]
        residue_ids = list(range(len(token_ids)))
        attention_mask = [1] * len(token_ids)
        
        pad_len = self.max_len - len(token_ids)
        token_ids = token_ids + [0] * pad_len
        residue_ids = residue_ids + [0] * pad_len
        attention_mask = attention_mask + [0] * pad_len
        
        return (
            torch.tensor(token_ids, dtype=torch.long),
            torch.tensor(attention_mask, dtype=torch.long),
            torch.tensor(residue_ids, dtype=torch.long)
        )
    
    def __getitem__(self, idx):
        pos = self.positives[idx]
        a_ids, a_mask, a_res = self._prepare(pos['token_ids'])
        p_ids, p_mask, p_res = self._prepare(pos['token_ids'])
        
        negs = random.sample(self.negatives, self.n_neg)
        n_ids, n_masks, n_res = [], [], []
        for neg in negs:
            ids, mask, res = self._prepare(neg['token_ids'])
            n_ids.append(ids)
            n_masks.append(mask)
            n_res.append(res)
        
        return {
            'anchor_ids': a_ids, 'anchor_mask': a_mask, 'anchor_res': a_res,
            'pos_ids': p_ids, 'pos_mask': p_mask, 'pos_res': p_res,
            'neg_ids': torch.stack(n_ids), 'neg_masks': torch.stack(n_masks), 'neg_res': torch.stack(n_res)
        }


class ProjectionHead(nn.Module):
    def __init__(self, in_dim=768, out_dim=256):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(in_dim, in_dim), nn.ReLU(), nn.Linear(in_dim, out_dim))
    def forward(self, x): return self.net(x)


def load_model(checkpoint_path, config_path, device):
    """Load model using EXACT config args from MultimodalGlycanBERTConfig."""
    with open(config_path) as f:
        cfg = yaml.safe_load(f)['model']
    
    seq = cfg['sequence']
    ms = cfg.get('mass_spectrometry', {})
    st = cfg.get('structure_3d', {})
    
    # EXACT args that MultimodalGlycanBERTConfig.__init__ accepts:
    config = MultimodalGlycanBERTConfig(
        # Sequence
        seq_vocab_size=seq.get('vocab_size', 2200),
        seq_hidden_size=seq.get('hidden_size', 768),
        seq_num_layers=seq.get('num_hidden_layers', 12),
        seq_num_heads=seq.get('num_attention_heads', 12),
        seq_max_length=seq.get('max_length', 256),
        # MS
        ms_vocab_size=ms.get('vocab_size', 242),
        ms_hidden_size=ms.get('hidden_size', 384),
        ms_num_layers=ms.get('num_hidden_layers', 6),
        ms_num_heads=ms.get('num_attention_heads', 6),
        ms_max_length=ms.get('max_length', 150),
        # Structure
        struct_vocab_size=st.get('vocab_size', 1024),
        struct_hidden_size=st.get('hidden_size', 512),
        struct_num_layers=st.get('num_hidden_layers', 8),
        struct_num_heads=st.get('num_attention_heads', 8),
        struct_max_length=st.get('max_length', 200),
        use_3d=st.get('enabled', True),
        # Cross-attention
        use_cross_attention=st.get('use_cross_attention', True),
        # Conv
        use_cnn_frontend=seq.get('use_cnn_frontend', True),
        cnn_kernel_size=seq.get('cnn_kernel_size', 3),
        # Dropout
        hidden_dropout_prob=seq.get('hidden_dropout_prob', 0.1),
        attention_probs_dropout_prob=seq.get('attention_probs_dropout_prob', 0.1),
    )
    
    logger.info("Creating model...")
    model = MultimodalGlycanBERT(config)
    
    logger.info(f"Loading checkpoint from {checkpoint_path}")
    ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False)
    
    if 'model_state_dict' in ckpt:
        model.load_state_dict(ckpt["model_state_dict"], strict=False)
        logger.info(f"Loaded epoch {ckpt.get('epoch', '?')}, best_val_loss={ckpt.get('best_val_loss', '?'):.4f}")
    else:
        model.load_state_dict(ckpt, strict=False)
    
    model.to(device)
    logger.info(f"Model loaded successfully! Params: {sum(p.numel() for p in model.parameters()):,}")
    return model, config


def main():
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--model_path', required=True)
    parser.add_argument('--config_path', required=True)
    parser.add_argument('--positives_path', required=True)
    parser.add_argument('--negatives_path', required=True)
    parser.add_argument('--output_dir', required=True)
    parser.add_argument('--epochs', type=int, default=50)
    parser.add_argument('--batch_size', type=int, default=16)
    parser.add_argument('--lr', type=float, default=2e-5)
    parser.add_argument('--n_neg', type=int, default=5)
    parser.add_argument('--temperature', type=float, default=0.1)
    parser.add_argument('--save_interval', type=int, default=5)
    args = parser.parse_args()
    
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    logger.info(f"Device: {device}")
    
    Path(args.output_dir).mkdir(parents=True, exist_ok=True)
    
    # Load model
    logger.info(f"Loading model from {args.model_path}")
    logger.info(f"Config from {args.config_path}")
    model, config = load_model(args.model_path, args.config_path, device)
    
    # Projection head
    proj_head = ProjectionHead(in_dim=config.seq_hidden_size, out_dim=256).to(device)
    
    # Data
    logger.info("Loading data...")
    with open(args.positives_path, 'rb') as f:
        positives = pickle.load(f)
    with open(args.negatives_path, 'rb') as f:
        negatives = pickle.load(f)
    logger.info(f"Positives: {len(positives)}, Negatives: {len(negatives)}")
    
    dataset = ContrastiveDataset(positives, negatives, args.n_neg)
    loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True, num_workers=4, pin_memory=True)
    
    # Optimizer
    optimizer = AdamW(list(model.parameters()) + list(proj_head.parameters()), lr=args.lr)
    scheduler = CosineAnnealingLR(optimizer, T_max=args.epochs)
    scaler = GradScaler()
    loss_fn = InfoNCELoss(args.temperature)
    
    best_loss = float('inf')
    logger.info(f"Starting {args.epochs} epochs, {len(loader)} batches per epoch...")
    
    for epoch in range(1, args.epochs + 1):
        model.train()
        proj_head.train()
        total_loss = 0
        
        pbar = tqdm(loader, desc=f"Epoch {epoch}/{args.epochs}")
        for batch in pbar:
            optimizer.zero_grad()
            
            B = batch['anchor_ids'].shape[0]
            N = batch['neg_ids'].shape[1]
            L = batch['neg_ids'].shape[2]
            
            with autocast():
                # Forward pass - model returns dict with various outputs
                anchor_out = model(
                    seq_token_ids=batch['anchor_ids'].to(device),
                    seq_attention_mask=batch['anchor_mask'].to(device),
                    seq_residue_ids=batch['anchor_res'].to(device),
                    compute_distance=False
                )
                pos_out = model(
                    seq_token_ids=batch['pos_ids'].to(device),
                    seq_attention_mask=batch['pos_mask'].to(device),
                    seq_residue_ids=batch['pos_res'].to(device),
                    compute_distance=False
                )
                
                # Get [CLS] embedding (first token of sequence hidden states)
                anchor_emb = anchor_out['seq_hidden'][:, 0, :]  # (B, hidden)
                pos_emb = pos_out['seq_hidden'][:, 0, :]
                
                # Negatives
                neg_ids = batch['neg_ids'].view(B * N, L).to(device)
                neg_masks = batch['neg_masks'].view(B * N, L).to(device)
                neg_res = batch['neg_res'].view(B * N, L).to(device)
                neg_out = model(
                    seq_token_ids=neg_ids,
                    seq_attention_mask=neg_masks,
                    seq_residue_ids=neg_res,
                    compute_distance=False
                )
                neg_emb = neg_out['seq_hidden'][:, 0, :].view(B, N, -1)
                
                # Project
                anchor_proj = proj_head(anchor_emb)
                pos_proj = proj_head(pos_emb)
                neg_proj = proj_head(neg_emb)
                
                loss = loss_fn(anchor_proj, pos_proj, neg_proj)
            
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
            
            total_loss += loss.item()
            pbar.set_postfix(loss=f"{loss.item():.4f}")
        
        avg_loss = total_loss / len(loader)
        scheduler.step()
        logger.info(f"Epoch {epoch}: avg_loss={avg_loss:.4f}")
        
        # Save
        if epoch % args.save_interval == 0 or avg_loss < best_loss:
            path = Path(args.output_dir) / f'checkpoint_epoch_{epoch}.pt'
            torch.save({
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
                'proj_head_state_dict': proj_head.state_dict(),
                'loss': avg_loss,
            }, path)
            
            if avg_loss < best_loss:
                best_loss = avg_loss
                best_path = Path(args.output_dir) / 'best_v51_contrastive_model.pt'
                torch.save({
                    'epoch': epoch,
                    'model_state_dict': model.state_dict(),
                    'proj_head_state_dict': proj_head.state_dict(),
                    'loss': best_loss,
                }, best_path)
                logger.info(f"New best! loss={best_loss:.4f}")
    
    logger.info("Training complete!")


if __name__ == '__main__':
    main()