Update app.py

app.py

@@ -87,196 +87,14 @@ Xs_te = scaler.transform(X_te).astype(np.float32)
 import joblib
 joblib.dump(scaler, "mtl_scaler.joblib")
 
-#-- 6) tensors --#
-
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-def tt(a, dtype=torch.float32): return torch.from_numpy(a).to(device).to(dtype)
-
-Xt_tr, Xt_va, Xt_te = tt(Xs_tr), tt(Xs_va), tt(Xs_te)
-yI_tr, yI_va, yI_te = tt(I_tr), tt(I_va), tt(I_te)
-yD_tr, yD_va, yD_te = tt(D_tr), tt(D_va), tt(D_te)
-
-if H_tr_all is not None:
-    H_tr_all_t = tt(H_tr_all)
-    H_va_all_t = tt(H_va_all)
-    H_te_all_t = tt(H_te_all)
-    mH_tr_t = torch.from_numpy(mH_tr.astype(bool)).to(device)
-    mH_va_t = torch.from_numpy(mH_va.astype(bool)).to(device)
-    mH_te_t = torch.from_numpy(mH_te.astype(bool)).to(device)
-else:
-    H_tr_all_t = H_va_all_t = H_te_all_t = None
-    mH_tr_t = mH_va_t = mH_te_t = None
-
-# ----------7) Model: Multi-Task MLP (shared trunk + 3 heads)----------
-class MTLNet(nn.Module):
-    def __init__(self, d_in=384, d_shared=384):
-        super().__init__()
-        self.trunk = nn.Sequential(
-            nn.Linear(d_in, d_shared), nn.ReLU(), nn.Dropout(0.2),
-            nn.Linear(d_shared, 192), nn.ReLU(), nn.Dropout(0.1)
-        )
-        self.head_imp = nn.Linear(192, 1)  # importance raw (clamped at inference)
-        self.head_dur = nn.Linear(192, 1)  # log-hours
-        self.head_hor = nn.Linear(192, 1)  # log-days
-
-        # learnable uncertainty weights (auto-balance task losses)
-        self.log_sigma_imp = nn.Parameter(torch.tensor(0.0))
-        self.log_sigma_dur = nn.Parameter(torch.tensor(0.0))
-        self.log_sigma_hor = nn.Parameter(torch.tensor(0.0))
-
-    def forward(self, x):
-        h = self.trunk(x)
-        return (
-            self.head_imp(h).squeeze(-1),
-            self.head_dur(h).squeeze(-1),
-            self.head_hor(h).squeeze(-1),
-        )
-
-    def multitask_loss(self, xb, yI, yD, yH=None, mH=None):
-        rI, rD, rH = self(xb)
-
-        # importance: SmoothL1 on raw scale
-        l_imp = nn.SmoothL1Loss()(rI, yI)
-        # duration: MSE on log1p(hours)
-        l_dur = nn.MSELoss()(rD, torch.log1p(yD))
-
-        loss = torch.exp(-self.log_sigma_imp)*l_imp + self.log_sigma_imp \
-             + torch.exp(-self.log_sigma_dur)*l_dur + self.log_sigma_dur
-
-        l_hor_val = None
-        if (yH is not None) and (mH is not None) and mH.any():
-            # horizon: only where label exists (mask True), MSE on log1p(days)
-            l_hor = nn.MSELoss()(rH[mH], torch.log1p(yH[mH]))
-            loss = loss + torch.exp(-self.log_sigma_hor)*l_hor + self.log_sigma_hor
-            l_hor_val = float(l_hor.item())
-        return loss, (float(l_imp.item()), float(l_dur.item()), l_hor_val)
-
-net = MTLNet(d_in=Xt_tr.shape[1]).to(device)
-opt = optim.AdamW(net.parameters(), lr=1e-3, weight_decay=1e-4)
-
-# ----------8) Helper Functions for NN ----------
-def predict_heads(Xt):
-    net.eval()
-    with torch.no_grad():
-        rI, rD, rH = net(Xt)
-        I = torch.clamp(rI, 1.0, 10.0)               # importance 1..10
-        Hh = torch.expm1(rD).clamp(0.25, 12.0)       # hours
-        Hd = torch.expm1(rH).clamp(0.0, 30.0)        # days; 0 allowed (today)
-    return I, Hh, Hd
-
-def eval_block(Xt, yI_true, yD_true, yH_true=None, mH=None):
-    I, Hh, Hd = predict_heads(Xt)
-    I_np, H_np, Hd_np = I.cpu().numpy(), Hh.cpu().numpy(), Hd.cpu().numpy()
-    maeI = mean_absolute_error(yI_true.cpu().numpy(), I_np)
-    maeD = mean_absolute_error(yD_true.cpu().numpy(), H_np)
-    rhoI = spearmanr(yI_true.cpu().numpy(), I_np).correlation if len(I_np) > 1 else float('nan')
-    rhoD = spearmanr(yD_true.cpu().numpy(), H_np).correlation if len(H_np) > 1 else float('nan')
-    out = {"maeI":maeI, "maeD":maeD, "rhoI":rhoI, "rhoD":rhoD}
-    if (yH_true is not None) and (mH is not None) and mH.any():
-        yH_np, mH_np = yH_true.cpu().numpy(), mH.cpu().numpy().astype(bool)
-        maeH = mean_absolute_error(yH_np[mH_np], Hd_np[mH_np])
-        rhoH = spearmanr(yH_np[mH_np], Hd_np[mH_np]).correlation if mH_np.sum()>1 else float('nan')
-        out.update({"maeH":maeH, "rhoH":rhoH})
-    return out
-
-# ------------------------------------------------------------
-# 6) Train (mini-batch) with LR schedule, AMP, clipping, early stop
-# ------------------------------------------------------------
-EPOCHS   = 100
-BATCH    = 32
-best_val = float("inf")
-patience = 20
-bad      = 0
-
-# (re)define optimizer if you like a lower LR for longer cosine cycles
-opt = optim.AdamW(net.parameters(), lr=3e-4, weight_decay=2e-4)
-
-# Cosine schedule with warm restarts (restart every ~40 epochs; doubles thereafter)
-sched = optim.lr_scheduler.CosineAnnealingWarmRestarts(opt, T_0=40, T_mult=2, eta_min=1e-6)
-
-# optional: also step down on plateaus (acts as a safety net)
-plateau = optim.lr_scheduler.ReduceLROnPlateau(opt, mode="min", factor=0.5, patience=6, min_lr=1e-6)
-
-scaler = torch.cuda.amp.GradScaler(enabled=torch.cuda.is_available())
-
-n_tr = Xt_tr.shape[0]
-for ep in range(1, EPOCHS + 1):
-    net.train()
-    order = torch.randperm(n_tr, device=device)
-    tot_loss = 0.0
-
-    for s in range(0, n_tr, BATCH):
-        e   = min(s + BATCH, n_tr)
-        idx = order[s:e]
-        xb, yi, yd = Xt_tr[idx], yI_tr[idx], yD_tr[idx]
-
-        if H_tr_all_t is not None:
-            yh, mh = H_tr_all_t[idx], mH_tr_t[idx]
-        else:
-            yh = mh = None
-
-        opt.zero_grad(set_to_none=True)
-
-        use_amp = torch.cuda.is_available()
-        amp_device = "cuda" if use_amp else "cpu"  # amp on CPU is a no-op fallback
-
-        with torch.amp.autocast(device_type=amp_device, enabled=use_amp):
-          loss, _ = net.multitask_loss(xb, yi, yd, yh, mh)
-
-
-        scaler.scale(loss).backward()
-        # gradient clipping for stability
-        scaler.unscale_(opt)
-        torch.nn.utils.clip_grad_norm_(net.parameters(), max_norm=1.0)
-        scaler.step(opt)
-        scaler.update()
-
-        tot_loss += float(loss.item())
-
-    # ---- validation ----
-    stats_va = eval_block(
-        Xt_va, yI_va, yD_va,
-        (H_va_all_t if H_va_all_t is not None else None),
-        (mH_va_t if mH_va_t is not None else None)
-    )
-    total_val = stats_va["maeI"] + stats_va["maeD"] + (stats_va.get("maeH", 0.0))
-
-    # step cosine scheduler every epoch
-    sched.step(ep)
-    # also step plateau on the combined val metric
-    plateau.step(total_val)
-
-    if ep % 5 == 0:
-        lr_now = opt.param_groups[0]["lr"]
-        extra = f" hor={stats_va.get('maeH', float('nan')):.3f}" if "maeH" in stats_va else ""
-
-    # ---- early stopping on summed MAE ----
-    if total_val < best_val - 1e-4:
-        best_val = total_val
-        bad = 0
-        torch.save(net.state_dict(), "mtl_net.pt")
-    else:
-        bad += 1
-        if bad >= patience:
-            break
-
-# ---- TEST with best checkpoint ----
-net.load_state_dict(torch.load("mtl_net.pt", map_location=device))
-stats_te = eval_block(
-    Xt_te, yI_te, yD_te,
-    (H_te_all_t if H_te_all_t is not None else None),
-    (mH_te_t    if mH_te_t    is not None else None)
-)
-
-# ============================================================
-# 0) Setup: device, seeds, helper
-# ============================================================
 import numpy as np, random
 import torch, torch.nn as nn, torch.optim as optim
 from scipy.stats import spearmanr
 from sklearn.preprocessing import StandardScaler
 from sklearn.metrics import mean_absolute_error
 
+# 0) Setup: device, seeds, helper
+
 torch.manual_seed(42); np.random.seed(42); random.seed(42)
 if torch.cuda.is_available(): torch.cuda.manual_seed_all(42)
 torch.backends.cudnn.benchmark = True
@@ -291,58 +109,19 @@ def safe_spearman(a, b):
     r = spearmanr(a, b).correlation
     return float('nan') if r is None else float(r)
 
-# ============================================================
-# 1) Deterministic split: TRAIN=350, VAL/TEST = remainder ~50/50
-#    NOTE: We assume your arrays named Xs_tr, I_tr, D_tr are the FULL dataset.
-#          Rename locally to avoid confusion.
-# ============================================================
-X_all = X_tr            # shape: [N, d]
-I_all = I_tr             # importance targets (raw 1..10)
-D_all = D_tr             # duration targets (hours, raw)
-H_all = H_tr_all if (H_tr_all is not None) else None   # horizon targets (days, optional)
-mH_all = mH_tr   if ('mH_tr' in globals() and mH_tr is not None) else None  # mask for horizon label
-
-N = X_all.shape[0]
-rng = np.random.RandomState(42)
-idx = np.arange(N)
-rng.shuffle(idx)
-
-# Adjust split sizes to ensure val and test sets are not empty
-n_train = 280 # Reduced training set size
-n_rem   = N - n_train
-n_val   = n_rem // 2
-n_test  = n_rem - n_val
-
-i_tr = idx[:n_train]
-i_va = idx[n_train:n_train+n_val]
-i_te = idx[n_train+n_val:]
-
-# ============================================================
-# 2) Standardize X using ONLY the train set, then tensorize
-# ============================================================
-scaler = StandardScaler().fit(X_all[i_tr])
-Xn = scaler.transform(X_all)
-
-Xt_tr, Xt_va, Xt_te = tt(Xn[i_tr]), tt(Xn[i_va]), tt(Xn[i_te])
-yI_tr, yI_va, yI_te = tt(I_all[i_tr]), tt(I_all[i_va]), tt(I_all[i_te])
-yD_tr, yD_va, yD_te = tt(D_all[i_tr]), tt(D_all[i_va]), tt(D_all[i_te])
-
-if H_all is not None:
-    H_tr_all_t = tt(H_all[i_tr])
-    H_va_all_t = tt(H_all[i_va])
-    H_te_all_t = tt(H_all[i_te])
-    # masks as bool tensors on device
-    mH_tr_t = torch.from_numpy(mH_all[i_tr].astype(bool)).to(device)
-    mH_va_t = torch.from_numpy(mH_all[i_va].astype(bool)).to(device)
-    mH_te_t = torch.from_numpy(mH_all[i_te].astype(bool)).to(device)
-else:
-    H_tr_all_t = H_va_all_t = H_te_all_t = None
-    mH_tr_t = mH_va_t = mH_te_t = None
+Xt_tr, Xt_va, Xt_te = tt(Xs_tr), tt(Xs_va), tt(Xs_te)
+yI_tr, yI_va, yI_te = tt(I_tr),    tt(I_va),    tt(I_te)
+yD_tr, yD_va, yD_te = tt(D_tr),    tt(D_va),    tt(D_te)
+
+if y_hor is not None:
+    H_tr_all_t, H_va_all_t, H_te_all_t = tt(H_tr_all), tt(H_va_all), tt(H_te_all)
+    mH_tr_t = torch.from_numpy(mH_tr.astype(bool)).to(device)
+    mH_va_t = torch.from_numpy(mH_va.astype(bool)).to(device)
+    mH_te_t = torch.from_numpy(mH_te.astype(bool)).to(device)
 
-# ============================================================
 # 3) Model: Multi-Task MLP (shared trunk + 3 heads)
 #    Slightly wider trunk; textbook uncertainty weighting (0.5 factor)
-# ============================================================
+
 class MTLNet(nn.Module):
     def __init__(self, d_in, d_hid=512):
         super().__init__()
@@ -390,9 +169,8 @@ class MTLNet(nn.Module):
 
 net = MTLNet(d_in=Xt_tr.shape[1]).to(device)
 
-# ============================================================
 # 4) Prediction + Eval helpers
-# ============================================================
+
 @torch.no_grad()
 def predict_heads(Xt):
     net.eval()
@@ -421,9 +199,8 @@ def eval_block(Xt, yI_true, yD_true, yH_true=None, mH=None):
             out.update({"maeH": maeH, "rhoH": rhoH})
     return out
 
-# ============================================================
 # 5) Train loop with per-batch cosine, AMP (new API), early stop
-# ============================================================
+
 EPOCHS   = 120
 BATCH    = 64
 best_val = float("inf")
@@ -484,9 +261,7 @@ for ep in range(1, EPOCHS + 1):
         if bad >= patience:
             break
 
-# ============================================================
 # 6) TEST with best checkpoint + final confirmation
-# ============================================================
 
 net.load_state_dict(torch.load("mtl_net.pt", map_location=device))
 stats_te = eval_block(
@@ -673,8 +448,6 @@ def reorder_tasks(tasks_string, user_due_iso=None):
 
     return task_lines, due_lines_out, duration_lines, checkbox_update
 
-import re
-
 import gradio as gr # For building the interface
 
 with gr.Blocks() as demo: