Upload neural_network.py

neural_network.py

@@ -0,0 +1,688 @@
+# -*- coding: utf-8 -*-
+"""neural network
+
+Automatically generated by Colab.
+
+Original file is located at
+    https://colab.research.google.com/drive/13Vym7d6JDkWLa9cv9p8h_amR_3uUnGp9
+"""
+
+# Cell A: Upload training dataset google sheets (CSV file)
+from google.colab import files
+import pandas as pd
+import io
+
+uploaded = files.upload()
+
+# Cell B: Define liability predictor model
+import torch
+import torch.nn as nn
+
+class LiabilityPredictor(nn.Module):
+    def __init__(
+        self,
+        input_dim: int = 640, #320 from VH + 320 from VL
+        output_dim: int = 4, #One output per liability
+        hidden_dims=(128, 64), #Two hidden layers. Layers between input and output.
+        dropout: float = 0.10, #Randomly turns off neurons during training (prevents overfitting)
+        activation: str = "gelu", #Smooth non-linearity (good for embeddings)
+        use_layernorm: bool = True, #Stabilises training
+    ):
+        super().__init__()
+
+#Choose activation function. Converts "gelu" string into actual PyTorch layer.
+        act_layer = {
+            "relu": nn.ReLU,
+            "gelu": nn.GELU,
+            "silu": nn.SiLU,
+        }.get(activation.lower())
+
+        if act_layer is None:
+            raise ValueError(f"Unknown activation='{activation}'. Use 'relu', 'gelu', or 'silu'.")
+
+        layers = []
+
+        if use_layernorm:
+            layers.append(nn.LayerNorm(input_dim))
+
+        prev = input_dim
+        for h in hidden_dims:
+            layers.append(nn.Linear(prev, h))
+            if use_layernorm:
+                layers.append(nn.LayerNorm(h))
+            layers.append(act_layer())
+            if dropout and dropout > 0:
+                layers.append(nn.Dropout(dropout))
+            prev = h
+
+        layers.append(nn.Linear(prev, output_dim))
+        self.net = nn.Sequential(*layers)
+
+        self._init_weights()
+
+    def _init_weights(self): #Xavier initialisation
+        # Stable init for small-data regression
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.xavier_uniform_(m.weight)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # Guardrails: ensure correct dtype/shape
+        if x.dim() == 1:
+            x = x.unsqueeze(0)  # (640,) -> (1, 640)
+        if x.dim() != 2:
+            raise ValueError(f"Expected x to have shape (batch, features). Got {tuple(x.shape)}")
+
+        return self.net(x.float())
+
+    def enable_mc_dropout(self): #Allows uncertainity estimation by turning dropout on during inference.
+        """
+        Optional: call before inference if you later want MC-dropout uncertainty.
+        Keeps BatchNorm/LayerNorm behavior in eval-like mode but enables Dropout layers.
+        """
+        for m in self.modules():
+            if isinstance(m, nn.Dropout):
+                m.train()
+
+# Cell C: Create dataset
+import torch
+from torch.utils.data import Dataset
+import pandas as pd
+from transformers import AutoModel, AutoTokenizer
+import numpy as np
+
+MODEL_NAME = "facebook/esm2_t6_8M_UR50D"
+CSV_PATH = "trainingdataset  - Sheet 1.csv"
+
+df = pd.read_csv(CSV_PATH)
+
+target_cols = ['polyreactivity', 'hydrophobicity', 'aggregation', 'charge_patch']
+for col in target_cols:
+    df[col] = pd.to_numeric(df[col], errors='coerce')
+
+df = df.dropna(subset=['VH','VL'] + target_cols).reset_index(drop=True)
+
+y = df[target_cols].values
+print("Target order:", target_cols)
+print("Rows kept:", len(df))
+
+#Load ESM-2
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+esm_model = AutoModel.from_pretrained(MODEL_NAME)
+esm_model.eval()
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+esm_model.to(device)
+
+hidden_size = esm_model.config.hidden_size
+
+#Embedding function
+def embed_sequences_meanpool_residues_only(seqs, batch_size=8):
+    """
+    Returns a dict: {seq_string: torch.Tensor(shape=(hidden_size,), on CPU)}
+    Mean-pools token embeddings over residues ONLY (excludes special tokens like CLS/EOS).
+    Uses attention_mask to ignore padding.
+    """
+    # Deduplicate while preserving order
+    unique_seqs = list(dict.fromkeys(seqs))
+
+    seq_to_vec = {}
+    for i in range(0, len(unique_seqs), batch_size):
+        batch_seqs = unique_seqs[i:i + batch_size]
+
+        tokenized = tokenizer(
+            batch_seqs,
+            return_tensors="pt",
+            padding=True,
+            truncation=False,
+        )
+        tokenized = {k: v.to(device) for k, v in tokenized.items()}
+
+        with torch.inference_mode():
+            out = esm_model(**tokenized)
+
+        token_emb = out.last_hidden_state
+        attn = tokenized["attention_mask"].long()
+
+        mask = attn.clone()
+
+        mask[:, 0] = 0
+
+        # Remove EOS at the last real token position for each sequence
+        lengths = attn.sum(dim=1)              # (B,) counts real tokens incl CLS/EOS
+        eos_idx = (lengths - 1).clamp(min=0)   # index of last real token
+        row_idx = torch.arange(mask.size(0), device=device)
+        mask[row_idx, eos_idx] = 0
+
+        # Mean pool over remaining (residue) tokens
+        denom = mask.sum(dim=1).clamp(min=1).unsqueeze(-1)          # (B, 1)
+        pooled = (token_emb * mask.unsqueeze(-1)).sum(dim=1) / denom  # (B, H)
+
+        pooled = pooled.detach().cpu()
+        for s, v in zip(batch_seqs, pooled):
+            seq_to_vec[s] = v
+
+    return seq_to_vec
+
+#Embed all VH and VL sequences. Embeds each unique sequence once.
+all_seqs = df["VH"].tolist() + df["VL"].tolist()
+seq_to_vec = embed_sequences_meanpool_residues_only(all_seqs, batch_size=8)
+
+X_tensors = []
+for _, row in df.iterrows():
+    vh_vec = seq_to_vec[row["VH"]]
+    vl_vec = seq_to_vec[row["VL"]]
+
+    assert vh_vec.shape == (hidden_size,), f"VH vec shape {vh_vec.shape} != ({hidden_size},)"
+    assert vl_vec.shape == (hidden_size,), f"VL vec shape {vl_vec.shape} != ({hidden_size},)"
+
+#Concatenate VH + VL
+    combined_vec = torch.cat([vh_vec, vl_vec], dim=0)  # (640,)
+    X_tensors.append(combined_vec)
+
+X = torch.stack(X_tensors, dim=0).numpy()
+assert X.shape[1] == 2 * hidden_size, f"Expected {2*hidden_size} features, got {X.shape[1]}"
+
+assert X.shape[0] == y.shape[0], f"X rows {X.shape[0]} != y rows {y.shape[0]}"
+
+#Create dataset object
+class AntibodyDataset(Dataset):
+    def __init__(self, X, y):
+        self.X = torch.tensor(X, dtype=torch.float32)
+        self.y = torch.tensor(y, dtype=torch.float32)
+
+    def __len__(self):
+        return len(self.X)
+
+    def __getitem__(self, idx):
+        return self.X[idx], self.y[idx]
+
+
+dataset = AntibodyDataset(X, y)
+
+print(
+    f"Dataset created: {len(dataset)} samples | "
+    f"X shape: {X.shape} | y shape: {y.shape}"
+)
+
+# double-check
+print("First name:", df["name"].iloc[0] if "name" in df.columns else "(no 'name' column)")
+print("First y row:", y[0])
+
+# Cell F (REPLACEMENT): 5-Fold Cross-Validation (with early stopping) + Baseline comparison
+!pip -q install scikit-learn
+
+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 KFold
+
+# ---- Dataset wrapper (raw y stored; z-scoring is done per fold) ----
+class AntibodyDatasetRaw(Dataset):
+    def __init__(self, X_np, y_np):
+        self.X = torch.tensor(X_np, dtype=torch.float32)
+        self.y = torch.tensor(y_np, dtype=torch.float32)
+    def __len__(self):
+        return self.X.shape[0]
+    def __getitem__(self, idx):
+        return self.X[idx], self.y[idx]
+
+def mae_rmse_r2(y_true, y_pred):
+    err = y_pred - y_true
+    mae = np.mean(np.abs(err), axis=0)
+    rmse = np.sqrt(np.mean(err**2, axis=0))
+    ss_res = np.sum((y_true - y_pred)**2, axis=0)
+    ss_tot = np.sum((y_true - np.mean(y_true, axis=0))**2, axis=0) + 1e-12
+    r2 = 1.0 - (ss_res / ss_tot)
+    return mae, rmse, r2
+
+def train_one_fold(X_train, y_train_raw, X_val, y_val_raw,
+                   hidden_dims=(128,64), dropout=0.10,
+                   batch_size=16, max_epochs=200,
+                   lr=3e-4, weight_decay=1e-4,
+                   patience=12, min_delta=1e-4):
+
+    # ----- z-score targets using TRAIN only (no leakage) -----
+    y_mean = y_train_raw.mean(axis=0)
+    y_std  = y_train_raw.std(axis=0) + 1e-8
+
+    y_train_z = (y_train_raw - y_mean) / y_std
+    y_val_z   = (y_val_raw   - y_mean) / y_std
+
+    train_ds = AntibodyDatasetRaw(X_train, y_train_z)
+    val_ds   = AntibodyDatasetRaw(X_val,   y_val_z)
+
+    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True)
+    val_loader   = DataLoader(val_ds,   batch_size=batch_size, shuffle=False)
+
+    # ----- model -----
+    model = LiabilityPredictor(
+        input_dim=X_train.shape[1],
+        hidden_dims=hidden_dims,
+        dropout=dropout
+    ).to(device)
+
+    loss_fn = nn.MSELoss()
+    optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)
+    scheduler = optim.lr_scheduler.ReduceLROnPlateau(
+        optimizer, mode="min", factor=0.5, patience=3, min_lr=1e-5
+    )
+
+    best_val = float("inf")
+    best_state = None
+    bad = 0
+
+    def epoch_loss(loader, train: bool):
+        model.train() if train else model.eval()
+        total, n = 0.0, 0
+        for xb, yb in loader:
+            xb = xb.to(device)
+            yb = yb.to(device)
+
+            if train:
+                optimizer.zero_grad()
+
+            with torch.set_grad_enabled(train):
+                pred = model(xb)
+                loss = loss_fn(pred, yb)
+                if train:
+                    loss.backward()
+                    optimizer.step()
+
+            bs = xb.size(0)
+            total += loss.item() * bs
+            n += bs
+        return total / max(n, 1)
+
+    @torch.no_grad()
+    def predict_val_raw():
+        model.eval()
+        preds_z = []
+        for xb, _ in val_loader:
+            xb = xb.to(device)
+            pz = model(xb).cpu().numpy()
+            preds_z.append(pz)
+        preds_z = np.vstack(preds_z)
+        return preds_z * y_std + y_mean
+
+    # ----- training loop -----
+    for ep in range(1, max_epochs + 1):
+        tr = epoch_loss(train_loader, True)
+        va = epoch_loss(val_loader, False)
+        scheduler.step(va)
+
+        if va < best_val - min_delta:
+            best_val = va
+            best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}
+            bad = 0
+        else:
+            bad += 1
+            if bad >= patience:
+                break
+
+    # load best
+    model.load_state_dict(best_state)
+
+    # predictions in raw units + metrics
+    y_pred_raw = predict_val_raw()
+    mae, rmse, r2 = mae_rmse_r2(y_val_raw, y_pred_raw)
+
+    # baseline: predict TRAIN mean in raw units
+    base_pred = np.tile(y_mean.reshape(1,-1), (y_val_raw.shape[0], 1))
+    b_mae, b_rmse, b_r2 = mae_rmse_r2(y_val_raw, base_pred)
+
+    return (mae, rmse, r2), (b_mae, b_rmse, b_r2)
+
+# -----------------------------
+# Run 5-fold CV
+# -----------------------------
+X_np = X.astype(np.float32)
+y_np = y_raw.astype(np.float32)
+
+kf = KFold(n_splits=5, shuffle=True, random_state=42)
+
+fold_metrics = []
+fold_baseline = []
+
+for fold, (tr_idx, va_idx) in enumerate(kf.split(X_np), start=1):
+    X_tr, X_va = X_np[tr_idx], X_np[va_idx]
+    y_tr, y_va = y_np[tr_idx], y_np[va_idx]
+
+    (mae, rmse, r2), (b_mae, b_rmse, b_r2) = train_one_fold(
+        X_tr, y_tr, X_va, y_va,
+        hidden_dims=(128,64),
+        dropout=0.10,
+        batch_size=16,
+        max_epochs=200,
+        lr=3e-4,
+        weight_decay=1e-4,
+        patience=12
+    )
+
+    fold_metrics.append((mae, rmse, r2))
+    fold_baseline.append((b_mae, b_rmse, b_r2))
+
+    print(f"\nFold {fold}/5")
+    print("  NN MAE :", dict(zip(target_cols, mae)))
+    print("  NN R2  :", dict(zip(target_cols, r2)))
+    print("  BASE MAE:", dict(zip(target_cols, b_mae)))
+    print("  BASE R2 :", dict(zip(target_cols, b_r2)))
+
+print("\nDone. Run Cell G for plots + summary + final training.")
+
+# Cell G: Post-CV plots + conclusion stats + Train final deployment model + Save
+import numpy as np
+import matplotlib.pyplot as plt
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import Dataset, DataLoader
+
+# -----------------------------
+# 1) CV summary plots + conclusions
+# -----------------------------
+K = len(fold_metrics)
+T = len(target_cols)
+
+nn_mae = np.stack([m[0] for m in fold_metrics], axis=0)   # (K,4)
+nn_rmse= np.stack([m[1] for m in fold_metrics], axis=0)
+nn_r2  = np.stack([m[2] for m in fold_metrics], axis=0)
+
+b_mae  = np.stack([m[0] for m in fold_baseline], axis=0)
+b_rmse = np.stack([m[1] for m in fold_baseline], axis=0)
+b_r2   = np.stack([m[2] for m in fold_baseline], axis=0)
+
+def mean_std(a):
+    return a.mean(axis=0), a.std(axis=0)
+
+nn_mae_m, nn_mae_s = mean_std(nn_mae)
+nn_r2_m,  nn_r2_s  = mean_std(nn_r2)
+b_mae_m,  b_mae_s  = mean_std(b_mae)
+b_r2_m,   b_r2_s   = mean_std(b_r2)
+
+x = np.arange(T)
+w = 0.35
+
+plt.figure()
+plt.bar(x - w/2, nn_mae_m, yerr=nn_mae_s, width=w, label="NN")
+plt.bar(x + w/2, b_mae_m,  yerr=b_mae_s,  width=w, label="Baseline")
+plt.xticks(x, target_cols, rotation=30, ha="right")
+plt.ylabel("MAE (raw units)")
+plt.title("5-Fold CV: MAE per target (mean ± std)")
+plt.legend()
+plt.show()
+
+plt.figure()
+plt.bar(x - w/2, nn_r2_m, yerr=nn_r2_s, width=w, label="NN")
+plt.bar(x + w/2, b_r2_m,  yerr=b_r2_s,  width=w, label="Baseline")
+plt.xticks(x, target_cols, rotation=30, ha="right")
+plt.ylabel("R²")
+plt.title("5-Fold CV: R² per target (mean ± std)")
+plt.legend()
+plt.show()
+
+# Worst-target MAE: because you need all four good
+nn_worst_mae = nn_mae.max(axis=1)
+b_worst_mae  = b_mae.max(axis=1)
+print("Worst-target MAE across folds:")
+print(f"  NN   worst-MAE mean ± std: {nn_worst_mae.mean():.4f} ± {nn_worst_mae.std():.4f}")
+print(f"  BASE worst-MAE mean ± std: {b_worst_mae.mean():.4f} ± {b_worst_mae.std():.4f}")
+
+print("\nPer-target summary (mean ± std):")
+for i, t in enumerate(target_cols):
+    print(f"{t:14s} | NN MAE {nn_mae_m[i]:.4f}±{nn_mae_s[i]:.4f}  R2 {nn_r2_m[i]:.4f}±{nn_r2_s[i]:.4f} "
+          f"|| BASE MAE {b_mae_m[i]:.4f}±{b_mae_s[i]:.4f}  R2 {b_r2_m[i]:.4f}±{b_r2_s[i]:.4f}")
+
+print("\nOverall (mean across targets):")
+print(f"  NN   MAE_mean  {nn_mae_m.mean():.4f} ± {nn_mae_s.mean():.4f} | R2_mean {nn_r2_m.mean():.4f} ± {nn_r2_s.mean():.4f}")
+print(f"  BASE MAE_mean  {b_mae_m.mean():.4f} ± {b_mae_s.mean():.4f} | R2_mean {b_r2_m.mean():.4f} ± {b_r2_s.mean():.4f}")
+
+# -----------------------------
+# 2) Train final model for deployment (on all data)
+# -----------------------------
+class AntibodyDatasetZ(Dataset):
+    def __init__(self, X_np, y_z_np):
+        self.X = torch.tensor(X_np, dtype=torch.float32)
+        self.y = torch.tensor(y_z_np, dtype=torch.float32)
+    def __len__(self): return len(self.X)
+    def __getitem__(self, idx): return self.X[idx], self.y[idx]
+
+y_mean_full = y_raw.mean(axis=0)
+y_std_full  = y_raw.std(axis=0) + 1e-8
+y_z_full = (y_raw - y_mean_full) / y_std_full
+
+ds_full = AntibodyDatasetZ(X.astype(np.float32), y_z_full.astype(np.float32))
+loader = DataLoader(ds_full, batch_size=16, shuffle=True)
+
+final_model = LiabilityPredictor(input_dim=640, hidden_dims=(128,64), dropout=0.10).to(device)
+loss_fn = nn.MSELoss()
+optimizer = optim.Adam(final_model.parameters(), lr=3e-4, weight_decay=1e-4)
+
+epochs = 80
+loss_hist = []
+
+final_model.train()
+for ep in range(1, epochs+1):
+    total, n = 0.0, 0
+    for xb, yb in loader:
+        xb, yb = xb.to(device), yb.to(device)
+        optimizer.zero_grad()
+        pred = final_model(xb)
+        loss = loss_fn(pred, yb)
+        loss.backward()
+        optimizer.step()
+        total += loss.item() * xb.size(0)
+        n += xb.size(0)
+    loss_epoch = total / max(n, 1)
+    loss_hist.append(loss_epoch)
+    if ep % 10 == 0 or ep == 1:
+        print(f"[FINAL] Epoch {ep:03d} | train_loss(zMSE) {loss_epoch:.4f}")
+
+plt.figure()
+plt.plot(np.arange(1, epochs+1), loss_hist)
+plt.xlabel("Epoch")
+plt.ylabel("Train MSE in z-space")
+plt.title("Final Model Training Curve (for deployment)")
+plt.show()
+
+# Save: model + normalization (critical for inference)
+final_artifacts = {
+    "state_dict": final_model.state_dict(),
+    "y_mean": y_mean_full,
+    "y_std": y_std_full,
+    "target_cols": target_cols,
+}
+torch.save(final_artifacts, "liability_predictor_final.pt")
+print("Saved: liability_predictor_final.pt")
+print("y_mean:", dict(zip(target_cols, y_mean_full)))
+print("y_std :", dict(zip(target_cols, y_std_full)))
+
+# Option A: Regression performance panel + baseline comparison
+!pip -q install scikit-learn
+
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+
+from sklearn.linear_model import Ridge
+from sklearn.ensemble import RandomForestRegressor
+from sklearn.multioutput import MultiOutputRegressor
+from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
+
+# -----------------------------
+# Helpers
+# -----------------------------
+def unz(y_z, y_mean, y_std):
+    return y_z * y_std + y_mean
+
+def regression_metrics(y_true_raw, y_pred_raw, target_cols):
+    mae  = mean_absolute_error(y_true_raw, y_pred_raw, multioutput='raw_values')
+    rmse = np.sqrt(mean_squared_error(y_true_raw, y_pred_raw, multioutput='raw_values'))
+    r2   = np.array([r2_score(y_true_raw[:,i], y_pred_raw[:,i]) for i in range(y_true_raw.shape[1])])
+
+    out = pd.DataFrame({
+        "target": target_cols,
+        "MAE": mae,
+        "RMSE": rmse,
+        "R2": r2
+    })
+    out.loc["mean"] = ["mean", mae.mean(), rmse.mean(), r2.mean()]
+    return out
+
+@torch.no_grad()
+def predict_nn_raw(loader, y_mean, y_std):
+    model.eval()
+    preds_z = []
+    trues_z = []
+    for xb, yb in loader:
+        xb = xb.to(device)
+        pred_z = model(xb).cpu().numpy()
+        preds_z.append(pred_z)
+        trues_z.append(yb.numpy())
+    preds_z = np.vstack(preds_z)
+    trues_z = np.vstack(trues_z)
+    return unz(trues_z, y_mean, y_std), unz(preds_z, y_mean, y_std)
+
+# -----------------------------
+# Prepare train/val arrays (raw y!)
+# -----------------------------
+# X is numpy (N,640); y_raw is numpy (N,4) from your Cell E
+X_train = X[train_idx]
+X_val   = X[val_idx]
+y_train_raw = y_raw[train_idx]
+y_val_raw   = y_raw[val_idx]
+
+# -----------------------------
+# Evaluate NN (your trained model already loaded best_state in Cell F)
+# -----------------------------
+y_val_true_nn, y_val_pred_nn = predict_nn_raw(val_loader, y_mean, y_std)
+nn_table = regression_metrics(y_val_true_nn, y_val_pred_nn, target_cols)
+print("\nNeural Network (val):")
+display(nn_table)
+
+# -----------------------------
+# Baselines
+# -----------------------------
+ridge = MultiOutputRegressor(Ridge(alpha=10.0, random_state=0))
+ridge.fit(X_train, y_train_raw)
+y_pred_ridge = ridge.predict(X_val)
+ridge_table = regression_metrics(y_val_raw, y_pred_ridge, target_cols)
+
+rf = MultiOutputRegressor(RandomForestRegressor(
+    n_estimators=600, random_state=0, min_samples_leaf=2
+))
+rf.fit(X_train, y_train_raw)
+y_pred_rf = rf.predict(X_val)
+rf_table = regression_metrics(y_val_raw, y_pred_rf, target_cols)
+
+# -----------------------------
+# Comparison summary (mean row only)
+# -----------------------------
+summary = pd.DataFrame({
+    "Model": ["NeuralNet", "Ridge", "RandomForest"],
+    "MAE_mean": [nn_table.loc["mean","MAE"], ridge_table.loc["mean","MAE"], rf_table.loc["mean","MAE"]],
+    "RMSE_mean": [nn_table.loc["mean","RMSE"], ridge_table.loc["mean","RMSE"], rf_table.loc["mean","RMSE"]],
+    "R2_mean": [nn_table.loc["mean","R2"], ridge_table.loc["mean","R2"], rf_table.loc["mean","R2"]],
+})
+print("\nModel comparison (val, mean across targets):")
+display(summary)
+
+# -----------------------------
+# Predicted vs True plots for NN (per target)
+# -----------------------------
+for i, t in enumerate(target_cols):
+    plt.figure()
+    plt.scatter(y_val_true_nn[:, i], y_val_pred_nn[:, i])
+    plt.xlabel(f"True {t} (raw)")
+    plt.ylabel(f"Predicted {t} (raw)")
+    plt.title(f"NN: Predicted vs True ({t})")
+    plt.show()
+
+# -----------------------------
+# Residual histogram (per target)
+# -----------------------------
+res = y_val_pred_nn - y_val_true_nn
+for i, t in enumerate(target_cols):
+    plt.figure()
+    plt.hist(res[:, i], bins=12)
+    plt.xlabel(f"Residual (Pred - True) for {t}")
+    plt.ylabel("Count")
+    plt.title(f"NN residuals ({t})")
+    plt.show()
+
+
+
+# Cell G: Plot graphs to visualise loss and accuracy
+import numpy as np
+import matplotlib.pyplot as plt
+import torch
+
+print("y_mean:", y_mean)
+print("y_std:", y_std)
+
+model.eval()
+
+y_true_z_list = []
+y_pred_z_list = []
+
+with torch.no_grad():
+    for xb, yb in val_loader:
+        xb = xb.to(device)
+
+        pred_z = model(xb).cpu().numpy()   # (batch, 4) in z-space
+        y_pred_z_list.append(pred_z)
+
+        y_true_z_list.append(yb.numpy())             # (batch, 4) in z-space
+
+y_true_z = np.vstack(y_true_z_list)
+y_pred_z = np.vstack(y_pred_z_list)
+
+# ---- Unscale HERE ----
+y_true = y_true_z * y_std + y_mean
+y_pred = y_pred_z * y_std + y_mean
+
+def pearsonr(a, b):
+    a = a - a.mean()
+    b = b - b.mean()
+    return float((a @ b) / (np.sqrt((a @ a) * (b @ b)) + 1e-12))
+
+def spearmanr(a, b):
+    ra = a.argsort().argsort().astype(float)
+    rb = b.argsort().argsort().astype(float)
+    return pearsonr(ra, rb)
+
+for j, name in enumerate(target_cols):
+    p = pearsonr(y_true[:, j], y_pred[:, j])
+    s = spearmanr(y_true[:, j], y_pred[:, j])
+
+    plt.figure()
+    plt.scatter(y_true[:, j], y_pred[:, j])
+    lo = min(y_true[:, j].min(), y_pred[:, j].min())
+    hi = max(y_true[:, j].max(), y_pred[:, j].max())
+    plt.plot([lo, hi], [lo, hi], linestyle="--")
+    plt.xlabel(f"True {name}")
+    plt.ylabel(f"Predicted {name}")
+    plt.title(f"{name} (val)  R={p:.2f}  ρ={s:.2f}")
+    plt.show()
+
+
+
+import torch
+from google.colab import files
+
+# Define the path where the model will be saved
+output_model_path = 'liability_predictor.pt'
+
+# Save the best model state dictionary
+torch.save(best_state, output_model_path)
+
+print(f"Model saved successfully to {output_model_path}")
+
+"""The model has been saved to `liability_predictor.pt` in your Colab environment. You can now download it to your local computer using the following code cell:"""
+
+# Download the saved model to your local computer
+files.download('liability_predictor.pt')