code/contrastive/contrastive_trainer_v5_SIMPLE.py

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

Simplified version that:
- Uses MultimodalGlycanBERT directly
- Only uses sequence modality (no MS, no 3D)
- Extracts [CLS] embedding for contrastive learning
"""

import os
import sys
import json
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

# Add path for custom model
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: (B, D), positive: (B, D), negatives: (B, N, D)
        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  # (B,)
        neg_sim = torch.bmm(negatives, anchor.unsqueeze(-1)).squeeze(-1) / self.temperature  # (B, N)
        
        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]
        # Create simple residue_ids (each token is its own residue for simplicity)
        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 get_cls_embedding(model, input_ids, attention_mask, residue_ids, device):
    """Get [CLS] token embedding from sequence encoder."""
    # Run forward with minimal inputs
    out = model(
        seq_token_ids=input_ids.to(device),
        seq_attention_mask=attention_mask.to(device),
        seq_residue_ids=residue_ids.to(device),
        compute_distance=False  # Save memory
    )
    # Extract [CLS] hidden state (first token)
    # The model returns dict with 'seq_hidden' or we use the logits
    if 'seq_pooled' in out:
        return out['seq_pooled']
    elif 'seq_hidden' in out:
        return out['seq_hidden'][:, 0, :]  # [CLS] token
    else:
        # Fallback: use the raw hidden states
        return out.get('hidden_states', out['seq_logits'])[:, 0, :]


def load_model(checkpoint_path, config_path, device):
    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', {})
    
    config = MultimodalGlycanBERTConfig(
        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_intermediate_size=seq.get('intermediate_size', 3072),
        seq_max_length=seq.get('max_length', 256),
        seq_hidden_dropout=seq.get('hidden_dropout_prob', 0.1),
        seq_attention_dropout=seq.get('attention_probs_dropout_prob', 0.1),
        use_cnn_frontend=seq.get('use_cnn_frontend', True),
        cnn_kernel_size=seq.get('cnn_kernel_size', 3),
        max_branch_depth=seq.get('max_branch_depth', 8),
        num_linkage_types=seq.get('num_linkage_types', 9),
        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_enabled=ms.get('enabled', True),
        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_enabled=st.get('enabled', True),
        use_cross_attention=st.get('use_cross_attention', True),
    )
    
    model = MultimodalGlycanBERT(config)
    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'])
        logger.info(f"Loaded epoch {ckpt.get('epoch', '?')}")
    else:
        model.load_state_dict(ckpt)
    
    model.to(device)
    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)
    logger.info(f"Model loaded successfully!")
    
    # 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...")
    
    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():
                # Get embeddings
                anchor_emb = get_cls_embedding(model, batch['anchor_ids'], batch['anchor_mask'], batch['anchor_res'], device)
                pos_emb = get_cls_embedding(model, batch['pos_ids'], batch['pos_mask'], batch['pos_res'], device)
                
                # Negatives: flatten, encode, reshape
                neg_ids = batch['neg_ids'].view(B * N, L)
                neg_masks = batch['neg_masks'].view(B * N, L)
                neg_res = batch['neg_res'].view(B * N, L)
                neg_emb = get_cls_embedding(model, neg_ids, neg_masks, neg_res, device)
                neg_emb = neg_emb.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()