Upload folder using huggingface_hub

README.md

@@ -4,28 +4,56 @@ emoji: 🧵
 colorFrom: indigo
 colorTo: gray
 sdk: gradio
-sdk_version: 6.17.3
+sdk_version: 6.29.1
 app_file: app.py
-short_description: Time series + language model for analysis and forecasting
+short_description: Understand and forecast time series with language
 python_version: "3.12"
-startup_duration_timeout: 30m
+startup_duration_timeout: 1h
+pinned: false
 ---
 
 # TimeBraid
 
-Interactive demo of **TimeBraid-2.5B** — a unified time-series + language model that
-interleaves a Qwen3-1.7B language backbone with TimesFM 2.5 200M time-series experts
-through Mixture-of-Transformers layers.
+Interactive demo for **TimeBraid-2.5B** — *Unifying Time Series and Language for
+Understanding and Forecasting*.
 
-Paste numeric series (one series per line, values comma- or space-separated) and ask a
-question, or set a forecast horizon to get a numeric point forecast plotted against the
-observed history.
+- Paper: <https://huggingface.co/papers/2609.29792>
+- Model: [`XinyueWangg/TimeBraid-2.5B`](https://huggingface.co/XinyueWangg/TimeBraid-2.5B)
+- Code: <https://github.com/CharonWangg/TimeBraid>
 
-- Model: [XinyueWangg/TimeBraid-2.5B](https://huggingface.co/XinyueWangg/TimeBraid-2.5B) (Apache-2.0)
-- Code: [CharonWangg/TimeBraid](https://github.com/CharonWangg/TimeBraid)
-- The `timebraid/` package in this Space is the authors' inference package, shipped
-  verbatim so `AutoModelForCausalLM` resolves the custom architecture with
-  `trust_remote_code=False`.
+TimeBraid braids a Qwen3-1.7B language backbone with a TimesFM 2.5 time-series
+expert through an interleaved Mixture-of-Transformers, so numeric series and text
+live in the same token stream. The same numbers can therefore be read very
+differently depending on the context you give them.
 
-Runs on ZeroGPU. Example inputs are drawn from the authors'
-`examples/inference_tasks.ipynb` notebook (Apache-2.0).
+## How to use it
+
+1. **Paste one series per line** — comma-separated numbers. One line = one
+   series; two lines = two related series the model can compare or choose
+   between.
+2. **Write your question or instruction.** Ask for an explanation of the
+   observed data, or ask for future values.
+3. **Set the forecast horizon.** `0` keeps the run in *understanding* mode and
+   the model answers in words. A positive `N` switches to *forecasting* and the
+   model returns `N` future values for one target series, plotted against the
+   history.
+4. Textual context matters for forecasts: describing an upcoming promotion, a
+   heatwave, or a school term materially changes the predicted numbers.
+
+Advanced options hold an optional system prompt, a 1-based selector for which
+input series to forecast, and the text-generation budget.
+
+## Implementation notes
+
+- Runs on ZeroGPU with the full BF16 checkpoint resident on the GPU.
+- The checkpoint vendors its own `timebraid` inference package, loaded through
+  `trust_remote_code=True`; no separate TimesFM checkpoint is required.
+- MoT mixed attention in TimeBraid is FlashAttention-2 only, and FA3/FA4 have no
+  sm_120 kernels, so the Space installs the prebuilt Blackwell FA2 wheel.
+- Decoding follows the released recipe: greedy (`do_sample=False`), one
+  returned sequence, one request at a time.
+
+## Attribution
+
+The example requests in the UI are the illustrative tasks from the TimeBraid
+model card and the authors' `examples/inference_tasks.ipynb` (Apache-2.0).

app.py

@@ -1,38 +1,64 @@
-"""TimeBraid: unified time series + language demo on ZeroGPU."""
+"""TimeBraid: unify time series and language for understanding and forecasting.
+
+A single-request demo around ``XinyueWangg/TimeBraid-2.5B``.  The model reads one
+or more numeric series plus a natural-language request and either explains what
+it sees or forecasts future values.
+"""
 
 import os
 
 os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
 
-import spaces  # MUST come before any torch / CUDA-touching import
-import time
+import spaces  # noqa: F401  # must be imported before torch / transformers
+
+from typing import Any
 
-import gradio as gr
 import matplotlib
 
 matplotlib.use("Agg")
-import matplotlib.pyplot as plt
-import torch
-from matplotlib.ticker import MaxNLocator
-from transformers import AutoModelForCausalLM, AutoProcessor
 
-from timebraid import TimeBraidProcessor  # noqa: F401 - registers HF AutoClasses
+import matplotlib.pyplot as plt  # noqa: E402
+import torch  # noqa: E402
+from matplotlib.ticker import MaxNLocator  # noqa: E402
+from transformers import AutoModelForCausalLM, AutoProcessor  # noqa: E402
+
+import gradio as gr  # noqa: E402
 
 MODEL_ID = "XinyueWangg/TimeBraid-2.5B"
 
-# Chart conventions from the authors' notebook (examples/inference_tasks.ipynb).
-_INPUT_COLOR = "#0072B2"
-_FORECAST_COLOR = "#D55E00"
-_BOUNDARY_COLOR = "#555555"
-_SERIES_COLORS = ("#0072B2", "#009E73", "#CC79A7", "#E69F00")
+MAX_SERIES = 8
+MAX_POINTS_PER_SERIES = 2048
+MAX_HORIZON = 64
+
+INPUT_COLOR = "#0072B2"
+OUTPUT_COLOR = "#D55E00"
+BOUNDARY_COLOR = "#555555"
+SERIES_COLORS = (
+    "#0072B2",
+    "#009E73",
+    "#CC79A7",
+    "#E69F00",
+    "#56B4E9",
+    "#F0E442",
+    "#000000",
+    "#999999",
+)
+
+# ---------------------------------------------------------------------------
+# Model (loaded once, at module scope, on the GPU)
+# ---------------------------------------------------------------------------
 
-processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=False)
+processor = AutoProcessor.from_pretrained(
+    MODEL_ID,
+    trust_remote_code=True,
+    fix_mistral_regex=False,
+)
 model = (
     AutoModelForCausalLM.from_pretrained(
         MODEL_ID,
+        trust_remote_code=True,
         dtype=torch.bfloat16,
         attn_implementation="flash_attention_2",
-        trust_remote_code=False,
         use_safetensors=True,
     )
     .eval()
@@ -40,290 +66,447 @@ model = (
 )
 
 
-def parse_series(text: str) -> list[list[float]]:
-    """Parse one series per line; values comma- or whitespace-separated."""
-    if text is None:
-        return []
+# ---------------------------------------------------------------------------
+# Input parsing / coercion helpers
+# ---------------------------------------------------------------------------
+
+
+def _coerce_text(value: Any, default: str = "") -> str:
+    return value if isinstance(value, str) else default
+
+
+def _coerce_int(value: Any, default: int, low: int, high: int) -> int:
+    if isinstance(value, bool) or not isinstance(value, (int, float)):
+        return default
+    try:
+        number = int(value)
+    except (TypeError, ValueError):
+        return default
+    return max(low, min(high, number))
+
+
+def parse_series_block(text: str) -> list[list[float]]:
+    """Parse one series per non-empty line, comma (or tab / space) separated."""
     series: list[list[float]] = []
-    for line in text.splitlines():
-        line = line.strip().rstrip(",")
+    for raw_line in (text or "").replace(";", "\n").replace("\t", ",").splitlines():
+        line = raw_line.strip()
         if not line:
             continue
-        tokens = line.replace(",", " ").split()
-        values = [float(token) for token in tokens]
+        chunks = line.replace(" ", ",") if "," not in line else line
+        values: list[float] = []
+        for chunk in chunks.split(","):
+            chunk = chunk.strip()
+            if not chunk:
+                continue
+            try:
+                number = float(chunk)
+            except ValueError as exc:
+                raise gr.Error(
+                    f"Could not read {chunk!r} as a number. Use one series per "
+                    "line with comma-separated values, e.g. ``1.0, 2.0, 3.0``."
+                ) from exc
+            values.append(number)
         if not values:
-            raise ValueError("A series line parsed to zero values.")
+            continue
+        if len(values) > MAX_POINTS_PER_SERIES:
+            raise gr.Error(
+                f"Each series can hold at most {MAX_POINTS_PER_SERIES} values "
+                f"(got {len(values)})."
+            )
         series.append(values)
+    if len(series) > MAX_SERIES:
+        raise gr.Error(f"At most {MAX_SERIES} input series are supported.")
     return series
 
 
-def render_plot(
-    history: list[list[float]],
-    forecast: list[float] | None,
-    target_index: int | None,
-) -> matplotlib.figure.Figure:
-    """Draw observed history and, when present, the forecast continuation."""
-    n = len(history)
-    fig, axes = plt.subplots(
-        n, 1, figsize=(9.5, 2.9 * n), constrained_layout=True, squeeze=False
-    )
-    for index, (ax, values) in enumerate(zip(axes[:, 0], history)):
-        color = _INPUT_COLOR if forecast is not None else _SERIES_COLORS[index % 4]
-        x = list(range(len(values)))
+# ---------------------------------------------------------------------------
+# Chart
+# ---------------------------------------------------------------------------
+
+
+def make_figure(
+    series: list[list[float]],
+    horizon: int,
+    target_index: int,
+    result: dict[str, Any],
+) -> plt.Figure:
+    """Draw the observed series and, when forecasting, the predicted values."""
+    fig, ax = plt.subplots(figsize=(7.8, 3.9), constrained_layout=False)
+    fig.patch.set_facecolor("#FFFFFF")
+
+    if not series:
+        ax.set_axis_off()
+        ax.text(
+            0.5,
+            0.5,
+            "No time series supplied\n(text-only completion)",
+            ha="center",
+            va="center",
+            fontsize=12,
+            color="#5F6368",
+        )
+        return fig
+
+    forecast = result.get("timeseries") if horizon > 0 else None
+    if forecast is not None and forecast.get("values"):
+        history = series[target_index]
+        values = list(forecast["values"])
+        x_hist = list(range(len(history)))
+        x_future = list(range(len(history), len(history) + len(values)))
         ax.plot(
-            x,
+            x_hist,
+            history,
+            "-o",
+            color=INPUT_COLOR,
+            markersize=3.5,
+            linewidth=2.0,
+            label=f"Observed (series {target_index + 1})",
+        )
+        ax.plot(
+            x_future,
             values,
-            color=color,
+            "-o",
+            color=OUTPUT_COLOR,
+            markersize=3.5,
             linewidth=2.0,
-            marker="o" if len(values) <= 64 else None,
-            markersize=4,
-            label="Observed history" if forecast is not None else f"Series {index + 1}",
+            label=f"Forecast (horizon {horizon})",
         )
-        title = f"INPUT TIMESERIES · Series {index + 1}"
-        if forecast is not None and target_index == index:
-            future_x = list(range(len(values), len(values) + len(forecast)))
-            boundary = len(values) - 0.5
-            ax.axvspan(boundary, future_x[-1] + 0.5, color=_FORECAST_COLOR, alpha=0.06)
-            ax.plot(
-                (len(values) - 1, *future_x),
-                (values[-1], *forecast),
-                color=_FORECAST_COLOR,
-                linewidth=2.0,
-                linestyle="--",
-                label="Point forecast",
-            )
-            ax.scatter(future_x, forecast, color=_FORECAST_COLOR, s=24, zorder=3)
-            ax.axvline(
-                boundary,
-                color=_BOUNDARY_COLOR,
-                linewidth=1.2,
-                linestyle=":",
-                label="Forecast boundary",
-            )
-            ax.text(
-                boundary,
-                1.01,
-                "forecast starts",
-                transform=ax.get_xaxis_transform(),
+        ax.axvline(
+            len(history) - 0.5,
+            color=BOUNDARY_COLOR,
+            linestyle="--",
+            linewidth=1.1,
+        )
+        for x_value, y_value in zip(x_future, values):
+            ax.annotate(
+                f"{y_value:.2f}",
+                (x_value, y_value),
+                textcoords="offset points",
+                xytext=(0, 7),
                 ha="center",
-                va="bottom",
-                fontsize=8,
-                color=_BOUNDARY_COLOR,
+                fontsize=7.5,
+                color=OUTPUT_COLOR,
             )
-            title = (
-                f"INPUT HISTORY + OUTPUT FORECAST · Series {index + 1} (forecast target)"
+        ax.set_title(
+            f"Forecast for series {target_index + 1} — {horizon} step(s), "
+            "values on the original scale",
+            fontsize=10.5,
+            color="#202124",
+        )
+    else:
+        for index, values in enumerate(series):
+            ax.plot(
+                range(len(values)),
+                values,
+                "-o",
+                color=SERIES_COLORS[index % len(SERIES_COLORS)],
+                markersize=3.2,
+                linewidth=1.9,
+                label=f"Series {index + 1}",
             )
-        ax.set_title(title, loc="left", fontsize=11, fontweight="bold")
-        ax.set_xlabel("Observation / forecast index" if forecast is not None else "Observation index")
-        ax.set_ylabel("Value (raw scale)")
-        ax.xaxis.set_major_locator(MaxNLocator(integer=True))
-        ax.grid(axis="y", color="#D9DCE1", linewidth=0.8, alpha=0.7)
-        ax.set_axisbelow(True)
-        ax.spines["top"].set_visible(False)
-        ax.spines["right"].set_visible(False)
-        ax.legend(loc="upper left", frameon=False, ncol=3, fontsize=9)
+        ax.set_title("Observed time series", fontsize=10.5, color="#202124")
+
+    ax.set_xlabel("Observation / forecast index")
+    ax.set_ylabel("Value (raw scale)")
+    ax.xaxis.set_major_locator(MaxNLocator(integer=True))
+    ax.grid(axis="y", color="#D9DCE1", linewidth=0.8, alpha=0.7)
+    ax.set_axisbelow(True)
+    ax.spines["top"].set_visible(False)
+    ax.spines["right"].set_visible(False)
+    ax.legend(loc="best", fontsize=8.5, frameon=False)
     return fig
 
 
-@spaces.GPU(duration=120)
-def run(
-    prompt: str,
+def format_answer(result: dict[str, Any], horizon: int) -> str:
+    """Render the model response plus a short run summary."""
+    lines: list[str] = [(result.get("content") or "(no text returned)").strip()]
+    forecast = result.get("timeseries") if horizon > 0 else None
+    if forecast is not None and forecast.get("values"):
+        values = ", ".join(f"{value:.3f}" for value in forecast["values"])
+        lines += ["", f"Forecast (original scale, {horizon} step(s)):", values]
+    lines += [
+        "",
+        "—",
+        f"finish_reason={result.get('finish_reason')}  "
+        f"prompt_tokens={result.get('prompt_tokens')}  "
+        f"completion_tokens={result.get('completion_tokens')}  "
+        f"decode={result.get('decode_impl')}",
+    ]
+    return "\n".join(lines)
+
+
+# ---------------------------------------------------------------------------
+# Inference
+# ---------------------------------------------------------------------------
+
+
+@spaces.GPU(duration=180)
+def run_timebraid(
     series_text: str,
-    horizon: float,
-    target_series_index: float,
-    max_new_tokens: float,
-    progress=gr.Progress(track_tqdm=True),
-):
-    """Analyze or forecast a time series with TimeBraid.
-
-    Args:
-        prompt: Question or instruction about the series (or a plain text question).
-        series_text: Numeric series, one per line; values comma- or space-separated.
-        horizon: Number of future points to forecast; 0 = no forecast (understanding only).
-        target_series_index: 1-based index of the series to forecast (multi-series only).
-        max_new_tokens: Maximum generated text tokens.
-    Returns:
-        (answer text, forecast values text, plot, elapsed seconds text).
+    prompt: str,
+    horizon: int = 0,
+    system_prompt: str = "",
+    target_number: int = 1,
+    max_new_tokens: int = 256,
+) -> tuple[str, Any]:
+    """Run one TimeBraid request and return the written answer plus a chart.
+
+    Supply one series per line in ``series_text`` (comma-separated numbers).
+    Leave ``horizon`` at 0 for an explanation of the observed series; set it to
+    a positive number to forecast that many future values for the target series.
     """
-    started = time.perf_counter()
+    series_text = _coerce_text(series_text)
+    prompt = _coerce_text(prompt)
+    system_prompt = _coerce_text(system_prompt)
+    horizon = _coerce_int(horizon, 0, 0, MAX_HORIZON)
+    target_number = _coerce_int(target_number, 1, 1, MAX_SERIES)
+    max_new_tokens = _coerce_int(max_new_tokens, 256, 32, 1024)
+
+    instruction = prompt.strip()
+    if not instruction:
+        raise gr.Error("Enter a question or instruction for the model.")
+
+    series = parse_series_block(series_text)
+
+    if horizon > 0 and not series:
+        raise gr.Error(
+            "Forecasting needs at least one input series — paste comma-separated "
+            "values, one series per line."
+        )
+
+    target_index = 0
+    if horizon > 0 and len(series) > 1:
+        if target_number > len(series):
+            raise gr.Error(
+                f"Only {len(series)} series were supplied, so the forecast target "
+                f"cannot be series {target_number}."
+            )
+        target_index = target_number - 1
+
+    messages: list[dict[str, str]] = []
+    if system_prompt.strip():
+        messages.append({"role": "system", "content": system_prompt.strip()})
+    messages.append({"role": "user", "content": instruction})
+
+    processor_kwargs: dict[str, Any] = {
+        "messages": messages,
+        "timeseries": series if series else None,
+        "horizon": horizon if horizon > 0 else None,
+        "return_tensors": "pt",
+    }
+    if horizon > 0 and len(series) > 1:
+        processor_kwargs["target_series_index"] = target_index
+
     try:
-        series = parse_series(series_text)
-    except ValueError as exc:
-        raise gr.Error(f"Could not parse the time series: {exc}")
-
-    horizon_i = int(horizon) if horizon and horizon > 0 else None
-    target_i: int | None = None
-    if horizon_i is not None:
-        if not series:
-            raise gr.Error("A forecast horizon needs at least one input series.")
-        if len(series) > 1:
-            if not target_series_index or target_series_index < 1:
-                raise gr.Error(
-                    "With multiple series, pick which one to forecast (Target series)."
-                )
-            if int(target_series_index) > len(series):
-                raise gr.Error(
-                    f"Target series is out of range: only {len(series)} series given."
-                )
-            target_i = int(target_series_index) - 1
-
-    messages = [{"role": "user", "content": prompt}]
-    inputs = processor.apply_chat_template(
-        messages,
-        timeseries=series if series else None,
-        horizon=horizon_i,
-        target_series_index=target_i,
-        add_generation_prompt=True,
-        tokenize=True,
-        return_dict=True,
-        return_tensors="pt",
-    ).to(model.device)
+        model_inputs = processor(**processor_kwargs)
+    except (ValueError, RuntimeError, TypeError) as exc:
+        raise gr.Error(f"Could not prepare the request: {exc}") from exc
+
+    device_inputs = {
+        key: value.to(model.device) if isinstance(value, torch.Tensor) else value
+        for key, value in model_inputs.items()
+    }
     with torch.inference_mode():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=int(max_new_tokens),
+        output = model.generate(
+            **device_inputs,
+            max_new_tokens=max_new_tokens,
             do_sample=False,
             num_beams=1,
             num_return_sequences=1,
         )
-    result = processor.post_process_generation(outputs, model_inputs=inputs)
-    elapsed = time.perf_counter() - started
-
-    answer = (result.get("content") or "").strip()
-    if not answer and horizon_i is None:
-        answer = "(empty response)"
-    ts = result.get("timeseries")
-    forecast_values = None
-    forecast_text = "No numeric forecast (understanding route)."
-    if ts is not None and ts.get("values"):
-        forecast_values = [float(v) for v in ts["values"]]
-        forecast_text = ", ".join(f"{v:.4g}" for v in forecast_values)
-
-    plot = render_plot(series, forecast_values, result.get("target_series_index"))
-    stats = (
-        f"route: {'forecast' if horizon_i is not None else 'understanding'} · "
-        f"prompt tokens: {result.get('prompt_tokens')} · "
-        f"completion tokens: {result.get('completion_tokens')} · "
-        f"finish: {result.get('finish_reason')} · {elapsed:.1f}s on GPU"
-    )
-    return answer, forecast_text, plot, stats
+    try:
+        result = processor.post_process_generation(output, model_inputs=model_inputs)
+    except (ValueError, RuntimeError, TypeError) as exc:
+        raise gr.Error(f"Could not decode the model output: {exc}") from exc
+
+    figure = make_figure(series, horizon, target_index, result)
+    return format_answer(result, horizon), figure
 
 
+# ---------------------------------------------------------------------------
+# Interface
+# ---------------------------------------------------------------------------
+
 CSS = """
-#col-container { max-width: 1150px; margin: 0 auto; }
 .dark .gradio-container { color: var(--body-text-color); }
+.dark .gradio-container .prose { color: var(--body-text-color); }
 """
 
+EXAMPLES = [
+    [
+        "2.0, 2.1, 2.2, 2.4, 2.8, 3.1, 3.0, 2.9, 3.4, 3.8, 4.1, 4.3",
+        "Describe the dominant trend, major turning points, and whether the "
+        "series becomes more volatile.",
+        0,
+        "",
+    ],
+    [
+        "20.1, 20.0, 20.2, 20.1, 20.3, 34.8, 20.2, 20.1, 20.0, 20.2",
+        "Does this series contain an isolated upward spike? Say only whether it "
+        "occurs near the beginning, middle, or end; do not give a numeric index. "
+        "Compare it with the neighboring level.",
+        0,
+        "",
+    ],
+    [
+        "100.0, 104.0, 107.0, 111.0, 116.0, 120.0, 125.0, 129.0\n"
+        "3.2, 3.1, 3.3, 3.4, 3.8, 3.7, 4.0, 4.2",
+        "Series 1 is website visits and Series 2 is conversion rate. Compare "
+        "their trends, volatility, and co-movement.",
+        0,
+        "",
+    ],
+    [
+        "100.0, 102.0, 105.0, 107.0, 103.0, 101.0, 99.0, 100.0, 103.0, 106.0, "
+        "108.0, 104.0, 102.0, 100.0, 101.0, 104.0",
+        "Forecast the next 8 values from the observed seasonal pattern.",
+        8,
+        "",
+    ],
+    [
+        "200.0, 208.0, 215.0, 205.0, 198.0, 210.0, 218.0, 207.0, 201.0, 212.0, "
+        "220.0, 209.0",
+        "A promotion starts at the first forecast step and is expected to lift "
+        "demand above the recent seasonal baseline. Forecast the next 6 values.",
+        6,
+        "You analyze weekly product demand.",
+    ],
+    [
+        "63.0, 62.0, 64.0, 63.0, 65.0, 64.0, 63.0, 72.0, 88.0, 91.0, 86.0, "
+        "74.0, 66.0, 64.0, 63.0, 62.0",
+        "This series is a patient's daily resting heart rate in beats per minute. "
+        "The patient reported a fever starting around day 8 that lasted a few "
+        "days. Is the heart-rate pattern consistent with that report, and does "
+        "the recovery look complete by the end of the series?",
+        0,
+        "",
+    ],
+    [
+        "420.0, 380.0, 395.0, 410.0, 430.0, 485.0, 560.0, 660.0, 790.0, 940.0, "
+        "1130.0, 1350.0",
+        "A new school term begins next week, which typically accelerates "
+        "transmission for several weeks. Forecast the next 6 weekly case counts.",
+        6,
+        "You analyze weekly influenza-like illness case counts for a regional "
+        "health authority.",
+    ],
+    [
+        "2.10, 2.05, 2.02, 2.00, 2.04, 2.15, 2.45, 2.80, 3.00, 3.10, 3.15, "
+        "3.20, 3.25, 3.30, 3.35, 3.30, 3.40, 3.55, 3.60, 3.45, 3.15, 2.80, "
+        "2.50, 2.25, 2.12, 2.06, 2.03, 2.02, 2.06, 2.18, 2.48, 2.83, 3.02, "
+        "3.12, 3.18, 3.24, 3.28, 3.34, 3.38, 3.33, 3.44, 3.58, 3.62, 3.48, "
+        "3.18, 2.83, 2.52, 2.28",
+        "The history covers two days of hourly load in gigawatts. A heatwave "
+        "arrives tomorrow and cooling demand is expected to lift the afternoon "
+        "and evening peak. Forecast the next 24 hourly values.",
+        24,
+        "You analyze hourly electricity load for a regional grid operator.",
+    ],
+    [
+        "",
+        "Explain in one sentence what a moving average reveals about a time series.",
+        0,
+        "",
+    ],
+]
+
+
 with gr.Blocks(title="TimeBraid") as demo:
-    with gr.Column(elem_id="col-container"):
-        gr.Markdown(
-            """
-            # 🧵 TimeBraid: Unifying Time Series and Language
-            Ask questions about numeric series, or request a point forecast.
-            One series per line, values comma- or space-separated. Leave the series
-            box empty for a pure text question. Set **Forecast horizon** to 0 for
-            analysis-only, or a positive number to get numeric forecast values.
-            **Model:** [XinyueWangg/TimeBraid-2.5B](https://huggingface.co/XinyueWangg/TimeBraid-2.5B) ·
-            **Code:** [GitHub](https://github.com/CharonWangg/TimeBraid) (Apache-2.0)
-            """
+    gr.Markdown(
+        "# TimeBraid\n"
+        "**Unifying time series and language for understanding and forecasting** — "
+        "[paper](https://huggingface.co/papers/2609.29792) · "
+        "[model](https://huggingface.co/XinyueWangg/TimeBraid-2.5B) · "
+        "[GitHub](https://github.com/CharonWangg/TimeBraid)\n\n"
+        "Paste one series per line (comma-separated numbers), write what you want "
+        "to know, and choose whether the model should *explain* the observed data "
+        "or *forecast* its future. TimeBraid-2.5B braids a Qwen3-1.7B language "
+        "backbone with a TimesFM 2.5 time-series expert, so the same numbers can "
+        "read very differently depending on the context you give them."
+    )
+
+    with gr.Row():
+        with gr.Column(scale=5):
+            series_box = gr.Textbox(
+                label="Time series — one series per line, comma-separated values",
+                placeholder="200.0, 208.0, 215.0, 205.0, 198.0, 210.0",
+                lines=7,
+            )
+            prompt_box = gr.Textbox(
+                label="Question or instruction",
+                placeholder=(
+                    "Describe the trend and turning points in this series."
+                ),
+                lines=4,
+            )
+            horizon_slider = gr.Slider(
+                minimum=0,
+                maximum=MAX_HORIZON,
+                value=0,
+                step=1,
+                precision=0,
+                label="Forecast horizon",
+                info="0 = explain the observed series; N > 0 = predict the next N values",
+            )
+            run_button = gr.Button("Run TimeBraid", variant="primary")
+
+        with gr.Column(scale=5):
+            answer_box = gr.Textbox(
+                label="Model response",
+                lines=12,
+                show_copy_button=True,
+            )
+            plot_box = gr.Plot(label="Series and forecast")
+
+    with gr.Accordion("Advanced options", open=False):
+        system_box = gr.Textbox(
+            label="System prompt (optional)",
+            placeholder="You analyze weekly product demand.",
+            lines=2,
+        )
+        target_number = gr.Number(
+            label="Forecast target series (1-based)",
+            value=1,
+            precision=0,
+            info="Used only when forecasting from more than one input series.",
         )
-        with gr.Row():
-            with gr.Column(scale=1):
-                prompt = gr.Textbox(
-                    label="Prompt",
-                    placeholder="e.g. Describe the dominant trend and turning points, or forecast the next 8 values.",
-                    lines=3,
-                )
-                series_box = gr.Textbox(
-                    label="Time series (one series per line)",
-                    placeholder="101.2, 101.8, 102.1, 102.5, 103.0, 103.4\n200, 208, 215, ...",
-                    lines=5,
-                )
-                horizon = gr.Number(
-                    label="Forecast horizon (0 = no forecast)",
-                    value=0,
-                    precision=0,
-                    minimum=0,
-                )
-                run_btn = gr.Button("Run", variant="primary")
-                with gr.Accordion("Advanced settings", open=False):
-                    target_idx = gr.Number(
-                        label="Target series (1-based; multi-series forecasts)",
-                        value=1,
-                        precision=0,
-                        minimum=1,
-                    )
-                    max_new_tokens = gr.Slider(
-                        label="Max new text tokens",
-                        minimum=64,
-                        maximum=1024,
-                        value=256,
-                        step=32,
-                    )
-            with gr.Column(scale=1):
-                answer_out = gr.Textbox(label="Answer", lines=8, buttons=["copy"])
-                forecast_out = gr.Textbox(label="Forecast values (raw scale)", buttons=["copy"])
-                plot_out = gr.Plot(label="Series and forecast")
-                stats_out = gr.Markdown()
-        gr.Examples(
-            examples=[
-                [
-                    "Forecast the next 8 values from the observed seasonal pattern.",
-                    "100, 102, 105, 107, 103, 101, 99, 100, 103, 106, 108, 104, 102, 100, 101, 104",
-                    8,
-                    1,
-                    256,
-                ],
-                [
-                    "A promotion starts at the first forecast step and is expected to lift demand above the recent seasonal baseline. Forecast the next 6 values.",
-                    "200, 208, 215, 205, 198, 210, 218, 207, 201, 212, 220, 209",
-                    6,
-                    1,
-                    256,
-                ],
-                [
-                    "Describe the dominant trend, major turning points, and whether the series becomes more volatile.",
-                    "2.0, 2.1, 2.2, 2.4, 2.8, 3.1, 3.0, 2.9, 3.4, 3.8, 4.1, 4.3",
-                    0,
-                    1,
-                    256,
-                ],
-                [
-                    "Does this series contain an isolated upward spike? Say only whether it occurs near the beginning, middle, or end; do not give a numeric index. Compare it with the neighboring level.",
-                    "20.1, 20.0, 20.2, 20.1, 20.3, 34.8, 20.2, 20.1, 20.0, 20.2",
-                    0,
-                    1,
-                    256,
-                ],
-                [
-                    "Series 1 is website visits and Series 2 is conversion rate. Compare their trends, volatility, and co-movement.",
-                    "100, 104, 107, 111, 116, 120, 125, 129\n3.2, 3.1, 3.3, 3.4, 3.8, 3.7, 4.0, 4.2",
-                    0,
-                    1,
-                    256,
-                ],
-                [
-                    "Explain in one sentence what a moving average reveals about a time series.",
-                    "",
-                    0,
-                    1,
-                    256,
-                ],
-            ],
-            inputs=[prompt, series_box, horizon, target_idx, max_new_tokens],
-            outputs=[answer_out, forecast_out, plot_out, stats_out],
-            fn=run,
-            cache_examples=True,
-            cache_mode="lazy",
-            label="Examples (from the authors' inference notebook)",
+        max_tokens_slider = gr.Slider(
+            minimum=64,
+            maximum=1024,
+            value=256,
+            step=32,
+            precision=0,
+            label="Maximum new text tokens",
         )
-    run_btn.click(
-        run,
-        inputs=[prompt, series_box, horizon, target_idx, max_new_tokens],
-        outputs=[answer_out, forecast_out, plot_out, stats_out],
-        api_name="run",
-        concurrency_limit=1,
+
+    inputs = [
+        series_box,
+        prompt_box,
+        horizon_slider,
+        system_box,
+        target_number,
+        max_tokens_slider,
+    ]
+    outputs = [answer_box, plot_box]
+
+    run_button.click(fn=run_timebraid, inputs=inputs, outputs=outputs)
+    prompt_box.submit(fn=run_timebraid, inputs=inputs, outputs=outputs)
+
+    gr.Examples(
+        examples=EXAMPLES,
+        inputs=[series_box, prompt_box, horizon_slider, system_box],
+        fn=run_timebraid,
+        cache_examples=True,
+        cache_mode="lazy",
+        examples_per_page=3,
+        label="Example requests (from the TimeBraid repository and model card)",
+    )
+
+    gr.Markdown(
+        "All series and forecasts are plotted on their original scale. Forecasts "
+        "come from greedy decoding on a fresh prompt, exactly as in the released "
+        "inference recipe."
     )
 
-if __name__ == "__main__":
-    demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)
+
+demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)

requirements.txt

@@ -1,8 +1,20 @@
+# TimeBraid-2.5B inference stack for HF Spaces ZeroGPU (NVIDIA RTX PRO 6000, sm_120).
+#
+# `spaces`, `gradio` and `huggingface_hub` are intentionally omitted: the platform
+# preinstalls and pins them.
+#
+# torch is pinned to 2.11.0 to match the prebuilt FlashAttention-2 wheel cell
+# below (`pt211-cu130-cp312`). ZeroGPU accepts 2.8.0 / 2.9.1 / 2.10.0 / 2.11.0.
 torch==2.11.0
-transformers==4.57.6
-huggingface_hub==0.36.2
-accelerate==1.11.0
 numpy==2.1.3
+einops
+accelerate==1.11.0
+transformers==4.57.6
 safetensors==0.5.3
+tokenizers==0.22.2
 matplotlib
-https://huggingface.co/datasets/multimodalart/zerogpu-blackwell-wheels/resolve/main/wheels/pt211-cu130-cp312/flash_attn-2.8.3-cp312-cp312-linux_x86_64.whl
+
+# MoT mixed attention in TimeBraid is FlashAttention-2 only, and FA3/FA4 cannot
+# run on sm_120, so we use the prebuilt Blackwell FA2 wheel rather than sdpa.
+# Requires `python_version: "3.12"` in README.md to match the cp312 tag.
+https://huggingface.co/datasets/multimodalart/zerogpu-blackwell-wheels/resolve/main/wheels/pt211-cu130-cp312/flash_attn-2.8.3-cp312-cp312-linux_x86_64.whl