Spaces:
Running on Zero
Running on Zero
| import spaces # must come before torch | |
| import os | |
| import tempfile | |
| import time | |
| import gradio as gr | |
| import numpy as np | |
| import pandas as pd | |
| import plotly.graph_objects as go | |
| from plotly.subplots import make_subplots | |
| import torch | |
| from timestral import TimestralForecaster | |
| MODEL_ID = "mistralai/Timestral-28M" | |
| MAX_HORIZON = 1024 | |
| MAX_TARGETS = 6 | |
| MAX_CONTEXT = 32768 | |
| forecaster = TimestralForecaster.from_pretrained(MODEL_ID, device="cpu") | |
| forecaster.model.to("cuda" if torch.cuda.is_available() else "cpu") | |
| LEVELS = forecaster.quantile_levels | |
| # ---------------------------------------------------------------- CSV helpers | |
| def _read_csv(path: str) -> tuple[pd.DataFrame, pd.Series | None]: | |
| """Read a CSV; detect an optional leading timestamp column.""" | |
| df = pd.read_csv(path) | |
| if df.shape[1] == 0 or len(df) == 0: | |
| raise gr.Error("The CSV is empty.") | |
| first = df.columns[0] | |
| timestamps = None | |
| if not pd.api.types.is_numeric_dtype(df[first]): | |
| parsed = pd.to_datetime(df[first], errors="coerce") | |
| if parsed.notna().mean() > 0.9: | |
| timestamps = parsed | |
| df = df.drop(columns=[first]) | |
| numeric = df.apply(pd.to_numeric, errors="coerce") | |
| numeric = numeric.loc[:, numeric.notna().any()] | |
| if numeric.shape[1] == 0: | |
| raise gr.Error("No numeric column found in the CSV.") | |
| return numeric, timestamps | |
| def _numeric_columns(path, current_targets=None, current_known=None): | |
| if not path: | |
| return gr.update(choices=[], value=[]), gr.update(choices=[], value=[]) | |
| try: | |
| df, _ = _read_csv(path) | |
| except Exception: | |
| return gr.update(choices=[], value=[]), gr.update(choices=[], value=[]) | |
| cols = list(df.columns) | |
| targets = [c for c in (current_targets or []) if c in cols] or [cols[0]] | |
| known = [c for c in (current_known or []) if c in cols and c not in targets] | |
| return gr.update(choices=cols, value=targets), gr.update(choices=cols, value=known) | |
| def _future_index(timestamps: pd.Series | None, n_ctx: int, horizon: int): | |
| if timestamps is None: | |
| return np.arange(n_ctx), np.arange(n_ctx, n_ctx + horizon) | |
| ts = pd.DatetimeIndex(timestamps) | |
| hist = ts[:n_ctx] | |
| freq = None | |
| try: | |
| freq = pd.infer_freq(hist[-min(len(hist), 200):]) | |
| except Exception: | |
| freq = None | |
| if freq is not None: | |
| fut = pd.date_range(hist[-1], periods=horizon + 1, freq=freq)[1:] | |
| else: | |
| step = pd.Series(hist).diff().median() | |
| fut = pd.DatetimeIndex([hist[-1] + step * (i + 1) for i in range(horizon)]) | |
| if len(ts) >= n_ctx + horizon: # CSV already holds the future timestamps | |
| fut = ts[n_ctx : n_ctx + horizon] | |
| return hist, fut | |
| # ---------------------------------------------------------------- plotting | |
| HIST_COLOR = "#3f3f46" | |
| FC_COLOR = "#e8590c" | |
| def _band(fig, x, lo, hi, color, name, row): | |
| fig.add_trace(go.Scatter(x=x, y=hi, mode="lines", line=dict(width=0), showlegend=False, | |
| hoverinfo="skip", legendgroup=name), row=row, col=1) | |
| fig.add_trace(go.Scatter(x=x, y=lo, mode="lines", line=dict(width=0), fill="tonexty", | |
| fillcolor=color, name=name, legendgroup=name, showlegend=row == 1, | |
| hoverinfo="skip"), row=row, col=1) | |
| def _make_figure(names, x_hist, history, x_fut, q, show_last): | |
| n = len(names) | |
| fig = make_subplots(rows=n, cols=1, shared_xaxes=True, subplot_titles=names, vertical_spacing=0.08) | |
| i01, i10, i50, i90, i99 = (LEVELS.index(v) for v in (0.01, 0.1, 0.5, 0.9, 0.99)) | |
| start = max(0, len(x_hist) - show_last) | |
| for v in range(n): | |
| row = v + 1 | |
| fig.add_trace(go.Scatter(x=x_hist[start:], y=history[v, start:], mode="lines", name="history", | |
| line=dict(color=HIST_COLOR, width=1.5), showlegend=row == 1, | |
| legendgroup="history"), row=row, col=1) | |
| _band(fig, x_fut, q[v, :, i01], q[v, :, i99], "rgba(232,89,12,0.15)", "1–99% interval", row) | |
| _band(fig, x_fut, q[v, :, i10], q[v, :, i90], "rgba(232,89,12,0.35)", "10–90% interval", row) | |
| fig.add_trace(go.Scatter(x=x_fut, y=q[v, :, i50], mode="lines", name="median forecast", | |
| line=dict(color=FC_COLOR, width=2), showlegend=row == 1, | |
| legendgroup="median"), row=row, col=1) | |
| fig.update_layout(height=max(380, 300 * n), margin=dict(l=40, r=20, t=50, b=40), | |
| legend=dict(orientation="h", y=1.02 + 0.04 / n, x=0, yanchor="bottom"), | |
| hovermode="x unified", template="plotly_white") | |
| return fig | |
| # ---------------------------------------------------------------- inference | |
| def forecast( | |
| csv_file: str, | |
| prediction_length: int = 96, | |
| target_columns: list[str] | None = None, | |
| known_covariate_columns: list[str] | None = None, | |
| context_length: int = 4096, | |
| symmetric_averaging: bool = True, | |
| show_last: int = 512, | |
| ): | |
| """Probabilistic zero-shot forecast of the series in a CSV file with Timestral-28M. | |
| The CSV may start with a timestamp column; every other numeric column can be a target. | |
| Trailing rows whose target cells are empty are treated as the future: their | |
| known-covariate values (e.g. calendar, weather forecast) condition the forecast, and | |
| their count sets the horizon when known covariates are used. | |
| Args: | |
| csv_file: path to a CSV file with one row per time step. | |
| prediction_length: forecast horizon in time steps (1-1024). | |
| target_columns: columns to forecast jointly (default: first numeric column). | |
| known_covariate_columns: columns known over both the past and the horizon. | |
| context_length: most recent time steps fed to the model. | |
| symmetric_averaging: average with the mirrored forecast of -context (author default). | |
| show_last: number of history points to show in the plot. | |
| Returns: | |
| A plot of history plus median and 10-90% / 1-99% intervals, a CSV of all 21 | |
| quantiles, and a short summary. | |
| """ | |
| if not csv_file: | |
| raise gr.Error("Upload a CSV file first.") | |
| df, timestamps = _read_csv(csv_file) | |
| targets = [c for c in (target_columns or []) if c in df.columns] or [df.columns[0]] | |
| targets = targets[:MAX_TARGETS] | |
| known = [c for c in (known_covariate_columns or []) if c in df.columns and c not in targets] | |
| # Trailing rows where every target is missing are the "future" part of the file. | |
| tgt = df[targets].to_numpy(dtype=np.float64) | |
| observed_rows = np.where(~np.isnan(tgt).all(axis=1))[0] | |
| if len(observed_rows) < 2: | |
| raise gr.Error("The target column(s) need at least 2 observed values.") | |
| n_ctx = int(observed_rows[-1]) + 1 | |
| n_future_rows = len(df) - n_ctx | |
| horizon = int(prediction_length) | |
| if known: | |
| if n_future_rows == 0: | |
| raise gr.Error("Known covariates need future rows: append rows with empty targets " | |
| "and filled covariate values for the forecast horizon.") | |
| horizon = n_future_rows | |
| horizon = int(np.clip(horizon, 1, MAX_HORIZON)) | |
| ctx_len = int(np.clip(context_length, 32, MAX_CONTEXT)) | |
| start = max(0, n_ctx - ctx_len) | |
| context = tgt[start:n_ctx].T[None] # (1, V, T) | |
| kwargs = {} | |
| if known: | |
| kc = df[known].to_numpy(dtype=np.float64)[start : n_ctx + horizon].T[None] | |
| if kc.shape[-1] < (n_ctx - start) + horizon: | |
| raise gr.Error("Known covariates must cover the full horizon.") | |
| kwargs["known_covariates"] = kc.astype(np.float32) | |
| forecaster.model.transformer.rope.rebuild_inv_freq() | |
| t0 = time.perf_counter() | |
| fc = forecaster.forecast(context.astype(np.float32), prediction_length=horizon, | |
| symmetric_averaging=bool(symmetric_averaging), **kwargs) | |
| elapsed = time.perf_counter() - t0 | |
| q = fc.quantiles[0].float().cpu().numpy() # (V, H, Q) | |
| x_hist, x_fut = _future_index(timestamps, n_ctx, horizon) | |
| fig = _make_figure(targets, x_hist, tgt[:n_ctx].T, x_fut, q, int(show_last)) | |
| rows = {"timestamp" if timestamps is not None else "step": list(x_fut)} | |
| for v, name in enumerate(targets): | |
| for k, level in enumerate(LEVELS): | |
| rows[f"{name}_q{level:g}"] = q[v, :, k] | |
| out = tempfile.NamedTemporaryFile(prefix="timestral_forecast_", suffix=".csv", delete=False) | |
| pd.DataFrame(rows).to_csv(out.name, index=False) | |
| summary = (f"Forecast **{horizon}** steps for **{', '.join(targets)}** from **{n_ctx - start}** context steps" | |
| + (f" with known covariates **{', '.join(known)}**" if known else "") | |
| + f" · model time {elapsed:.2f}s") | |
| return fig, out.name, summary | |
| # ---------------------------------------------------------------- UI | |
| CSS = """ | |
| #col-container { max-width: 1200px; margin: 0 auto; } | |
| .dark .gradio-container { color: var(--body-text-color); } | |
| """ | |
| with gr.Blocks(title="Timestral-28M") as demo: | |
| with gr.Column(elem_id="col-container"): | |
| gr.Markdown( | |
| "# ⏳ Timestral-28M · zero-shot time series forecasting\n" | |
| "Upload a CSV (optional timestamp column first, one numeric column per series) and get a " | |
| "probabilistic forecast with 21 quantiles from Mistral AI's " | |
| "[Timestral-28M](https://huggingface.co/mistralai/Timestral-28M). " | |
| "Pick several targets to forecast them jointly (multivariate). For **known-future covariates**, append " | |
| "future rows with empty target cells and filled covariate values. " | |
| "[Code](https://github.com/mistralai/timestral) · [Blog](https://mistral.ai/news/timestral)" | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=1, min_width=300): | |
| csv_file = gr.File(label="Time series CSV", file_types=[".csv"], type="filepath") | |
| target_columns = gr.Dropdown(label="Target column(s)", multiselect=True, choices=[], | |
| allow_custom_value=True, | |
| info="Default: first numeric column") | |
| prediction_length = gr.Slider(1, MAX_HORIZON, value=96, step=1, label="Forecast horizon (steps)", | |
| info="Ignored when known covariates are used (set by future rows)") | |
| run = gr.Button("Forecast", variant="primary") | |
| with gr.Accordion("Advanced options", open=False): | |
| known_covariate_columns = gr.Dropdown(label="Known-future covariate column(s)", | |
| multiselect=True, choices=[], allow_custom_value=True) | |
| context_length = gr.Slider(32, MAX_CONTEXT, value=4096, step=32, | |
| label="Context length (most recent steps)") | |
| symmetric_averaging = gr.Checkbox(value=True, label="Symmetric averaging (test-time augmentation)") | |
| show_last = gr.Slider(16, 4096, value=512, step=16, label="History points shown in plot") | |
| with gr.Column(scale=2): | |
| plot = gr.Plot(label="Forecast") | |
| summary = gr.Markdown() | |
| quantile_file = gr.File(label="All 21 quantiles (CSV)") | |
| csv_file.change(_numeric_columns, [csv_file, target_columns, known_covariate_columns], [target_columns, known_covariate_columns], api_visibility="private") | |
| run.click( | |
| forecast, | |
| inputs=[csv_file, prediction_length, target_columns, known_covariate_columns, | |
| context_length, symmetric_averaging, show_last], | |
| outputs=[plot, quantile_file, summary], | |
| api_name="forecast", | |
| ) | |
| gr.Examples( | |
| examples=[ | |
| ["examples/air_passengers.csv", 36, ["passengers"], []], | |
| ["examples/mauna_loa_co2.csv", 120, ["co2_ppm"], []], | |
| ["examples/bike_sharing_hourly.csv", 48, ["rentals"], ["temp", "workingday"]], | |
| ], | |
| inputs=[csv_file, prediction_length, target_columns, known_covariate_columns], | |
| outputs=[plot, quantile_file, summary], | |
| fn=forecast, | |
| cache_examples=True, | |
| cache_mode="lazy", | |
| label="Examples: monthly airline passengers · Mauna Loa CO₂ · hourly bike rentals with known covariates", | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True) | |