Added pytorch lightning version

model_pl.py

@@ -0,0 +1,52 @@
+from typing import Any
+
+import pytorch_lightning as pl
+import torch
+import torch.nn as nn
+from pytorch_lightning.utilities.types import STEP_OUTPUT
+from torch.optim import Adam
+
+from model import LigandGNN, ProteinTransformer
+
+class BindingAffinityModelPL(pl.LightningModule):
+    def __init__(self, num_node_features, hidden_channels_gnn, lr):
+        super().__init__()
+        self.save_hyperparameters() # Save hyperparameters for easy access
+        self.lr = lr
+
+        self.ligand_gnn = LigandGNN(input_dim=num_node_features, hidden_channels=hidden_channels_gnn)
+        self.protein_transformer = ProteinTransformer(vocab_size=26)
+        self.head = nn.Sequential(
+            nn.Linear(128 + 128, 256),
+            nn.ReLU(),
+            nn.Dropout(0.2),
+            nn.Linear(256, 1)
+        )
+        self.criterion = nn.MSELoss()
+
+    def forward(self, x, edge_index, batch, protein_seq):
+        ligand_vec = self.ligand_gnn(x, edge_index, batch)
+        batch_size = batch.max().item() + 1
+        protein_seq = protein_seq.view(batch_size, -1)
+
+        protein_vec = self.protein_transformer(protein_seq)
+        combined = torch.cat([ligand_vec, protein_vec], dim=1)
+        return self.head(combined)
+
+    def training_step(self, batch, batch_idx):
+        # We don't need .to(device), zero_grad, backward, PL handles that
+        out = self(batch.x, batch.edge_index, batch.batch, batch.protein_seq)
+        loss = self.criterion(out.squeeze(), batch.y.squeeze())
+
+        self.log("train_loss", loss, prog_bar=True, on_step=False, on_epoch=True)
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        out = self(batch.x, batch.edge_index, batch.batch, batch.protein_seq)
+        loss = self.criterion(out.squeeze(), batch.y.squeeze())
+
+        self.log("val_loss", loss, prog_bar=True, on_step=False, on_epoch=True)
+        return loss
+
+    def configure_optimizers(self):
+        return Adam(self.parameters(), lr=self.lr, weight_decay=1e-4)

train_pl.py

@@ -0,0 +1,50 @@
+import pytorch_lightning as pl
+from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping
+from torch_geometric.loader import DataLoader
+from dataset import BindingDataset
+from torch.utils.data import random_split
+from model_pl import BindingAffinityModelPL
+import pandas as pd
+
+def main():
+    lr = 0.0005
+    # Load dataset
+    dataframe = pd.read_csv('pdbbind_refined_dataset.csv')
+    dataframe.dropna(inplace=True)
+    print("Dataset loaded with {} samples".format(len(dataframe)))
+    dataset = BindingDataset(dataframe)
+    print("Dataset transformed with {} samples".format(len(dataset)))
+
+    if len(dataset) == 0:
+        print("Dataset is empty")
+        return
+
+    train_size = int(0.8 * len(dataset))
+    test_size = len(dataset) - train_size
+    train_dataset, test_dataset = random_split(dataset, [train_size, test_size])
+
+    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
+    val_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)
+    num_features = train_dataset[0].x.shape[1]
+    print("Number of node features:", num_features)
+
+    model = BindingAffinityModelPL(num_node_features=84, hidden_channels_gnn=128, lr=lr)
+    checkpoint_callback = ModelCheckpoint(
+        monitor='val_loss',
+        dirpath='checkpoints/',
+        filename='best-checkpoint',
+        save_top_k=3,
+        mode='min'
+    )
+    early_stop_callback = EarlyStopping(monitor="val_loss", patience=5)
+
+    trainer = pl.Trainer(
+        max_epochs=20,
+        accelerator="auto", # Use GPU if available
+        devices=1,
+        callbacks=[checkpoint_callback, early_stop_callback]
+    )
+    trainer.fit(model, train_loader, val_loader)
+
+if __name__ == "__main__":
+    main()