Update app.py

app.py

@@ -169,7 +169,14 @@ def run_forecast(
     nixtla_api_key,
     xcols,
     future_x_file,
+    last_eval,
+    last_fig,
+    last_future,
+    last_files,
 ):
+    def keep(msg):
+        return last_eval, last_fig, last_future, msg, last_files, last_eval, last_fig, last_future, last_files
+
     try:
         season_length = _to_int(season_length, "season_length")
         cv_h = _to_int(cv_h, "cv_h")
@@ -179,7 +186,7 @@ def run_forecast(
 
         df, candidate_xcols, msg = load_training_data(train_file)
         if df is None:
-            return None, None, None, msg, None
+            return keep(msg)
 
         if use_timegpt and (not xcols) and candidate_xcols:
             xcols = list(candidate_xcols)
@@ -220,15 +227,17 @@ def run_forecast(
                 uni_models.append(m)
 
         needs_future_x = use_exog and (bool(exog_models) or use_timegpt)
-        X_future = None
         X_future_h = None
         df_exog = None
 
         if needs_future_x:
             X_future, xmsg = load_future_x(future_x_file, xcols)
             if X_future is None:
-                return None, None, None, xmsg, None
-            X_future_h = align_future_x_to_horizon(df_y, X_future, xcols, future_h)
+                return keep(xmsg)
+            try:
+                X_future_h = align_future_x_to_horizon(df_y, X_future, xcols, future_h)
+            except Exception as e:
+                return keep(f"Future predictors invalid: {type(e).__name__}: {e}")
             df_exog = df[["unique_id", "ds", "y"] + list(xcols)].copy()
             for c in xcols:
                 if not pd.api.types.is_numeric_dtype(df_exog[c]):
@@ -247,18 +256,26 @@ def run_forecast(
             future_parts.append(fc_u)
 
         if exog_models:
+            if not needs_future_x or X_future_h is None:
+                return keep("Predictors selected for an exogenous model, but X_df is missing/invalid.")
             sf_x = StatsForecast(models=exog_models, freq=freq, n_jobs=-1)
             fc_x = sf_x.forecast(df=df_exog, h=future_h, X_df=X_future_h)
             future_parts.append(fc_x)
 
         if use_timegpt:
             if not nixtla_api_key:
-                return None, None, None, "TimeGPT selected but no Nixtla API key provided.", None
-            from nixtla import NixtlaClient
-            client = NixtlaClient(api_key=nixtla_api_key)
+                return keep("TimeGPT selected but no Nixtla API key provided.")
+            try:
+                from nixtla import NixtlaClient
+            except Exception as e:
+                return keep(f"TimeGPT selected but nixtla import failed: {type(e).__name__}: {e}")
 
+            client = NixtlaClient(api_key=nixtla_api_key)
             df_tgpt = df_exog if use_exog else df_y
 
+            if use_exog and X_future_h is None:
+                return keep("TimeGPT predictors detected/selected, but X_df is missing/invalid.")
+
             tgpt_fc = client.forecast(
                 df=df_tgpt,
                 h=future_h,
@@ -279,7 +296,7 @@ def run_forecast(
             future_parts.append(tgpt)
 
         if not future_parts and not eval_parts:
-            return None, None, None, "No models selected.", None
+            return keep("No models selected.")
 
         eval_df = pd.concat(eval_parts, ignore_index=True) if eval_parts else None
         cv_df = pd.concat(cv_parts, ignore_index=True) if cv_parts else None
@@ -291,21 +308,19 @@ def run_forecast(
         future_df = future_df.sort_values(["unique_id", "ds"]).reset_index(drop=True)
 
         fig = create_future_plot(future_df, df_y, title="Forecast")
-
         files = export_results(eval_df, cv_df, future_df)
 
         status = msg
         if use_timegpt and use_exog:
             status += f" | TimeGPT using predictors: {', '.join(xcols)}"
         elif use_exog:
-            status += f" | Predictors selected: {', '.join(xcols)}"
+            status += f" | Predictors used: {', '.join(xcols)}"
 
-        return eval_df, fig, future_df, status, files
+        return eval_df, fig, future_df, status, files, eval_df, fig, future_df, files
 
     except Exception as e:
-        tb = traceback.format_exc()
-        print(tb)
-        return None, None, None, f"Run failed: {type(e).__name__}: {e}", None
+        print(traceback.format_exc())
+        return keep(f"Run failed: {type(e).__name__}: {e}")
 
 
 def on_train_upload(file_obj):
@@ -353,6 +368,11 @@ Training CSV required columns: unique_id, ds, y
         use_timegpt = gr.Checkbox(value=False, label="Enable TimeGPT")
         nixtla_api_key = gr.Textbox(label="Nixtla API Key", type="password")
 
+    last_eval = gr.State(None)
+    last_fig = gr.State(None)
+    last_future = gr.State(None)
+    last_files = gr.State([])
+
     run_btn = gr.Button("Run Forecast", variant="primary")
 
     eval_df_out = gr.Dataframe(label="Evaluation metrics")
@@ -384,8 +404,22 @@ Training CSV required columns: unique_id, ds, y
             nixtla_api_key,
             xcols,
             future_x_file,
+            last_eval,
+            last_fig,
+            last_future,
+            last_files,
+        ],
+        outputs=[
+            eval_df_out,
+            plot_out,
+            future_df_out,
+            status,
+            files_out,
+            last_eval,
+            last_fig,
+            last_future,
+            last_files,
         ],
-        outputs=[eval_df_out, plot_out, future_df_out, status, files_out],
     )
 
 demo.queue().launch()