Update app.py

app.py

@@ -1,51 +1,128 @@
 import os
-import tempfile
+import traceback
+from pathlib import Path
+
+import gradio as gr
 import numpy as np
-import torch
-import torch.nn as nn
-import onnx
-from onnx import external_data_helper
-
-class OnnxWrapper(nn.Module):
-    def __init__(self, net: nn.Module, mu: np.ndarray, sd: np.ndarray):
-        super().__init__()
-        self.net = net
-        self.register_buffer("mu", torch.tensor(mu, dtype=torch.float32))
-        self.register_buffer("sd", torch.tensor(sd, dtype=torch.float32))
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = (x - self.mu) / self.sd
-        logits = self.net(x).squeeze(-1)
-        return torch.sigmoid(logits)
-
-def export_onnx_model(trained_model, mu: np.ndarray, sd: np.ndarray, n_features: int) -> str:
-    # trained_model is your LightningModule; we export its .net
-    wrapper = OnnxWrapper(trained_model.net.cpu().eval(), mu=mu, sd=sd).eval()
-    dummy = torch.zeros(1, n_features, dtype=torch.float32)
-
-    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".onnx", prefix="model_")
-    onnx_path = tmp.name
-
-    torch.onnx.export(
-        wrapper,
-        dummy,
-        onnx_path,
-        input_names=["features"],
-        output_names=["p_up"],
-        dynamic_axes={"features": {0: "batch"}, "p_up": {0: "batch"}},
-        opset_version=17,
-        do_constant_folding=True,
+import onnxruntime as ort
+
+MODEL_PATH = Path("model.onnx")
+
+# Lazy-loaded ORT session cache
+_SESSION = None
+_INPUT_NAME = None
+
+
+def _load_session():
+    """Load ONNX Runtime session only when needed (prevents startup crash)."""
+    global _SESSION, _INPUT_NAME
+
+    if _SESSION is not None:
+        return _SESSION, _INPUT_NAME
+
+    if not MODEL_PATH.exists():
+        raise FileNotFoundError(
+            "model.onnx not found in the Space root. "
+            "Upload your ONNX file and name it exactly: model.onnx"
+        )
+
+    # CPU provider is the most compatible on Spaces
+    sess = ort.InferenceSession(str(MODEL_PATH), providers=["CPUExecutionProvider"])
+    inp_name = sess.get_inputs()[0].name
+
+    _SESSION = sess
+    _INPUT_NAME = inp_name
+    return _SESSION, _INPUT_NAME
+
+
+def _parse_vector(text: str) -> np.ndarray:
+    """
+    Parse a comma/space separated vector like:
+      "0.1, 0.2, 0.3"
+      "0.1 0.2 0.3"
+    Returns shape (1, n_features)
+    """
+    if not text or not text.strip():
+        raise ValueError("Vector input is empty.")
+    parts = [p for p in text.replace(",", " ").split() if p.strip()]
+    vals = [float(p) for p in parts]
+    x = np.array([vals], dtype=np.float32)
+    return x
+
+
+def predict_5(ret_1, ret_5, sma_ratio, rsi, vol):
+    """
+    For your 5-feature model wrapper:
+      [ret_1, ret_5, sma_ratio, rsi, vol]
+    """
+    try:
+        sess, inp_name = _load_session()
+        x = np.array([[ret_1, ret_5, sma_ratio, rsi, vol]], dtype=np.float32)
+        y = sess.run(None, {inp_name: x})[0]
+        y = np.array(y).reshape(-1)[0]
+        return float(y), "OK"
+    except Exception:
+        return None, traceback.format_exc()
+
+
+def predict_vector(vec_text: str):
+    """Generic vector inference for any ONNX model expecting [batch, features]."""
+    try:
+        sess, inp_name = _load_session()
+        x = _parse_vector(vec_text)
+        y = sess.run(None, {inp_name: x})[0]
+        y = np.array(y).reshape(-1)
+        # show first value for convenience, but also return full output
+        first = float(y[0]) if y.size else None
+        return first, y.tolist(), "OK"
+    except Exception:
+        return None, None, traceback.format_exc()
+
+
+with gr.Blocks(title="ONNX Inference Only") as demo:
+    gr.Markdown(
+        """
+# ONNX Inference Only (No training / no data)
+
+Place your model in the Space root as **`model.onnx`**.
+
+⚠️ If your ONNX was exported with external weights, you must also upload the referenced
+`*.onnx.data` file into the same folder — OR re-export as a single-file ONNX.
+"""
     )
 
-    # Merge external data back into one .onnx if needed
-    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
+    with gr.Tab("5-Feature Input (recommended for your stock model)"):
+        with gr.Row():
+            ret_1 = gr.Number(label="ret_1", value=0.001)
+            ret_5 = gr.Number(label="ret_5", value=0.01)
+            sma_ratio = gr.Number(label="sma_ratio", value=0.02)
+            rsi = gr.Number(label="rsi", value=55.0)
+            vol = gr.Number(label="vol", value=0.012)
+        btn1 = gr.Button("Run ONNX", variant="primary")
+        out1 = gr.Number(label="Model output (e.g., p_up)")
+        status1 = gr.Textbox(label="Status / Error", lines=10)
+
+        btn1.click(
+            fn=predict_5,
+            inputs=[ret_1, ret_5, sma_ratio, rsi, vol],
+            outputs=[out1, status1],
+        )
+
+    with gr.Tab("Vector Input (any feature size)"):
+        vec = gr.Textbox(
+            label="Input vector (comma or space separated)",
+            value="0.001, 0.01, 0.02, 55.0, 0.012",
+        )
+        btn2 = gr.Button("Run ONNX (vector)", variant="primary")
+        out2_first = gr.Number(label="First output value")
+        out2_full = gr.JSON(label="Full output")
+        status2 = gr.Textbox(label="Status / Error", lines=10)
+
+        btn2.click(
+            fn=predict_vector,
+            inputs=[vec],
+            outputs=[out2_first, out2_full, status2],
+        )
+
+if __name__ == "__main__":
+    demo.launch()