import pandas as pd import numpy as np import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import Dataset, DataLoader from sklearn.model_selection import StratifiedKFold from sklearn.metrics import precision_recall_curve, auc, f1_score, roc_auc_score import os # --- Configurations --- DATA_PATH = "./data/cda_qm_spectra_pre2008277_lvl3_train.parquet" MODELS_DIR = "./hmoc_binary_classifier/models" os.makedirs(MODELS_DIR, exist_ok=True) BATCH_SIZE = 64 EPOCHS = 100 PATIENCE = 15 LEARNING_RATE = 1e-3 WEIGHT_DECAY = 1e-4 # --- Focal Loss --- class FocalLoss(nn.Module): def __init__(self, alpha=0.85, gamma=2.0, reduction='mean'): super(FocalLoss, self).__init__() self.alpha = alpha self.gamma = gamma self.reduction = reduction def forward(self, inputs, targets): BCE_loss = nn.functional.binary_cross_entropy_with_logits(inputs, targets, reduction='none') targets = targets.type(torch.float32) pt = torch.exp(-BCE_loss) F_loss = self.alpha * (1-pt)**self.gamma * BCE_loss if self.reduction == 'mean': return torch.mean(F_loss) else: return F_loss # --- Model Architecture --- class Spectra1DCNN(nn.Module): def __init__(self, seq_len): super(Spectra1DCNN, self).__init__() self.conv = nn.Sequential( nn.Conv1d(1, 32, kernel_size=7, padding=3), nn.BatchNorm1d(32), nn.ReLU(), nn.Conv1d(32, 64, kernel_size=5, padding=2), nn.BatchNorm1d(64), nn.ReLU(), nn.Conv1d(64, 128, kernel_size=3, padding=1), nn.BatchNorm1d(128), nn.ReLU(), nn.AdaptiveAvgPool1d(1) ) self.fc = nn.Sequential( nn.Linear(128, 64), nn.ReLU(), nn.Dropout(0.4), nn.Linear(64, 1) ) def forward(self, x): # x shape: (batch, 1, seq_len) x = self.conv(x) x = x.squeeze(-1) # (batch, 128) x = self.fc(x) return x class SpectraDataset(Dataset): def __init__(self, X, y): self.X = torch.tensor(X, dtype=torch.float32).unsqueeze(1) # Add channel dim self.y = torch.tensor(y, dtype=torch.float32).unsqueeze(1) def __len__(self): return len(self.X) def __getitem__(self, idx): return self.X[idx], self.y[idx] def pad_or_truncate(x, target_len=1000): if len(x) > target_len: return x[:target_len] elif len(x) < target_len: return np.pad(x, (0, target_len - len(x))) return x def main(): print("Loading data...") df = pd.read_parquet(DATA_PATH) # Filter negatives exclude_classes = ['?', 'Noise'] df_filtered = df[~df['class'].isin(exclude_classes)].copy() # Map classes df_filtered['label'] = (df_filtered['class'] == 'HMOC').astype(int) print(f"Dataset shape after filtering: {df_filtered.shape}") print(f"Class distribution:\n{df_filtered['label'].value_counts()}") # Process sequences print("Processing spectra arrays...") TARGET_LEN = 1000 X = np.stack(df_filtered['spectrum'].apply(lambda s: pad_or_truncate(np.array(s), TARGET_LEN)).values) # Simple Min-Max normalization per spectrum X_min = X.min(axis=1, keepdims=True) X_max = X.max(axis=1, keepdims=True) X = (X - X_min) / (X_max - X_min + 1e-8) y = df_filtered['label'].values print(f"X shape: {X.shape}, y shape: {y.shape}") device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') print(f"Using device: {device}") skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42) ensemble_seeds = [42, 100, 2026, 777, 888] for seed_idx, seed in enumerate(ensemble_seeds): print(f"\\n--- Training Ensemble Model {seed_idx+1}/{len(ensemble_seeds)} (Seed {seed}) ---") torch.manual_seed(seed) np.random.seed(seed) # We will just train on one fold per seed for the ensemble to maximize data efficiency, # or we could do 5 models total, each being a full train on 4/5ths of the data. # Let's use the fold corresponding to the seed index fold = seed_idx % 5 train_idx, val_idx = list(skf.split(X, y))[fold] X_train, y_train = X[train_idx], y[train_idx] X_val, y_val = X[val_idx], y[val_idx] train_dataset = SpectraDataset(X_train, y_train) val_dataset = SpectraDataset(X_val, y_val) # Use simple augmentation (Label smoothing handled in loss conceptually, but let's just train standard PyTorch) # To handle imbalance in batches, we use a balanced sampler class_sample_count = np.array([len(np.where(y_train == t)[0]) for t in np.unique(y_train)]) weight = 1. / class_sample_count samples_weight = np.array([weight[t] for t in y_train]) samples_weight = torch.from_numpy(samples_weight) sampler = torch.utils.data.WeightedRandomSampler(samples_weight.type('torch.DoubleTensor'), len(samples_weight)) train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, sampler=sampler, num_workers=2) val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False) model = Spectra1DCNN(TARGET_LEN).to(device) criterion = FocalLoss(alpha=0.85, gamma=2.0) optimizer = optim.AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY) best_auprc = 0.0 patience_counter = 0 for epoch in range(EPOCHS): model.train() train_loss = 0.0 for batch_x, batch_y in train_loader: batch_x, batch_y = batch_x.to(device), batch_y.to(device) optimizer.zero_grad() outputs = model(batch_x) loss = criterion(outputs, batch_y) loss.backward() optimizer.step() train_loss += loss.item() model.eval() val_preds = [] val_targs = [] with torch.no_grad(): for batch_x, batch_y in val_loader: batch_x = batch_x.to(device) outputs = torch.sigmoid(model(batch_x)) val_preds.extend(outputs.cpu().numpy().flatten()) val_targs.extend(batch_y.numpy().flatten()) val_preds = np.array(val_preds) val_targs = np.array(val_targs) # Metrics precision, recall, _ = precision_recall_curve(val_targs, val_preds) auprc = auc(recall, precision) if auprc > best_auprc: best_auprc = auprc torch.save(model.state_dict(), os.path.join(MODELS_DIR, f"hmoc_model_seed_{seed}.pth")) patience_counter = 0 else: patience_counter += 1 if np.isnan(auprc) or (epoch + 1) % 5 == 0: print(f"Epoch {epoch+1:03d} | Train Loss: {train_loss/len(train_loader):.4f} | Val AUPRC: {auprc:.4f}") if patience_counter >= PATIENCE: print(f"Early stopping at epoch {epoch+1}") break print(f"Best Val AUPRC for seed {seed}: {best_auprc:.4f}") if __name__ == "__main__": main()