Delete neural_network.py

neural_network.py

@@ -1,688 +0,0 @@
-# -*- 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')