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