Update app.py

app.py

@@ -1,4 +1,3 @@
-import io
 import numpy as np
 import pandas as pd
 import streamlit as st
@@ -7,102 +6,131 @@ from chronos import BaseChronosPipeline
 
 st.set_page_config(page_title="Chronos-Bolt Zero-Shot Forecast", layout="centered")
 st.title("Chronos-Bolt Zero-Shot Forecast")
-st.caption("Paste a series or pull a ticker and get a probabilistic forecast (q10 / q50 / q90).")
+st.caption("Zero-shot probabilistic forecasting (q10/q50/q90) using amazon/chronos-bolt-* models.")
 
-# --------- SETTINGS ---------
+# -------------------- Model options --------------------
 MODEL_CHOICES = {
-    "Bolt Mini (fast, CPU-friendly)": "amazon/chronos-bolt-mini",
-    "Bolt Small (better, may need GPU)": "amazon/chronos-bolt-small",
+    "Bolt Mini (CPU-friendly)": "amazon/chronos-bolt-mini",
+    "Bolt Small (better; GPU if available)": "amazon/chronos-bolt-small",
 }
-DEFAULT_MODEL = "Bolt Mini (fast, CPU-friendly)"
 
-# --------- HELPERS ---------
 @st.cache_resource(show_spinner=True)
 def load_pipeline(model_id: str):
     device = "cuda" if torch.cuda.is_available() else "cpu"
     dtype = torch.bfloat16 if device == "cuda" else torch.float32
-    return BaseChronosPipeline.from_pretrained(
-        model_id, device_map=device, torch_dtype=dtype
-    )
+    return BaseChronosPipeline.from_pretrained(model_id, device_map=device, torch_dtype=dtype)
+
+# -------------------- Data loaders (always return 1-D) --------------------
+def _force_1d(a):
+    a = pd.Series(a, dtype="float32").replace([np.inf, -np.inf], np.nan).dropna()
+    return a.to_numpy().reshape(-1)
 
 @st.cache_data(show_spinner=False)
 def load_ticker_series(ticker: str, period: str = "2y"):
     import yfinance as yf
-    df = yf.download(ticker, period=period, interval="1d", progress=False)
-    close = df["Close"].dropna().astype("float32").values
-    return close
+    df = yf.download(ticker, period=period, interval="1d", auto_adjust=True, progress=False)
+    if df.empty:
+        return np.asarray([], dtype="float32")
+    close = df["Close"]
+    if isinstance(close, pd.DataFrame):  # handle rare multi-index cases
+        close = close.iloc[:, 0]
+    return _force_1d(close)
 
 def parse_pasted_series(txt: str):
+    import re
+    toks = re.split(r"[,\s]+", txt.strip())
     vals = []
-    for line in txt.splitlines():
-        line = line.strip().replace(",", "")
-        if line:
-            try:
-                vals.append(float(line))
-            except:
-                pass
-    return np.asarray(vals, dtype="float32")
-
-# --------- UI ---------
-col1, col2 = st.columns(2)
-with col1:
+    for t in toks:
+        if not t: 
+            continue
+        try:
+            vals.append(float(t))
+        except:
+            pass
+    return _force_1d(vals)
+
+def load_csv_series(file, column=None):
+    df = pd.read_csv(file)
+    if column is None:
+        num_cols = [c for c in df.columns if np.issubdtype(df[c].dtype, np.number)]
+        column = num_cols[0] if num_cols else None
+    if column is None:
+        return np.asarray([], dtype="float32"), df, None
+    return _force_1d(df[column]), df, column
+
+# -------------------- UI --------------------
+c1, c2 = st.columns(2)
+with c1:
     model_label = st.selectbox("Model", list(MODEL_CHOICES.keys()), index=0)
-with col2:
-    pred_len = st.number_input("Prediction length (steps)", 1, 200, 30)
+with c2:
+    pred_len = st.number_input("Prediction length (steps)", 1, 365, 30)
 
-src = st.radio("Data source", ["Ticker (yfinance)", "Paste numbers"], horizontal=True)
+src = st.radio("Data source", ["Ticker (yfinance)", "Paste numbers", "Upload CSV"], horizontal=True)
 
 series = None
 if src == "Ticker (yfinance)":
-    t1, t2 = st.columns([2,1])
+    t1, t2 = st.columns([2, 1])
     with t1:
         ticker = st.text_input("Ticker (e.g., AAPL, SPY, BTC-USD)", "AAPL")
     with t2:
         period = st.selectbox("History window", ["6mo", "1y", "2y", "5y"], index=2)
     if st.button("Load data"):
-        with st.spinner("Downloading..."):
-            series = load_ticker_series(ticker.strip(), period)
-            if series.size == 0:
-                st.error("No series returned. Try a different ticker or window.")
-                series = None
+        series = load_ticker_series(ticker.strip(), period)
+        if series.size == 0:
+            st.error("No data returned. Try another ticker/window.")
+elif src == "Paste numbers":
+    txt = st.text_area("One value per line (or comma/space separated)", "1\n2\n3\n4\n5\n6\n7\n8\n9\n10")
+    if st.button("Use pasted data"):
+        series = parse_pasted_series(txt)
 else:
-    txt = st.text_area("One value per line", "1\n2\n3\n4\n5\n6\n7\n8\n9\n10")
-    series = parse_pasted_series(txt)
+    uploaded = st.file_uploader("Upload CSV", type=["csv"])
+    if uploaded is not None:
+        df = pd.read_csv(uploaded)
+        numeric_cols = [c for c in df.columns if np.issubdtype(df[c].dtype, np.number)]
+        col = st.selectbox("Pick numeric column", numeric_cols) if numeric_cols else None
+        if st.button("Load CSV column") and col:
+            series, _, _ = load_csv_series(uploaded, column=col)
+        elif uploaded and not numeric_cols:
+            st.error("No numeric columns found in CSV.")
 
+# -------------------- Plot + Forecast --------------------
 if series is not None and series.size > 5:
     st.write(f"Loaded {series.size} points.")
-    st.line_chart(pd.DataFrame({"value": series}))
+    st.line_chart(pd.DataFrame(series, columns=["value"]))  # always 1-D -> no error
 
     if st.button("Forecast"):
         with st.spinner("Running Chronos-Bolt..."):
-            model_id = MODEL_CHOICES[model_label]
-            pipe = load_pipeline(model_id)
-
+            pipe = load_pipeline(MODEL_CHOICES[model_label])
             ctx = torch.tensor(series, dtype=torch.float32)
             q_levels = [0.10, 0.50, 0.90]
-            # predict_quantiles returns (quantiles, mean)
+
             quantiles, mean = pipe.predict_quantiles(
                 context=ctx,
                 prediction_length=int(pred_len),
                 quantile_levels=q_levels,
             )
 
-            q = quantiles[0].cpu().numpy()  # shape [pred_len, 3]
-            lo, med, hi = q[:, 0], q[:, 1], q[:, 2]
+            q_np = quantiles[0].cpu().numpy()  # shape [pred_len, 3]
+            lo, med, hi = q_np[:, 0], q_np[:, 1], q_np[:, 2]
 
-            # Plot
             import matplotlib.pyplot as plt
             hist_x = np.arange(len(series))
-            fut_x = np.arange(len(series), len(series) + pred_len)
+            fut_x = np.arange(len(series), len(series) + int(pred_len))
 
             fig = plt.figure(figsize=(9, 4.5))
             plt.plot(hist_x, series, label="history")
             plt.plot(fut_x, med, label="median forecast")
-            plt.fill_between(fut_x, lo, hi, alpha=0.3, label="q10–q90 interval")
+            plt.fill_between(fut_x, lo, hi, alpha=0.3, label="q10–q90 band")
             plt.legend()
             plt.grid(True, alpha=0.3)
             st.pyplot(fig)
 
-            st.success("Done.")
+            out = pd.DataFrame({"t": fut_x, "q10": lo, "q50": med, "q90": hi})
+            st.download_button(
+                "Download forecast CSV",
+                out.to_csv(index=False).encode("utf-8"),
+                file_name="chronos_forecast.csv",
+                mime="text/csv",
+            )
 else:
-    st.info("Load a ticker or paste at least ~20 numbers to begin.")
+    st.info("Load a ticker, paste values, or upload a CSV to begin.")