import os import sys import time import json import numpy as np import torch import torch.nn as nn import torch.optim as optim from torch.amp import autocast, GradScaler from sklearn.metrics import (accuracy_score, precision_score, recall_score, f1_score, roc_auc_score, matthews_corrcoef, confusion_matrix) from collections import Counter from model import PeptEdge, count_parameters from data_utils import load_genpept_data, get_dataloaders def evaluate(model, loader, device, multilabel=False): model.eval() all_preds, all_labels, all_probs = [], [], [] with torch.no_grad(): for x, y in loader: x, y = x.to(device), y.to(device) logits = model(x) if multilabel: probs = torch.sigmoid(logits) preds = (probs > 0.5).long() else: probs = torch.softmax(logits, dim=1) preds = logits.argmax(dim=1) all_preds.append(preds.cpu()) all_labels.append(y.cpu()) all_probs.append(probs.cpu()) preds = torch.cat(all_preds).numpy() labels = torch.cat(all_labels).numpy() probs = torch.cat(all_probs).numpy() if multilabel: from sklearn.metrics import average_precision_score ap_per_class = [] for i in range(probs.shape[1]): ap_per_class.append(average_precision_score(labels[:, i], probs[:, i])) mAP = np.mean(ap_per_class) f1 = f1_score(labels, preds, average='macro', zero_division=0) acc = accuracy_score(labels.flatten(), preds.flatten()) return { 'accuracy': float(acc), 'f1_macro': float(f1), 'mAP': float(mAP), } else: if probs.shape[1] == 2: auc = roc_auc_score(labels, probs[:, 1]) else: auc = roc_auc_score(labels, probs, multi_class='ovr') mcc = matthews_corrcoef(labels, preds) sens = recall_score(labels, preds, pos_label=1, zero_division=0) spec = recall_score(labels, preds, pos_label=0, zero_division=0) return { 'accuracy': float(accuracy_score(labels, preds)), 'precision': float(precision_score(labels, preds, zero_division=0)), 'recall': float(recall_score(labels, preds, zero_division=0)), 'specificity': float(spec), 'f1': float(f1_score(labels, preds, zero_division=0)), 'auc': float(auc), 'mcc': float(mcc), } def train_epoch(model, loader, criterion, optimizer, scaler, device, use_amp=True): model.train() total_loss = 0 for x, y in loader: x, y = x.to(device), y.to(device) optimizer.zero_grad() if use_amp: with autocast(device_type='cuda'): logits = model(x) loss = criterion(logits, y) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() else: logits = model(x) loss = criterion(logits, y) loss.backward() optimizer.step() total_loss += loss.item() return total_loss / len(loader) def train(config=None): if config is None: config = { 'batch_size': 64, 'lr': 5e-4, 'weight_decay': 1e-4, 'epochs': 150, 'd_model': 128, 'n_heads': 4, 'num_layers': 3, 'ff_dim': 256, 'dropout': 0.2, 'max_len': 200, 'use_amp': True, 'patience': 20, 'label_smoothing': 0.05, } device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') print(f'Device: {device}') print(f'Config: {json.dumps(config, indent=2)}') sequences, labels = load_genpept_data() print(f'Dataset: {len(sequences)} sequences, {Counter(labels)}') train_loader, val_loader, test_loader = get_dataloaders( sequences, labels, batch_size=config['batch_size'], max_len=config['max_len'] ) model = PeptEdge( vocab_size=21, max_len=config['max_len'], d_model=config['d_model'], n_heads=config['n_heads'], num_layers=config['num_layers'], ff_dim=config['ff_dim'], num_classes=2, dropout=config['dropout'], ).to(device) total_params = count_parameters(model) print(f'Model params: {total_params:,}') criterion = nn.CrossEntropyLoss(label_smoothing=config['label_smoothing']) optimizer = optim.AdamW(model.parameters(), lr=config['lr'], weight_decay=config['weight_decay']) scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=config['epochs']) scaler = GradScaler('cuda') if device.type == 'cuda' else None best_val_f1 = 0 best_state = None patience_counter = 0 history = [] print(f'\n{"Epoch":>5} | {"Train Loss":>10} | {"Val Acc":>8} | {"Val F1":>8} | {"Val AUC":>8} | {"Best":>5} | {"LR":>10}') print('-' * 60) for epoch in range(config['epochs']): train_loss = train_epoch(model, train_loader, criterion, optimizer, scaler, device, config['use_amp']) val_metrics = evaluate(model, val_loader, device) scheduler.step() current_lr = scheduler.get_last_lr()[0] is_best = val_metrics['f1'] > best_val_f1 if is_best: best_val_f1 = val_metrics['f1'] best_state = model.state_dict().copy() patience_counter = 0 else: patience_counter += 1 history.append({ 'epoch': epoch + 1, 'train_loss': train_loss, **val_metrics }) if (epoch + 1) % 2 == 0 or epoch == 0: print(f'{epoch + 1:>5} | {train_loss:>10.4f} | {val_metrics["accuracy"]:>8.4f} | ' f'{val_metrics["f1"]:>8.4f} | {val_metrics["auc"]:>8.4f} | ' f'{"*" if is_best else "":>5} | {current_lr:>10.2e}') if patience_counter >= config['patience']: print(f'Early stopping at epoch {epoch + 1}') break model.load_state_dict(best_state) test_metrics = evaluate(model, test_loader, device) print('\n' + '=' * 50) print('TEST SET RESULTS') print('=' * 50) for k, v in test_metrics.items(): print(f' {k}: {v:.4f}') print(f' params: {total_params:,}') results = { 'config': config, 'total_params': total_params, 'best_val_f1': best_val_f1, 'test_metrics': test_metrics, 'history': history, } os.makedirs('results', exist_ok=True) torch.save({ 'model_state_dict': best_state, 'config': config, 'test_metrics': test_metrics, 'total_params': total_params, }, 'results/peptedge_genpept.pt') with open('results/training_results.json', 'w') as f: json.dump(results, f, indent=2, default=str) print(f'\nModel saved to results/peptedge_genpept.pt') print(f'Results saved to results/training_results.json') return results if __name__ == '__main__': results = train() print(f"\nSOTA Comparison:") print(f" ESM-2 LoRA (650M params): 88.3% F1") print(f" PeptEdge ({results['total_params']:,} params): {results['test_metrics']['f1']:.2%} F1") improvement = results['test_metrics']['f1'] - 0.883 print(f" Δ: {improvement:+.2%}")