Metadata-Version: 2.4
Name: relational-transformer
Version: 1.8.0
Requires-Dist: beartype
Requires-Dist: einops
Requires-Dist: huggingface-hub
Requires-Dist: ml-dtypes
Requires-Dist: numpy
Requires-Dist: orjson
Requires-Dist: pandas
Requires-Dist: polars
Requires-Dist: pyarrow
Requires-Dist: pyyaml
Requires-Dist: safetensors
Requires-Dist: scikit-learn
Requires-Dist: sentence-transformers
Requires-Dist: torch
Requires-Dist: wandb
Requires-Dist: matplotlib ; extra == 'dev'
Requires-Dist: pre-commit ; extra == 'dev'
Requires-Dist: pytest ; extra == 'dev'
Requires-Dist: tqdm ; extra == 'dev'
Requires-Dist: wandb-workspaces ; extra == 'dev'
Requires-Dist: xgboost ; extra == 'dev'
Provides-Extra: dev
License-File: LICENSE
Summary: Relational Transformer (RT): a foundation model for relational data, bundling a native Rust data-processing engine (rt.rustler).
Requires-Python: >=3.11
Description-Content-Type: text/markdown; charset=UTF-8; variant=GFM

# Relational Transformer (RT)

The official implementation of the **Relational Transformer (RT)**, an
architecture for **Relational Foundation Models (RFMs)** that predict directly
over relational databases (tables linked by foreign keys) and generalize
zero-shot to new databases, tasks, and schemas.

| Paper | Venue | Implementation |
|---|---|---|
| [Relational Transformer: Toward Zero-Shot Foundation Models for Relational Data](https://arxiv.org/abs/2510.06377) | ICLR 2026 | [`rt-v1`](https://github.com/stanford-star/relational-transformer/tree/rt-v1) |
| [PluRel: Synthetic Data unlocks Scaling Laws for Relational Foundation Models](https://arxiv.org/abs/2602.04029) | ICML 2026 | [`stanford-star/plurel`](https://github.com/stanford-star/plurel) |
| RT-J: Large-Scale Pretraining of Relational Transformers for Context-Efficient Predictions | In progress | [`main`](https://github.com/stanford-star/relational-transformer) |

## Installation

Install the `rt` package (the model plus the native Rust data engine) from
GitHub. It builds the extension from source, so you need a
[Rust toolchain](https://rustup.rs) and Python 3.12+:

```bash
pip install "git+https://github.com/stanford-star/relational-transformer.git"
```

## Quickstart

The quickest way to try a released checkpoint is on a RelBench
database already preprocessed into RT's tensor format on the Hub. The example below predicts whether an F1 driver fails to finish a race (`driver-dnf`) with a released RT-J checkpoint:

```python
import os

# flex_attention's compiled kernel is CUDA-only; run it eager on CPU/MPS
os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")

import torch
from huggingface_hub import snapshot_download

from rt import RelationalTransformer
from rt.eval import build_evaluator
from rt.data import get_tasks

device = (
    "cuda" if torch.cuda.is_available()
    else "mps" if torch.backends.mps.is_available()
    else "cpu"
)

# 1. download one RelBench database, already preprocessed into RT's tensor format
pre_dir = snapshot_download(
    "stanford-star/relbench-preprocessed",
    repo_type="dataset",
    allow_patterns="rel-f1/*",
)

# 2. load a pretrained checkpoint (RT-J here)
model = RelationalTransformer.from_pretrained(
    "stanford-star/rt-j/classification", device=device
).to(torch.bfloat16)
cfg = model.config

# 3. build an evaluator for one task and predict zero-shot for 5 test rows.
#    the evaluator samples each row's context from the preprocessed DB;
tasks = get_tasks(pre_dir, [("rel-f1", "driver-dnf")], ("test",),
                  embedder=cfg["embedder"])
ev = build_evaluator(
    tasks, pre_dir,
    embedder=cfg["embedder"], d_text=cfg["d_text"], device=device,
    ctx_size=128, local_ctx_size=64, bfs_width=32, prefer_latest=True,
    num_walks=10_000, walk_length=20, tokens_per_gpu=2**18,
    items_per_task=5, num_workers=0, prefetch_factor=2,
    shuffle_seed=0, context_seed=0, mmap_populate=True, vector_db_path=None,
)

# evaluate_raw yields one (task, ctx, labels, preds, n) per task
results = ev.evaluate_raw([(model, "")], [128])
_task, _ctx, _labels, out, _n = next(iter(results))
preds = torch.sigmoid(torch.tensor(out[""], dtype=torch.float32))
print("driver-dnf probability:", [round(p, 3) for p in preds.tolist()])
```

> [!NOTE]
> `items_per_task=5` and `ctx_size=128` keep this demo quick (most of the runtime
> is one-time warmup). On a GPU, raise `ctx_size` toward RT-J's training context of
> 8192 (with `local_ctx_size <= ctx_size`) for full accuracy over the whole test
> split.

## Bring your own database

Point RT at your **own** database, define a
prediction task, and infer with a released checkpoint:

- **Colab, no setup**: the [fully worked notebook](byod/colab.ipynb)
  ([open in Colab](https://colab.research.google.com/github/stanford-star/relational-transformer/blob/main/byod/colab.ipynb))
  runs the whole flow end-to-end on your database (or the bundled demo).

## Development

We use [pixi](https://pixi.sh) to manage one self-contained
environment (Python, PyTorch + CUDA, Rust, and all dependencies), built on first
use. There is nothing to build past `pixi install`: the native Rust extension is
compiled as part of the project's own editable install.

```bash
git clone https://github.com/stanford-star/relational-transformer.git
cd relational-transformer
pixi run pytest                        # the test suite
pixi run python examples/train.py      # or eval.py, preprocess.py, ...
```

## Documentation

| Guide | Description |
|---|---|
| [Downloads](docs/downloads.md) | Bulk-download raw data, preprocessed data, and checkpoints from HuggingFace |
| [Preprocess](docs/preprocess.md) | Convert RelBench-format databases into RT's on-disk format |
| [Inference](docs/inference.md) | Run a trained checkpoint; evaluate, engineer, tune, and ensemble contexts |
| [Pretrain](docs/train.md) | Train RT from scratch, single-GPU to multi-node |
| [Experiments](expts/README.md) | How experiments in `expts/` are laid out and submitted to slurm |

There is no CLI: RT is a library, and a run is a script that calls it. Copy
something from [`examples/`](examples/) and edit it — every entry point takes
its arguments explicitly, so nothing is hidden in a default you did not choose.

