train.py

import os
import sys

# Enable importing from parent package
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import json
import torch

import torch.nn as nn
import torch.optim as optim
from torch.utils.data import TensorDataset, DataLoader
import numpy as np
from sklearn.metrics import roc_auc_score, accuracy_score, f1_score
from gemmasight.config import (
    DATA_DIR, MODELS_DIR, DIM_FUSED, RANDOM_SEED,
    EPOCHS, BATCH_SIZE, LEARNING_RATE, WEIGHT_DECAY,
    FAISS_INDEX_PATH
)
from gemmasight.models.classifier import MSIClassifier
from gemmasight.models.retriever import CaseRetriever

# Set random seeds for reproducibility
torch.manual_seed(RANDOM_SEED)
np.random.seed(RANDOM_SEED)

def generate_synthetic_data(num_samples=100):
    """
    Generates realistic synthetic 1536-dim embeddings and binary labels for simulation training.
    """
    print(f"Generating {num_samples} synthetic histopathology samples...")
    # Simulate a structure where some features are correlated with label 1 (MSI-High)
    embeddings = np.random.randn(num_samples, DIM_FUSED).astype(np.float32)
    labels = np.random.randint(0, 2, size=(num_samples,)).astype(np.float32)
    
    # Inject signal correlating with labels to allow classifier learning
    for i in range(num_samples):
        if labels[i] == 1:
            embeddings[i, :200] += 0.5 # Shift some features positively for MSI-High
        else:
            embeddings[i, :200] -= 0.5 # Shift some features negatively for MSS
            
    # L2 normalize embeddings to match frozen encoder output format
    norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
    norms[norms == 0] = 1e-12
    embeddings = embeddings / norms
    
    return embeddings, labels

def main():
    print("=== GemmaSight Classifier Training Pipeline ===")
    
    train_emb_path = os.path.join(DATA_DIR, "train_embeddings.npy")
    train_lbl_path = os.path.join(DATA_DIR, "train_labels.npy")
    val_emb_path = os.path.join(DATA_DIR, "val_embeddings.npy")
    val_lbl_path = os.path.join(DATA_DIR, "val_labels.npy")
    
    # 1. Load or Generate Dataset
    if os.path.exists(train_emb_path) and os.path.exists(train_lbl_path):
        print("Loading pre-extracted training embeddings from disk...")
        X_train = np.load(train_emb_path)
        y_train = np.load(train_lbl_path)
        
        if os.path.exists(val_emb_path) and os.path.exists(val_lbl_path):
            X_val = np.load(val_emb_path)
            y_val = np.load(val_lbl_path)
        else:
            print("Validation files not found. Splitting training data 80/20.")
            indices = np.random.permutation(len(X_train))
            split_idx = int(0.8 * len(X_train))
            X_val = X_train[indices[split_idx:]]
            y_val = y_train[indices[split_idx:]]
            X_train = X_train[indices[:split_idx]]
            y_train = y_train[indices[:split_idx]]
    else:
        print("No pre-extracted embeddings found. Launching Synthetic Dataset Generator baseline...")
        X_train, y_train = generate_synthetic_data(num_samples=160)
        X_val, y_val = generate_synthetic_data(num_samples=40)
        
        # Save so they are available for other pipelines
        np.save(train_emb_path, X_train)
        np.save(train_lbl_path, y_train)
        np.save(val_emb_path, X_val)
        np.save(val_lbl_path, y_val)
        print("Synthetic dataset successfully cached in data/ directory.")

    # 2. Build and Save FAISS Index (Once from Training Cohort)
    print("\nBuilding CaseRetriever FAISS database from training cohort...")
    retriever = CaseRetriever(index_dim=DIM_FUSED)
    
    # Generate metadata for training cohorts
    patient_ids = [f"GS-{1000 + i}" for i in range(len(X_train))]
    visual_descriptions = []
    for lbl in y_train:
        if lbl == 1:
            visual_descriptions.append("poorly differentiated, medullary architecture, abundant tumor-infiltrating lymphocytes (TILs)")
        else:
            visual_descriptions.append("well-formed glands, regular tubular structures, preserved crypt morphology")
            
    retriever.build_index(X_train, y_train, visual_descriptions, patient_ids)

    # 3. Create PyTorch DataLoaders
    train_dataset = TensorDataset(torch.tensor(X_train), torch.tensor(y_train).unsqueeze(1))
    val_dataset = TensorDataset(torch.tensor(X_val), torch.tensor(y_val).unsqueeze(1))
    
    train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False)

    # 4. Initialize MLP Classifier
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"\nInitializing 3-layer MSI-Classifier on device: {device}")
    model = MSIClassifier(input_dim=DIM_FUSED).to(device)
    
    criterion = nn.BCELoss()
    optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY)
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=EPOCHS)

    # 5. Training Loop
    best_val_auroc = 0.0
    best_weights_path = os.path.join(MODELS_DIR, "best_classifier.pt")
    
    print("\nStarting training loop...")
    for epoch in range(1, EPOCHS + 1):
        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() * batch_x.size(0)
            
        train_loss /= len(train_loader.dataset)
        scheduler.step()
        
        # Validation evaluation
        model.eval()
        val_preds = []
        val_targets = []
        
        with torch.no_grad():
            for batch_x, batch_y in val_loader:
                batch_x = batch_x.to(device)
                outputs = model(batch_x)
                val_preds.extend(outputs.cpu().numpy())
                val_targets.extend(batch_y.numpy())
                
        val_preds = np.array(val_preds)
        val_targets = np.array(val_targets)
        
        # Compute metrics
        val_auc = roc_auc_score(val_targets, val_preds)
        val_preds_bin = (val_preds >= 0.5).astype(int)
        val_acc = accuracy_score(val_targets, val_preds_bin)
        val_f1 = f1_score(val_targets, val_preds_bin)
        
        print(f"Epoch {epoch}/{EPOCHS} | Train Loss: {train_loss:.4f} | Val AUROC: {val_auc:.4f} | Val Acc: {val_acc:.4f} | Val F1: {val_f1:.4f}")
        
        # Save best checkpoint
        if val_auc > best_val_auroc:
            best_val_auroc = val_auc
            torch.save(model.state_dict(), best_weights_path)
            print(f"--> Saved new best checkpoint with Val AUROC: {val_val_auc:.4f}" if False else "--> Saved new best model checkpoint.")

    print(f"\nTraining Complete! Best model saved to: {best_weights_path}")
    print(f"Best Validation AUROC achieved: {best_val_auroc:.4f}")

if __name__ == "__main__":
    main()