Update app.py

app.py

@@ -1,3 +1,4 @@
+import os
 import io
 import math
 import tempfile
@@ -20,6 +21,8 @@ import torch.nn.functional as F
 import lightning.pytorch as pl
 from torch.utils.data import DataLoader, TensorDataset
 
+import onnx
+from onnx import external_data_helper
 import onnxruntime as ort
 
 
@@ -230,8 +233,12 @@ class OnnxWrapper(nn.Module):
         return torch.sigmoid(logits)
 
 
-def export_onnx_model(trained_model: LitClassifier, mu: np.ndarray, sd: np.ndarray, n_features: int) -> str:
-    # Build wrapper on CPU for export
+def export_onnx_model(trained_model, mu: np.ndarray, sd: np.ndarray, n_features: int) -> str:
+    """
+    Exports a SINGLE-FILE ONNX.
+    If PyTorch writes external data (onnx_path + '.data'), we merge it into the .onnx
+    so you do NOT need a separate weights file for inference.
+    """
     wrapper = OnnxWrapper(trained_model.net.cpu().eval(), mu=mu, sd=sd).eval()
 
     dummy = torch.zeros(1, n_features, dtype=torch.float32)
@@ -246,12 +253,24 @@ def export_onnx_model(trained_model: LitClassifier, mu: np.ndarray, sd: np.ndarr
         output_names=["p_up"],
         dynamic_axes={"features": {0: "batch"}, "p_up": {0: "batch"}},
         opset_version=17,
+        do_constant_folding=True,
     )
+
+    # Merge external data into the ONNX (if created)
+    data_path = onnx_path + ".data"
+    if os.path.exists(data_path):
+        m = onnx.load_model(onnx_path, load_external_data=True)
+        external_data_helper.convert_model_from_external_data(m)
+        onnx.save_model(m, onnx_path)
+        try:
+            os.remove(data_path)
+        except OSError:
+            pass
+
     return onnx_path
 
 
 def onnx_predict_probs(onnx_path: str, X: np.ndarray) -> np.ndarray:
-    # CPU provider is the most compatible for Spaces
     sess = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
     input_name = sess.get_inputs()[0].name
     out = sess.run(None, {input_name: X.astype(np.float32)})
@@ -326,7 +345,7 @@ def run_app(
     )
     trainer.fit(model, train_dataloaders=train_loader, val_dataloaders=val_loader)
 
-    # ---- Export ONNX (includes preprocessing + sigmoid)
+    # ---- Export ONNX (single-file; includes preprocessing + sigmoid)
     onnx_path = export_onnx_model(model, mu=mu, sd=sd, n_features=n_features)
 
     # ---- Inference: latest row per ticker (compare Torch vs ONNX)
@@ -335,16 +354,14 @@ def run_app(
     torch_probs_for_onnx_compare = []
     onnx_inputs = []
 
-    # Use RAW (unstandardized) latest feature row for ONNX input
     for t in tickers:
         dft_raw = df[df["ticker"] == t].sort_values("date")
         if dft_raw.empty:
             continue
         last_raw = dft_raw.iloc[-1]
-        x_raw = last_raw[feature_cols].values.astype(np.float32)  # raw features
+        x_raw = last_raw[feature_cols].values.astype(np.float32)  # raw features (ONNX expects raw)
         onnx_inputs.append(x_raw)
 
-        # Torch probability (do same preprocessing here for comparison)
         x_std = (x_raw - mu) / sd
         x_t = torch.tensor(x_std, dtype=torch.float32).unsqueeze(0)
         with torch.no_grad():
@@ -354,10 +371,9 @@ def run_app(
 
     onnx_probs = np.array([])
     if len(onnx_inputs) > 0:
-        X_onnx = np.stack(onnx_inputs, axis=0)  # shape (n_tickers, n_features)
+        X_onnx = np.stack(onnx_inputs, axis=0)
         onnx_probs = onnx_predict_probs(onnx_path, X_onnx)
 
-    # Build final table
     idx = 0
     for t in tickers:
         dft_raw = df[df["ticker"] == t].sort_values("date")
@@ -367,8 +383,12 @@ def run_app(
 
         p_torch = float(torch_probs_for_onnx_compare[idx])
         p_onnx = float(onnx_probs[idx]) if len(onnx_probs) else float("nan")
-        sig = signal_from_prob(p_onnx if not math.isnan(p_onnx) else p_torch,
-                               float(buy_threshold), float(sell_threshold))
+
+        sig = signal_from_prob(
+            p_onnx if not math.isnan(p_onnx) else p_torch,
+            float(buy_threshold),
+            float(sell_threshold),
+        )
 
         out_rows.append(
             {
@@ -390,12 +410,8 @@ def run_app(
     # Toy backtest for first ticker (val split only)
     backtest_img = None
     t0 = tickers[0]
-    d0 = df_std[(df_std["ticker"] == t0) & (df_std["split"] == "val")].sort_values("date").copy()
-    if len(d0) >= 30:
-        X0_std = d0[feature_cols].values.astype(np.float32)
-
-        # Use ONNX for backtest probability (feed RAW features to ONNX wrapper)
-        d0_raw = df[(df["ticker"] == t0) & (df["split"] == "val")].sort_values("date").copy()
+    d0_raw = df[(df["ticker"] == t0) & (df["split"] == "val")].sort_values("date").copy()
+    if len(d0_raw) >= 30:
         X0_raw = d0_raw[feature_cols].values.astype(np.float32)
         p = onnx_predict_probs(onnx_path, X0_raw)
 
@@ -417,15 +433,16 @@ def run_app(
     # Data preview + download
     export_df = df.copy()
     export_df["date"] = export_df["date"].dt.date.astype(str)
-    export_df = export_df[["date", "ticker", "split", "close", "ret_1", "ret_5", "sma_ratio", "rsi", "vol", "ret_next", "target"]]
+    export_df = export_df[
+        ["date", "ticker", "split", "close", "ret_1", "ret_5", "sma_ratio", "rsi", "vol", "ret_next", "target"]
+    ]
     preview_df = export_df.head(25).round(6)
     csv_path = save_df_to_temp_csv(export_df.round(8), prefix="signals_dataset_")
 
-    # ONNX download + simple inference snippet for students
     inference_snippet = """import numpy as np
 import onnxruntime as ort
 
-onnx_path = "signals_model.onnx"  # downloaded file
+onnx_path = "model.onnx"
 sess = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
 inp = sess.get_inputs()[0].name
 
@@ -434,7 +451,7 @@ inp = sess.get_inputs()[0].name
 x = np.array([[0.001, 0.01, 0.02, 55.0, 0.012]], dtype=np.float32)
 
 p_up = sess.run(None, {inp: x})[0]
-print("p_up:", float(p_up[0]))
+print("p_up:", float(np.array(p_up).reshape(-1)[0]))
 """
     snippet_path = save_bytes_to_temp_file(inference_snippet.encode("utf-8"), suffix=".py", prefix="onnx_inference_example_")
 
@@ -443,7 +460,7 @@ print("p_up:", float(p_up[0]))
         f"Tickers requested (max 10): {', '.join(tickers)}",
         f"Rows: {len(export_df)} | train={int((export_df['split']=='train').sum())} | val={int((export_df['split']=='val').sum())}",
         f"BUY if p_up >= {buy_threshold:.2f} | SELL if p_up <= {sell_threshold:.2f}",
-        "ONNX export: wrapper includes preprocessing + sigmoid, so ONNX input is RAW features.",
+        "ONNX export: wrapper includes preprocessing + sigmoid; exported ONNX is SINGLE-FILE (no .onnx.data).",
     ]
     if failed:
         summary_lines.append(f"Tickers with no data / error: {', '.join(failed)}")
@@ -481,18 +498,18 @@ with gr.Blocks(title="Educational Stock Signals (Lightning + ONNX)") as demo:
     run_btn = gr.Button("Train + Export ONNX + Infer", variant="primary")
 
     with gr.Tab("Signals (Torch vs ONNX)"):
-        signals_out = gr.Dataframe(label="Signals (educational) + Torch/ONNX comparison", wrap=True)
+        signals_out = gr.Dataframe(label="Signals + Torch/ONNX comparison", wrap=True)
 
     with gr.Tab("Backtest (toy)"):
         backtest_out = gr.Image(label="Toy equity curve (val only; first ticker) using ONNX probs", type="numpy")
 
     with gr.Tab("Data"):
         preview_out = gr.Dataframe(label="Feature dataset preview", wrap=True)
-        download_out = gr.File(label="Download full dataset CSV (features + target + split)")
+        download_out = gr.File(label="Download dataset CSV (features + target + split)")
         summary_out = gr.Textbox(label="Run summary", lines=10)
 
     with gr.Tab("ONNX Export"):
-        onnx_file = gr.File(label="Download ONNX model (.onnx)")
+        onnx_file = gr.File(label="Download ONNX model (.onnx) — single-file")
         onnx_example = gr.File(label="Download ONNX inference example (.py)")
 
     run_btn.click(