perf: Add cache for predictions

src/leaderboard.py

@@ -191,6 +191,36 @@ def _load_metrics_cached(model_name, dataset_term, horizon):
     return result
 
 
+@lru_cache(maxsize=10)
+def _load_predictions_cached(model_name, dataset_term, horizon):
+    """
+    Load and cache predictions.npz for a specific (model, dataset, horizon).
+
+    predictions.npz is the heaviest file loaded in the Per Test Window tab
+    (contains quantile predictions for all series / windows / variates).
+    On HF Space the disk I/O is much slower than a local SSD, so caching
+    avoids redundant reads when cascading Gradio change events re-trigger
+    plot_window_series with identical parameters.
+
+    maxsize=10 keeps memory bounded — at most 10 (model, dataset, horizon)
+    combos stay in RAM at a time.
+
+    Returns:
+        dict with keys from the npz file (typically "predictions_quantiles"
+        and "quantile_levels"), or None if the file does not exist.
+    """
+    results_root = str(RESULTS_ROOT)
+    predictions_path = os.path.join(
+        results_root, model_name, dataset_term, horizon, "predictions.npz"
+    )
+    if not os.path.exists(predictions_path):
+        return None
+    npz = np.load(predictions_path)
+    result = {k: npz[k] for k in npz.files}
+    npz.close()
+    return result
+
+
 def load_test_windows(display_name, horizon, model_name="moirai_small", series=None, variate=None, window_id=None, parse_series=False):
     """
     Load test window results from TIME NPZ files.

src/tab.py

@@ -32,7 +32,7 @@ import gradio as gr
 from src.about import DATASET_CHOICES, ALL_MODELS, RESULTS_ROOT, FEATURES_DF, FEATURES_BOOL_DF, PATTERN_MAP
 from src.leaderboard import (get_overall_leaderboard, get_dataset_multilevel_leaderboard,
                             get_window_leaderboard, get_pattern_leaderboard, resolve_dataset_id,
-                            _get_dataset_metadata)
+                            _get_dataset_metadata, _load_predictions_cached)
 from src.about import DATASETS_DF, ALL_HORIZONS
 # get_datasets_root, get_config_root no longer needed here — handled by _get_dataset_metadata
 import numpy as np
@@ -276,20 +276,16 @@ def plot_window_series(display_name, series, variate, window_id, horizon, select
         fig.update_layout(title="Dataset not found")
         return fig, ""
 
-    predictions_path = os.path.join(results_root, model, dataset_term, horizon, "predictions.npz")
-    print(f"📁 predictions_path: {predictions_path}, exists: {os.path.exists(predictions_path)}")
-
-    if not os.path.exists(predictions_path):
-        print("❌ Predictions file not found")
+    # --- Cached predictions loading (biggest I/O in Per Test Window) ---
+    pred_data = _load_predictions_cached(model, dataset_term, horizon)
+    if pred_data is None:
+        print(f"❌ Predictions file not found for {model}/{dataset_term}/{horizon}")
         fig = go.Figure()
-        fig.update_layout(title="Predictions file not found")
+        fig.update_layout(title="Predictions file not found for this horizon")
         return fig, ""
 
-
-    predictions = np.load(predictions_path)
-    # Load pre-computed quantiles (new format only)
-    predictions_quantiles = predictions["predictions_quantiles"]  # (num_series, num_windows, 9, num_variates, prediction_length)
-    quantile_levels = predictions["quantile_levels"]  # [0.1, 0.2, ..., 0.9]
+    predictions_quantiles = pred_data["predictions_quantiles"]  # (num_series, num_windows, 9, num_variates, prediction_length)
+    quantile_levels = pred_data["quantile_levels"]  # [0.1, 0.2, ..., 0.9]
 
     # Load prediction scale factor from config.json (for float16 overflow prevention)
     model_config_path = os.path.join(results_root, model, dataset_term, horizon, "config.json")
@@ -300,6 +296,7 @@ def plot_window_series(display_name, series, variate, window_id, horizon, select
             prediction_scale_factor = model_config.get("prediction_scale_factor", 1.0)
     if prediction_scale_factor != 1.0:
         print(f"📊 Applying inverse scale factor: {prediction_scale_factor}")
+        # Copy to avoid mutating the cached array
         predictions_quantiles = predictions_quantiles.astype(np.float32) * prediction_scale_factor
 
     # Use cached metadata for name-to-index mappings and Dataset object
@@ -923,7 +920,20 @@ def init_per_window_tab(demo):
 
     table_window = gr.DataFrame(elem_classes="custom-table", interactive=False)
 
-    # When dataset changes: first update horizon choices, then update dropdowns
+    # ── Shared input / output lists ────────────────────────────────────
+    _plot_in  = [dataset_dropdown, series_dropdown, variate_dropdown,
+                 window_dropdown, horizons, quantiles, model]
+    _plot_out = [ts_visualization, prediction_info]
+    _tbl_in   = [dataset_dropdown, series_dropdown, variate_dropdown,
+                 window_dropdown, horizons]
+    _tbl_out  = table_window
+
+    # ── dataset changes ─────────────────────────────────────────────────
+    # Chain: update horizons → update dropdowns → refresh plot → refresh table.
+    # The chain already calls plot & table at the end, so we do NOT bind
+    # separate .change() on series/variate/window for this trigger path —
+    # otherwise updating the 3 dropdowns cascades into 3 extra duplicate
+    # plot_window_series calls (the #1 cause of slowness on HF Space).
     dataset_dropdown.change(
         fn=update_horizon_choices,
         inputs=[dataset_dropdown],
@@ -933,62 +943,44 @@ def init_per_window_tab(demo):
         inputs=[dataset_dropdown, horizons],
         outputs=[series_dropdown, variate_dropdown, window_dropdown],
     ).then(
-        # After dropdowns are updated, refresh the visualization and table
-        fn=plot_window_series,
-        inputs=[dataset_dropdown, series_dropdown, variate_dropdown, window_dropdown, horizons, quantiles, model],
-        outputs=[ts_visualization, prediction_info]
+        fn=plot_window_series,  inputs=_plot_in, outputs=_plot_out,
     ).then(
-        fn=get_window_leaderboard,
-        inputs=[dataset_dropdown, series_dropdown, variate_dropdown, window_dropdown, horizons],
-        outputs=table_window
+        fn=get_window_leaderboard, inputs=_tbl_in, outputs=_tbl_out,
     )
 
-    # When horizon changes: update dropdowns, then refresh visualization
+    # ── horizon changes ─────────────────────────────────────────────────
     horizons.change(
         fn=update_series_variate_and_window,
         inputs=[dataset_dropdown, horizons],
         outputs=[series_dropdown, variate_dropdown, window_dropdown],
     ).then(
-        fn=plot_window_series,
-        inputs=[dataset_dropdown, series_dropdown, variate_dropdown, window_dropdown, horizons, quantiles, model],
-        outputs=[ts_visualization, prediction_info]
+        fn=plot_window_series,  inputs=_plot_in, outputs=_plot_out,
     ).then(
-        fn=get_window_leaderboard,
-        inputs=[dataset_dropdown, series_dropdown, variate_dropdown, window_dropdown, horizons],
-        outputs=table_window
+        fn=get_window_leaderboard, inputs=_tbl_in, outputs=_tbl_out,
     )
 
-    # For series, variate, window changes - update visualization and table
+    # ── series / variate / window manual changes ────────────────────────
+    # Use a single .then() chain per dropdown so each user-initiated
+    # change fires plot + table exactly ONCE instead of 2 separate events.
     for comp in [series_dropdown, variate_dropdown, window_dropdown]:
         comp.change(
-            fn=get_window_leaderboard,
-            inputs=[dataset_dropdown, series_dropdown, variate_dropdown, window_dropdown, horizons],
-            outputs=table_window
-        )
-        comp.change(
-            fn=plot_window_series,
-            inputs=[dataset_dropdown, series_dropdown, variate_dropdown, window_dropdown, horizons, quantiles, model],
-            outputs=[ts_visualization, prediction_info]
+            fn=plot_window_series,  inputs=_plot_in, outputs=_plot_out,
+        ).then(
+            fn=get_window_leaderboard, inputs=_tbl_in, outputs=_tbl_out,
         )
 
-    # For quantiles and model changes - only update visualization (no table change needed)
+    # ── quantiles / model changes ───────────────────────────────────────
     for comp in [quantiles, model]:
         comp.change(
-            fn=plot_window_series,
-            inputs=[dataset_dropdown, series_dropdown, variate_dropdown, window_dropdown, horizons, quantiles, model],
-            outputs=[ts_visualization, prediction_info]
+            fn=plot_window_series, inputs=_plot_in, outputs=_plot_out,
         )
 
-    # Load initial visualization and table on page load
+    # ── initial page load ───────────────────────────────────────────────
     demo.load(
-        fn=plot_window_series,
-        inputs=[dataset_dropdown, series_dropdown, variate_dropdown, window_dropdown, horizons, quantiles, model],
-        outputs=[ts_visualization, prediction_info]
+        fn=plot_window_series,  inputs=_plot_in, outputs=_plot_out,
     )
     demo.load(
-        fn=get_window_leaderboard,
-        inputs=[dataset_dropdown, series_dropdown, variate_dropdown, window_dropdown, horizons],
-        outputs=table_window
+        fn=get_window_leaderboard, inputs=_tbl_in, outputs=_tbl_out,
     )
 
     # CSV Export