import os
import time
import inspect
from typing import Any, Dict, List, Optional, Tuple
import numpy as np
import pandas as pd
import gradio as gr
import torch
import plotly.graph_objects as go
from chronos import Chronos2Pipeline
MODEL_ID_DEFAULT = os.getenv("CHRONOS_MODEL_ID", "amazon/chronos-2")
DATA_DIR = "data"
OUT_DIR = "/tmp"
# -------------------------
# Data
# -------------------------
def available_test_csv() -> List[str]:
if not os.path.isdir(DATA_DIR):
return []
return sorted([f for f in os.listdir(DATA_DIR) if f.lower().endswith(".csv")])
def pick_device(ui_choice: str) -> str:
return "cuda" if (ui_choice or "").startswith("cuda") and torch.cuda.is_available() else "cpu"
def make_sample_series(n: int, seed: int, trend: float, season_period: int, season_amp: float, noise: float) -> np.ndarray:
rng = np.random.default_rng(int(seed))
t = np.arange(int(n), dtype=np.float32)
y = (trend * t + season_amp * np.sin(2 * np.pi * t / max(1, int(season_period))) + rng.normal(0, noise, size=int(n))).astype(np.float32)
if float(np.min(y)) < 0:
y -= float(np.min(y))
return y
def load_series_from_csv(csv_path: str, column: Optional[str]) -> Tuple[np.ndarray, str]:
df = pd.read_csv(csv_path)
col = (column or "").strip()
if not col:
numeric_cols = [c for c in df.columns if pd.api.types.is_numeric_dtype(df[c])]
if not numeric_cols:
# try coercion
for c in df.columns:
coerced = pd.to_numeric(df[c], errors="coerce")
if coerced.notna().sum() > 0:
numeric_cols.append(c)
if not numeric_cols:
raise ValueError("Non trovo colonne numeriche nel CSV.")
col = numeric_cols[0]
if col not in df.columns:
raise ValueError(f"Colonna '{col}' non trovata. Disponibili: {list(df.columns)}")
y = pd.to_numeric(df[col], errors="coerce").dropna().astype(np.float32).to_numpy()
if len(y) < 10:
raise ValueError("Serie troppo corta.")
return y, col
# -------------------------
# Model cache
# -------------------------
_PIPE = None
_META = {"model_id": None, "device": None}
def get_pipeline(model_id: str, device: str) -> Chronos2Pipeline:
global _PIPE, _META
model_id = (model_id or MODEL_ID_DEFAULT).strip()
device = "cuda" if device == "cuda" and torch.cuda.is_available() else "cpu"
if _PIPE is None or _META["model_id"] != model_id or _META["device"] != device:
_PIPE = Chronos2Pipeline.from_pretrained(model_id, device_map=device)
_META = {"model_id": model_id, "device": device}
return _PIPE
# -------------------------
# Predict (STABLE)
# -------------------------
def _to_numpy(x: Any) -> np.ndarray:
if isinstance(x, np.ndarray):
return x
if torch.is_tensor(x):
return x.detach().cpu().numpy()
return np.asarray(x)
def _extract_samples(raw: Any) -> np.ndarray:
if isinstance(raw, dict):
for k in ["samples", "predictions", "prediction", "output"]:
if k in raw:
return _to_numpy(raw[k])
if len(raw) > 0:
return _to_numpy(next(iter(raw.values())))
return np.asarray([], dtype=np.float32)
return _to_numpy(raw)
def chronos2_predict(pipe: Chronos2Pipeline, y: np.ndarray, horizon: int, requested_samples: int) -> Tuple[np.ndarray, bool, str]:
"""
Returns:
samples: (S, H)
multi: whether S>1 is real (not replicated)
note: debug note
"""
sig = inspect.signature(pipe.predict)
params = sig.parameters
# input format: ALWAYS batch = [series]
inputs = [y.tolist()]
# kw for horizon
horizon_kw = None
for cand in ["prediction_length", "horizon", "steps", "n_steps", "pred_len"]:
if cand in params:
horizon_kw = cand
break
# kw for samples count (many versions don't have it!)
sample_kw = None
for cand in ["n_samples", "num_return_sequences", "num_samples"]:
if cand in params:
sample_kw = cand
break
kwargs: Dict[str, Any] = {}
if horizon_kw:
kwargs[horizon_kw] = int(horizon)
else:
# worst case: try positional horizon if supported (rare)
kwargs["prediction_length"] = int(horizon)
if sample_kw:
kwargs[sample_kw] = int(requested_samples)
# call
raw = pipe.predict(inputs=inputs, **kwargs) if "inputs" in params else pipe.predict(inputs, **kwargs)
arr = _extract_samples(raw).astype(np.float32, copy=False)
# normalize shape -> (S,H)
arr = np.squeeze(arr)
if arr.ndim == 1:
# could be (H,) or (S,) - assume horizon if length == H
arr = arr[None, :]
# Sometimes output is (B,S,H) or (B,H). If batch dim exists, take first
if arr.ndim == 3:
# assume (B,S,H) or (S,B,H); safest: pick first on axis=0
arr = arr[0]
if arr.ndim == 1:
arr = arr[None, :]
# ensure horizon length
if arr.shape[-1] != horizon:
if arr.shape[-1] > horizon:
arr = arr[..., :horizon]
else:
pad = horizon - arr.shape[-1]
last = arr[..., -1:]
arr = np.concatenate([arr, np.repeat(last, pad, axis=-1)], axis=-1)
# If we got only 1 sample, we can still plot median but band is not meaningful
real_multi = arr.shape[0] > 1
note = f"predict_signature={sig} | used_horizon_kw={horizon_kw} | used_sample_kw={sample_kw} | got_shape={tuple(arr.shape)}"
return arr, real_multi, note
# -------------------------
# Plotly
# -------------------------
def plot_forecast(y, median, low, high, title, show_band: bool, band_label: str) -> go.Figure:
t_hist = np.arange(len(y))
t_fcst = np.arange(len(y), len(y) + len(median))
fig = go.Figure()
fig.add_trace(go.Scatter(x=t_hist, y=y, mode="lines", name="History"))
fig.add_trace(go.Scatter(x=t_fcst, y=median, mode="lines", name="Forecast (median)"))
fig.add_vline(x=len(y) - 1, line_width=1, line_dash="dash", opacity=0.6)
if show_band:
fig.add_trace(go.Scatter(x=t_fcst, y=high, mode="lines", line=dict(width=0),
showlegend=False, hoverinfo="skip"))
fig.add_trace(go.Scatter(
x=t_fcst, y=low, mode="lines", fill="tonexty",
line=dict(width=0), name=band_label
))
fig.update_layout(
title=title,
hovermode="x unified",
margin=dict(l=10, r=10, t=55, b=10),
legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="left", x=0),
xaxis_title="t",
yaxis_title="value",
)
return fig
def kpi_card(label: str, value: str, hint: str = "") -> str:
hint_html = f"{hint}
" if hint else ""
return f"""
{label}
{value}
{hint_html}
"""
def kpi_grid(cards: List[str]) -> str:
return f"{''.join(cards)}
"
def explain(y, median, low, high, band_enabled: bool, q_low: float, q_high: float, extra: str) -> str:
horizon = len(median)
base = float(np.mean(y))
delta = float(median[-1] - median[0])
pct = (delta / max(1e-6, base)) * 100.0
if abs(pct) < 2:
trend_txt = "sostanzialmente stabile"
elif pct > 0:
trend_txt = "in crescita"
else:
trend_txt = "in calo"
txt = f"""
### 🧠 Spiegazione
Nei prossimi **{horizon} step** la previsione mediana è **{trend_txt}** (variazione ≈ **{pct:+.1f}%** rispetto al livello medio storico).
- **Ultimo valore mediano previsto:** **{median[-1]:.2f}**
"""
if band_enabled:
txt += f"- **Banda [{q_low:.0%}–{q_high:.0%}] (ultimo step):** **[{low[-1]:.2f} – {high[-1]:.2f}]**\n"
else:
txt += "- **Banda di incertezza:** disattivata (questa versione di Chronos2 non restituisce campioni multipli con i parametri disponibili).\n"
txt += f"\nDebug
\n\n`{extra}`\n\n