add prob_% column (probability of predicted class)

app.py

@@ -106,14 +106,21 @@ def _run(train_df, test_df, target_col, task, n_estimators, progress=gr.Progress
         clf = TabFMClassifier(model=model, random_state=42, n_estimators=int(n_estimators))
         clf.fit(Xtr, ytr)
         preds = clf.predict(Xte)
+        proba = clf.predict_proba(Xte)
+        classes = list(clf.classes_)
+        pred_idx = [classes.index(p) for p in preds]
+        prob_pred = np.array([proba[i, pred_idx[i]] for i in range(len(preds))])
     else:
         reg = TabFMRegressor(model=model, random_state=42, n_estimators=int(n_estimators))
         reg.fit(Xtr, ytr)
         preds = reg.predict(Xte)
+        prob_pred = None
     infer_s = time.time() - t0
 
     out = test_df.copy()
     out["tabfm_prediction"] = preds
+    if prob_pred is not None:
+        out["prob_%"] = (prob_pred * 100).round(1)
 
     summary = f"Loaded weights in {load_s:.1f}s (cached after first run). Inference took {infer_s:.1f}s for {len(Xte)} test rows.\n"
     if has_target_in_test:
@@ -187,14 +194,17 @@ def _run_stock_backtest(dataset_key, ticker, n_estimators, progress=gr.Progress(
 
     progress(0.6, desc=f"Predicting {len(Xva)} test rows for {ticker}...")
     clf.fit(data["Xtr"], ytr_bin)
-    pred_bin = clf.predict(Xva)
-    pred = np.where(pred_bin == 1, "BUY", "SELL")
+    proba   = clf.predict_proba(Xva)
+    pred_bin = (proba[:, 1] > proba[:, 0]).astype(int)
+    pred    = np.where(pred_bin == 1, "BUY", "SELL")
+    prob_pred = np.where(pred_bin == 1, proba[:, 1], proba[:, 0])
 
     acc = (pred == yva_str).mean()
     out = pd.DataFrame({
         "date":                    [m["date"] for m in meta],
         "actual_label":            yva_str,
         "tabfm_prediction":        pred,
+        "prob_%":                  (prob_pred * 100).round(1),
         "actual_5d_fwd_return_%":  [round(m["fwd_return"] * 100, 2) for m in meta],
         "correct":                 pred == yva_str,
     }).sort_values("date")