"""Interactive demo of ci-proto, a zero-shot probabilistic time-series forecasting model. Upload a CSV with a timestamp column and a value column (or try the bundled bicycle-rentals series) and get a forecast: the median trajectory with nested quantile bands from 0.1 to 0.9, plus the raw quantile values for every step. The forecasting logic mirrors the authors' ``examples/forecast.py`` reference: - A ``split`` column marks ``history`` rows as observed; the ``forecast`` rows that follow them set the horizon and its timestamps. - Without one, every row is history, the frequency is inferred with ``pd.infer_freq`` and future timestamps continue it for the chosen horizon. """ import os from pathlib import Path import spaces # ZeroGPU: must come before torch and any CUDA-touching import import matplotlib matplotlib.use("Agg") import gradio as gr import pandas as pd import torch # ZeroGPU's torch hijack reports a minimal fake CUDA device (a SimpleNamespace # with name/major/minor/total_memory/multi_processor_count). The vendored # ci-proto backend selector reads .shared_memory_per_block_optin at load time # to pick a kernel chunk budget, which the fake namespace lacks. Wrap it once # here so the read works under ZeroGPU; inside @spaces.GPU the real function # is restored, so this is a no-op there. _original_get_device_properties = torch.cuda.get_device_properties def _get_device_properties_with_optin(*args, **kwargs): """Return device properties that always report ``shared_memory_per_block_optin``.""" properties = _original_get_device_properties(*args, **kwargs) if not hasattr(properties, "shared_memory_per_block_optin"): # Blackwell (sm_120) GPUs have 227 KB opt-in shared memory per block. properties.shared_memory_per_block_optin = 232448 return properties torch.cuda.get_device_properties = _get_device_properties_with_optin import continuousintelligence as ci MODEL_ID = "ContinuousIntelligence/ci-proto" LEVELS = torch.arange(1, 10, dtype=torch.float32) / 10 LABELS = {"rentals": "Bike rentals per hour"} BIKE_CSV = Path(ci.__file__).parent / "_data" / "bike_sharing.csv" torch.set_num_threads(min(4, os.cpu_count() or 1)) model = ci.load_model(MODEL_ID, device="auto") def read_series( dataframe: pd.DataFrame, time_column: str, value_column: str, horizon: int | None, ) -> tuple[pd.DatetimeIndex, torch.Tensor, pd.DatetimeIndex]: """Return history timestamps, history values ``[1, T]`` and forecast timestamps.""" missing = {time_column, value_column} - set(dataframe.columns) if missing: raise ValueError(f"the CSV has no column(s) {sorted(missing)}") times = pd.DatetimeIndex(pd.to_datetime(dataframe[time_column])) if "split" in dataframe.columns: if horizon is not None: raise ValueError("horizon applies only to CSVs without a split column") split = dataframe["split"] if not split.isin(["history", "forecast"]).all(): raise ValueError("split values must be 'history' or 'forecast'") observed = split.eq("history").to_numpy() count = int(observed.sum()) if not observed[:count].all(): raise ValueError("history rows must come before forecast rows") history_times, future_times = times[:count], times[count:] values = dataframe[value_column].iloc[:count] else: if horizon is None or horizon < 1: raise ValueError("pass a horizon of at least 1 for a CSV without a split column") freq = pd.infer_freq(times) if len(times) >= 3 else None if freq is None: raise ValueError("timestamps have no regular frequency; add a split column") history_times = times future_times = pd.date_range(times[-1], periods=horizon + 1, freq=freq)[1:] values = dataframe[value_column] if history_times.empty or len(future_times) == 0: raise ValueError("the CSV needs at least one history row and one forecast step") history = torch.tensor(values.to_numpy(dtype=float), dtype=torch.float32)[None] return history_times, history, future_times @spaces.GPU(duration=20) def forecast_series( csv_file, horizon: int, time_column: str, value_column: str, ylabel: str, progress=gr.Progress(), ): """Forecast a time series from CSV history and plot quantile bands. Args: csv_file: CSV with a timestamp column and a value column; an optional ``split`` column marks the rows to forecast. horizon: steps to forecast when the CSV has no ``split`` column. time_column: name of the timestamp column. value_column: name of the column to forecast. ylabel: label for the plot's y-axis; empty uses the value column. Returns: A tuple of (forecast plot figure, per-step quantile table, summary). """ if csv_file is None: dataframe = pd.read_csv(BIKE_CSV) else: dataframe = pd.read_csv(csv_file if isinstance(csv_file, str) else csv_file.name) history_times, history, future_times = read_series( dataframe, time_column, value_column, horizon if "split" not in dataframe.columns else None ) horizon_steps = len(future_times) progress(0.3, desc=f"Forecasting {horizon_steps} steps from {history.shape[-1]} observations") prediction = model.forecast(history, horizon=horizon_steps) columns = (model.quantiles[:, None] - LEVELS).abs().argmin(dim=0) if not torch.allclose(model.quantiles[columns], LEVELS, atol=1e-6, rtol=0): raise ValueError("the checkpoint must provide quantiles 0.1 to 0.9 in increments of 0.1") prediction = prediction[:, columns] if not torch.isfinite(prediction).all(): raise RuntimeError("the model returned nonfinite forecasts") axis_label = ylabel or LABELS.get(value_column, value_column) figure = ci.plot_forecast( history, prediction, model.quantiles[columns], timestamps=history_times.append(future_times), ylabel=axis_label, ) levels = [f"q{int(round(float(level * 100)))}" for level in LEVELS] table = pd.DataFrame( prediction[0].T.numpy(), index=[str(ts) for ts in future_times], columns=levels, ) summary = ( f"Device: {model.device} | History: {history.shape[-1]} steps | " f"Horizon: {horizon_steps} steps | Backend: {model.backend}" ) return figure, table, summary BIKE_HISTORY_CSV = Path(__file__).parent / "bike_history.csv" EXAMPLES = [ [str(BIKE_CSV), 24, "timestamp", "rentals", "Bike rentals per hour"], [str(BIKE_HISTORY_CSV), 48, "timestamp", "rentals", "Bike rentals per hour"], ] custom_css = """ .dark .gradio-container { color: var(--body-text-color); } """ with gr.Blocks() as demo: gr.Markdown( """ # ci-proto — zero-shot probabilistic forecasting Upload a CSV time series (or try the bundled example) and forecast its next steps with **ci-proto**, a 167M-parameter time-series foundation model. The plot shows the median forecast and nested 10–90% quantile bands; the table lists every quantile per step. """ ) with gr.Row(): with gr.Column(): csv_input = gr.File( label="CSV time series", file_types=[".csv"], type="filepath", ) horizon_input = gr.Slider( minimum=1, maximum=320, value=24, step=1, label="Horizon (steps to forecast, when no split column)", info="Ignored when the CSV has a split column", ) with gr.Accordion("Column options", open=False): time_column_input = gr.Text( value="timestamp", label="Timestamp column", ) value_column_input = gr.Text( value="rentals", label="Value column", ) ylabel_input = gr.Text( value="Bike rentals per hour", label="Y-axis label (empty uses the value column)", ) forecast_button = gr.Button("Forecast", variant="primary") with gr.Column(): plot_output = gr.Plot(label="Forecast with quantile bands") table_output = gr.Dataframe(label="Forecast quantiles by step") summary_output = gr.Textbox(label="Summary", lines=1) gr.Examples( examples=EXAMPLES, inputs=[csv_input, horizon_input, time_column_input, value_column_input, ylabel_input], fn=forecast_series, outputs=[plot_output, table_output, summary_output], cache_examples=True, cache_mode="lazy", label="Try an example — hourly bike rentals", ) forecast_button.click( fn=forecast_series, inputs=[csv_input, horizon_input, time_column_input, value_column_input, ylabel_input], outputs=[plot_output, table_output, summary_output], api_name="forecast", ) demo.launch(theme=gr.themes.Citrus(), css=custom_css, mcp_server=True)