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.

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 and Python 3.12+:

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:

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 (open in Colab) runs the whole flow end-to-end on your database (or the bundled demo).

Development

We use pixi 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.

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

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