big refactoring

.example.env

@@ -0,0 +1 @@
+TOKEN=example_token

.gitignore

@@ -0,0 +1,8 @@
+tox21_test.csv
+results.csv
+predict copy.py
+hp_search/logs/*
+hp_search/models/*
+__pycache__
+.env
+notes.txt

checkpoints/model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:39696815c68d20cae615f3e271f5d406d75b866a930851a43f6506f3a593282c
+size 628746

config/config.json

@@ -0,0 +1,12 @@
+{
+  "lr": 0.0001,
+  "dropout": 0.1,
+  "hidden_dim": 128,
+  "num_layers": 5,
+  "patience": 100,
+  "max_epochs": 200,
+  "batch_size": 64,
+  "seed": 0,
+  "add_or_mean": "mean",
+  "window_size": 15
+}

predict.py

@@ -2,9 +2,9 @@ from torch_geometric.data import Batch
 from torch_geometric.utils import from_rdmol
 import torch
 
-from src.model import GIN
-from src.preprocess import create_clean_mol_objects
-from src.seed import set_seed
+from utils.model import GIN
+from utils.preprocess import create_clean_mol_objects
+from utils.seed import set_seed
 
 def predict(smiles_list):
     """
@@ -26,7 +26,7 @@ def predict(smiles_list):
 
     # setup model
     model = GIN(num_features=9, num_classes=12, dropout=0.1, hidden_dim=128, num_layers=5, add_or_mean="mean")
-    model_path = "./assets/best_gin_model.pt"
+    model_path = "./checkpoints/model.pt"
     model.load_state_dict(torch.load(model_path, map_location=DEVICE))
     print(f"Loaded model from {model_path}")
     model.to(DEVICE)
@@ -54,4 +54,4 @@ def predict(smiles_list):
             pred_dict = {t: 0.5 for t in TARGET_NAMES}
             predictions[smiles] = pred_dict
 
-    return predictions
+    return predictions

train.py

@@ -0,0 +1,91 @@
+import torch
+from torch_geometric.loader import DataLoader
+import torch_geometric
+import numpy as np
+import json
+import os
+from dotenv import load_dotenv
+
+from utils.model import GIN
+from utils.preprocess import get_graph_datasets
+from utils.train_evaluate import train_model, evaluate, compute_roc_auc_avg_and_per_class
+from utils.seed import set_seed
+
+
+def train(config):
+    SEED=config["seed"]
+    set_seed(SEED)
+    best_model_path = "./checkpoints/model.pt"
+
+     # get dataloaders
+    print("Loading Datasets...")
+    torch_geometric.seed_everything(SEED)
+    token = os.getenv("TOKEN")
+    train_dataset, val_dataset = get_graph_datasets(token)
+    train_loader = DataLoader(train_dataset, batch_size=config["batch_size"], shuffle=True)
+    val_loader   = DataLoader(val_dataset, batch_size=config["batch_size"])
+
+    # initialize
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = GIN(num_features=9, num_classes=12, dropout=config["dropout"], hidden_dim=config["hidden_dim"], num_layers=config["num_layers"], add_or_mean=config["add_or_mean"]).to(device)
+    optimizer = torch.optim.Adam(model.parameters(), lr=config["lr"])
+
+    # training loop
+    best_mean_auc = -float("inf") 
+    best_mean_epoch = 0
+    aucs = []
+    window_size = config["window_size"]
+    epoch_checkpoints = {}
+    print("Starting Training...")
+
+    for epoch in range(0, config["max_epochs"]):
+        train_loss = train_model(model, train_loader, optimizer, device)
+        val_loss = evaluate(model, val_loader, device)
+        val_auc_per_class, val_auc_avg = compute_roc_auc_avg_and_per_class(model, val_loader, device)
+        aucs.append(val_auc_avg)
+
+        # log
+        if epoch % 10 == 0:
+            print(f"Epoch {epoch:03d} | "
+                f"Train Loss: {train_loss:.4f} | "
+                f"Val Loss: {val_loss:.4f} | "
+                f"Val ROC-AUC: {val_auc_avg:.4f}")
+        
+        # store model parameters for this epoch in cache (on CPU to save GPU memory)
+        epoch_checkpoints[epoch] = {k: v.cpu() for k, v in model.state_dict().items()}
+
+        # keep cache size limited
+        if len(epoch_checkpoints) > window_size + 2:
+            oldest = min(epoch_checkpoints.keys())
+            del epoch_checkpoints[oldest]
+
+        # once we have enough epochs, compute rolling mean
+            if len(aucs) >= window_size:
+                current_window = aucs[-window_size:]
+                current_mean_auc = np.mean(current_window)
+                middle_epoch = epoch - window_size // 2
+
+                # check if current mean beats the best so far
+                if current_mean_auc > best_mean_auc:
+                    best_mean_auc = current_mean_auc
+                    best_mean_epoch = middle_epoch
+
+                    # save only the best middle model
+                    if middle_epoch in epoch_checkpoints:
+                        torch.save(epoch_checkpoints[middle_epoch], best_model_path)
+                        print(f"🟢 New best mean AUC = {best_mean_auc:.4f} "
+                            f"(center epoch {best_mean_epoch}) — model saved!")
+
+            # early stopping based on best mean epoch
+            if epoch - best_mean_epoch >= config["patience"]:
+                print(f"⛔ Early stopping at epoch {epoch}. "
+                    f"Best mean AUC = {best_mean_auc:.4f} (center epoch {best_mean_epoch})")
+                break
+
+    print("best_smoothed_val_auc" + str(best_mean_auc) + ", best_middle_epoch" + str(best_mean_epoch))
+
+if __name__ == "__main__":
+    with open("./config/config.json", "r") as f:
+        config = json.load(f)
+    load_dotenv()
+    train(config)

{src → utils}/model.py

@@ -7,7 +7,7 @@ import numpy as np
 
 
 class GIN(torch.nn.Module):
-    def __init__(self, num_features, num_classes, dropout, hidden_dim=64, num_layers=5, add_or_mean="add"):
+    def __init__(self, num_features, num_classes, dropout, hidden_dim=128, num_layers=5, add_or_mean="add"):
         super().__init__()
         self.num_layers = num_layers
         self.hidden_dim = hidden_dim

{src → utils}/preprocess.py

@@ -7,6 +7,33 @@ from rdkit.Chem.MolStandardize import rdMolStandardize
 from rdkit import Chem
 from torch_geometric.data import InMemoryDataset
 from torch_geometric.utils import from_rdmol
+from datasets import load_dataset
+
+def get_tox21_split(token, cvfold=None):
+    ds = load_dataset("tschouis/tox21", token=token)
+    
+    train_df = ds["train"].to_pandas()
+    val_df = ds["validation"].to_pandas()
+
+    if cvfold is None:
+        return {
+            "train": train_df,
+            "validation": val_df
+        }
+    
+    combined_df = pd.concat([train_df, val_df], ignore_index=True)
+    cvfold = float(cvfold)
+
+    # create new splits
+    cvfold = float(cvfold)
+    train_df = combined_df[combined_df.CVfold != cvfold]
+    val_df = combined_df[combined_df.CVfold == cvfold]
+
+    # exclude train mols that occur in the validation split
+    val_inchikeys = set(val_df["inchikey"])
+    train_df = train_df[~train_df["inchikey"].isin(val_inchikeys)]
+
+    return {"train": train_df.reset_index(drop=True), "validation": val_df.reset_index(drop=True)}
 
 def create_clean_mol_objects(smiles: list[str]) -> tuple[list[Chem.Mol], np.ndarray]:
     """Create cleaned RDKit Mol objects from SMILES.
@@ -87,7 +114,7 @@ class Tox21Dataset(InMemoryDataset):
         self.data, self.slices = self.collate(data_list)
 
 
-def get_graph_dataset(filepath:str):
+def get_graph_datasets(token):
     """returns an InMemoryDataset that can be used in dataloaders
 
     Args:
@@ -96,6 +123,8 @@ def get_graph_dataset(filepath:str):
     Returns:
         Tox21Dataset: dataset for dataloaders
     """
-    df = pd.read_csv(filepath)
-    dataset = Tox21Dataset(df)
-    return dataset
+    datasets = get_tox21_split(token, cvfold=4)
+    train_df, val_df = datasets["train"], datasets["validation"]
+    train_dataset = Tox21Dataset(train_df)
+    val_dataset = Tox21Dataset(val_df)
+    return train_dataset, val_dataset

utils/train_evaluate.py

@@ -0,0 +1,116 @@
+import torch
+import torch.nn as nn
+import numpy as np
+from sklearn.metrics import roc_auc_score
+
+def masked_bce_loss(logits, labels, mask):
+    """
+    logits: [batch_size, num_classes] (raw outputs)
+    labels: [batch_size, num_classes] (0/1 with filler)
+    mask:   [batch_size, num_classes] (True if label is valid)
+    """
+    criterion = nn.BCEWithLogitsLoss(reduction="none")
+    loss_raw = criterion(logits, labels)
+    loss = (loss_raw * mask.float()).sum() / mask.float().sum()
+    return loss
+
+def train_model(model, loader, optimizer, device):
+    model.train()
+    total_loss = 0
+    for batch in loader:
+        batch = batch.to(device)
+
+        optimizer.zero_grad()
+        out = model(batch.x, batch.edge_index, batch.batch)  # [num_graphs, num_classes]
+
+        loss = masked_bce_loss(out, batch.y, batch.mask)
+        loss.backward()
+        optimizer.step()
+
+        total_loss += loss.item() * batch.num_graphs
+    return total_loss / len(loader.dataset)
+
+
+@torch.no_grad()
+def evaluate(model, loader, device):
+    model.eval()
+    total_loss = 0
+    for batch in loader:
+        batch = batch.to(device)
+        out = model(batch.x, batch.edge_index, batch.batch)
+        loss = masked_bce_loss(out, batch.y, batch.mask)
+        total_loss += loss.item() * batch.num_graphs
+    return total_loss / len(loader.dataset)
+
+
+@torch.no_grad()
+def compute_roc_auc(model, loader, device):
+    model.eval()
+    y_true, y_pred, y_mask = [], [], []
+
+    for batch in loader:
+        batch = batch.to(device)
+        out = model(batch.x, batch.edge_index, batch.batch)
+
+        # Store predictions (sigmoid → probabilities)
+        y_pred.append(torch.sigmoid(out).cpu())
+        y_true.append(batch.y.cpu())
+        y_mask.append(batch.mask.cpu())
+
+    # Concatenate across all batches
+    y_true = torch.cat(y_true, dim=0).numpy()
+    y_pred = torch.cat(y_pred, dim=0).numpy()
+    y_mask = torch.cat(y_mask, dim=0).numpy()
+
+    auc_list = []
+    for i in range(y_true.shape[1]):  # per label
+        mask_i = y_mask[:, i].astype(bool)
+        if mask_i.sum() > 0:  # at least one valid label
+            try:
+                auc = roc_auc_score(y_true[mask_i, i], y_pred[mask_i, i])
+                auc_list.append(auc)
+            except ValueError:
+                # happens if only one class present (all 0 or all 1)
+                pass
+
+    return np.mean(auc_list) if len(auc_list) > 0 else float("nan")
+
+@torch.no_grad()
+def compute_roc_auc_avg_and_per_class(model, loader, device):
+    model.eval()
+    y_true, y_pred, y_mask = [], [], []
+
+    with torch.no_grad():
+        for batch in loader:
+            batch = batch.to(device)
+            out = model(batch.x, batch.edge_index, batch.batch)
+
+            # Store predictions (sigmoid → probabilities)
+            y_pred.append(torch.sigmoid(out).cpu())
+            y_true.append(batch.y.cpu())
+            y_mask.append(batch.mask.cpu())
+
+    # Concatenate across all batches
+    y_true = torch.cat(y_true, dim=0).numpy()
+    y_pred = torch.cat(y_pred, dim=0).numpy()
+    y_mask = torch.cat(y_mask, dim=0).numpy()
+
+    # Compute AUC per class
+    auc_list = []
+    for i in range(y_true.shape[1]):
+        mask_i = y_mask[:, i].astype(bool)
+        if mask_i.sum() > 0:
+            try:
+                auc = roc_auc_score(y_true[mask_i, i], y_pred[mask_i, i])
+            except ValueError:
+                auc = np.nan  # in case only one class present
+        else:
+            auc = np.nan
+        auc_list.append(auc)
+
+    # Convert to numpy array for easier manipulation
+    auc_array = np.array(auc_list, dtype=np.float32)
+    mean_auc = np.nanmean(auc_array)  # overall mean ignoring NaNs
+
+    # Return both per-class and mean
+    return auc_array, mean_auc