Update src/streamlit_app.py

src/streamlit_app.py

@@ -20,48 +20,25 @@ if str(ROOT) not in sys.path:
     sys.path.insert(0, str(ROOT))
 
 from src.survival_utils import prepare_cox_df, fit_cox, make_patient_design_row, predict_patient_survival
-
-COX_TRAIN_PATH = ROOT / "assets" / "cox_training.parquet"
+from src.cox_persist import load_cox_artifacts, ensure_cox_artifacts_available
 
 @st.cache_resource
-def get_or_train_cox():
+def get_or_load_cox():
     """
-    Train Cox once per app start (cached). No persistence needed across restarts.
+    Load Cox artifacts trained in Bulk page.
+    If running on Spaces without persistence, download from HF dataset repo when missing.
     """
-    if not COX_TRAIN_PATH.exists():
-        raise FileNotFoundError(f"Missing Cox training file: {COX_TRAIN_PATH}")
-
-    df_train = pd.read_parquet(COX_TRAIN_PATH)
-
-    # REQUIRED columns must exist in df_train:
-    # OS_time_days, Event_clean, Predicted_GVHD_Risk + your covariates
-    covariates = [
-        "R_Age_at_transplant_cutoff18",
-        "HLA_total_mismatch",
-        "Hematological Diagnosis_Grouped",
-        "Donor_relation to recepient",
-        "Source of cells",
-        "Donor_type",
-        "Conditioning_intensity",
-        "GVHD_Prophylaxis_Cat",
-    ]
-
-    # df_train already should contain Predicted_GVHD_Risk.
-    # If it doesn't, you must generate it offline and store it in this parquet.
-    preds = df_train["Predicted_GVHD_Risk"].values
-
-    df_cox, design_cols, cat_cols = prepare_cox_df(df_train, preds, covariates)
-
-    cph = fit_cox(df_cox)
-    meta = {
-        "n": len(df_cox),
-        "n_events": int(df_cox["Event_clean"].sum()),
-        "c_index": float(cph.concordance_index_),
-        "covariates": covariates,
-    }
+    # Ensure artifacts exist locally (downloads to /tmp/saved_models by default)
+    ensure_cox_artifacts_available()  # downloads cox_os_model.joblib + cox_os_meta.json
+
+    payload, meta = load_cox_artifacts(prefix="cox_os")  # loads from DEFAULT_DIR (/tmp/saved_models)
+    cph = payload["cph"]
+    design_cols = payload["design_cols"]
+    cat_cols = payload["cat_cols"]
     return cph, design_cols, cat_cols, meta
 
 
+
 # --- Country options for UI (pycountry backed) ---
 try:
     import pycountry
@@ -447,7 +424,8 @@ if submitted:
     
     # ---- load saved Cox artifacts (trained once in Bulk page) ----
     try:
-        cph, design_cols, cat_cols, meta = get_or_train_cox()
+        cph, design_cols, cat_cols, meta = get_or_load_cox()
+        st.success("Cox model ready (loaded from saved artifacts).")
         covariates = meta["covariates"]
         st.success("Cox model ready (trained on app start).")
         st.caption(f"N={meta['n']} | events={meta['n_events']} | C-index={meta['c_index']:.3f}")