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 @spaces.GPU(duration=10) 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)