diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml new file mode 100644 index 0000000..2bac91d --- /dev/null +++ b/.github/workflows/docs.yml @@ -0,0 +1,38 @@ +name: Deploy Documentation + +on: + push: + branches: + - main + paths: + - "docs/**" + - "src/tirex2/**" + - "mkdocs.yml" + - ".github/workflows/docs.yml" + workflow_dispatch: + +permissions: + contents: write + +jobs: + build-and-deploy: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v5 + + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - name: Install documentation dependencies + run: pip install -r docs/requirements.txt + + - name: Build documentation + run: mkdocs build --strict + + - name: Deploy to GitHub Pages + uses: peaceiris/actions-gh-pages@v4 + with: + github_token: ${{ secrets.GITHUB_TOKEN }} + publish_dir: ./site + publish_branch: gh-pages diff --git a/.github/workflows/test-inference.yaml b/.github/workflows/test-inference.yaml index cdcc49e..7959406 100644 --- a/.github/workflows/test-inference.yaml +++ b/.github/workflows/test-inference.yaml @@ -9,6 +9,21 @@ on: jobs: inference-tests: runs-on: ubuntu-latest + env: + # Local Broker + TEST_MQTT_BROKER_HOST: localhost + TEST_MQTT_BROKER_PORT: "1883" + services: + mqtt: + image: emqx/emqx:5.10.4 + ports: + - 1883:1883 + options: >- + --health-cmd "emqx ctl status" + --health-interval 5s + --health-timeout 5s + --health-retries 20 + --health-start-period 10s steps: - uses: actions/checkout@v5 - name: Setup Python @@ -28,11 +43,4 @@ jobs: - name: Run inference tests working-directory: ./inference - env: - HF_TOKEN: ${{ secrets.HF_TOKEN }} - run: | - if [ -z "$HF_TOKEN" ]; then - echo "HF_TOKEN secret is not configured; skipping gated inference tests." - exit 0 - fi - pytest tests -s + run: pytest tests -s diff --git a/.gitignore b/.gitignore index 75510fb..85d17e4 100644 --- a/.gitignore +++ b/.gitignore @@ -11,6 +11,7 @@ output/ *.csv build/ dist/ +/site/ .ipynb_checkpoints/ .idea/ diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 0000000..96385ba --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,58 @@ +# Contributing + +Development and documentation contribution guidelines. + +## Environment Setup + +- **Pixi** (recommended): install [Pixi](https://pixi.prefix.dev/latest/), then run `pixi install`. + Environments are defined in [`pyproject.toml`](pyproject.toml) under `tool.pixi.environments` + (e.g. `cuda128`, `cuda126`, `test-cu128`, `test-cu126`, `example`, `example-cu128`, `example-cu126`). +- **pip**: create a virtual environment and install the package in editable mode: + `python -m venv .venv && source .venv/bin/activate && pip install -e ".[examples,fev,gluonts]"` +- **Tooling**: run `pre-commit install` once, then `pre-commit run --all-files` and `pixi run test` + (or `pytest test/` in a pip environment) before opening a PR. + +## Workflow Overview + +1. Branch from `main` and keep changes focused (docs versus code versus tooling). +2. Run pre-commit and tests locally before pushing. +3. Open a PR with a clear summary, test notes, and any follow-up TODOs. +4. Address CI feedback — red checks block review. + +Commit messages are linted by [`conventional-pre-commit`](.pre-commit-config.yaml) and must +follow `type(scope): summary` with one of `chore`, `ci`, `docs`, `feat`, `fix`, `test` +(a scope is required). + +## Documentation Specifics + +The documentation site lives under [`docs/`](docs/) and is built with +[MkDocs](https://www.mkdocs.org/) + [Material for MkDocs](https://squidfunk.github.io/mkdocs-material/) ++ [mkdocstrings](https://mkdocstrings.github.io/). + +- Install docs dependencies: `pip install -r docs/requirements.txt` +- Preview locally: `mkdocs serve` +- Build (as CI does): `mkdocs build --strict` +- The [API reference](docs/api/) is generated automatically from docstrings in + `src/tirex2/` via mkdocstrings — update the docstring, not the generated page, and add a + runnable usage example to any public function or class that doesn't already have one. +- Add new guides under `docs/` and register them in the `nav` section of [`mkdocs.yml`](mkdocs.yml). + +## Commit & Review Etiquette + +- Avoid committing generated artifacts (e.g. `.pixi/`, `output/`, `model`, `*.csv`, + `*.egg-info`, `site/` — see [`.gitignore`](.gitignore)) unless they are intended changes. +- Rebase (don't merge) when syncing from `main`. +- Respond to every review comment; clarify disagreements rather than ignoring them. + +## Getting Help + +- Open a draft PR early for directional feedback. +- Use GitHub Issues/Discussions for larger proposals. + +## NXAI Contributor License Agreement + +Read the full CLA for Individual Contributors here: [CLA](https://github.com/NX-AI/CLA/blob/main/CLA.md) + +### Contact + +If you have any question about the CLA, feel free to reach out to [contact@nx-ai.com](mailto:contact@nx-ai.com) diff --git a/README.md b/README.md index db20d69..f44c844 100644 --- a/README.md +++ b/README.md @@ -8,10 +8,13 @@ [![Hugging Face](https://img.shields.io/badge/HuggingFace-TiRex--2-yellow?logo=huggingface)](https://huggingface.co/NX-AI/TiRex-2) [![PyPI](https://img.shields.io/pypi/v/tirex-2?color=blue)](https://pypi.org/project/tirex-2/) [![PyPI Downloads](https://static.pepy.tech/personalized-badge/tirex-2?period=total&units=INTERNATIONAL_SYSTEM&left_color=GREY&right_color=BLUE&left_text=downloads)](https://pepy.tech/projects/tirex-2) +[![Docker](https://img.shields.io/badge/GHCR-tirex2--cpu%20%2F%20tirex2--gpu-2496ED?logo=docker&logoColor=white)](https://github.com/NX-AI/tirex-2/pkgs/container/tirex2-cpu) [![License](https://img.shields.io/badge/license-Apache--2.0-green)](LICENSE) +[![docs](https://img.shields.io/website?url=https%3A%2F%2Fnx-ai.github.io%2Ftirex-2%2F&label=docs&up_message=online&up_color=green&down_message=offline&down_color=red)](https://nx-ai.github.io/tirex-2/) [![Tests](https://github.com/NX-AI/tirex-2/actions/workflows/test.yaml/badge.svg)](https://github.com/NX-AI/tirex-2/actions/workflows/test.yaml) [![pre-commit](https://img.shields.io/badge/pre--commit-enabled-brightgreen?logo=pre-commit)](https://github.com/pre-commit/pre-commit) [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/NX-AI/tirex-2/blob/main/examples/getting_started.ipynb) +[![TiRex-2 Demo](https://img.shields.io/badge/HuggingFace-TiRex--2%20Demo-yellow?logo=huggingface)](https://huggingface.co/spaces/NX-AI/TiRex-2-demo) @@ -120,22 +123,6 @@ fig.show() ``` ![output of plot_multivariate function visualizing context and forecast of multivariate input](/resources/multivariate-prediction.png) -### Optional FlexAttention for large multivariate batches - -The default dense attention backend avoids compilation overhead and is suitable for small batches. For large -CUDA batches containing many independent multivariate series, opt into the block-sparse FlexAttention backend -when loading the checkpoint: - -```python -model = load_model("NX-AI/TiRex-2", device="cuda", use_flex_attention=True) -forecasts = model.forecast(timeseries, prediction_length=64, batch_size=64) -``` - -FlexAttention compiles its CUDA kernel on first use and can be slower for small batches, so benchmark both the -backend and `batch_size` on the target GPU. Leave `use_flex_attention` unset to preserve the checkpoint setting -and package default, or pass `False` to force dense attention. The two CUDA kernels are numerically close but -not bit-identical; re-evaluate forecast metrics when changing backends in reproducible benchmarks. - ### Benchmarking diff --git a/docs/about.md b/docs/about.md new file mode 100644 index 0000000..6b1ad89 --- /dev/null +++ b/docs/about.md @@ -0,0 +1,21 @@ +# About NXAI + +![NXAI](images/nxai-logo-light.svg#only-light){ width="160" } +![NXAI](images/nxai-logo-dark.svg#only-dark){ width="160" } + +TiRex-2 is developed by [NXAI](https://nx-ai.com) — *"Minds for Machines."* NXAI builds AI +that masters any cognitive task, developing proprietary xLSTM memory architectures designed +to outperform Transformer architectures. TiRex-2 is part of their time-series product line, +focused on deploying efficient, edge-optimized models for real-world industrial +applications. + +- Website: [nx-ai.com](https://nx-ai.com) +- GitHub: [github.com/NX-AI](https://github.com/NX-AI) +- Hugging Face: [huggingface.co/NX-AI](https://huggingface.co/NX-AI) +- LinkedIn: [linkedin.com/company/nxai](https://www.linkedin.com/company/nxai/) +- Contact: [contact@nx-ai.com](mailto:contact@nx-ai.com) + +TiRex-2 builds on [xLSTM](https://arxiv.org/abs/2405.04517) and generalizes NXAI's original +univariate forecasting model, [TiRex](https://github.com/NX-AI/tirex) (see the +[Introduction](introduction.md)). For additional capabilities beyond this open-source +release, see [TiRex-2 Pro](pro.md). diff --git a/docs/api/demo.md b/docs/api/demo.md new file mode 100644 index 0000000..9fbe162 --- /dev/null +++ b/docs/api/demo.md @@ -0,0 +1,8 @@ +# Demo utilities + +Used to build the runnable examples in the [Quickstart](../getting-started/quickstart.md) +and [Covariates](../how-to/covariates.md) pages. + +::: tirex2.demo.Demo + +::: tirex2.demo.plot_demo_forecast diff --git a/docs/api/forecast.md b/docs/api/forecast.md new file mode 100644 index 0000000..030d5e7 --- /dev/null +++ b/docs/api/forecast.md @@ -0,0 +1,3 @@ +# Forecasting API + +::: tirex2.ForecastModel diff --git a/docs/api/index.md b/docs/api/index.md new file mode 100644 index 0000000..2cc9d60 --- /dev/null +++ b/docs/api/index.md @@ -0,0 +1,14 @@ +# API reference + +This reference is generated from the docstrings in +[`src/tirex2/`](https://github.com/NX-AI/tirex-2/tree/main/src/tirex2). It covers the public +API surface exported from the top-level `tirex2` package. + +- [Forecasting](forecast.md) — `load_model`, `TimeseriesType`, `ForecastModel`. +- [Demo](demo.md) — `Demo` and `plot_demo_forecast`, used in the + [Quickstart](../getting-started/quickstart.md). +- [Plotting](plotting.md) — `plot_multivariate`, `plot_forecast`, `plot_covariate`. + +::: tirex2.load_model + +::: tirex2.TimeseriesType diff --git a/docs/api/plotting.md b/docs/api/plotting.md new file mode 100644 index 0000000..0323d0e --- /dev/null +++ b/docs/api/plotting.md @@ -0,0 +1,10 @@ +# Plotting + +Requires either `matplotlib` or `plotly` to be installed, e.g. via +`pip install "tirex-2[examples]"`. + +::: tirex2.plotting.plot_multivariate + +::: tirex2.plotting.plot_forecast + +::: tirex2.plotting.plot_covariate diff --git a/docs/benchmarks.md b/docs/benchmarks.md new file mode 100644 index 0000000..4869295 --- /dev/null +++ b/docs/benchmarks.md @@ -0,0 +1,59 @@ +# Benchmarks + +TiRex-2's reported results are reproducible against two public leaderboards. Full +instructions live in the repository next to the benchmark code: + +- [GIFT-Eval](https://github.com/NX-AI/tirex-2/blob/main/examples/gifteval/README.md) — + reproduce results on the + [GIFT-Eval](https://huggingface.co/spaces/Salesforce/GIFT-Eval) leaderboard. +- [fev-bench](https://github.com/NX-AI/tirex-2/blob/main/examples/fevbench/README.md) — + reproduce results on the [fev-bench](https://huggingface.co/spaces/autogluon/fev-bench) + leaderboard. + +## GIFT-Eval + +Download the dataset once: + +```bash +pixi run -e example-cu128 huggingface-cli download Salesforce/GiftEval --repo-type=dataset --local-dir PATH_TO_SAVE +``` + +Run the benchmark, choosing a model type: + +```bash +pixi run -e example-cu128 python examples/gifteval/run_gifteval.py pretrained +``` + +- `pretrained` loads `NX-AI/TiRex-2-gifteval-pretrain`. +- `zero-shot` loads `NX-AI/TiRex-2-gifteval-zs`. + +By default (`--eval-mode multivariate`) the native multivariate target is kept intact and +scored jointly, which exercises TiRex-2's cross-variate path but is **not** directly +comparable to the public GIFT-Eval leaderboard. Pass `--eval-mode univariate` to split every +multivariate dataset into independent univariate channels, matching the official leaderboard +protocol: + +```bash +pixi run -e example-cu126 python examples/gifteval/run_gifteval.py \ + --eval-mode univariate +``` + +An interactive notebook is also available: start `pixi run notebook` and open +`examples/gifteval/gifteval.ipynb`. + +## fev-bench + +Optionally pre-download the data: + +```bash +pixi run -e example-cu128 huggingface-cli download autogluon/fev_datasets --repo-type=dataset --local-dir +``` + +Run the benchmark — this always loads `NX-AI/TiRex-2-fevbench` from Hugging Face: + +```bash +pixi run fevbench [/path/to/fevbench_storage] [--tasks examples/fevbench/tasks.yaml] +``` + +If the storage path is omitted, the dataset is downloaded at runtime and cached under +`$HOME/.cache`. diff --git a/docs/contributing.md b/docs/contributing.md new file mode 100644 index 0000000..d5330f1 --- /dev/null +++ b/docs/contributing.md @@ -0,0 +1,7 @@ +# Contributing + +Development and documentation contribution guidelines live in +[`CONTRIBUTING.md`](https://github.com/NX-AI/tirex-2/blob/main/CONTRIBUTING.md) at the +repository root — covering environment setup (Pixi/pip), pre-commit and testing, how to +preview and build this documentation site, commit/review etiquette, and the NXAI Contributor +License Agreement. diff --git a/docs/deployment.md b/docs/deployment.md new file mode 100644 index 0000000..e455d18 --- /dev/null +++ b/docs/deployment.md @@ -0,0 +1,253 @@ +# Deployment + +TiRex-2 ships a Docker-based inference server that exposes the model over three APIs: + +- **HTTP API** (FastAPI) +- **MQTT** (request/reply over MQTT v5) +- **MCP** (Model Context Protocol, for tools like Claude Desktop) + +Source: [`inference/`](https://github.com/NX-AI/tirex-2/tree/main/inference) in the repository +— this page documents what is actually implemented there. + +## Images + +Two container images are published: + +- `ghcr.io/nx-ai/tirex2-cpu` — Linux image for `linux/amd64` and `linux/arm64`. Runs on Linux, + macOS, or Windows via Docker Desktop's Linux container backend. +- `ghcr.io/nx-ai/tirex2-gpu` — CUDA Linux image for `linux/amd64`. Runs on Linux with the + NVIDIA Container Toolkit, or on Windows via Docker Desktop's WSL2 backend with NVIDIA WSL + GPU support. + +Both images run a warmup forecast on startup so the model is `torch.compile`d (C++ kernels on +CPU, Triton on GPU) before the first real request; this download-and-warmup step can take up +to ~10-20 seconds. + +### Run the CPU image + +```bash +docker run -it -p 8000:8000 ghcr.io/nx-ai/tirex2-cpu +``` + +PowerShell: + +```powershell +docker run -it -p 8000:8000 ghcr.io/nx-ai/tirex2-cpu +``` + +### Run the GPU image + +```bash +docker run -it --gpus 1 -p 8000:8000 ghcr.io/nx-ai/tirex2-gpu +``` + +PowerShell: + +```powershell +docker run -it --gpus 1 -p 8000:8000 ghcr.io/nx-ai/tirex2-gpu +``` + +Once running, the HTTP API is at `http://localhost:8000/`, with Swagger docs at +[http://localhost:8000/docs](http://localhost:8000/docs). + +## HTTP API + +Every request is batched — pass a list of series even for a single forecast. There is no +internal batching, so choose a batch size appropriate for your hardware; larger batches are +more efficient but too-large batches can cause out-of-memory errors. + +### Univariate endpoints + +`POST /univariate/forecast/mean` and `POST /univariate/forecast/quantiles` take a batch of +plain 1D series: + +```bash +# Univariate series +curl -s -X POST "http://localhost:8000/univariate/forecast/mean" \ + -H 'Content-Type: application/json' \ + -d '{ + "context": [[1, 2, 3, 4, 5, 6, 7, 8]], + "prediction_length": 5 + }' + +# Univariate batch (batch size 2): two independent series forecast in a single request +curl -s -X POST "http://localhost:8000/univariate/forecast/mean" \ + -H 'Content-Type: application/json' \ + -d '{ + "context": [[1, 2, 3, 4, 5, 6, 7, 8], [8, 7, 6, 5, 4, 3, 2, 1]], + "prediction_length": 5 + }' +``` + +### Multivariate endpoints + +`POST /multivariate/forecast/mean` and `POST /multivariate/forecast/quantiles` take a batch +of objects, each with a multi-row `target` and optional `future_covariates`: + +```bash +# Multivariate (multi-target) series +curl -s -X POST "http://localhost:8000/multivariate/forecast/mean" \ + -H 'Content-Type: application/json' \ + -d '{ + "context": [{"target": [[1, 2, 3, 4, 5, 6, 7, 8], [10, 20, 30, 40, 50, 60, 70, 80]]}], + "prediction_length": 5 + }' + +# Multivariate with future covariates +curl -s -X POST "http://localhost:8000/multivariate/forecast/mean" \ + -H 'Content-Type: application/json' \ + -d '{ + "context": [{ + "target": [[1, 2, 3, 4, 5, 6, 7, 8]], + "future_covariates": [[0, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0]] + }], + "prediction_length": 5 + }' +``` + +Batching multiple multivariate series works the same way, as a list under `context`. See +[inference/README.md](https://github.com/NX-AI/tirex-2/blob/main/inference/README.md) for the +full set of curl/Python examples, including batched multivariate-with-covariates requests. + +### `/quantiles` vs `/mean` + +The `/mean` endpoints return the median forecast. The `/quantiles` endpoints return all 9 +quantiles (10, 20, 30, 40, 50, 60, 70, 80, 90%) for the same inputs. + +### Python client example + +```python +import requests + +resp = requests.post( + "http://localhost:8000/univariate/forecast/mean", + json={"context": [[1, 2, 3, 4, 5, 6, 7, 8]], "prediction_length": 5}, +) +print(resp.json()) +``` + +## MQTT API + +The MQTT integration uses **MQTT v5** with a request/reply pattern: TiRex-2 subscribes to +fixed forecast *request* topics and publishes each result back to the +**response topic the requester specifies on the request** (the MQTT v5 `Response Topic` +property). Every client receives only its own results — there is **no shared result topic**, +unlike a design where all clients read from one common response topic. + +Requests without a Response Topic are rejected. Optionally set `Correlation Data` to match a +reply back to its request. + +You need a **v5-capable** MQTT broker. For quick testing, a public broker like +[broker.emqx.io](https://broker.emqx.io) works (don't send sensitive data to a public +broker). The [MQTTX CLI](https://mqttx.app/cli) is convenient for testing: + +```bash +# Linux x86_64 — standalone binary +curl -sL https://github.com/emqx/MQTTX/releases/latest/download/mqttx-cli-linux-x64 -o mqttx && sudo install mqttx /usr/local/bin/mqttx +``` + +Start the container with MQTT enabled: + +```bash +docker run -p 8000:8000 -it -e MQTT_ENABLED=1 -e MQTT_BROKER_HOST=broker.emqx.io -e MQTT_BROKER_PORT=1883 ghcr.io/nx-ai/tirex2-cpu +``` + +Subscribe to your own reply topic first, over MQTT v5: + +```bash +mqttx sub -V 5 -t 'tirex/my-client/result' -h 'broker.emqx.io' -p 1883 +``` + +Then send a forecast request, pointing its Response Topic at that reply topic: + +```bash +mqttx pub -V 5 \ + -t 'tirex/univariate/forecast/request' \ + --response-topic 'tirex/my-client/result' \ + --correlation-data '1234' \ + -h 'broker.emqx.io' -p 1883 \ + -m '{"id": "1234", "context": [[0, 1, 2, 3]], "prediction_length": 4}' +``` + +The result is published to your Response Topic, with the Correlation Data echoed back. +Successful results contain `mean` and `quantiles`; if an error occurs during processing, the +message published to the same Response Topic contains an `error` field instead. + +## MCP + +Start the container as in the HTTP API section above, then connect a tool like Claude +Desktop by following its +[guide for connecting local servers](https://modelcontextprotocol.io/docs/develop/connect-local-servers). +Add the following to `claude_desktop_config.json` under `mcpServers`: + +```json +{ + "mcpServers": { + "tirex": { "command": "npx", "args": ["-y", "mcp-remote", "http://127.0.0.1:8000/mcp"] } + } +} +``` + +Two MCP tools are exposed: a univariate `tirex_model(context, prediction_length)` and a +multivariate `tirex_model_multivariate(target, prediction_length, past_covariates, +future_covariates)`. Unlike the HTTP and MQTT APIs, MCP is **not batched** — each call +forecasts a single series. + +## Configuration options + +Set these as environment variables via `-e`, e.g. +`docker run -e MODEL_DEVICE=cuda ghcr.io/nx-ai/tirex2-cpu`: + +| Environment Variable | Default Value | Description | +| :-------------------- | :------------- | :----------- | +| `MODEL_PATH` | `NX-AI/TiRex-2` | The Hugging Face model id. | +| `MODEL_DEVICE` | `cpu` | Device to run the model on (`cpu` or `cuda`). | +| `HTTP_HOST` | `0.0.0.0` | Host the HTTP server binds to. | +| `HTTP_PORT` | `8000` | Port the HTTP server binds to. | +| `MQTT_ENABLED` | `0` | Enable MQTT client functionality (`1`=True, `0`=False). | +| `MQTT_BROKER_HOST` | `None` | Hostname or IP address of the MQTT broker. | +| `MQTT_BROKER_PORT` | `None` | Port of the MQTT broker. | +| `MQTT_BROKER_USERNAME` | `None` | Username for authenticating with the MQTT broker (if required). | +| `MQTT_BROKER_PASSWORD` | `None` | Password for authenticating with the MQTT broker (if required). | +| `MQTT_CLIENT_ID` | `tirex-worker` | Stable, unique client id so the broker can resume the session on reconnect. | +| `MQTT_SESSION_EXPIRY` | `3600` | Seconds the broker retains the session (and queued requests) while disconnected. | +| `MQTT_TOPIC_UNIVARIATE_FORECAST` | `tirex/univariate/forecast/request` | Topic to subscribe to for univariate forecast requests. | +| `MQTT_TOPIC_MULTIVARIATE_FORECAST` | `tirex/multivariate/forecast/request` | Topic to subscribe to for multivariate forecast requests. | + +## Building the images yourself + +```bash +cd inference +docker build -f Dockerfile.cpu -t tirex2-inference-cpu . +docker run --rm -p 8000:8000 tirex2-inference-cpu +``` + +```bash +docker build -f Dockerfile.gpu -t tirex2-inference-gpu . +docker run --rm --gpus 1 -p 8000:8000 tirex2-inference-gpu +``` + +## Development setup + +```bash +pip install -r requirements.txt -r requirements-dev.txt +python -m app.main +``` + +Run the tests against a locally-started server: + +```bash +pytest tests +``` + +Or against an already-running container: + +```bash +TEST_START_SERVER=0 TEST_PORT=8000 pytest tests -s +``` + +## License + +The inference server (this Docker image and the `inference/` directory) is licensed under the +same [Apache License 2.0](https://github.com/NX-AI/tirex-2/blob/main/LICENSE) as the rest of +TiRex-2. diff --git a/docs/getting-started/install.md b/docs/getting-started/install.md new file mode 100644 index 0000000..1ecc7db --- /dev/null +++ b/docs/getting-started/install.md @@ -0,0 +1,33 @@ +# Installation + +## Via pip + +```bash +pip install tirex-2 +``` + +Install with additional dependencies: + +```bash +pip install "tirex-2[examples,fev,gluonts]" +``` + +The Python package installation is currently only tested on Linux and macOS. Docker usage +(covered in [Deployment](../deployment.md)) additionally supports Windows via Docker Desktop. + +## Via Pixi + +[Pixi](https://pixi.prefix.dev/latest/) is used for the development and benchmarking +environment, to ensure it is set up correctly. Install it with: + +```bash +curl -fsSL https://pixi.sh/install.sh | sh +``` + +Environments (e.g. `example-cu128`) are defined in +[`pyproject.toml`](https://github.com/NX-AI/tirex-2/blob/main/pyproject.toml) under +`tool.pixi.environments`; pick the one matching your CUDA version and use case. + +## Next steps + +Continue with the [Quickstart](quickstart.md) for a first forecast. diff --git a/docs/getting-started/quickstart.md b/docs/getting-started/quickstart.md new file mode 100644 index 0000000..4285811 --- /dev/null +++ b/docs/getting-started/quickstart.md @@ -0,0 +1,80 @@ +# Quickstart + +The easiest way to get started is the +["Getting Started" notebook](https://github.com/NX-AI/tirex-2/blob/main/examples/getting_started.ipynb), +which you can also run directly in Google Colab: + +[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/NX-AI/tirex-2/blob/main/examples/getting_started.ipynb) + +If you have cloned the repository, start it locally via Pixi: + +```bash +pixi run notebook +``` + +Depending on your CUDA version and use case you may need a different environment, e.g. +`-e example-cu128`, as defined in +[`pyproject.toml`](https://github.com/NX-AI/tirex-2/blob/main/pyproject.toml) under +`tool.pixi.environments`. + +## Minimal usage: predicting a simple sine wave + +```python +import torch +from tirex2 import TimeseriesType, load_model +from tirex2.plotting import plot_multivariate # requires matplotlib to be installed + +# load model +model = load_model("NX-AI/TiRex-2", device="cpu") # use `device="cuda"` if cuda is available + +# generate data - target expects time series of shape (n_targets, context_length) +context = torch.sin(torch.arange(128).float() / 8) +ts = TimeseriesType(target=context.unsqueeze(0), past_covariates=None, future_covariates=None) + +# perform forecast - each forecast is of shape (n_targets, 9 quantiles, prediction_length) +forecast = model.forecast([ts], prediction_length=32, output_type="numpy")[0] + +# visualize result +fig = plot_multivariate(ts, forecast, engine="matplotlib") +fig.show() +``` + +![output of plot_multivariate function visualizing context and forecast](../images/sine-wave-prediction.png) + +## Covariate example + +This example originates from the "Getting Started" notebook, showing the value of additional +covariates. + +```python +from tirex2 import load_model +from tirex2.demo import Demo, plot_demo_forecast + +# load model +model = load_model("NX-AI/TiRex-2", device="cpu") # use `device="cuda"` if cuda is available + +# load data +demo = Demo.create_nonstationary_demo() +ts_univariate = demo.to_timeseries_type(include_covariates=False) +ts_multivariate = demo.to_timeseries_type(include_covariates=True) + +# perform forecast - each forecast is of shape (n_targets, 9 quantiles, prediction_length) +forecasts = model.forecast( + timeseries=[ts_univariate, ts_multivariate], + prediction_length=demo.horizon, + output_type="numpy", +) + +# visualize result +fig = plot_demo_forecast(demo, *forecasts, engine="matplotlib") +fig.show() +``` + +![output of plot_multivariate function visualizing context and forecast of multivariate input](../images/multivariate-prediction.png) + +## Next steps + +- [Forecasting](../how-to/forecasting.md) — univariate and multivariate forecasting in depth. +- [Covariates](../how-to/covariates.md) — past and future-known covariates. +- [Streaming](../how-to/streaming.md) — what's open-source and what's Pro-only. +- [API reference](../api/index.md) — full signatures and parameters. diff --git a/docs/how-to/covariates.md b/docs/how-to/covariates.md new file mode 100644 index 0000000..d7ea8c4 --- /dev/null +++ b/docs/how-to/covariates.md @@ -0,0 +1,103 @@ +# Covariates + +TiRex-2 natively conditions its forecast on two kinds of covariates, in addition to the +target history. Both are optional and independent of each other and of whether the target +itself is univariate or multivariate. + +## Past covariates + +`past_covariates` are known only up to the current time — like the target itself, they stop +at the end of the context window. Shape: `(num_past_covariates, context_length)`, matching +the target's `context_length` exactly. + +```python +import torch +from tirex2 import TimeseriesType, load_model + +context_length = 512 +prediction_length=42 +target = torch.randn(1, context_length) + +past_covariates = torch.randn(2, context_length) # 2 past-only covariates + +ts = TimeseriesType(target=target, past_covariates=past_covariates, future_covariates=None) + +model = load_model("NX-AI/TiRex-2", device="cpu") +forecast = model.forecast([ts], prediction_length=prediction_length, output_type="numpy")[0] +``` + +## Future-known covariates + +`future_covariates` are known ahead of time for the whole forecast horizon — calendar +features, holidays, promotions, or scheduled interventions are typical examples. Shape: +`(num_future_covariates, context_length + prediction_length)`; if you pass more steps than +`context_length + prediction_length`, the extra trailing steps are ignored. + +```python +import torch +from tirex2 import TimeseriesType, load_model + +context_length = 512 +prediction_length=42 +target = torch.randn(1, context_length) + +future_covariates = torch.zeros(1, context_length + prediction_length) +future_covariates[0, 100::7] = 1.0 # e.g. a weekly recurring event flag + +ts = TimeseriesType(target=target, past_covariates=None, future_covariates=future_covariates) + +model = load_model("NX-AI/TiRex-2", device="cpu") +forecast = model.forecast([ts], prediction_length=42, output_type="numpy")[0] +``` + +## Combining both + +Past and future covariates can be combined freely on the same series: + +```python +ts = TimeseriesType( + target=target, + past_covariates=past_covariates, + future_covariates=future_covariates, +) +``` + +## Worked example: a non-stationary series with two covariates + +The [`Demo`][tirex2.demo.Demo] class used in the [Quickstart](../getting-started/quickstart.md) +builds exactly this kind of input — a continuous future-known driver that sets a wandering +baseline level, plus a binary future-known promotion flag that adds spikes: + +```python +import torch +from tirex2 import TimeseriesType, load_model +from tirex2.demo import Demo +import numpy as np + +demo_nonstationary = Demo.create_nonstationary_demo() +# univariate target shape: (1, context_length) +target = torch.from_numpy(demo_nonstationary.target_context).unsqueeze(0) + +# future-known covariates shape: (n_covariates, context_length + horizon) +future_covariates = torch.from_numpy( + np.stack([np.concatenate([c.context, c.future]) for c in demo_nonstationary.covariates]).astype(np.float32) +) + +# multivariate forecast conditioning: the target plus future-known covariates. +multivariate_nonstationary = TimeseriesType( + target=target, + past_covariates=None, + future_covariates=future_covariates, +) + +model = load_model("NX-AI/TiRex-2", device="cpu") + +forecast = model.forecast( + timeseries=[multivariate_nonstationary], + prediction_length=42, + output_type="numpy", +)[0] + +# forecast.shape == (1, 9, 42) -> (num_target_variates, num_quantiles, prediction_length) +``` +![Multivariate context and forecast, with future-known covariates plotted below](../images/multivariate-prediction.png) diff --git a/docs/how-to/forecasting.md b/docs/how-to/forecasting.md new file mode 100644 index 0000000..c72b91b --- /dev/null +++ b/docs/how-to/forecasting.md @@ -0,0 +1,177 @@ +# Forecasting + +## The `TimeseriesType` input + +Every forecast call takes a list of [`TimeseriesType`][tirex2.model.types.TimeseriesType] +objects — one per series in the batch. Each holds: + +- `target`: tensor of shape `(num_target_variates, context_length)`. +- `past_covariates`: `None`, or a tensor of shape `(num_past_covariates, context_length)`. +- `future_covariates`: `None`, or a tensor of shape + `(n_future_covariates, context_length + prediction_length)` (extra trailing steps + beyond what's needed are ignored). See [Covariates](covariates.md) for a full walkthrough. + +A single-variate `target` (a plain 1D series) is still passed as a 2D tensor with +`num_target_variates == 1` — the same model and API path serve both univariate and +multivariate forecasting. + +## Loading the model + +```python +from tirex2 import load_model + +model = load_model("NX-AI/TiRex-2", device="cpu") # or device="cuda" +``` + +`load_model` returns a [`ForecastModel`][tirex2.api_adapter.forecast.ForecastModel] wrapping +the backbone; unknown attributes fall through to the underlying model, so +`model.quantiles` and similar backbone attributes remain reachable directly. + +`use_flex_attention` overrides every variate mixer's checkpoint setting. `True` enables block-sparse FlexAttention, which can reduce the cost of large grouped multivariate batches on CUDA but adds first-call compilation overhead. `False` forces dense attention. Leave it unset to keep the checkpoint configuration. + +```python +model = load_model("NX-AI/TiRex-2", device="cuda", use_flex_attention=True) +``` + +## Univariate forecasting + +```python +import torch +from tirex2 import TimeseriesType, load_model + +# (num_target_variates=1, context_length) +context = torch.sin(torch.arange(128).float() / 8) +ts_univariate = TimeseriesType(target=context, past_covariates=None, future_covariates=None) + +model = load_model("NX-AI/TiRex-2", device="cpu") + +forecast = model.forecast([ts_univariate], prediction_length=64, output_type="numpy")[0] +# forecast.shape == (1, 9, 64) -> (num_target_variates, num_quantiles, prediction_length) +``` + +![Sine-wave context and forecast produced by plot_multivariate](../images/sine-wave-prediction.png) + +## Multivariate forecasting + +The primary goal of multitarget forecasting is to model complex systems where multiple signals interact jointly, allowing the model to capture both the temporal structure within each individual time series and the cross-variate dependencies among them. + +Pass a target with more than one row to forecast several variates jointly from a single +checkpoint — no separate model or per-variate training is needed: + + +```python +import torch +from tirex2 import TimeseriesType, load_model +from tirex2.demo import Demo + +demo_nonstationary = Demo.create_nonstationary_demo() +demo_holidays = Demo.create_holidays_demo() + +# Stack variates together +multi_target = torch.stack( + [ + torch.from_numpy(demo_holidays.target_context), + torch.from_numpy(demo_nonstationary.target_context), + ] +) + +multi_target_ts = TimeseriesType( + target=multi_target, + past_covariates=None, + future_covariates=None, +) + +model = load_model("NX-AI/TiRex-2", device="cpu") + +multi_target_forecast = model.forecast( + [multi_target_ts], + prediction_length=42, + output_type="numpy", + batch_size=1, +)[0] + +# forecast.shape == (2, 9, 42) -> (num_target_variates, num_quantiles, prediction_length) +``` + +See [Covariates](covariates.md) for past vs. future-known covariates. + + +## Batching multiple series + +`forecast` accepts a list of `TimeseriesType` — each entry can have a different number of +variates and a different context length, and covariates are optional per series: + +```python +import torch +from tirex2 import TimeseriesType, load_model + +context_length_a = 128 +ts_a = torch.randn(1, context_length_a) +ts_a = TimeseriesType( + target=ts_a, + past_covariates=None, + future_covariates=None, +) + +context_length_b = 128 +ts_b = torch.randn(1, context_length_b) +ts_b = TimeseriesType( + target=ts_b, + past_covariates=None, + future_covariates=None, +) + +context_length_c = 128 +ts_c = torch.randn(1, context_length_c) +ts_c = TimeseriesType( + target=ts_c, + past_covariates=None, + future_covariates=None, +) + +model = load_model("NX-AI/TiRex-2", device="cpu") +forecasts = model.forecast([ts_a, ts_b, ts_c], prediction_length=64, output_type="numpy") +# forecasts is a list, one entry per input series +``` + +Internally, series are grouped into batches of at most `batch_size` (default `512`); on a +CUDA out-of-memory error the batch size is automatically halved and the failing batch +retried, without affecting the rest of the call. + +## Output types + +`output_type` controls the returned format: + +| `output_type` | Returns | Requires | +| :------------ | :------ | :------- | +| `"torch"` (default) | list of `torch.Tensor`, shape `(V, 9, H)` | — | +| `"numpy"` | list of `numpy.ndarray`, shape `(V, 9, H)` | — | +| `"gluonts"` | list of GluonTS `QuantileForecast` | `pip install "tirex-2[gluonts]"` | +| `"fev"` | a `datasets.DatasetDict` for `fev.Task.evaluation_summary` | `pip install "tirex-2[fev]"` | + +The 9 quantiles are the levels `0.1, 0.2, ..., 0.9`, with index `4` being the median. + +## Test-time augmentation options + +Extra keyword arguments passed to `forecast(...)` are forwarded to the backbone's +`predict`: + +- `tta_sign_flip: bool` — opt-in sign-flip test-time augmentation: the model is also run on + the sign-flipped input and the two passes are averaged in level space. Roughly doubles + inference cost. Defaults to the checkpoint's configured setting when omitted. +- `tta_diff: bool` — opt-in differencing inside the postprocessor. Defaults to the + checkpoint's configured setting when omitted. + +```python +forecast = model.forecast([ts], prediction_length=64, tta_sign_flip=True) +``` + +## GluonTS and FEV integration + +For GluonTS datasets, use +[`forecast_gluon`][tirex2.api_adapter.forecast.ForecastModel.forecast_gluon] instead of +building `TimeseriesType` objects by hand; it extracts targets, covariates, and item metadata +from each dataset entry directly. For FEV evaluation windows, use +[`forecast_fev`][tirex2.api_adapter.forecast.ForecastModel.forecast_fev]. See +[Benchmarks](../benchmarks.md) for the full GIFT-Eval and fev-bench reproduction paths that +build on these. diff --git a/docs/how-to/streaming.md b/docs/how-to/streaming.md new file mode 100644 index 0000000..c508909 --- /dev/null +++ b/docs/how-to/streaming.md @@ -0,0 +1,38 @@ +# Streaming + +TiRex-2's README describes it as "built on a recurrent architecture designed for efficient +streaming settings." That describes the *architecture choice*, not a capability exposed by +this open-source release — read carefully, since it's easy to conflate the two. + +## What this open-source release actually does + +Every call to +[`ForecastModel.forecast`][tirex2.api_adapter.forecast.ForecastModel.forecast] (and the +underlying `TiRex2.predict`) recomputes the forecast **from scratch over the full context +array you pass in**. There is no stateful, incremental call that lets you feed only the new +observations since your last call: + +- `TiRex2.forward` initializes fresh block state on every call + (`state = {i: None for i in range(len(self.stack))}`) — nothing carries over between calls. +- `TiRex2._predict_once` pads/truncates its input to the model's fixed + `context_len + future_len` window and runs the full stack over it every time. + +So to get an updated forecast as new data points arrive, you re-call `forecast` with the +target tensor extended by the new observations (and re-run over the whole, now-longer or +truncated, context window) — there is no way to avoid recomputing over the full history in +this release. + +```python +# Every call below recomputes over its full context; nothing is cached between calls. +forecast_t1 = model.forecast([ts_up_to_t1], prediction_length=32, output_type="numpy")[0] +# ... new observations arrive ... +forecast_t2 = model.forecast([ts_up_to_t2], prediction_length=32, output_type="numpy")[0] +``` + +## What's Pro-only + +**Streaming** — incremental forecast updates as new observations arrive, without recomputing +over the full history — is listed explicitly as a [TiRex-2 Pro](../pro.md) capability. The +recurrent architecture used in this open-source release is what makes that incremental mode +possible in principle, but the incremental, no-recompute code path itself is not part of this +release. If you need it, see [TiRex-2 Pro](../pro.md). diff --git a/docs/images/multivariate-prediction.png b/docs/images/multivariate-prediction.png new file mode 100644 index 0000000..e018702 Binary files /dev/null and b/docs/images/multivariate-prediction.png differ diff --git a/docs/images/nxai-logo-dark.svg b/docs/images/nxai-logo-dark.svg new file mode 100644 index 0000000..1cd7edc --- /dev/null +++ b/docs/images/nxai-logo-dark.svg @@ -0,0 +1,21 @@ + + + + + + + + + diff --git a/docs/images/nxai-logo-light.svg b/docs/images/nxai-logo-light.svg new file mode 100644 index 0000000..07be137 --- /dev/null +++ b/docs/images/nxai-logo-light.svg @@ -0,0 +1,15 @@ + + + + + + + + + + diff --git a/docs/images/sine-wave-prediction.png b/docs/images/sine-wave-prediction.png new file mode 100644 index 0000000..c2bd918 Binary files /dev/null and b/docs/images/sine-wave-prediction.png differ diff --git a/docs/index.md b/docs/index.md new file mode 100644 index 0000000..c16009b --- /dev/null +++ b/docs/index.md @@ -0,0 +1,75 @@ +# TiRex-2 + +**Generalizing TiRex to Multivariate Data and Streaming** + +[![Paper](https://img.shields.io/static/v1?label=Paper&message=2607.01204&color=B31B1B&logo=arXiv)](https://arxiv.org/abs/2607.01204) +[![Hugging Face](https://img.shields.io/badge/HuggingFace-TiRex--2-yellow?logo=huggingface)](https://huggingface.co/NX-AI/TiRex-2) +[![GitHub](https://img.shields.io/badge/GitHub-NX--AI%2Ftirex--2-181717?logo=github)](https://github.com/NX-AI/tirex-2) +[![PyPI](https://img.shields.io/pypi/v/tirex-2?color=blue)](https://pypi.org/project/tirex-2/) +[![Docker](https://img.shields.io/badge/GHCR-tirex2--cpu%20%2F%20tirex2--gpu-2496ED?logo=docker&logoColor=white)](https://github.com/NX-AI/tirex-2/pkgs/container/tirex2-cpu) +[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/NX-AI/tirex-2/blob/main/examples/getting_started.ipynb) +[![TiRex-2 Demo](https://img.shields.io/badge/HuggingFace-TiRex--2%20Demo-yellow?logo=huggingface)](https://huggingface.co/spaces/NX-AI/TiRex-2-demo) + +TiRex-2 is a **pretrained time series foundation model** that forecasts one or many target +variates directly from their history, optionally conditioned on past and future-known +covariates. A single checkpoint serves both univariate and multivariate forecasting, built +on a recurrent architecture designed for efficient streaming settings — all zero-shot, with +no task-specific training or fine-tuning. + +TiRex-2 generalizes the original univariate model, [TiRex](https://github.com/NX-AI/tirex), to +multivariate forecasting with past and future covariates. See the [Introduction](introduction.md) +for background and the [paper](https://arxiv.org/pdf/2607.01204) for details. + +!!! info "Looking for TiRex-1?" + This site documents **TiRex-2**. Documentation for the original univariate TiRex model + lives at [nx-ai.github.io/tirex](https://nx-ai.github.io/tirex/) — the two are separate + projects and this site does not attempt to unify them. + +## Key facts + +- **Zero-shot multivariate forecasting** — TiRex-2 forecasts multiple target variates out of + the box, without training or fine-tuning on your data. +- **Past and future-known covariates** — TiRex-2 natively conditions on past covariates and + future-known covariates, such as calendar features, holidays, promotions, or scheduled + interventions. +- **Small active footprint** — TiRex-2 activates 38.4M parameters in univariate mode and an + additional 44.1M parameters for multivariate forecasting. + +## Where to go next + +- [Installation](getting-started/install.md) — `pip install tirex-2`, Pixi setup, and gated + Hugging Face weight access. +- [Quickstart](getting-started/quickstart.md) — minimal sine-wave and covariate examples. +- [How-to guides](how-to/forecasting.md) — univariate/multivariate forecasting, covariates in + depth, and what streaming does (and doesn't) mean in this open-source release. +- [Deployment](deployment.md) — the Docker-based HTTP/MQTT/MCP inference server. +- [API reference](api/index.md) — generated reference for the public `tirex2` API. + +## TiRex-2 Pro + +This repository is NXAI's open-source release. A Pro version extends TiRex-2 with: + +- **Streaming**: incremental forecast updates as new observations arrive, without recomputing + over the full history. +- **Speed**: performance-optimized inference, including optimization for dedicated hardware + such as edge, embedded, and industrial PC deployments. +- **Finetuning**: models fine-tuned on your data or with different pretraining. +- **Classification & Regression**: TiRex-2 adapted for classification and regression tasks. + +See [TiRex-2 Pro](pro.md) for details, or contact [contact@nx-ai.com](mailto:contact@nx-ai.com). + +## Cite Our Work + +If you use TiRex-2 in your research, please cite our work: + +```bibtex +@misc{podest2026tirex2generalizingtirexmultivariate, + title={TiRex-2: Generalizing TiRex to Multivariate Data and Streaming}, + author={Patrick Podest and Marco Pichler and Elias Bürger and Levente Zólyomi and Bernhard Voggenberger and Wilhelm Berghammer and Daniel Klotz and Sebastian Böck and Günter Klambauer and Sepp Hochreiter}, + year={2026}, + eprint={2607.01204}, + archivePrefix={arXiv}, + primaryClass={cs.LG}, + url={https://arxiv.org/abs/2607.01204}, +} +``` diff --git a/docs/introduction.md b/docs/introduction.md new file mode 100644 index 0000000..99254a6 --- /dev/null +++ b/docs/introduction.md @@ -0,0 +1,46 @@ +# Introduction + +TiRex-2 is introduced in the paper +[TiRex-2: Generalizing TiRex to Multivariate Data and Streaming](https://arxiv.org/pdf/2607.01204) +(arXiv:2607.01204). + +## From TiRex to TiRex-2 + +The original [TiRex](https://github.com/NX-AI/tirex) model is a univariate, zero-shot time +series forecasting model built on the [xLSTM](https://arxiv.org/abs/2405.04517) architecture. +TiRex-2 generalizes it along two axes: + +- **Multivariate forecasting**: a single checkpoint forecasts one or many target variates + jointly, and can condition on past covariates and future-known covariates (e.g. calendar + features, holidays, promotions, or scheduled interventions) alongside the target history. +- **Streaming-oriented architecture**: TiRex-2 is built on a recurrent architecture (extending + the xLSTM-based design) chosen for efficient streaming settings. + +Both univariate and multivariate forecasting are served zero-shot, without any task-specific +training or fine-tuning, from the same pretrained checkpoint published on +[Hugging Face](https://huggingface.co/NX-AI/TiRex-2). + +## What "streaming-oriented" means in this release + +The recurrent architecture is what makes efficient incremental inference possible in +principle, but this open-source release does not itself expose an incremental, +state-carrying forecast API — every call to +[`forecast`][tirex2.api_adapter.forecast.ForecastModel.forecast] recomputes over the full +context array you pass in. Incremental (no-recompute) streaming updates are part of +[TiRex-2 Pro](pro.md). See [How-to: Streaming](how-to/streaming.md) for the full explanation. + +## Citation + +If you use TiRex-2 in your research, please cite: + +```bibtex +@misc{podest2026tirex2generalizingtirexmultivariate, + title={TiRex-2: Generalizing TiRex to Multivariate Data and Streaming}, + author={Patrick Podest and Marco Pichler and Elias Bürger and Levente Zólyomi and Bernhard Voggenberger and Wilhelm Berghammer and Daniel Klotz and Sebastian Böck and Günter Klambauer and Sepp Hochreiter}, + year={2026}, + eprint={2607.01204}, + archivePrefix={arXiv}, + primaryClass={cs.LG}, + url={https://arxiv.org/abs/2607.01204}, +} +``` diff --git a/docs/pro.md b/docs/pro.md new file mode 100644 index 0000000..c805589 --- /dev/null +++ b/docs/pro.md @@ -0,0 +1,19 @@ +# TiRex-2 Pro + +TiRex-2 already provides state-of-the-art performance for zero-shot prediction, so this +open-source release can be used as-is without training on your own data. + +A Pro version extends TiRex-2 with additional capabilities, including: + +- **Streaming**: incremental forecast updates as new observations arrive, without recomputing + over the full history. (This open-source release recomputes over the full context on every + call — see [How-to: Streaming](how-to/streaming.md) for the exact distinction.) +- **Speed**: performance-optimized inference, including optimization for dedicated hardware + such as edge, embedded, and industrial PC deployments. +- **Finetuning**: models fine-tuned on your data or with different pretraining. +- **Classification & Regression**: TiRex-2 adapted for classification and regression tasks. + +These are Pro-exclusive additions — this documentation does not cover them as usable APIs of +the open-source package, since they aren't part of it. + +If you are interested in any of these, contact [contact@nx-ai.com](mailto:contact@nx-ai.com). diff --git a/docs/requirements.txt b/docs/requirements.txt new file mode 100644 index 0000000..9927536 --- /dev/null +++ b/docs/requirements.txt @@ -0,0 +1,4 @@ +mkdocs>=1.6 +mkdocs-material>=9.5 +mkdocstrings[python]>=0.26 +mkdocs-material-extensions>=1.3 diff --git a/examples/getting_started.ipynb b/examples/getting_started.ipynb index e90037c..d4dce8b 100644 --- a/examples/getting_started.ipynb +++ b/examples/getting_started.ipynb @@ -138,31 +138,6 @@ "```" ] }, - { - "cell_type": "markdown", - "id": "0c8d63e3", - "metadata": {}, - "source": [ - "## Setup\n", - "
\n", - " ℹ️ Note: To gain access to our model weights, please generate yourself a Hugging Face access token, and insert it into the cell below. \n", - " Click here \n", - " to generate the token and ensure that you enable \"Read access to contents of all public gated repos you can access\".\n", - "
" - ] - }, - { - "cell_type": "code", - "execution_count": 2, - "id": "f398dfe4-a1b2-48f0-a658-280111914f3a", - "metadata": {}, - "outputs": [], - "source": [ - "import os\n", - "\n", - "os.environ[\"HF_TOKEN\"] = \"\"" - ] - }, { "cell_type": "code", "execution_count": 3, @@ -4261,7 +4236,7 @@ ], "metadata": { "kernelspec": { - "display_name": "venv (3.13.11)", + "display_name": ".venv-docs (3.12.3)", "language": "python", "name": "python3" }, @@ -4275,7 +4250,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.13.11" + "version": "3.12.3" } }, "nbformat": 4, diff --git a/inference/Dockerfile.cpu b/inference/Dockerfile.cpu index 4bbca2b..b5cc3bb 100644 --- a/inference/Dockerfile.cpu +++ b/inference/Dockerfile.cpu @@ -1,5 +1,5 @@ # Copyright (c) NXAI GmbH. -# This software may be used and distributed according to the terms of the NXAI Community License Agreement. +# Licensed under the Apache License, Version 2.0; see LICENSE for details. # ------------------------- # Stage 1: Build diff --git a/inference/Dockerfile.gpu b/inference/Dockerfile.gpu index d223971..09451c1 100644 --- a/inference/Dockerfile.gpu +++ b/inference/Dockerfile.gpu @@ -1,5 +1,5 @@ # Copyright (c) NXAI GmbH. -# This software may be used and distributed according to the terms of the NXAI Community License Agreement. +# Licensed under the Apache License, Version 2.0; see LICENSE for details. # ------------------------- # Builder (CUDA toolkit for JIT-compatible builds) diff --git a/inference/README.md b/inference/README.md index eec66ef..6efd63c 100644 --- a/inference/README.md +++ b/inference/README.md @@ -13,22 +13,22 @@ These images are Linux containers. The CPU image targets `linux/amd64` and `linu Run the CPU container on Linux/macOS Bash: ```bash -docker run -it -p 8000:8000 -e HF_TOKEN=$YOUR_HF_TOKEN ghcr.io/nx-ai/tirex2-cpu +docker run -it -p 8000:8000 ghcr.io/nx-ai/tirex2-cpu ``` Run the CPU container on Windows PowerShell: ```powershell -docker run -it -p 8000:8000 -e HF_TOKEN=$env:HF_TOKEN ghcr.io/nx-ai/tirex2-cpu +docker run -it -p 8000:8000 ghcr.io/nx-ai/tirex2-cpu ``` Run the GPU container on a Linux host with NVIDIA GPU support: ```bash -docker run -it --gpus 1 -p 8000:8000 -e HF_TOKEN=$YOUR_HF_TOKEN ghcr.io/nx-ai/tirex2-gpu +docker run -it --gpus 1 -p 8000:8000 ghcr.io/nx-ai/tirex2-gpu ``` Run the GPU container on Windows PowerShell: ```powershell -docker run -it --gpus 1 -p 8000:8000 -e HF_TOKEN=$env:HF_TOKEN ghcr.io/nx-ai/tirex2-gpu +docker run -it --gpus 1 -p 8000:8000 ghcr.io/nx-ai/tirex2-gpu ``` Both the CPU and GPU containers run a warmup forecast on startup so the model is compiled before the first request. torch.compile generates kernels for parts of the model (C++ on CPU, Triton on GPU). Download of the model and warmup can take up to ~10-20 seconds. @@ -178,7 +178,7 @@ print(resp.json()) Every request is batched, so provide a list of timeseries as context, even when you only forecast a single timeseries. Bigger batch sizes are more efficient for the hardware, but too big batch sizes can lead to out of memory errors. There isn't any internal batching done, so the consumer of the API is responsible to call with an appropriate batch size for the hardware. -The HTTP API also provides `/univariate/forecast/quantiles` and `/multivariate/forecast/quantiles`, where the 10, 20, 30, 50 (mean), 60, 70, 80 and 90% quantiles are returned, using the same arguments as the `/univariate/forecast/mean` and `/multivariate/forecast/mean` endpoints respectively. +The HTTP API also provides `/univariate/forecast/quantiles` and `/multivariate/forecast/quantiles`, where the 10, 20, 30, 40, 50 (mean), 60, 70, 80 and 90% quantiles are returned, using the same arguments as the `/univariate/forecast/mean` and `/multivariate/forecast/mean` endpoints respectively. ### MQTT API The MQTT integration uses **MQTT v5** with a request/reply pattern. TiRex subscribes to the forecast **request** topics and publishes each result back to the **response topic the requester specifies on the request** (the MQTT v5 `Response Topic` property). Every client therefore receives only its own results — there is no shared result topic. @@ -193,7 +193,7 @@ curl -sL https://github.com/emqx/MQTTX/releases/latest/download/mqttx-cli-linux- Start the container with MQTT: ``` -docker run -p 8000:8000 -it -e HF_TOKEN=$YOUR_HF_TOKEN -e MQTT_ENABLED=1 -e MQTT_BROKER_HOST=broker.emqx.io -e MQTT_BROKER_PORT=1883 ghcr.io/nx-ai/tirex2-cpu +docker run -p 8000:8000 -it -e MQTT_ENABLED=1 -e MQTT_BROKER_HOST=broker.emqx.io -e MQTT_BROKER_PORT=1883 ghcr.io/nx-ai/tirex2-cpu ``` Each request must be sent over MQTT v5 and set a **Response Topic** telling TiRex where to publish the result. Optionally set **Correlation Data** to match the reply back to the request. Requests without a Response Topic are rejected. @@ -261,16 +261,12 @@ docker build -f Dockerfile.cpu -t tirex2-inference-cpu . Run the CPU container on Linux/macOS Bash: ```bash -docker run --rm -p 8000:8000 \ - -e HF_TOKEN=$YOUR_HF_TOKEN \ - tirex2-inference-cpu +docker run --rm -p 8000:8000 tirex2-inference-cpu ``` Run the CPU container on Windows PowerShell: ```powershell -docker run --rm -p 8000:8000 ` - -e HF_TOKEN=$env:HF_TOKEN ` - tirex2-inference-cpu +docker run --rm -p 8000:8000 tirex2-inference-cpu ``` ### GPU Container @@ -282,16 +278,12 @@ docker build -f Dockerfile.gpu -t tirex2-inference-gpu . Run the GPU container on a Linux host with NVIDIA GPU support: ```bash -docker run --rm --gpus 1 -p 8000:8000 \ - -e HF_TOKEN=$YOUR_HF_TOKEN \ - tirex2-inference-gpu +docker run --rm --gpus 1 -p 8000:8000 tirex2-inference-gpu ``` Run the GPU container on Windows PowerShell: ```powershell -docker run --rm --gpus 1 -p 8000:8000 ` - -e HF_TOKEN=$env:HF_TOKEN ` - tirex2-inference-gpu +docker run --rm --gpus 1 -p 8000:8000 tirex2-inference-gpu ``` ## Development Setup @@ -326,4 +318,4 @@ pytest tests -s ## License -TiRex is licensed under the [NXAI community license](../LICENSE). +TiRex-2 is licensed under the [Apache License 2.0](../LICENSE). diff --git a/inference/app/__init__.py b/inference/app/__init__.py index a6e436d..da2c990 100644 --- a/inference/app/__init__.py +++ b/inference/app/__init__.py @@ -1,2 +1,2 @@ # Copyright (c) NXAI GmbH. -# This software may be used and distributed according to the terms of the NXAI Community License Agreement. +# Licensed under the Apache License, Version 2.0; see LICENSE for details. diff --git a/inference/app/config.py b/inference/app/config.py index 06d5e31..633a457 100644 --- a/inference/app/config.py +++ b/inference/app/config.py @@ -1,5 +1,5 @@ # Copyright (c) NXAI GmbH. -# This software may be used and distributed according to the terms of the NXAI Community License Agreement. +# Licensed under the Apache License, Version 2.0; see LICENSE for details. from pydantic_settings import BaseSettings diff --git a/inference/app/http_server.py b/inference/app/http_server.py index 51edd30..2d36353 100644 --- a/inference/app/http_server.py +++ b/inference/app/http_server.py @@ -1,5 +1,5 @@ # Copyright (c) NXAI GmbH. -# This software may be used and distributed according to the terms of the NXAI Community License Agreement. +# Licensed under the Apache License, Version 2.0; see LICENSE for details. import torch from fastapi import FastAPI, Request @@ -123,8 +123,8 @@ def multivariate_forecast_quantiles(req: MultivariateForecastRequest) -> list[li disclaimer = ( "Disclaimer: NXAI is not responsible for any incorrect interpretations of the " - "forecasted values by LLMs. Check the TiRex license for more details: " - "https://github.com/NX-AI/tirex\n\n" + "forecasted values by LLMs. Check the TiRex-2 license for more details: " + "https://github.com/NX-AI/tirex-2\n\n" ) diff --git a/inference/app/main.py b/inference/app/main.py index c0bc01a..b8c5d2f 100644 --- a/inference/app/main.py +++ b/inference/app/main.py @@ -1,5 +1,5 @@ # Copyright (c) NXAI GmbH. -# This software may be used and distributed according to the terms of the NXAI Community License Agreement. +# Licensed under the Apache License, Version 2.0; see LICENSE for details. import atexit import multiprocessing as mp diff --git a/inference/app/model.py b/inference/app/model.py index eab570a..6e33a87 100644 --- a/inference/app/model.py +++ b/inference/app/model.py @@ -1,5 +1,5 @@ # Copyright (c) NXAI GmbH. -# This software may be used and distributed according to the terms of the NXAI Community License Agreement. +# Licensed under the Apache License, Version 2.0; see LICENSE for details. import numpy as np import torch diff --git a/inference/app/mqtt_server.py b/inference/app/mqtt_server.py index fff046a..605132a 100644 --- a/inference/app/mqtt_server.py +++ b/inference/app/mqtt_server.py @@ -1,5 +1,5 @@ # Copyright (c) NXAI GmbH. -# This software may be used and distributed according to the terms of the NXAI Community License Agreement. +# Licensed under the Apache License, Version 2.0; see LICENSE for details. import json import time diff --git a/inference/tests/conftest.py b/inference/tests/conftest.py index 9fdc90e..32d1a17 100644 --- a/inference/tests/conftest.py +++ b/inference/tests/conftest.py @@ -1,5 +1,5 @@ # Copyright (c) NXAI GmbH. -# This software may be used and distributed according to the terms of the NXAI Community License Agreement. +# Licensed under the Apache License, Version 2.0; see LICENSE for details. import json import os diff --git a/inference/tests/test_http.py b/inference/tests/test_http.py index 9025d61..4e5768b 100644 --- a/inference/tests/test_http.py +++ b/inference/tests/test_http.py @@ -1,5 +1,5 @@ # Copyright (c) NXAI GmbH. -# This software may be used and distributed according to the terms of the NXAI Community License Agreement. +# Licensed under the Apache License, Version 2.0; see LICENSE for details. import requests diff --git a/inference/tests/test_mcp.py b/inference/tests/test_mcp.py index 7a129dc..ba8464a 100644 --- a/inference/tests/test_mcp.py +++ b/inference/tests/test_mcp.py @@ -1,5 +1,5 @@ # Copyright (c) NXAI GmbH. -# This software may be used and distributed according to the terms of the NXAI Community License Agreement. +# Licensed under the Apache License, Version 2.0; see LICENSE for details. import json import re diff --git a/inference/tests/test_mqtt.py b/inference/tests/test_mqtt.py index 0d42014..d8e9462 100644 --- a/inference/tests/test_mqtt.py +++ b/inference/tests/test_mqtt.py @@ -1,5 +1,5 @@ # Copyright (c) NXAI GmbH. -# This software may be used and distributed according to the terms of the NXAI Community License Agreement. +# Licensed under the Apache License, Version 2.0; see LICENSE for details. import json import time diff --git a/mkdocs.yml b/mkdocs.yml new file mode 100644 index 0000000..e4a4d62 --- /dev/null +++ b/mkdocs.yml @@ -0,0 +1,99 @@ +site_name: TiRex-2 +site_description: Documentation for TiRex-2, a pretrained multivariate time series forecasting model. +site_url: https://nx-ai.github.io/tirex-2/ +repo_name: NX-AI/tirex-2 +repo_url: https://github.com/NX-AI/tirex-2 +edit_uri: edit/main/docs/ +copyright: Copyright © NXAI GmbH + +theme: + name: material + logo: images/tirex.svg + favicon: images/tirex.svg + palette: + - media: "(prefers-color-scheme: light)" + scheme: default + primary: black + accent: indigo + toggle: + icon: material/brightness-7 + name: Switch to dark mode + - media: "(prefers-color-scheme: dark)" + scheme: slate + primary: black + accent: indigo + toggle: + icon: material/brightness-4 + name: Switch to light mode + features: + - content.code.copy + - content.code.annotate + - content.tabs.link + - navigation.instant + - navigation.tracking + - navigation.sections + - navigation.tabs + - navigation.tabs.sticky + - navigation.path + - navigation.top + - search.suggest + - toc.follow + +nav: + - Home: index.md + - Introduction: introduction.md + - Getting started: + - Installation: getting-started/install.md + - Quickstart: getting-started/quickstart.md + - How-to guides: + - Forecasting: how-to/forecasting.md + - Covariates: how-to/covariates.md + - Streaming: how-to/streaming.md + - Benchmarks: benchmarks.md + - Deployment: deployment.md + - TiRex-2 Pro: pro.md + - Contributing: contributing.md + - About NXAI: about.md + - API reference: + - Overview: api/index.md + - Forecasting: api/forecast.md + - Demo: api/demo.md + - Plotting: api/plotting.md + +plugins: + - search + - mkdocstrings: + handlers: + python: + paths: [src] + options: + docstring_style: numpy + show_source: true + show_root_heading: true + members_order: source + separate_signature: true + show_signature_annotations: true + filters: ["!^_"] + +markdown_extensions: + - attr_list + - admonition + - pymdownx.details + - pymdownx.superfences + - pymdownx.tabbed: + alternate_style: true + - pymdownx.highlight + - pymdownx.inlinehilite + - pymdownx.snippets + - tables + - toc: + permalink: true + +extra: + social: + - icon: fontawesome/brands/github + link: https://github.com/NX-AI/tirex-2 + - icon: fontawesome/brands/python + link: https://pypi.org/project/tirex-2/ + - icon: simple/huggingface + link: https://huggingface.co/NX-AI/TiRex-2 diff --git a/src/tirex2/api_adapter/forecast.py b/src/tirex2/api_adapter/forecast.py index 2ae5ace..b04cd81 100644 --- a/src/tirex2/api_adapter/forecast.py +++ b/src/tirex2/api_adapter/forecast.py @@ -341,6 +341,16 @@ def forecast( (roughly doubles inference cost), and ``tta_diff`` controls postprocessor differencing; when omitted, the checkpoint's configured defaults (``model-config.yaml``) are used. Pass ``True``/``False`` to override. + + Examples + -------- + >>> import torch + >>> from tirex2 import TimeseriesType, load_model + >>> model = load_model("NX-AI/TiRex-2", device="cpu") + >>> ts = TimeseriesType(target=torch.randn(1, 128), past_covariates=None, future_covariates=None) + >>> forecasts = model.forecast([ts], prediction_length=32, output_type="numpy") + >>> forecasts[0].shape + (1, 9, 32) """ assert batch_size >= 1, "Batch size must be >= 1" return _gen_forecast( diff --git a/src/tirex2/base.py b/src/tirex2/base.py index e3df807..b529c89 100644 --- a/src/tirex2/base.py +++ b/src/tirex2/base.py @@ -81,6 +81,16 @@ def load_model( The instantiated backbone (with the checkpoint weights loaded, set to evaluation mode) wrapped in a :class:`ForecastModel` that exposes the high-level ``forecast`` / ``forecast_gluon`` API. + + Examples + -------- + >>> import torch + >>> from tirex2 import TimeseriesType, load_model + >>> model = load_model("NX-AI/TiRex-2", device="cpu") + >>> ts = TimeseriesType(target=torch.randn(1, 128), past_covariates=None, future_covariates=None) + >>> forecast = model.forecast([ts], prediction_length=32, output_type="numpy")[0] + >>> forecast.shape + (1, 9, 32) """ if device.startswith("cuda") and not torch.cuda.is_available(): raise RuntimeError("Execution on CUDA was requested but is not available.") diff --git a/src/tirex2/demo.py b/src/tirex2/demo.py index 4955bcf..e590f42 100644 --- a/src/tirex2/demo.py +++ b/src/tirex2/demo.py @@ -28,6 +28,29 @@ class Covariate: @dataclass class Demo: + """A synthetic forecasting scenario for showcasing TiRex-2, built by + :meth:`create_nonstationary_demo` or :meth:`create_holidays_demo`. + + Bundles a target series (split into context and held-out future) with the + covariates that explain it, and converts to a :class:`~tirex2.model.types.TimeseriesType` + ready to pass to :meth:`~tirex2.api_adapter.forecast.ForecastModel.forecast`. + + Examples + -------- + >>> from tirex2 import load_model + >>> from tirex2.demo import Demo, plot_demo_forecast + >>> model = load_model("NX-AI/TiRex-2", device="cpu") + >>> demo = Demo.create_nonstationary_demo() + >>> ts_univariate = demo.to_timeseries_type(include_covariates=False) + >>> ts_multivariate = demo.to_timeseries_type(include_covariates=True) + >>> forecasts = model.forecast( + ... timeseries=[ts_univariate, ts_multivariate], + ... prediction_length=demo.horizon, + ... output_type="numpy", + ... ) + >>> fig = plot_demo_forecast(demo, *forecasts, engine="matplotlib") + """ + title: str description: str target_context: np.ndarray # (context_length,) diff --git a/src/tirex2/model/component/patch_tokenizer.py b/src/tirex2/model/component/patch_tokenizer.py index 0b277e8..e3a70e3 100644 --- a/src/tirex2/model/component/patch_tokenizer.py +++ b/src/tirex2/model/component/patch_tokenizer.py @@ -7,10 +7,14 @@ class Tokenizer: """Tokenizer that applies sliding-window patching with optional left padding. - Args: - output_patch_size: Number of timesteps produced per output patch. - input_patch_size: Window size used when unfolding the input sequence. - input_patch_stride: Stride between successive input patches. + Parameters + ---------- + output_patch_size : int + Number of timesteps produced per output patch. + input_patch_size : int + Window size used when unfolding the input sequence. + input_patch_stride : int + Stride between successive input patches. """ def __init__( @@ -40,10 +44,14 @@ def output_transform(self, data: torch.Tensor, tokenizer_state: dict): class Patch(nn.Module): """Utility module that unfolds 1D tensors into patches with optional padding. - Args: - patch_size: Size of each extracted patch. - patch_stride: Step between adjacent patches; must divide ``patch_size``. - left_pad: If ``True``, pad on the left; otherwise pad on the right. + Parameters + ---------- + patch_size : int + Size of each extracted patch. + patch_stride : int + Step between adjacent patches; must divide ``patch_size``. + left_pad : bool + If ``True``, pad on the left; otherwise pad on the right. """ def __init__(self, patch_size: int, patch_stride: int, left_pad: bool = False) -> None: diff --git a/src/tirex2/model/component/scaler.py b/src/tirex2/model/component/scaler.py index aa540a0..27bced6 100644 --- a/src/tirex2/model/component/scaler.py +++ b/src/tirex2/model/component/scaler.py @@ -6,10 +6,14 @@ class Scaler: """Normalizes using mean/variance computed over the full sequence. - Args: - eps: Small constant added to avoid division by zero. - use_arcsinh: If ``True``, apply an arcsinh squashing after standardizing. - binaryaware: If ``True``, detect binary variates and bypass scaling for them. + Parameters + ---------- + eps : float + Small constant added to avoid division by zero. + use_arcsinh : bool + If ``True``, apply an arcsinh squashing after standardizing. + binaryaware : bool + If ``True``, detect binary variates and bypass scaling for them. """ def __init__(self, eps: float = 1e-8, use_arcsinh: bool = False, binaryaware: bool = False, **kwargs) -> None: diff --git a/src/tirex2/model/tirex2.py b/src/tirex2/model/tirex2.py index 15a5c86..3925f7a 100644 --- a/src/tirex2/model/tirex2.py +++ b/src/tirex2/model/tirex2.py @@ -285,26 +285,31 @@ def predict( ): """Return quantile forecasts aligned with the input sequence length. - Args: - timeseries: A batch of multivariate timeseries. Each - :class:`TimeseriesType` holds the target and optional covariate - tensors; the target has shape ``(num_variates, sequence_length)``. - prediction_length: The forecast horizon - tta_sign_flip: Opt-in sign-flip test-time augmentation. Leave as - ``None`` (the default) to use the checkpoint's configured setting - (``self.tta_sign_flip``, from ``model-config.yaml``); pass an - explicit ``True``/``False`` to override it for this call. When - enabled, the model is run a second time on the sign-flipped input - (target *and* every covariate negated); the flipped forecast is - mapped back to level space (values negated and the quantile axis - reversed via the ``q -> 1-q`` complement map) and the two passes - are averaged in level space. Requires a symmetric quantile set - and roughly doubles inference cost. When it resolves to disabled, - the output is byte-identical to a single pass. - tta_diff: Opt-in differencing path inside the postprocessor. Leave as - ``None`` (the default) to use the checkpoint's configured setting - (``self.tta_diff``, from ``model-config.yaml``); pass an explicit - ``True``/``False`` to override trend differencing for this call. + Parameters + ---------- + timeseries : list[TimeseriesType] + A batch of multivariate timeseries. Each :class:`TimeseriesType` holds + the target and optional covariate tensors; the target has shape + ``(num_variates, sequence_length)``. + prediction_length : int + The forecast horizon. + tta_sign_flip : bool, optional + Opt-in sign-flip test-time augmentation. Leave as ``None`` (the + default) to use the checkpoint's configured setting + (``self.tta_sign_flip``, from ``model-config.yaml``); pass an + explicit ``True``/``False`` to override it for this call. When + enabled, the model is run a second time on the sign-flipped input + (target *and* every covariate negated); the flipped forecast is + mapped back to level space (values negated and the quantile axis + reversed via the ``q -> 1-q`` complement map) and the two passes + are averaged in level space. Requires a symmetric quantile set + and roughly doubles inference cost. When it resolves to disabled, + the output is byte-identical to a single pass. + tta_diff : bool, optional + Opt-in differencing path inside the postprocessor. Leave as ``None`` + (the default) to use the checkpoint's configured setting + (``self.tta_diff``, from ``model-config.yaml``); pass an explicit + ``True``/``False`` to override trend differencing for this call. """ if tta_sign_flip is None: tta_sign_flip = self.tta_sign_flip diff --git a/src/tirex2/model/types.py b/src/tirex2/model/types.py index 336792b..5f7cf52 100644 --- a/src/tirex2/model/types.py +++ b/src/tirex2/model/types.py @@ -5,6 +5,35 @@ @dataclass class TimeseriesType: + """A single (possibly multivariate) series with optional covariates, as passed to + :meth:`~tirex2.api_adapter.forecast.ForecastModel.forecast`. + + Parameters + ---------- + target : torch.Tensor + Target history, shape ``[V_t, T]`` (``V_t`` target variates, context length ``T``). + A univariate series is still 2D, with ``V_t == 1``. + past_covariates : torch.Tensor or None + Covariates known only up to the current time, shape ``[V_p, T]``, matching the + target's context length. ``None`` if there are no past covariates. + future_covariates : torch.Tensor or None + Covariates known ahead of time for the whole forecast horizon, shape + ``[V_f, >=T+H]`` (``H`` is the requested ``prediction_length``); extra trailing + steps beyond ``T+H`` are ignored. ``None`` if there are no future covariates. + + Examples + -------- + >>> import torch + >>> from tirex2 import TimeseriesType + >>> ts = TimeseriesType( + ... target=torch.randn(1, 128), + ... past_covariates=None, + ... future_covariates=None, + ... ) + >>> ts.past_length + 128 + """ + target: torch.Tensor # [V_t, T] past_covariates: torch.Tensor | None # [V_p, T] future_covariates: torch.Tensor | None # [V_f, >=T+H]; extra future steps are ignored diff --git a/src/tirex2/plotting.py b/src/tirex2/plotting.py index 9ab71d1..51d69c6 100644 --- a/src/tirex2/plotting.py +++ b/src/tirex2/plotting.py @@ -397,13 +397,20 @@ def _plot_covariates_plotly( """ Plots multivariate covariates into a specified Plotly figure/subplot. - Args: - covariates (dict): Format -> {'covariate_name': {'x': x_array, 'y': y_array}} - fig (go.Figure): Plotly figure object. - row (int): Row index for the subplot. - col (int): Column index for the subplot. - cutoff_x (float/datetime): The x-axis value where the forecast starts. - color_cutoff_line (str): Color of the vertical cutoff line. + Parameters + ---------- + covariates : dict + Format -> ``{'covariate_name': {'x': x_array, 'y': y_array}}`` + fig : go.Figure + Plotly figure object. + row : int + Row index for the subplot. + col : int + Column index for the subplot. + cutoff_x : float or datetime + The x-axis value where the forecast starts. + color_cutoff_line : str + Color of the vertical cutoff line. """ try: import plotly.express as px diff --git a/test/test_references.py b/test/test_references.py index a20ca4b..e367404 100644 --- a/test/test_references.py +++ b/test/test_references.py @@ -6,12 +6,14 @@ import torch from tirex2 import TimeseriesType, load_model +from tirex2.model.component.attention_block import is_flex_attention_available pytestmark = pytest.mark.skipif(sys.platform != "linux", reason="reference outputs were recorded on Linux") REFERENCES = Path(__file__).parent / "references" CHECKPOINT = "NX-AI/TiRex-2" RTOL, ATOL = 1e-5, 1e-5 +FLEX_RTOL, FLEX_ATOL = 1e-4, 1e-4 def _reference(name: str) -> dict: @@ -35,14 +37,28 @@ def _reference(name: str) -> dict: } -@pytest.fixture(scope="module") -def model(): +def _load_reference_model(**kwargs): try: - return load_model(CHECKPOINT, device="cpu") + return load_model(CHECKPOINT, device="cpu", **kwargs) except Exception as exc: pytest.skip(f"reference checkpoint {CHECKPOINT} is unavailable: {exc}") +@pytest.fixture(scope="module") +def model(): + return _load_reference_model() + + +@pytest.fixture(scope="module") +def flex_model(): + if not is_flex_attention_available(): + pytest.skip("FlexAttention is not available in this PyTorch installation") + + model = _load_reference_model(use_flex_attention=True) + assert all(block.variate_mixer.attn.use_flex_attention for block in model.model.stack) + return model + + @pytest.mark.parametrize("scenario", list(SCENARIOS)) def test_forecast_matches_cpu_reference(model, scenario): forecast = model.forecast([SCENARIOS[scenario]], prediction_length=H, output_type="torch")[0] @@ -50,3 +66,12 @@ def test_forecast_matches_cpu_reference(model, scenario): assert forecast.shape == expected.shape torch.testing.assert_close(forecast, expected, rtol=RTOL, atol=ATOL) + + +@pytest.mark.parametrize("scenario", list(SCENARIOS)) +def test_flex_attention_forecast_matches_cpu_reference(flex_model, scenario): + forecast = flex_model.forecast([SCENARIOS[scenario]], prediction_length=H, output_type="torch")[0] + expected = torch.tensor(REFERENCE_OUTPUT[scenario], dtype=torch.float32) + + assert forecast.shape == expected.shape + torch.testing.assert_close(forecast, expected, rtol=FLEX_RTOL, atol=FLEX_ATOL)