Update app.py

app.py

@@ -1,31 +1,34 @@
-import os
 import warnings
 warnings.filterwarnings("ignore")
 
 import gradio as gr
-import pandas as pd
 import numpy as np
+import pandas as pd
 import yfinance as yf
 import matplotlib.pyplot as plt
 
-import torch
-from gluonts.dataset.common import ListDataset
 from pandas.tseries.frequencies import to_offset
+from gluonts.dataset.common import ListDataset
 
-# Moirai 2.0 via Uni2TS
+# --- Moirai 2.0 via Uni2TS ---
+# Make sure your requirements install Uni2TS from GitHub:
+# git+https://github.com/SalesforceAIResearch/uni2ts.git
 try:
     from uni2ts.model.moirai2 import Moirai2Forecast, Moirai2Module
 except Exception as e:
     raise ImportError(
-        "Moirai 2.0 not found in your Uni2TS install. "
-        "Make sure requirements.txt installs Uni2TS from GitHub: "
-        "git+https://github.com/SalesforceAIResearch/uni2ts.git\n"
+        "Moirai 2.0 not found in your Uni2TS install.\n"
+        "Ensure requirements.txt includes:\n"
+        "  git+https://github.com/SalesforceAIResearch/uni2ts.git\n"
         f"Original error: {e}"
     )
 
 MODEL_ID = "Salesforce/moirai-2.0-R-small"
+DEFAULT_CONTEXT = 1680  # from Moirai examples, but we clamp to series length
 
-# ---- Model loader (one-time) ----
+# ----------------------------
+# Model loader (single instance)
+# ----------------------------
 _MODULE = None
 def load_module():
     global _MODULE
@@ -33,18 +36,25 @@ def load_module():
         _MODULE = Moirai2Module.from_pretrained(MODEL_ID)
     return _MODULE
 
-# ---- Utilities ----
+# ----------------------------
+# Shared forecasting core
+# ----------------------------
 def _future_index(last_idx: pd.Timestamp, freq: str, horizon: int) -> pd.DatetimeIndex:
-    """Create future timestamps continuing the given freq."""
     off = to_offset(freq)
     start = last_idx + off
     return pd.date_range(start=start, periods=horizon, freq=freq)
 
-def _run_forecast_on_series(y: pd.Series, freq: str, horizon: int, context_hint: int, title: str):
-    """Core forecasting routine on an indexed univariate series y with pandas freq string."""
+def _run_forecast_on_series(
+    y: pd.Series,
+    freq: str,
+    horizon: int,
+    context_hint: int,
+    title: str,
+):
     if len(y) < 50:
         raise gr.Error("Need at least 50 points to forecast.")
-    ctx = int(np.clip(context_hint or 1680, 32, len(y)))
+
+    ctx = int(np.clip(context_hint or DEFAULT_CONTEXT, 32, len(y)))
     target = y.values[-ctx:].astype(np.float32)
     start_idx = y.index[-ctx]
 
@@ -59,7 +69,7 @@ def _run_forecast_on_series(y: pd.Series, freq: str, horizon: int, context_hint:
         feat_dynamic_real_dim=0,
         past_feat_dynamic_real_dim=0,
     )
-    predictor = model.create_predictor(batch_size=32)  # device managed internally
+    predictor = model.create_predictor(batch_size=32)  # device handled internally
 
     forecast = next(iter(predictor.predict(ds)))
     if hasattr(forecast, "mean"):
@@ -73,7 +83,7 @@ def _run_forecast_on_series(y: pd.Series, freq: str, horizon: int, context_hint:
 
     yhat = np.asarray(yhat).ravel()[:horizon]
     future_idx = _future_index(y.index[-1], freq, horizon)
-    pred = pd.Series(yhat, index=future_idx, name="predicted")
+    pred = pd.Series(yhat, index=future_idx, name="prediction")
 
     # Plot
     fig = plt.figure(figsize=(10, 5))
@@ -85,9 +95,11 @@ def _run_forecast_on_series(y: pd.Series, freq: str, horizon: int, context_hint:
     out_df = pd.DataFrame({"date": pred.index, "prediction": pred.values})
     return fig, out_df
 
-# ---- Ticker path ----
+# ----------------------------
+# Ticker helpers
+# ----------------------------
 def fetch_series(ticker: str, years: int) -> pd.Series:
-    """Fetch daily close price and align to business-day frequency."""
+    """Fetch daily close prices and align to business-day frequency."""
     data = yf.download(
         ticker,
         period=f"{years}y",
@@ -115,7 +127,7 @@ def fetch_series(ticker: str, years: int) -> pd.Series:
     y.name = ticker
     y.index = pd.DatetimeIndex(y.index).tz_localize(None)
 
-    # Business-day index, forward-fill market holidays
+    # Business-day index; forward-fill holidays
     bidx = pd.bdate_range(y.index.min(), y.index.max())
     y = y.reindex(bidx).ffill()
 
@@ -132,13 +144,13 @@ def forecast_ticker(ticker: str, horizon: int, lookback_years: int, context_hint
     y = fetch_series(ticker, lookback_years)
     return _run_forecast_on_series(y, "B", horizon, context_hint, f"{ticker} — forecast (Moirai 2.0 R-small)")
 
-# ---- CSV path ----
+# ----------------------------
+# CSV helpers
+# ----------------------------
 def _read_csv_columns(file_path: str) -> pd.DataFrame:
-    # Try very tolerant CSV read
     try:
         df = pd.read_csv(file_path)
     except Exception:
-        # if it’s actually TSV or weird delimiter, try python engine
         df = pd.read_csv(file_path, sep=None, engine="python")
     return df
 
@@ -146,16 +158,16 @@ def _coerce_numeric_series(s: pd.Series) -> pd.Series:
     s = pd.to_numeric(s, errors="coerce")
     return s.dropna().astype(np.float32)
 
-def build_series_from_csv(file, value_col: str, date_col: str, freq_choice: str) -> tuple[pd.Series, str]:
+def build_series_from_csv(file, value_col: str, date_col: str, freq_choice: str):
     """
     Returns (series y with DateTimeIndex, freq string).
-    - If date_col provided: parse dates and (optionally) infer freq.
-    - If no date_col: require freq_choice != 'auto'; build synthetic dates from 2000-01-01.
+    - If date_col is provided: parse dates and infer/align frequency.
+    - If NO date_col: create a synthetic date index using freq_choice (default to 'D' if auto/blank).
     """
     if file is None:
         raise gr.Error("Please upload a CSV file.")
 
-    # Gradio v4/v5 file object compatibility
+    # Gradio file object handling (v4/v5)
     path = getattr(file, "name", None) or getattr(file, "path", None) or (file if isinstance(file, str) else None)
     if path is None:
         raise gr.Error("Could not read the uploaded file path.")
@@ -164,72 +176,75 @@ def build_series_from_csv(file, value_col: str, date_col: str, freq_choice: str)
     if df.empty:
         raise gr.Error("Uploaded file is empty.")
 
-    # Pick value column
+    # Value column selection
+    value_col = (value_col or "").strip()
     if value_col:
         if value_col not in df.columns:
             raise gr.Error(f"Value column '{value_col}' not found. Available: {list(df.columns)}")
         vals = _coerce_numeric_series(df[value_col])
     else:
-        # Try the first numeric-looking column
         numeric_cols = [c for c in df.columns if pd.api.types.is_numeric_dtype(df[c])]
-        if not numeric_cols:
-            # Coerce first column
-            vals = _coerce_numeric_series(df.iloc[:, 0])
-        else:
+        if numeric_cols:
             vals = _coerce_numeric_series(df[numeric_cols[0]])
+        else:
+            vals = _coerce_numeric_series(df.iloc[:, 0])
 
     if vals.empty or len(vals) < 10:
         raise gr.Error("Not enough numeric values after parsing (need at least 10).")
 
-    # With datetime column
+    date_col = (date_col or "").strip()
+    freq_choice_norm = (freq_choice or "").strip().upper()
+
     if date_col:
         if date_col not in df.columns:
             raise gr.Error(f"Date column '{date_col}' not found. Available: {list(df.columns)}")
         dt = pd.to_datetime(df[date_col], errors="coerce")
         mask = dt.notna() & vals.notna()
-        dt = pd.DatetimeIndex(dt[mask])
+        dt = pd.DatetimeIndex(dt[mask]).tz_localize(None)
         vals = vals[mask]
+
         if len(vals) < 10:
             raise gr.Error("Too few valid rows after parsing date/value columns.")
-        # sort by date
+
+        # Sort & dedupe index BEFORE inferring/aligning freq
         order = np.argsort(dt.values)
         dt = dt[order]
         vals = vals.iloc[order].reset_index(drop=True)
+
         y = pd.Series(vals.values, index=dt, name=value_col or "value").copy()
-        y.index = y.index.tz_localize(None)
+        y = y[~y.index.duplicated(keep="last")].sort_index()
 
-        # Determine frequency
-        freq = None
-        if freq_choice and freq_choice != "auto":
-            freq = freq_choice
-            y = y.asfreq(freq, method="ffill")
+        # Choose frequency
+        if freq_choice_norm and freq_choice_norm != "AUTO":
+            freq = freq_choice_norm
         else:
-            # try to infer; if None, fallback to 'D'
-            freq = pd.infer_freq(y.index)
-            if freq is None:
-                # try business day if looks like weekdays only
+            inferred = pd.infer_freq(y.index)
+            if inferred:
+                freq = inferred
+            else:
                 weekday_ratio = (y.index.dayofweek < 5).mean()
                 freq = "B" if weekday_ratio > 0.95 else "D"
-            y = y.asfreq(freq, method="ffill")
+
+        # Align to chosen frequency
+        y = y.asfreq(freq, method="ffill")
 
     else:
-        # No date column: require explicit freq
-        if not freq_choice or freq_choice == "auto":
-            raise gr.Error("No date column given. Please choose a frequency (e.g., D, B, H).")
-        freq = freq_choice
+        # No date column: build synthetic index
+        freq = "D" if (not freq_choice_norm or freq_choice_norm == "AUTO") else freq_choice_norm
         idx = pd.date_range(start="2000-01-01", periods=len(vals), freq=freq)
         y = pd.Series(vals.values, index=idx, name=value_col or "value").copy()
 
-    # Final sanity
     if y.isna().all():
         raise gr.Error("Series is all-NaN after processing.")
     return y, freq
 
 def forecast_csv(file, value_col: str, date_col: str, freq_choice: str, horizon: int, context_hint: int):
-    y, freq = build_series_from_csv(file, value_col.strip(), date_col.strip(), freq_choice.strip())
+    y, freq = build_series_from_csv(file, value_col, date_col, freq_choice)
     return _run_forecast_on_series(y, freq, horizon, context_hint, f"Uploaded series — forecast (freq={freq})")
 
-# ---- UI ----
+# ----------------------------
+# UI
+# ----------------------------
 with gr.Blocks(title="Moirai 2.0 — Time Series Forecast (Research)") as demo:
     gr.Markdown(
         """
@@ -246,7 +261,7 @@ Use **Salesforce/moirai-2.0-R-small** (via Uni2TS) to forecast either a stock ti
             horizon_t = gr.Slider(5, 120, value=30, step=1, label="Forecast horizon (steps)")
         with gr.Row():
             lookback = gr.Slider(1, 10, value=5, step=1, label="Lookback window (years of history)")
-            ctx_t = gr.Slider(64, 2000, value=1680, step=16, label="Context length")
+            ctx_t = gr.Slider(64, 5000, value=1680, step=16, label="Context length")
         run_t = gr.Button("Run forecast", variant="primary")
         plot_t = gr.Plot(label="History + Forecast")
         table_t = gr.Dataframe(label="Forecast table", interactive=False)
@@ -255,7 +270,7 @@ Use **Salesforce/moirai-2.0-R-small** (via Uni2TS) to forecast either a stock ti
     with gr.Tab("Upload CSV"):
         gr.Markdown(
             "Upload a CSV with either (1) a **date/time column** and a **value column**, "
-            "or (2) just a numeric value column (then choose a frequency)."
+            "or (2) just a numeric value column (then choose a frequency, or leave **auto** to default to **D**)."
         )
         with gr.Row():
             file = gr.File(label="CSV file", file_types=[".csv"])
@@ -267,7 +282,7 @@ Use **Salesforce/moirai-2.0-R-small** (via Uni2TS) to forecast either a stock ti
                 label="Frequency",
                 value="auto",
                 choices=["auto", "B", "D", "H", "W", "M", "MS"],
-                info="If no date column, pick a freq (e.g., D)."
+                info="If no date column, 'auto' defaults to D (daily)."
             )
         with gr.Row():
             horizon_u = gr.Slider(1, 500, value=60, step=1, label="Forecast horizon (steps)")
@@ -275,7 +290,11 @@ Use **Salesforce/moirai-2.0-R-small** (via Uni2TS) to forecast either a stock ti
         run_u = gr.Button("Run forecast on CSV", variant="primary")
         plot_u = gr.Plot(label="History + Forecast (CSV)")
         table_u = gr.Dataframe(label="Forecast table (CSV)", interactive=False)
-        run_u.click(forecast_csv, inputs=[file, value_col, date_col, freq_choice, horizon_u, ctx_u], outputs=[plot_u, table_u])
+        run_u.click(
+            forecast_csv,
+            inputs=[file, value_col, date_col, freq_choice, horizon_u, ctx_u],
+            outputs=[plot_u, table_u],
+        )
 
 if __name__ == "__main__":
     demo.launch()