#!/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()